import MNN
import cv2
import numpy as np
import logging

class Face_Recognizer(object):
    def __init__(self, model_path):
        logging.info('******** Start Init Face ID ********')
        self.reg_interpreter = MNN.Interpreter(model_path)
        self.reg_session = self.reg_interpreter.createSession()
        self.reg_input_tensor = self.reg_interpreter.getSessionInput(self.reg_session)
        logging.info('******** Success Init Face ID  ********')

    def recognize(self, imgs):
        
        feats = []
        for i in range(len(imgs)):
            img = imgs[i] 

            img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
            img = img.astype(np.float)
            img = (img / 255. - 0.5) / 0.5
            img = img.transpose(2, 0, 1)
            img = np.expand_dims(img, axis=0)
        
            input_tensor = MNN.Tensor((1, 3, 112, 112), MNN.Halide_Type_Float, img, MNN.Tensor_DimensionType_Caffe)
            self.reg_input_tensor.copyFrom(input_tensor)
            self.reg_interpreter.runSession(self.reg_session)
            output_tensor = self.reg_interpreter.getSessionOutput(self.reg_session, 'output0')
            output = output_tensor.getData()

            feats.append(output)

        feats_np = np.array(feats)
        return feats_np