from emotion import EmotionRekognition
from face import FaceRekognition
from utils import draw_bounding_box, draw_text, read_image, show_image

# 人脸检测模型
detection_model = 'models/shape_predictor_68_face_landmarks.dat'

# 表情分类模型
classification_model = 'models/emotion_predictor_FERPlus_mini_XCEPTION.h5'
# 表情分类标签
classification_labels = ['neutral', 'happiness', 'surprise', 'sadness', 'anger', 'disgust', 'fear', 'contempt']


def rekognition(image):
    # 人脸检测
    face_rekognition = FaceRekognition(detection_model)
    # 表情分类
    emotion_rekognition = EmotionRekognition(classification_model, classification_labels)
    # 遍历人脸，识别表情
    face_details = [{
        'BoundingBox': face['BoundingBox'],
        'Emotions': emotion_rekognition.predict(face['FaceImage'])
    } for face in face_rekognition.predict(image)]

    return face_details


if __name__ == '__main__':
    image_path = 'images/test.jpg'

    image_array = read_image(image_path)
    result = rekognition(image_array)
    for item in result:
        box = list(item['BoundingBox'].values())
        text = list(item['Emotions'].keys())[0]
        draw_bounding_box(box, image_array, (0, 0, 255))
        draw_text(box, image_array, text, (0, 0, 255), 0, -20, 1, 2)
        print(item)
    show_image(image_array)