#coding=utf-8
from keras.engine.saving import load_model
import cv2
import numpy as np
import rospy

from load_dataset import resize_image, IMAGE_SIZE
from train_model import labels_name
rospy.init_node("keras_img_class",anonymous=True)
from std_msgs.msg import String
from geometry_msgs.msg import Twist
cmd_vel_pub=rospy.Publisher("cmd_vel",Twist,queue_size=10)
cmd_vel_data=Twist()

MODEL_PATH = './model/model.h5'
# 加载模型
model = load_model(MODEL_PATH)
cap = cv2.VideoCapture(2)
# 实时识别
while True:
    # 从摄像头读取图片
    success, img = cap.read()
    # img = cv2.flip(img, 1)
    img2=img.copy()
    img = cv2.rectangle(img, (212, 112), (468, 368), (0, 255, 0), 2)
    traffic_sign = img2[112:368,212:468]

    image = resize_image(traffic_sign, IMAGE_SIZE, IMAGE_SIZE)
    image = image.reshape((1, IMAGE_SIZE, IMAGE_SIZE, 3))

    # 浮点并归一化
    image = image.astype('float32')
    image /= 255.0
    # 给出输入属于各个类别的概率，我们是二值类别，则该函数会给出输入图像属于0和1的概率各为多少
    rst = model.predict(image)
    # 给出类别预测
    result = model.predict_classes(image)
    cls = labels_name[int(result)]
    if cls=="left":
        cmd_vel_data.angular.z=0.6
        cmd_vel_data.linear.x=0
        cmd_vel_pub.publish(cmd_vel_data)
    if cls=="right":
        cmd_vel_data.angular.z=-0.6
        cmd_vel_data.linear.x=0
        cmd_vel_pub.publish(cmd_vel_data)
    if cls=="stop":
        cmd_vel_data.angular.z=0
        cmd_vel_data.linear.x=0
        cmd_vel_pub.publish(cmd_vel_data)
    if cls=="forward":
        cmd_vel_data.angular.z=0
        cmd_vel_data.linear.x=0.15
        cmd_vel_pub.publish(cmd_vel_data)
    # if cls=="null":
    #     cmd_vel_data.angular.z=0
    #     cmd_vel_data.linear.x=0
    #     cmd_vel_pub.publish(cmd_vel_data)

    img = cv2.putText(img, cls, (212, 100), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 2)
    cv2.imshow('traffic_sign', img)   # imshow文件不能是中文，处理报错

    k = cv2.waitKey(1)
    if k == ord('q'):
        break

cap.release()
cv2.destroyAllWindows()