import cv2
import threading
from math import *
import numpy as np
from PyQt5.QtCore import QFile
from PyQt5.QtWidgets import QFileDialog, QMessageBox
from PyQt5.QtGui import QImage, QPixmap
import time
import shutil
import os
import torch
from torchvision import datasets, models, transforms
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader
import numpy as np
import matplotlib.pyplot as plt
import models
import config
from PIL import Image
import scipy
import scipy.io

model_path = r'trained_models\last.pth'
test_valid_transforms = transforms.Compose([
    # transforms.RandomResizedCrop(size=224, scale=(0.82, 1.0)),#随机裁剪到256*256
    transforms.Resize((224,224)),
    # transforms.CenterCrop(224),
    # transforms.FiveCrop(224),
    transforms.ToTensor(),
    # transforms.Normalize([0.485, 0.456, 0.406],
    #  [0.229, 0.224, 0.225])
    transforms.Normalize([0.44540259, 0.47563922, 0.5588167],#归一化
                         [0.21878081, 0.23083013, 0.25116897])
])

def rotateImage(img,degree,pt1,pt2,pt3,pt4):
    height,width=img.shape[:2]
    heightNew = int(width * fabs(sin(radians(degree))) + height * fabs(cos(radians(degree))))
    widthNew = int(height * fabs(sin(radians(degree))) + width * fabs(cos(radians(degree))))
    matRotation=cv2.getRotationMatrix2D((width/2,height/2),degree,1)
    matRotation[0, 2] += (widthNew - width) / 2
    matRotation[1, 2] += (heightNew - height) / 2
    imgRotation = cv2.warpAffine(img, matRotation, (widthNew, heightNew), borderValue=(255, 255, 255))
    pt1 = list(pt1)
    pt3 = list(pt3)
    [[pt1[0]], [pt1[1]]] = np.dot(matRotation, np.array([[pt1[0]], [pt1[1]], [1]]))
    [[pt3[0]], [pt3[1]]] = np.dot(matRotation, np.array([[pt3[0]], [pt3[1]], [1]]))
    imgOut=imgRotation[int(pt1[1]):int(pt3[1]),int(pt1[0]):int(pt3[0])]
    # cv2.imshow("imgOut",imgOut) #裁减得到的旋转矩形框
    # cv2.imwrite("imgOut.jpg",imgOut)
    return imgOut


class Display:
    def __init__(self, ui, mainWnd):
        self.ui = ui
        self.mainWnd = mainWnd
        # 默认视频源为相机
        self.isCamera = True
        #信号槽设置
        ui.detect.clicked.connect(self.Detect)
        ui.pushButton_3.clicked.connect(self.update_model)
        #显示模板图片
        self._img_model = cv2.imread("img_model.jpg")
        if self._img_model.shape[1] > self._img_model.shape[0]:
            self._img_model = cv2.rotate(self._img_model, cv2.cv2.ROTATE_90_CLOCKWISE)
        img_model = cv2.cvtColor(self._img_model, cv2.COLOR_RGB2BGR)#BGR转RGB
        Q_img_model = QImage(img_model.data, img_model.shape[1], img_model.shape[0],
                             img_model.shape[1]*3, QImage.Format_RGB888)
        self.ui.img_model.setPixmap(QPixmap.fromImage(Q_img_model))

        # 创建一个关闭事件并设为未触发
        self.stopEvent = threading.Event()
        self.stopEvent.clear()
        self.Open()
        #载入网络
        self.lode_net()

    def lode_net(self):#导入网络
        self.resnet50 = torch.load(model_path).cuda()
        self.resnet50.eval()  # test
    def classify(self,fream):#神经网络对芯片进行分类的函数
        device = torch.device(
            "cuda:0" if torch.cuda.is_available() else "cpu")  # 若有gpu可用则用gpu
        img = Image.fromarray(cv2.cvtColor(fream, cv2.COLOR_BGR2RGB))
        img = test_valid_transforms(img)  # 图片预处理函数
        input = img.cuda().view(1, 3, 224, 224)  #
        outputs = self.resnet50(input)
        # 类别输出:predictions 0:err 1:true
        ret, predictions = torch.max(outputs.data, 1)
        if predictions.data == 1:
            predictions1 = 1
        if predictions.data == 0:
            predictions1 = 0
        print(str(predictions1) + '\n')
        return predictions1
    def Open(self):#打开摄像头
        self.cap = cv2.VideoCapture(1)
        self.cap.set(15, -11)#
        # self.cap.set(3, 1280)
        # self.cap.set(4, 1024)
        # 创建视频显示线程
        th = threading.Thread(target=self.Display)
        th.start()

    def Close(self):
        # 关闭事件设为触发，关闭视频播放
        self.stopEvent.set()

    def Display(self):#图像实时显示
        while self.cap.isOpened():
            success, frame = self.cap.read()
            # RGB转BGR
            frame = cv2.cvtColor(frame, cv2.COLOR_RGB2BGR)
            img = QImage(frame.data, frame.shape[1], frame.shape[0], QImage.Format_RGB888)
            self.ui.DispalyLabel.setPixmap(QPixmap.fromImage(img))

            if self.isCamera:
                cv2.waitKey(1)
            else:
                cv2.waitKey(int(1000 / self.frameRate))
            # 判断关闭事件是否已触发
            if True == self.stopEvent.is_set():
                # 关闭事件置为未触发，清空显示label
                self.stopEvent.clear()
                self.ui.DispalyLabel.clear()
                print("close")
                break
    def Detect(self):#检测函数，先用神经网络进行分类，判断是否残缺品，若是残缺品则通过机器视觉检测出残缺处并显示
        self.ui.result.setText("检测中")
        success, fream = self.cap.read()#读取一帧视频图像
        if self.classify(fream) == 0:
            # 读取制作模板和待检测的原图
            kernel = np.ones((3, 3), np.uint8)
            img_model = self._img_model
            # fream = cv2.imread("imgOut.jpg")
            fream = cv2.rotate(fream, cv2.cv2.ROTATE_90_CLOCKWISE)
            fream_copy = fream.copy()
            # 灰度化
            gray_img_fream = cv2.cvtColor(fream, cv2.COLOR_BGR2GRAY)
            # 二值化
            ret_, gray_img_model_thresh_fream = cv2.threshold(
                gray_img_fream, 155, 255, cv2.THRESH_BINARY_INV)
            cv2.imshow("ss", gray_img_model_thresh_fream)
            # 腐蚀膨胀去除噪点
            image_contours_mask_fream = cv2.erode(gray_img_model_thresh_fream, kernel, iterations=1)
            image_contours_mask_fream = cv2.dilate(image_contours_mask_fream, np.ones((7, 7), np.uint8), iterations=1)

            #寻找包围框
            image_contours_fream, det_image_fream, hierarchy_ = cv2.findContours(
                image_contours_mask_fream, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)  # 找方框
            #寻找最大包围框
            if len(det_image_fream) != 0:
                area_fream = []
                for k in range(len(det_image_fream)):
                    det_img_new_ = cv2.drawContours(fream_copy, det_image_fream[k], -1, (0, 0, 255), 3)
                    area_fream.append(cv2.contourArea(det_image_fream[k]))
                # cv2.imshow("all", fream_copy)
                max_area_contours_fream = np.argmax(area_fream)
                rect_fream = cv2.minAreaRect(det_image_fream[max_area_contours_fream])

                center_fream, (h_fream, w_fream), degree_fream = rect_fream  # （中心(x,y), (宽,高), 旋转角度）
                box_fream = cv2.boxPoints(rect_fream)  # 获取最小外接矩形的4个顶点坐标
                box_fream = np.int0(box_fream)
                cv2.drawContours(fream_copy, [box_fream], 0, (0, 0, 255), 5)#在图像上画框，框出芯片
                # cv2.imshow("max", fream_copy)

                imgOut = self._img_model#获取模板
                #旋转通过外接芯片矩形的四个顶点和旋转角度裁剪出芯片
                imgOut_fream = rotateImage(fream, -degrees(atan2(abs(box_fream[0][1] - box_fream[1][1]),
                                                                 abs(box_fream[0][0] - box_fream[1][0]))), box_fream[0],
                                                                         box_fream[1], box_fream[2], box_fream[3])
                #通过长宽判断芯片裁剪方向是否正确，不正确则旋转90度
                if imgOut_fream.shape[0] < imgOut_fream.shape[1]:
                    imgOut_fream = cv2.rotate(imgOut_fream, cv2.cv2.ROTATE_90_CLOCKWISE)
                #把芯片resize到模板的尺寸
                high_imgOut, wigth_imgOut, _ = imgOut.shape
                imgOut_fream = cv2.resize(imgOut_fream, (wigth_imgOut,high_imgOut))
                #将模板和裁剪的芯片图像灰度化
                imgOut_gray = cv2.cvtColor(imgOut, cv2.COLOR_BGR2GRAY)
                imgOut_fream_gray = cv2.cvtColor(imgOut_fream, cv2.COLOR_BGR2GRAY)
                #二值化
                _, imgOut_gray = cv2.threshold(
                    imgOut_gray, 120, 255, cv2.THRESH_BINARY)  # 二值化
                _, imgOut_fream_gray = cv2.threshold(
                    imgOut_fream_gray, 120, 255, cv2.THRESH_BINARY)  # 二值化
                #模板与图像相减
                result = imgOut_gray - imgOut_fream_gray
                #腐蚀膨胀去除噪点
                result = cv2.erode(result, np.ones((7, 7), np.uint8), iterations=1)
                result = cv2.dilate(result, np.ones((7, 7), np.uint8), iterations=1)
                # _, result = cv2.threshold(
                #     result, 200, 255, cv2.THRESH_BINARY)  # 二值化
                #寻找残缺处的包围圈
                image_contours, det_image_result, hierarchy = cv2.findContours(
                    result, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_NONE)  # 找方框
                area_fream = []
                if len(det_image_result) != 0:
                    for k in range(len(det_image_result)):
                        if cv2.contourArea(det_image_result[k]) > 40:
                            det_img_new = cv2.drawContours(imgOut_fream, det_image_result[k], -1, (0, 0, 255), 3)
                            area_fream.append(cv2.contourArea(det_image_result[k]))
                    imgOut_fream = cv2.cvtColor(imgOut_fream, cv2.COLOR_RGB2BGR)
                    img = QImage(imgOut_fream.data, imgOut_fream.shape[1], imgOut_fream.shape[0], imgOut_fream.shape[1]*3, QImage.Format_RGB888)
                    self.ui.img_model_2.setPixmap(QPixmap.fromImage(img))
                    # cv2.imshow("result", result)  # 裁减得到的旋转矩形框
                    # cv2.imshow("imgOut", imgOut)  # 裁减得到的旋转矩形框
                    # cv2.imshow("imgOut_fream_gray", imgOut_fream_gray)  # 裁减得到的旋转矩形框
                    # cv2.imshow("1111", fream_copy)
                    # cv2.imshow("imgOut_fream", imgOut_fream)
                if len(area_fream)==0:
                    self.ui.result.setText("标准品")
                else:
                    self.ui.result.setText("瑕疵品")
        else:
            kernel = np.ones((3, 3), np.uint8)
            fream_copy = fream.copy()
            # 灰度化并保存
            gray_img_fream = cv2.cvtColor(fream, cv2.COLOR_BGR2GRAY)

            ret_, gray_img_model_thresh_fream = cv2.threshold(
                gray_img_fream, 150, 255, cv2.THRESH_BINARY_INV)  # 二值化
            image_contours_mask_fream = cv2.erode(gray_img_model_thresh_fream, kernel, iterations=1)
            image_contours_mask_fream = cv2.dilate(image_contours_mask_fream, np.ones((7, 7), np.uint8), iterations=1)

            image_contours_fream, det_image_fream, hierarchy_ = cv2.findContours(
                image_contours_mask_fream, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)  # 找方框
            if len(det_image_fream) != 0:
                area_fream = []
                for k in range(len(det_image_fream)):
                    det_img_new_ = cv2.drawContours(fream_copy, det_image_fream[k], -1, (0, 0, 255), 3)
                    area_fream.append(cv2.contourArea(det_image_fream[k]))
                # print(area)

                max_area_contours_fream = np.argmax(area_fream)
                rect_fream = cv2.minAreaRect(det_image_fream[max_area_contours_fream])

                center_fream, (h_fream, w_fream), degree_fream = rect_fream  # （中心(x,y), (宽,高), 旋转角度）
                box_fream = cv2.boxPoints(rect_fream)  # 获取最小外接矩形的4个顶点坐标
                box_fream = np.int0(box_fream)
                cv2.drawContours(fream_copy, [box_fream], 0, (0, 0, 255), 5)
                imgOut_fream = rotateImage(fream, -degrees(atan2(abs(box_fream[0][1] - box_fream[1][1]),
                                                                 abs(box_fream[0][0] - box_fream[1][0]))), box_fream[0],
                                           box_fream[1], box_fream[2], box_fream[3])
                # cv2.imshow("imgOut_fream1", imgOut_fream)  # 裁减得到的旋转矩形框
                if imgOut_fream.shape[0] < imgOut_fream.shape[1]:
                    imgOut_fream = cv2.rotate(imgOut_fream, cv2.cv2.ROTATE_90_CLOCKWISE)
                # cv2.imshow("imgOut_fream", imgOut_fream)  # 裁减得到的旋转矩形框
                # self._img_model = imgOut_fream
                img_show = cv2.cvtColor(imgOut_fream, cv2.COLOR_RGB2BGR)
                Q_img_model = QImage(img_show.data, img_show.shape[1], img_show.shape[0],
                                     img_show.shape[1] * 3, QImage.Format_RGB888)
                self.ui.img_model_2.setPixmap(QPixmap.fromImage(Q_img_model))
                self.ui.result.setText("标准品")

    def update_model(self):
        success, fream = self.cap.read()
        # 读取制作模板和待检测的原图
        kernel = np.ones((3, 3), np.uint8)
        # fream = cv2.imread("1.jpg")
        fream_copy = fream.copy()
        # 灰度化并保存
        gray_img_fream = cv2.cvtColor(fream, cv2.COLOR_BGR2GRAY)

        ret_, gray_img_model_thresh_fream = cv2.threshold(
            gray_img_fream, 155, 255, cv2.THRESH_BINARY_INV)  # 二值化
        image_contours_mask_fream = cv2.erode(gray_img_model_thresh_fream, kernel, iterations=1)
        image_contours_mask_fream = cv2.dilate(image_contours_mask_fream, np.ones((7, 7), np.uint8), iterations=1)

        image_contours_fream, det_image_fream, hierarchy_ = cv2.findContours(
            image_contours_mask_fream, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)  # 找方框
        if len(det_image_fream) != 0:
            area_fream = []
            for k in range(len(det_image_fream)):
                det_img_new_ = cv2.drawContours(fream_copy, det_image_fream[k], -1, (0, 0, 255), 3)
                area_fream.append(cv2.contourArea(det_image_fream[k]))
            # print(area)
            max_area_contours_fream = np.argmax(area_fream)
            rect_fream = cv2.minAreaRect(det_image_fream[max_area_contours_fream])

            center_fream, (h_fream, w_fream), degree_fream = rect_fream  # （中心(x,y), (宽,高), 旋转角度）
            box_fream = cv2.boxPoints(rect_fream)  # 获取最小外接矩形的4个顶点坐标
            box_fream = np.int0(box_fream)
            cv2.drawContours(fream_copy, [box_fream], 0, (0, 0, 255), 5)
            imgOut_fream = rotateImage(fream, -degrees(atan2(abs(box_fream[0][1] - box_fream[1][1]),
                                                             abs(box_fream[0][0] - box_fream[1][0]))), box_fream[0],
                                       box_fream[1], box_fream[2], box_fream[3])
            cv2.imshow("imgOut_fream1", imgOut_fream)  # 裁减得到的旋转矩形框
            if imgOut_fream.shape[0] < imgOut_fream.shape[1]:
                imgOut_fream = cv2.rotate(imgOut_fream, cv2.cv2.ROTATE_90_CLOCKWISE)
            cv2.imshow("imgOut_fream", imgOut_fream)  # 裁减得到的旋转矩形框
            self._img_model = imgOut_fream
            img_model = cv2.cvtColor(self._img_model, cv2.COLOR_RGB2BGR)
            Q_img_model = QImage(img_model.data, img_model.shape[1], img_model.shape[0],
                                 img_model.shape[1] * 3, QImage.Format_RGB888)
            self.ui.img_model.setPixmap(QPixmap.fromImage(Q_img_model))
            cv2.imwrite("img_model.jpg", self._img_model)