from tensorflow.keras.preprocessing.image import img_to_array
from tensorflow.keras.models import load_model
import numpy as np
import pickle
import cv2
import time
import os


def resize(image, width=None, height=None, inter=cv2.INTER_AREA):
    dim = None
    (h, w) = image.shape[:2]
    # 如果高和宽为None则直接返回
    if width is None and height is None:
        return image
    # 检查宽是否是None
    if width is None:
        # 计算高度的比例并并按照比例计算宽度
        r = height / float(h)
        dim = (int(w * r), height)
    # 高为None
    else:
        # 计算宽度比例，并计算高度
        r = width / float(w)
        dim = (width, int(h * r))
    resized = cv2.resize(image, dim, interpolation=inter)
    # return the resized image
    return resized

# 构建脸部/活体识别检测器
class Face_inference():
    def __init__(self):
        self.face_detector_path = "./models/face_detector/"
        self.recognize_embed_path = "./models/nn4.small2.v1.t7"
        self.liveness_model_path = "./models/liveness.model"
        self.recognize_pickle = "./pickles/face_label/recognizer.pickle"
        self.recognize_class_pickle = "./pickles/face_label/le.pickle"
        self.liveness_class_pickle = "./pickles/liveness/le.pickle"
        
        # 人脸检测
        self.face_detector = cv2.dnn.readNetFromCaffe(
            self.face_detector_path + "deploy.prototxt",
            self.face_detector_path + "res10_300x300_ssd_iter_140000.caffemodel"
        )

        # 人脸embedding
        self.face_embed_model = cv2.dnn.readNetFromTorch(self.recognize_embed_path)

        # embedding SVM 分类
        with open(self.recognize_pickle,"rb") as f:
            self.face_recog_model = pickle.loads(f.read())
         
        # 活体分类
        self.liveness_model = load_model(self.liveness_model_path)

        # 对应类别标签

        # 人脸类别
        with open(self.recognize_class_pickle,"rb") as f:
            self.face_class = pickle.loads(f.read()).classes_
        
        # 活体类别
        with open(self.liveness_class_pickle,"rb") as f:
            self.liveness_class = pickle.loads(f.read()).classes_

    def face_detect(self,img,conf_threshold = 0.1,is_show = False):
        if img is None:
            return img
        frame = resize(img,width = 600)
        h,w = frame.shape[:2]
        blob = cv2.dnn.blobFromImage(
            cv2.resize(frame,(300,300)),1.0,(300,300),
            (104.0,177.0,123.0),swapRB = False,crop = False
        )

        self.face_detector.setInput(blob)
        detections = self.face_detector.forward()

        for i in range(detections.shape[2]):
            conf = detections[0,0,i,2]
            if conf > conf_threshold:
                iou = detections[0,0,i,3:7] * np.array([w,h,w,h])
                (start_x,start_y,end_x,end_y) = iou.astype('int')
                face = frame[start_y:end_y,start_x:end_x]

                fh,fw = face.shape[:2]
                if fh < 20 or fw < 20:
                    continue
                face_blob = cv2.dnn.blobFromImage(
                    face,1.0/255.,(96,96),
                    (0,0,0),swapRB = True , crop = False
                )

                self.face_embed_model.setInput(face_blob)
                vec = self.face_embed_model.forward()

                preds = self.face_recog_model.predict_proba(vec)[0]
                pred_index = np.argmax(preds)
                pred_class,pred_prob = self.face_class[pred_index],preds[pred_index]
                pred_class = pred_class if pred_prob > conf_threshold else "err"

                if is_show:
                    text = f"{pred_class}: {pred_prob*100:.2f}%"
                    cv2.rectangle(frame,(start_x,start_y),(end_x,end_y),
                                  (0,0,255,2))
                    y = start_y - 10 if start_y - 10 > 10 else start_y + 10
                    cv2.putText(frame,text,(start_x,y),cv2.FONT_HERSHEY_SIMPLEX, 0.45, (0, 0, 255), 2)

                    return pred_class,pred_prob,frame
                else:
                    return pred_class,pred_prob
                

        
    def live_detect(self,img,name = "",conf_threshold = 0.1,is_show = False):
        if img is None:
            return img
        frame = resize(img,width=600)
        h,w = frame.shape[:2]
        blob = cv2.dnn.blobFromImage(
            cv2.resize(frame,(300,300)),1.0,(300,300),
            (104.0,177.0,123.0),swapRB = False,crop = False
        )

        self.face_detector.setInput(blob)
        detections = self.face_detector.forward()
                
        for i in range(detections.shape[2]):
            conf = detections[0,0,i,2]
            if conf > conf_threshold:
                iou = detections[0,0,i,3:7] * np.array([w,h,w,h])
                (start_x,start_y,end_x,end_y) = iou.astype('int')
                face = frame[start_y:end_y,start_x:end_x]

                face = cv2.resize(face,(32,32))
                face = face.astype('float')/255.
                face = img_to_array(face)
                face = np.expand_dims(face,axis = 0)

                preds = self.liveness_model(face)[0]
                index = np.argmax(preds)
                pred_class,pred_prob = self.liveness_class[index],preds[index]

                if is_show:
                    text = f"{name} , {pred_class}: {pred_prob*100:.4f}%"
                    cv2.rectangle(frame,(start_x,start_y),(end_x,end_y),
                                  (0,0,255,2))
                    y = start_y - 10 if start_y - 10 > 10 else start_y + 10
                    cv2.putText(frame,text,(start_x,y),cv2.FONT_HERSHEY_SIMPLEX, 0.45, (0, 0, 255), 2)

                    return pred_class,pred_prob,frame
                else:
                    return pred_class,pred_prob


if __name__ == "__main__":
    model = Face_inference()
    img = cv2.imread("R.jpg")
    name,prob,res_frame = model.face_detect(img,is_show=True)
    cv2.imshow("face",res_frame)
    label,prob,res_frame = model.live_detect(img,name,is_show=True)
    cv2.imshow("live",res_frame)
    cv2.waitKey(0)