import os
import cv2
import MNN
import logging
import numpy as np

class Abnormal_Face(object):

    def __init__(self, model_path):
        logging.info('******** Start Init Abnormal Face ********')
        self.interpreter = MNN.Interpreter(model_path)
        self.session = self.interpreter.createSession()
        self.input_tensor = self.interpreter.getSessionInput(self.session)
        logging.info('******** Success Init Abnormal Face ********')

    def reg(self, image, bbox):

        image_h, image_w, image_c = image.shape
        face_area = image[bbox[1]:bbox[3], bbox[0]:bbox[2], ]

        face_h, face_w, _ = face_area.shape
        image_resize = cv2.resize(face_area, (128, 128), interpolation=cv2.INTER_CUBIC)
        image_transpose = image_resize.astype(np.float32).transpose((2, 0, 1)) / 255
        input_tensor = MNN.Tensor((1, 3, 128, 128), MNN.Halide_Type_Float, image_transpose, MNN.Tensor_DimensionType_Caffe)
        self.input_tensor.copyFrom(input_tensor)
        self.interpreter.runSession(self.session)
        output_tensor = self.interpreter.getSessionOutput(self.session)
        output_data = output_tensor.getData()
        output_data = output_data[0]
        output_data = output_data.tolist()
        predict_res = output_data.index(max(output_data))
       
        return predict_res