from PyQt5.QtGui import *
from PyQt5.QtCore import *
from PyQt5.QtWidgets import *
import threading
import sys
import torch.backends.cudnn as cudnn
from models.experimental import *
from utils.datasets_notprint import *
from utils.utils import *


# 窗口主类
class MainWindow(QTabWidget):
    # 基本配置不动，然后只动第三个界面
    def __init__(self):
        # 初始化设置
        super().__init__()
        self.setWindowTitle('红外车辆检测系统')
        self.resize(1100, 700)
        self.setWindowIcon(QIcon("UI_images/logo.png"))
        # 图像检索
        self.input_fname = "UI_images/about2.jpg"

        # 视频读取相关设置
        self.source = 'images/cam2.mkv'
        self.single_source = ""
        self.stopEvent = threading.Event()
        self.stopEvent.clear()
        self.isCamera = True
        self.initUI()
        self.set_down()

    def initUI(self):
        # 主布局
        main_widget = QWidget()
        generally_font = QFont('楷体', 15)
        main_layout = QVBoxLayout()
        font_v = QFont('楷体', 14)

        '''
        *** 3.图片检索界面 ***
        '''
        up_widget = QWidget()
        up_layout = QHBoxLayout()
        # 图片检索左上区域
        up_left_widget = QWidget()
        up_left_layout = QVBoxLayout()
        title_0 = QLabel("Query Image")
        self.input_img_label = QLabel()
        self.output_imge_label = QLabel()
        input_src_img = cv2.imread(self.input_fname)
        input_img = cv2.resize(input_src_img, (360, 400))
        cv2.imwrite("UI_images/input.jpg", input_img)
        self.input_img_label.setPixmap(QPixmap("UI_images/input.jpg"))
        self.output_imge_label.setPixmap(QPixmap("UI_images/input.jpg"))
        btn_upload = QPushButton("上传图片📂")
        btn_upload.setFont(generally_font)
        btn_upload.setStyleSheet("QPushButton{color:white}"
                                 "QPushButton:hover{background-color: rgb(2,110,180);}"
                                 "QPushButton{background-color:rgb(48,124,208)}"
                                 "QPushButton{border:2px}"
                                 "QPushButton{border-radius:5px}"
                                 "QPushButton{padding:5px 5px}"
                                 "QPushButton{margin:5px 5px}")
        btn_upload.setCursor(QCursor(Qt.PointingHandCursor))
        btn_upload.clicked.connect(self.get_img_file)
        up_left_layout.addWidget(title_0)
        up_left_layout.addWidget(self.input_img_label, 0, Qt.AlignCenter)
        up_left_layout.addWidget(btn_upload)
        up_left_widget.setLayout(up_left_layout)
        # 图片检索右上区域
        up_right_widget = QWidget()
        up_right_layout = QVBoxLayout()
        title_1 = QLabel("Search Result")
        btn_search = QPushButton("开始检索🔭")
        btn_search.setFont(generally_font)
        btn_search.setStyleSheet("QPushButton{color:white}"
                                 "QPushButton:hover{background-color: rgb(2,110,180);}"
                                 "QPushButton{background-color:rgb(48,124,208)}"
                                 "QPushButton{border:2px}"
                                 "QPushButton{border-radius:5px}"
                                 "QPushButton{padding:5px 5px}"
                                 "QPushButton{margin:5px 5px}")
        btn_search.setCursor(QCursor(Qt.PointingHandCursor))
        btn_search.clicked.connect(self.to_search)
        up_right_layout.addWidget(title_1)
        up_right_layout.addWidget(self.output_imge_label, 0, Qt.AlignCenter)
        up_right_layout.addWidget(btn_search)
        up_right_widget.setLayout(up_right_layout)
        up_layout.addWidget(up_left_widget)
        up_layout.addWidget(up_right_widget)
        up_widget.setLayout(up_layout)
        # 图片检索下方区域

        '''
        *** 4. 视频识别界面 ***
        '''
        video_widget = QWidget()
        video_layout = QVBoxLayout()
        self.video_title2 = QLabel("车辆检测区")
        self.video_title2.setFont(font_v)
        self.video_title2.setAlignment(Qt.AlignCenter)
        self.video_title2.setFont(font_v)
        self.DisplayLabel = QLabel()
        self.DisplayLabel.setPixmap(QPixmap(""))

        self.btn_open = QPushButton("开始识别（选择文件）")
        self.btn_open.setFont(font_v)
        self.btn_open.setStyleSheet("QPushButton{color:white}"
                                    "QPushButton:hover{background-color: rgb(2,110,180);}"
                                    "QPushButton{background-color:rgb(48,124,208)}"
                                    "QPushButton{border:2px}"
                                    "QPushButton{border-radius:5px}"
                                    "QPushButton{padding:5px 5px}"
                                    "QPushButton{margin:5px 5px}")
        self.btn_close = QPushButton("结束演示")
        self.btn_close.setFont(font_v)
        self.btn_close.setStyleSheet("QPushButton{color:white}"
                                     "QPushButton:hover{background-color: rgb(2,110,180);}"
                                     "QPushButton{background-color:rgb(48,124,208)}"
                                     "QPushButton{border:2px}"
                                     "QPushButton{border-radius:5px}"
                                     "QPushButton{padding:5px 5px}"
                                     "QPushButton{margin:5px 5px}")
        # self.btn_open_sure.clicked.connect(self.open_sure)
        self.btn_open.clicked.connect(self.open)
        self.btn_close.clicked.connect(self.close)
        video_layout.setAlignment(Qt.AlignCenter)
        video_layout.addWidget(self.video_title2)
        video_layout.addWidget(self.DisplayLabel)
        # video_layout.addWidget(self.btn_open_sure)
        video_layout.addWidget(self.btn_open)
        video_layout.addWidget(self.btn_close)
        video_widget.setLayout(video_layout)

        '''
        *** 5. 关于界面 ***
        '''
        about_widget = QWidget()
        about_layout = QVBoxLayout()
        about_title = QLabel('欢迎使用红外车辆检测系统😀')
        about_title.setFont(QFont('楷体', 18))
        about_title.setAlignment(Qt.AlignCenter)
        about_img = QLabel()
        about_img.setPixmap(QPixmap('UI_images/about3.jpg'))
        about_img.setAlignment(Qt.AlignCenter)
        label_super = QLabel()
        label_super.setText("<a href='https://www.baidu.com/'>我的个人主页</a>")
        label_super.setFont(QFont('楷体', 12))
        label_super.setOpenExternalLinks(True)
        label_super.setAlignment(Qt.AlignRight)
        about_layout.addWidget(about_title)
        about_layout.addStretch()
        about_layout.addWidget(about_img)
        about_layout.addStretch()
        about_layout.addWidget(label_super)
        about_widget.setLayout(about_layout)

        # 主页面设置
        main_layout.addWidget(up_widget)
        main_widget.setLayout(main_layout)

        self.addTab(main_widget, '图片检测')
        self.addTab(video_widget, '视频检测')
        self.addTab(about_widget, '关于')
        self.setTabIcon(0, QIcon('UI_images/图片.png'))
        self.setTabIcon(1, QIcon('UI_images/直播.png'))
        self.setTabIcon(3, QIcon('UI_images/logo_about.png'))

    def closeEvent(self, event):
        reply = QMessageBox.question(self,
                                     '退出',
                                     "是否要退出程序？",
                                     QMessageBox.Yes | QMessageBox.No,
                                     QMessageBox.No)
        if reply == QMessageBox.Yes:
            self.close()
            event.accept()
        else:
            event.ignore()

    '''
    ### 4. 视频识别相关功能 ### 
    '''
    def open_sure(self):
        # 也可以设置成摄像头
        # TODO 修改文件夹位置
        self.fileName = 'images/cam2.mkv'
        self.cap = cv2.VideoCapture(self.fileName)
        th = threading.Thread(target=self.display_video)
        th.start()

    def open(self):
        # 选择录像文件进行读取
        mp4_fileName, fileType = QFileDialog.getOpenFileName(self, 'Choose file', '', '*.mp4, *.mkv')
        if mp4_fileName:
            self.source = mp4_fileName
            th = threading.Thread(target=self.display_video)
            th.start()

    def close(self):
        self.stopEvent.set()
        self.set_down()

    def display_video(self):
        # 首先把打开按钮关闭
        self.btn_open.setEnabled(False)
        # self.btn_open_sure.setEnabled(False)
        self.btn_close.setEnabled(True)

        # todo 这里执行显示的逻辑
        webcam =False
        agnostic_nms = False
        augment = False
        classes = None
        conf_thres = 0.4
        device = ''
        img_size = 640
        iou_thres = 0.5

        view_img = False
        weights = './weights/best.pt'
        # Initialize
        device = torch_utils.select_device(device)
        # if os.path.exists(out):
        #     shutil.rmtree(out)  # delete output folder
        # os.makedirs(out)  # make new output folder
        half = device.type != 'cpu'  # half precision only supported on CUDA

        # Load model
        model = attempt_load(weights, map_location=device)  # load FP32 model
        imgsz = check_img_size(img_size, s=model.stride.max())  # check img_size
        if half:
            model.half()  # to FP16

        # Second-stage classifier
        classify = False
        if classify:
            modelc = torch_utils.load_classifier(name='resnet101', n=2)  # initialize
            modelc.load_state_dict(torch.load('weights/resnet101.pt', map_location=device)['model'])  # load weights
            modelc.to(device).eval()

        # Set Dataloader
        vid_path, vid_writer = None, None
        if webcam:
            view_img = True
            cudnn.benchmark = True  # set True to speed up constant image size inference
            dataset = LoadStreams(self.source, img_size=imgsz)
        else:
            save_img = True
            dataset = LoadImages(self.source, img_size=imgsz)

        # Get names and colors
        names = model.module.names if hasattr(model, 'module') else model.names
        colors = [[random.randint(0, 255) for _ in range(3)] for _ in range(len(names))]

        # Run inference
        t0 = time.time()
        img = torch.zeros((1, 3, imgsz, imgsz), device=device)  # init img
        _ = model(img.half() if half else img) if device.type != 'cpu' else None  # run once
        for path, img, im0s, vid_cap in dataset:
            img = torch.from_numpy(img).to(device)
            img = img.half() if half else img.float()  # uint8 to fp16/32
            img /= 255.0  # 0 - 255 to 0.0 - 1.0
            if img.ndimension() == 3:
                img = img.unsqueeze(0)

            # Inference
            t1 = torch_utils.time_synchronized()
            pred = model(img, augment=augment)[0]

            # Apply NMS
            pred = non_max_suppression(pred, conf_thres, iou_thres, classes=classes, agnostic=agnostic_nms)
            t2 = torch_utils.time_synchronized()

            # Apply Classifier
            if classify:
                pred = apply_classifier(pred, modelc, img, im0s)

            # Process detections
            for i, det in enumerate(pred):  # detections per image
                if webcam:  # batch_size >= 1
                    p, s, im0 = path[i], '%g: ' % i, im0s[i].copy()
                else:
                    p, s, im0 = path, '', im0s

                # save_path = str(Path(out) / Path(p).name)
                # txt_path = str(Path(out) / Path(p).stem) + ('_%g' % dataset.frame if dataset.mode == 'video' else '')
                # s += '%gx%g ' % img.shape[2:]  # print string
                gn = torch.tensor(im0.shape)[[1, 0, 1, 0]]  # normalization gain whwh
                if det is not None and len(det):
                    # Rescale boxes from img_size to im0 size
                    det[:, :4] = scale_coords(img.shape[2:], det[:, :4], im0.shape).round()

                    # Print results
                    for c in det[:, -1].unique():
                        n = (det[:, -1] == c).sum()  # detections per class
                        s += '%g %ss, ' % (n, names[int(c)])  # add to string
                    # Write results
                    for *xyxy, conf, cls in det:
                        # print(names[int(cls)])
                        # if names[int(cls)] == "person":
                            # if save_img or view_img:  # Add bbox to image
                        label = '%s %.2f' % (names[int(cls)], conf)
                        plot_one_box(xyxy, im0, label=label, color=(0, 0, 255), line_thickness=3)
                frame = im0
                frame_height = frame.shape[0]
                frame_width = frame.shape[1]
                frame_scale = 500 / frame_height
                frame_resize = cv2.resize(frame, (int(frame_width * frame_scale), int(frame_height * frame_scale)))
                # frame_final = cv2.cvtColor(frame_resize, cv2.COLOR_RGB2BGR)
                # cv2.imshow("x", frame_resize)
                cv2.imwrite("tmp.jpg", frame_resize)
                # frame_final = frame_resize
                # img = QImage(frame_final.data, frame_final.shape[1], frame_final.shape[0], QImage.Format_RGB888)
                self.DisplayLabel.setPixmap(QPixmap("tmp.jpg"))
                # print(self.stopEvent.is_set())
            if cv2.waitKey(25) & self.stopEvent.is_set() == True:
                self.stopEvent.clear()
                self.DisplayLabel.clear()
                self.btn_close.setEnabled(False)
                self.btn_open.setEnabled(True)
                # self.btn_open_sure.setEnabled(True)
                self.set_down()
                break
        self.btn_open.setEnabled(True)
        # self.btn_open_sure.setEnabled(True)
        self.btn_close.setEnabled(False)
        self.set_down()

    def set_down(self):
        self.DisplayLabel.setPixmap(QPixmap("UI_images/ae862.jpg"))


    def get_img_file(self):
        # 上传图片
        openfile_name = QFileDialog.getOpenFileName(self, '选择文件', '', 'Image files(*.jpg , *.png)')
        img_name = openfile_name[0]
        if img_name == '':
            pass
        else:
            src_img = cv2.imread(img_name)
            src_img_height = src_img.shape[0]
            src_img_width = src_img.shape[1]
            target_img_height = 400
            ratio = target_img_height / src_img_height
            target_img_width = int(src_img_width * ratio)
            target_img = cv2.resize(src_img, (target_img_width, target_img_height))
            cv2.imwrite("UI_images/input_single.jpg", target_img)
            self.input_img_label.setPixmap(QPixmap("UI_images/input_single.jpg"))
            self.single_source = img_name

    def to_search(self):
        if self.single_source == "":
            QMessageBox.information(self, "警告", "请先上传图片！")
        else:
            # todo 这里执行显示的逻辑
            webcam = False
            agnostic_nms = False
            augment = False
            classes = None
            conf_thres = 0.4
            device = ''
            img_size = 640
            iou_thres = 0.5

            view_img = False
            weights = './weights/best.pt'
            # Initialize
            device = torch_utils.select_device(device)
            # if os.path.exists(out):
            #     shutil.rmtree(out)  # delete output folder
            # os.makedirs(out)  # make new output folder
            half = device.type != 'cpu'  # half precision only supported on CUDA

            # Load model
            model = attempt_load(weights, map_location=device)  # load FP32 model
            imgsz = check_img_size(img_size, s=model.stride.max())  # check img_size
            if half:
                model.half()  # to FP16

            # Second-stage classifier
            classify = False
            if classify:
                modelc = torch_utils.load_classifier(name='resnet101', n=2)  # initialize
                modelc.load_state_dict(torch.load('weights/resnet101.pt', map_location=device)['model'])  # load weights
                modelc.to(device).eval()

            # Set Dataloader
            vid_path, vid_writer = None, None
            if webcam:
                view_img = True
                cudnn.benchmark = True  # set True to speed up constant image size inference
                dataset = LoadStreams(self.single_source, img_size=imgsz)
            else:
                save_img = True
                dataset = LoadImages(self.single_source, img_size=imgsz)

            # Get names and colors
            names = model.module.names if hasattr(model, 'module') else model.names
            colors = [[random.randint(0, 255) for _ in range(3)] for _ in range(len(names))]

            # Run inference
            t0 = time.time()
            img = torch.zeros((1, 3, imgsz, imgsz), device=device)  # init img
            _ = model(img.half() if half else img) if device.type != 'cpu' else None  # run once
            for path, img, im0s, vid_cap in dataset:
                img = torch.from_numpy(img).to(device)
                img = img.half() if half else img.float()  # uint8 to fp16/32
                img /= 255.0  # 0 - 255 to 0.0 - 1.0
                if img.ndimension() == 3:
                    img = img.unsqueeze(0)

                # Inference
                t1 = torch_utils.time_synchronized()
                pred = model(img, augment=augment)[0]

                # Apply NMS
                pred = non_max_suppression(pred, conf_thres, iou_thres, classes=classes, agnostic=agnostic_nms)
                t2 = torch_utils.time_synchronized()

                # Apply Classifier
                if classify:
                    pred = apply_classifier(pred, modelc, img, im0s)

                # Process detections
                for i, det in enumerate(pred):  # detections per image
                    if webcam:  # batch_size >= 1
                        p, s, im0 = path[i], '%g: ' % i, im0s[i].copy()
                    else:
                        p, s, im0 = path, '', im0s

                    gn = torch.tensor(im0.shape)[[1, 0, 1, 0]]  # normalization gain whwh
                    if det is not None and len(det):
                        # Rescale boxes from img_size to im0 size
                        det[:, :4] = scale_coords(img.shape[2:], det[:, :4], im0.shape).round()

                        # Print results
                        for c in det[:, -1].unique():
                            n = (det[:, -1] == c).sum()  # detections per class
                            s += '%g %ss, ' % (n, names[int(c)])  # add to string
                        # Write results
                        for *xyxy, conf, cls in det:
                            # print(names[int(cls)])
                            # if names[int(cls)] == "person":
                                # if save_img or view_img:  # Add bbox to image
                            label = '%s %.2f' % (names[int(cls)], conf)
                            plot_one_box(xyxy, im0, label=label, color=(0, 0, 255), line_thickness=3)
                    frame = im0
                    frame_height = frame.shape[0]
                    frame_width = frame.shape[1]
                    frame_scale = 400 / frame_height
                    frame_resize = cv2.resize(frame, (int(frame_width * frame_scale), int(frame_height * frame_scale)))
                    frame_final = cv2.cvtColor(frame_resize, cv2.COLOR_RGB2BGR)
                    cv2.imwrite("images/single_detect_result.jpg", frame_final)
                    self.output_imge_label.setPixmap(QPixmap("images/single_detect_result.jpg"))


if __name__ == "__main__":
    app = QApplication(sys.argv)
    mainWindow = MainWindow()
    mainWindow.show()
    sys.exit(app.exec_())