#!/usr/bin/env python
# coding: utf-8

from PyQt5 import QtWidgets
from PyQt5.QtWidgets import *
from PyQt5.QtCore import *
from PyQt5.QtGui import *

# import Ui_MainWindow as ui
import ui_test as ui

from uvc_thermal import Thermal
from camera import Camera
# from detector import FaceDetector
from mask_detector import FaceMaskDetector
from anti import AntiSpoofing
from heater import Heater
from aligner import BoxAligner
from correcter import Correcter
from osd import OSD

import cv2
import numpy as np
import time

class Main(QMainWindow, ui.Ui_MainWindow):
    def __init__(self):
        super().__init__()
        self.setupUi(self)
        self.resize(640,480)
        self.rgb_frame.setScaledContents(True)
        self.thermal_frame.setScaledContents(True)
        self.thermal = Thermal(width=160, height=120, framerate=9, frame_width=160, frame_height=120, log=None)
        self.camera = Camera(width=640, height=480, framerate=30, log=None, channel=0, flip=1)
#         self.detector = FaceDetector(model_path="../Ultra-Light-Fast-Generic-Face-Detector-1MB/MNN/model/version-RFB/RFB-320.mnn" , input_size=(320,240))
        self.detector = FaceMaskDetector(model_path="/home/allen/RetinaFaceMask/mnn_models/RetinaFaceMask_mobileNet0.25.mnn" , input_size=(640,640))
        # self.anti = AntiSpoofing(model_path="../Ultra-Light-Fast-Generic-Face-Detector-1MB/MNN/model/4_0_0_80x80_MiniFASNetV1SE.mnn")
        self.aligner = BoxAligner(img_h = 480,img_w = 640)
        self.heater = Heater(pwm_pin=12 , target_temp=35.0)
        self.correcter = Correcter(model_path="../thermal-tk/20210421_correcter.sav")
        self.osd = OSD(local='Asia/Taipei', str_format='%Y-%m-%d %H:%M:%S')
        self.thermal.start()
        time.sleep(2)
        self.camera.start()
        time.sleep(2)
        self.heater.start()
        time.sleep(2)
        self.osd.start()
        time.sleep(2)
        self.faces = []
        self.timer = QTimer(self)
        self.timer.timeout.connect(self.capture)
        self.timer.start(10)
#         self.painter = QPainter(self)
    def capture(self):
        ret, frame = self.camera.capture()
        thermal_ret,thermal_frame, thermal_row = self.thermal.capture()
        if(thermal_ret and ret):
            blackbody_max = thermal_row[:,:34].max()
            heater_temp = self.heater.last_temp
            # 偵測人臉
            str_time = self.osd.str_time
            boxes, labels, probs, mask_probs = self.detector.predict(frame)
            # 取出所有偵測的結果
            frame=cv2.cvtColor(frame,cv2.COLOR_BGR2RGB)
            thermal_frame=cv2.cvtColor(thermal_frame,cv2.COLOR_BGR2RGB)
            for i in range(boxes.shape[0]):
                box = boxes[i, :]
                score = probs[i]
                mask_score = mask_probs[i]
                if(score > 0.9):
#                     color = (255, 0, 0)
                    thermal_box = self.aligner.box_aligment([box])[0]
                    face_frame = frame[box[1]:box[3],box[0]:box[2]].copy()
                    Qface_frame=QImage(face_frame.data,face_frame.shape[1],face_frame.shape[0],face_frame.shape[1]*3,QImage.Format_RGB888)
                    face_pixMap=QPixmap.fromImage(Qface_frame)
                    
                    face_temp = self.max_temperature(thermal_box,thermal_row,black_h=32)
                    corrected_face_temp = self.correcter.predict(x_detect_k= face_temp, x_black_k= blackbody_max, x_black_c= heater_temp)
                    
                    if(mask_score > 0.7):
                        bbox_color = (0, 255, 0)
                    else:
                        bbox_color = (255, 0, 0)
                    
                    if(corrected_face_temp > 37.0):
                        color = (255, 0, 0)
                    elif(corrected_face_temp < 35.0):
                        color = (255, 255, 0)
                    else:
                        color = (0, 255, 0)
                    
                    cv2.rectangle(thermal_frame, (thermal_box[0]//4, thermal_box[1]//4), (thermal_box[2]//4, thermal_box[3]//4), bbox_color, 2)
                    cv2.rectangle(frame, (box[0], box[1]), (box[2], box[3]), bbox_color, 2)
                    cv2.putText(frame, "{:.2f}->{:.2f}".format(face_temp*0.0092,corrected_face_temp), ((box[0]+box[2])//2, (box[1]+box[3])//2), cv2.FONT_HERSHEY_COMPLEX_SMALL,1, color, 1, cv2.LINE_AA)
                    
                    cv2.putText(thermal_frame, "{:.2f}".format(heater_temp), (320//4, 64//4), cv2.FONT_HERSHEY_COMPLEX_SMALL,1, (0, 0, 0), 1, cv2.LINE_AA)
                    
                    self.face1.setPixmap(face_pixMap)
                    self.temp1.setStyleSheet("color: rgb{};".format(color))
                    self.temp1.setText("{:.2f}".format(corrected_face_temp))
                    self.time1.setText(str_time)
            # 顯示至UI
#             frame=cv2.cvtColor(frame,cv2.COLOR_BGR2RGB)
#             thermal_frame=cv2.cvtColor(thermal_frame,cv2.COLOR_BGR2RGB)
            Qframe=QImage(frame.data,frame.shape[1],frame.shape[0],frame.shape[1]*3,QImage.Format_RGB888)
            Qthermal_frame=QImage(thermal_frame.data,thermal_frame.shape[1],thermal_frame.shape[0],thermal_frame.shape[1]*3,QImage.Format_RGB888)
            pixMap=QPixmap.fromImage(Qframe)
            thermal_pixMap=QPixmap.fromImage(Qthermal_frame)
            self.rgb_frame.setPixmap(pixMap)
            self.thermal_frame.setPixmap(thermal_pixMap)
            self.update()
    def closeEvent(self,event):
        # 關閉程式後執行
        self.thermal.stop()
        self.camera.stop()
        self.heater.stop()
        self.osd.stop()
    def max_temperature(self,box,thermal_row,black_h=0,thermal_height=120,thermal_width=160,image_height=480,image_width=640):
        scale_ratio_h = thermal_height/image_height
        scale_ratio_w = thermal_width/image_width
        x1 = max(int(box[0]*scale_ratio_w),0)
        y1 = max(int(box[1]*scale_ratio_h),black_h)
        x2 = min(int(box[2]*scale_ratio_w),thermal_width)
        y2 = min(int(box[3]*scale_ratio_h),thermal_height)
        box_temperature = thermal_row[y1:y2,x1:x2]
        if(box_temperature.size != 0):
            return box_temperature.max()
        else:
            return 0
if __name__ == '__main__':
    import sys
    import os
    app = QtWidgets.QApplication(sys.argv)
    window = Main()
    window.show()
#     sys.exit(app.exec_())
    status = app.exec_()
    os._exit(status)






