# -*- coding: UTF-8 -*-
import aidlite_gpu
import cv2
import os
import time
import numpy as np
from PIL import Image

import matplotlib.pyplot as plt
def mask_to_image(mask: np.ndarray):
    if mask.ndim == 2:
        return Image.fromarray((mask * 255).astype(np.uint8))
    elif mask.ndim == 3:
        return Image.fromarray((np.argmax(mask, axis=0) * 255 / mask.shape[0]).astype(np.uint8))


def aidlux_tflite_infer(model_path, img_path, save_path):
    # step1: 初始化aidlite类并创建aidlite对象
    aidlite = aidlite_gpu.aidlite()
    print('model initial success!!')

    # step2: 加载模型
    inp_shape = [256*256*1*4]
    out_shape = [256*256*2*4]
    value = aidlite.ANNModel(model_path, inp_shape, out_shape, 4, 0) 
    # step3: 传入模型输入数据
    img = cv2.imread(img_path, 0)
    img = cv2.resize(img, (256, 256))
    img = img[np.newaxis, ...]
    img = img / 255.0
    img = np.expand_dims(img, axis=0)
    img = img.astype(dtype=np.float32)
    print("image shape is ", img.shape)
    aidlite.setInput_Float32(img)
    
    # step4: 执行推理
    start = time.time()
    aidlite.invoke()
    end = time.time()
    print("infer time(ms):{0}", 1000 * (end - start))
    # step5: 获取输出
    pred = aidlite.getOutput_Float32(0)
    # step6: 后处理
    pred = np.array(pred)
    pred = np.reshape(pred,(2,256,256))
    mask_img = mask_to_image(pred)
    
    mask_img.save(save_path) 
    # mask_img = np.array(mask_img)  
    # cv2.imshow('mask_img', mask_img)
    # cv2.waitKey(0)
    # cv2.destroyAllWindows() 
    
if __name__ == '__main__':
    model_path = "/home/dataset2aidlux/unetmodel_fp32.tflite"
    img_path = "/home/dataset2aidlux/test_imgs/0597.PNG"
    save_path = '/home/dataset2aidlux/test_imgs/result_0597.png'
    aidlux_tflite_infer(model_path, img_path, save_path)