import sys
import cv2
import numpy as np

sys_path = sys.path[0].replace("\\", "/")

# 加载人脸检测模型
tensorflowWeightFile = sys_path + "/opencv_face_detector.pbtxt"
tensorflowModelFile = sys_path + "/opencv_face_detector_uint8.pb"

conf_threshold = 0.7

net = cv2.dnn.readNetFromTensorflow(tensorflowModelFile, tensorflowWeightFile)

def process_face_detection(image : cv2.typing.MatLike, fps : int = 10) :
    # 读取图片
    #image = cv2.imread(sys_path + "/example.jpg")
    height, width, _ = image.shape
    
    # # 准备输入数据
    # blob = cv2.dnn.blobFromImage(cv2.resize(image, (300, 300)), 1.0,
    #                             (300, 300), (104.0, 177.0, 123.0))
    
    # 设置参数
    scalefactor = 1.0
    size = (300, 300)
    mean = (104.0, 177.0, 123.0)
    swapRB = False
    crop = False
    ddepth = cv2.CV_32F
    
    # 将图像转换为blob
    blob = cv2.dnn.blobFromImage(cv2.resize(image, (320, 240)), scalefactor, size, mean, swapRB, crop, ddepth)
 
    net.setInput(blob)

    # 运行网络
    detections = net.forward()

    text = "fps : %d"%(fps)
    cv2.putText(image, text, org = (50, 50), fontFace = cv2.FONT_HERSHEY_SIMPLEX, fontScale = 1, color = (0, 255, 0), thickness = 1)

    # 解析检测结果
    for i in range(detections.shape[2]):
        confidence = detections[0, 0, i, 2]
        if confidence > 0.5:
            box = detections[0, 0, i, 3:7] * np.array([width, height, width, height])
            (startX, startY, endX, endY) = box.astype("int")
            
            # 画出检测到的人脸
            cv2.rectangle(image, (startX, startY), (endX, endY), (0, 255, 0), 2)
            
            
if __name__ == "__main__":
    # 显示图片
    image = cv2.imread(sys_path + "/example.jpg")
    process_face_detection(image)

    cv2.imshow("Face Detection", image)
    cv2.waitKey(0)
    cv2.destroyAllWindows()