from pathlib import Path
from typing import Union

import cv2
import numpy
from keras.models import load_model

from utils import read_image


class EmotionRekognition:

    def __init__(self, model: Union[str, Path], labels: Union[list, tuple]):
        self.model = None
        self.labels = None
        self.load_model(model, labels)

    def load_model(self, model: Union[str, Path], labels: Union[list, tuple]):
        """ 加载表情识别模型
        :param model: 模型的路径（字符串路径或路径对象）
        :param labels: 表情标签（字符串列表或字符串元组）
        """
        self.model = load_model(model)
        self.labels = labels

    def _preprocess_image(self, image: Union[str, Path, bytes, numpy.ndarray]):
        """ 读取并预处理图像
        :param image: 图像的路径（字符串路径或路径对象）或图像的字节数据或NumPy数组
        :return: 预处理后的图像（NumPy数组）
        """
        if not isinstance(image, numpy.ndarray):
            image = read_image(image)

        if self.model.input_shape[-1] == 1:
            image = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)

        image = cv2.resize(image, self.model.input_shape[1:-1])
        image = numpy.array(image).reshape(*self.model.input_shape[1:])
        image = image / 255.0
        image = numpy.expand_dims(image, axis=0)

        return image

    def predict(self, image: Union[str, Path, bytes, numpy.ndarray]):
        """ 使用加载的模型对图像进行表情预测
        :param image: 图像的路径（字符串路径或路径对象）或图像的字节数据或NumPy数组
        :return: 表情预测结果，按置信度从高到低排序的字典
        :raises RuntimeError: 如果未加载模型
        """
        if self.model is None:
            raise RuntimeError('未加载模型')

        result = {}
        predictions = self.model.predict(self._preprocess_image(image))

        for prediction in predictions:
            for index, value in enumerate(prediction):
                result[self.labels[index]] = value
            result = dict(sorted(result.items(), key=lambda item: item[1], reverse=True))
            break

        return result


if __name__ == '__main__':
    classification_model = 'models/emotion_predictor_FERPlus_mini_XCEPTION.h5'
    classification_labels = ['neutral', 'happiness', 'surprise', 'sadness', 'anger', 'disgust', 'fear', 'contempt']

    emotion_rekognition = EmotionRekognition(classification_model, classification_labels)
    rekognition_result = emotion_rekognition.predict('images/happy.jpg')
    print(rekognition_result)