import cv2
import numpy as np
import MNN
import os

def image_infer_mnn(mnn_model_path, image_path, class_list):
    image = cv2.imread(image_path)
    input_image = cv2.resize(image,(128,128))
    input_data = input_image.astype(np.float32).transpose((2, 0, 1)) / 255
    interpreter = MNN.Interpreter(mnn_model_path)
    session = interpreter.createSession()
    input_tensor = interpreter.getSessionInput(session)
    tmp_input = MNN.Tensor((1, 3, 128, 128), MNN.Halide_Type_Float, input_data, MNN.Tensor_DimensionType_Caffe)
    input_tensor.copyFrom(tmp_input)
    interpreter.runSession(session)
    infer_result = interpreter.getSessionOutput(session)
    output_data = infer_result.getData()
    out = output_data.tolist()
    out = out.index(max(out))
    print(out)
    cv2.putText(image,class_list[int(out)],(50, 50),cv2.FONT_HERSHEY_SIMPLEX,2,(0,0,255))
    return image


image_path = r'input/0000/0.jpg'
mnn_model_path = r'models/cls_abnormal_face_mnn_1.0.0_v0.0.1.mnn'
class_list = ['mask', 'no_mask']

image = image_infer_mnn(mnn_model_path, image_path, class_list)