from PIL import Image
from util import *
import json

detector = FaceDetector('./weights/detection.onnx')
fer = HSEmotionRecognizer('./weights/emotion.onnx')

def detect_face(frame):
    boxes = detector.detect(frame, (640, 640))
    return boxes if boxes is not None and len(boxes) else None


def porcess_video(input_video_path):
    stream = cv2.VideoCapture(input_video_path)
    frame_width = int(stream.get(3))
    frame_height = int(stream.get(4))
    fps = stream.get(cv2.CAP_PROP_FPS)
    
    idx_class = fer.idx_to_class
    frame_counter = 0
    emotion_reg_list = []
    emotion_score_list = []
    
    while True:
        ret, frame = stream.read()
        if ret:
            frame_counter += 1
            ## 每30帧进行处理一次
            if frame_counter % 30 == 0:
                image = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
                boxes = detect_face(image)
                if boxes is not None:
                    for box in boxes.astype('int32'):
                        x1, y1, x2, y2 = box[:4]
                        face_image = image[y1:y2, x1:x2]
                        pil_image = Image.fromarray(face_image).convert('RGB')
                        pil_image = pil_image.resize((224, 224))  # Adjust size as per model requirements
                        emotion, scores,pred = fer.predict_emotions(np.array(pil_image), logits=False)
                        ###### 做应用逻辑
                        print("frame_id:",frame_counter)
                        print("emotion",emotion)
                        print("scores",scores[pred])
                        emotion_reg_list.append(emotion)
                        emotion_score_list.append(float(scores[pred]))
                        print("idx_class",idx_class)
                        ######
                        draw_emotion_bars(fer, frame, scores, (x2 + 10, y1), bar_height=15, width=100)
        else:
            print("退出")
            break

    stream.release()
    
    
    data = {
        "video_url":input_video_path,
        "emotion_reg_list":emotion_reg_list,
        "emotion_score_list":emotion_score_list
    }
    json_str = json.dumps(data)
    print("json_str:",json_str)
    return json_str

if __name__ == '__main__':
    #自己录制一段视频
    video_path = "/Users/apple/Downloads/教学代码/111111.mp4"
    porcess_video(video_path)
    
#结果如下
# json_str: {"video_url": "/Users/apple/Downloads/\u6559\u5b66\u4ee3\u7801/111111.mp4", "emotion_reg_list": ["NEUTRAL", "HAPPINESS"], "emotion_score_list": [0.5329781174659729, 0.43945103883743286]}