import sys
import cv2
import numpy as np
import mediapipe as mp

sys_path = sys.path[0].replace("\\", "/")
image = cv2.imread(sys_path + "/example.jpg")

# 初始化BlazeFace模型
mp_face_detection = mp.solutions.face_detection
face_detection = mp_face_detection.FaceDetection(min_detection_confidence=0.5, model_selection=2)

def process_face_detection(image : cv2.typing.MatLike, fps : int = 10) :
    image_width  = image.shape[1]
    image_height = image.shape[0]

    # 检测人脸
    results = face_detection.process(image)
    total_score = 0
    
    if results.detections:
        for face in results.detections:
            relative_box = face.location_data.relative_bounding_box
            
            # print(f'Detection score: {detection.score}')
            # print(f'Box location: {relative_box}')

            if face.score[0] < 0.55 :
                continue
            
            total_score += face.score[0]
            
            # 在图片上绘制方框
            xmin = relative_box.xmin
            ymin = relative_box.ymin
            ymax = relative_box.height + ymin
            xmax = relative_box.width + xmin
            
            # ymin, xmin, ymax, xmax = detection.location_data.relative_bounding_box
            cv2.rectangle(image, (int(xmin * image_width), int(ymin * image_height)),
                        (int(xmax * image_width), int(ymax * image_height)), (0, 255, 0), 2)
            
            key_points = np.array([(p.x, p.y) for p in face.location_data.relative_keypoints])
            key_points_coords = np.multiply(key_points,[image_width,image_height]).astype(int)
            for p in key_points_coords:
                cv2.circle(image, p, 4, (255, 255, 255), 2)
                cv2.circle(image, p, 2, (0, 0, 0), -1)

        average_score = total_score / len(results.detections)
        text = "fps : %d, score = %f"%(fps, average_score)
    else :
        text = "fps :%d"%(fps)
    cv2.putText(image, text, org = (50, 50), fontFace = cv2.FONT_HERSHEY_SIMPLEX, fontScale = 1, color = (255, 255, 0), thickness = 1)
 
if __name__ == "__main__":
    # 显示图片
    image = cv2.imread(sys_path + "/example.jpg")
    process_face_detection(image)

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