import torch
from PIL import Image
import torchvision.transforms as transforms
from torch.autograd import Variable
from torch import nn
import matplotlib.pyplot as plt
import cv2
import time
import os
from model import resnet34
import json
import serial
import Find_COM


def get_picture(cap):  # 获取照片
    # 捕获一帧的数据
    ret, frame = cap.read()
    if frame is None:
        print(frame)
    if ret:
        # 默认不阻塞
        cv2.imshow("picture", frame)
        cv2.waitKey(1)
    # 数据帧写入图片中
    label = "1"
    timeStamp = 1381419600
    image_name = str(int(time.time())) + ".jpg"
    # 照片存储位置
    filepath = "Input/" + image_name  # 改成跟上面一样的位置
    str_name = filepath.replace('%s', label)
    cv2.imwrite(str_name, frame)    # 将照片保存起来
    return image_name


def start_move_1(port, baudrate):  # 快速加酸程序
    ser = serial.Serial(port, baudrate)
    data = b"q1h15d"  # 每分钟转15圈，一个比慢滴略高的旋转速度
    ser.write(data)
    time.sleep(0.01)
    data = b"q5h1d"  # 转1秒
    ser.write(data)
    time.sleep(0.01)
    data = b"q6h3d"  # 逆时针
    ser.write(data)
    time.sleep(20)  # 等待20秒，为快滴时间，注意，这里等待时间不能少于1秒（阀门开启的时间）
    data = b"q6h2d"  # 顺时针
    ser.write(data)
    time.sleep(1) # 转回去的时间
    ser.close()


def start_move_2(port, baudrate):  # 缓慢加酸程序
    ser = serial.Serial(port, baudrate)
    data = b"q1h14d"  # 每分钟转14圈，每秒转14/60=0.233圈，可以根据实际情况调整
    ser.write(data)
    time.sleep(0.01)
    # 注意：实际上我们也可以修改速度的小数部分，如下列注释所示
    # data = b"q2h50d"  # 结合q1指令，每分钟转14.5圈，可以根据实际情况调整
    # ser.write(data)
    # time.sleep(0.01)
    data = b"q5h1d"  # 转1秒
    ser.write(data)
    time.sleep(0.01)
    data = b"q6h3d"  # 逆时针
    ser.write(data)
    time.sleep(1)  # 转阀门的时间
    # 注意，这里没有将阀门转回去，而是持续几秒钟滴加一次的状态
    ser.close()


def start_move_3(port, baudrate):  # 停止加酸程序
    ser = serial.Serial(port, baudrate)
    data = b"q1h14d"  # 每分钟转14圈，需要与move2的速度保持一致
    ser.write(data)
    time.sleep(0.01)
    data = b"q5h1d"  # 转1秒
    ser.write(data)
    time.sleep(0.01)
    data = b"q6h2d"  # 顺时针
    ser.write(data)
    time.sleep(1)  # 转阀门的时间
    # 将阀门转回去
    ser.close()


def main():
    # port = "COM6"  # 串口名，根据实际情况修改
    port = Find_COM.list_ch340_ports()[0]  # 串口名，根据实际情况修改
    baudrate = 9600  # 波特率，根据实际情况修改
    # # 快速滴加过程，这里请自己根据滴加量优化
    # start_move_1(port, baudrate)
    # time.sleep(15)
    videoSourceIndex = 0  # 摄像机编号，请根据自己的情况调整
    cap = cv2.VideoCapture(videoSourceIndex, cv2.CAP_DSHOW)  # 打开摄像头
    # 是否用GPU
    device = torch.device( "cuda:0" if torch.cuda.is_available() else "cpu")
    start_move_2(port, baudrate)  # 开启慢滴状态
    while True:
        # 读取图片
        name = get_picture(cap)
        # 图片完整路径
        im_file = 'Input/' + name
        # 使用PIL库打开图片
        image = Image.open(im_file)
        # print(type(image)) # 打印图片的类型
        # 定义图片预处理流程
        data_transform = transforms.Compose(
            [
                # 调整图片大小为256x256
                transforms.Resize(256),
                # 从中心裁剪出224x224大小的图片
                transforms.CenterCrop(224),
                # 将图片转换为PyTorch的Tensor格式
                transforms.ToTensor(),
                # 对图片进行归一化，使用ImageNet的均值和标准差
                transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
            ])
        # [N, C, H, W]
        img = data_transform(image)
        # expand batch dimension
        img = torch.unsqueeze(img, dim=0)
        # 定义模型权重文件的路径
        json_path = './class_indices.json'

        with open(json_path, "r") as f:
            class_indict = json.load(f)
        # create model
        model = resnet34(num_classes=2).to(device)

        # load model weights
        weights_path = "./resnet34-1Net.pth"
        assert os.path.exists(weights_path), "file: '{}' dose not exist.".format(weights_path)
        model.load_state_dict(torch.load(weights_path, map_location=device, weights_only=True))

        # prediction
        model.eval()
        with torch.no_grad():
            # predict class
            output = torch.squeeze(model(img.to(device))).cpu()
            # 对预测结果进行softmax，得到每个类别的概率
            predict = torch.softmax(output, dim=0)
            # 找到概率最大的类别的索引
            predict_cla = torch.argmax(predict).numpy()
        # 根据索引从类别字典中获取类别名称
        class_a = "{}".format(class_indict[str(predict_cla)])
        # 格式化概率值，保留三位小数
        prob_a = "{:.3}".format(predict[predict_cla].numpy())
        # 将概率值转换为浮点数
        prob_b = float(prob_a)
        # 打印预测的类别和概率
        print(class_a)
        print(prob_b)
        if class_a == "orange" and prob_b >= 0.5:  # 到达滴定终点
            # 关闭阀门
            start_move_3(port, baudrate)
            print('----->>End<<-----')
            print(im_file)
            time.sleep(1)
            # 释放摄像头
            cap.release()
            # 关闭所有OpenCV窗口
            cv2.destroyAllWindows()

            break
        time.sleep(1)  # 拍照间隔


if True:
    main()