import numpy as np
from PIL import Image
import cv2
import migraphx

def Preprocess(pil_img, newW, newH):
    assert newW > 0 and newH > 0, 'Scale is too small'
    pil_img = pil_img.resize((newW, newH))
    img_nd = np.array(pil_img)
    if len(img_nd.shape) == 2:
        img_nd = np.expand_dims(img_nd, axis=2)

    # HWC to CHW
    img_print = pil_img
    img_trans = img_nd.transpose((2, 0, 1))
    if img_trans.max() > 1:
        img_trans = img_trans / 255
    img_trans = np.expand_dims(img_trans, 0)

    return img_trans, img_print

def Sigmoid(x):
  return 1 / (1 + np.exp(-x))

if __name__ == '__main__':
    
    # 加载模型
    model = migraphx.parse_onnx("./model/unet_13_256.onnx")
    inputName = model.get_parameter_names()
    inputShape = model.get_parameter_shapes()
    print("inputName:{0} \ninputShape:{1}".format(inputName, inputShape))

    # 编译模型
    model.compile(migraphx.get_target("gpu"), device_id=0)   # device_id: 设置GPU设备，默认为0号设备

    # 图像预处理
    img = Image.open("./car.jpeg")
    img, imPrint = Preprocess(img, 256, 256)
    input_img = np.zeros((1,3,256,256),dtype='float32') 
    np.lib.stride_tricks.as_strided(input_img, shape=img.shape, strides=input_img.strides)[:] = img

    # 模型推理
    mask = model.run({'inputs':input_img})      
    output_mask = np.array(mask[0])[0]                        # 获取推理结果，shape为(1,256,256)
    probs = Sigmoid(output_mask)                              # 计算sigmoid值

     # 0/1像素值，当大于0.996时，值为255，小于等于0.996时，值为0
    output_mask[probs > 0.996] = 255
    output_mask[probs <= 0.996] = 0

    output = output_mask.astype(np.uint8)[0]                  # 将浮点型转换为uint8整型，shape为（256，256）
    cv2.imwrite("output.jpg", output)                         # 保存图像分割结果