# -*- coding: utf-8 -*-
"""
分类器示例
"""

import cv2
import numpy as np
import migraphx

def Preprocessing(pathOfImage):
    image = cv2.imread(pathOfImage,cv2.IMREAD_COLOR)             # cv2.IMREAD_COLOR:彩色图，cv2.IMREAD_GRAYSCALE:灰度图
    image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
    image = cv2.resize(image, (224,224))
    image = image.transpose(2, 0, 1)
    image = np.expand_dims(image, 0)
    image = np.ascontiguousarray(image)
    image = image.astype(np.float32)
    input = image / 255
    return input

if __name__ == '__main__':
    # 加载模型
    model = migraphx.parse_onnx("../Resource/Models/resnet50-v2-7.onnx")
    inputName=model.get_parameter_names()[0]
    inputShape=model.get_parameter_shapes()[inputName].lens()
    print("inputName:{0} \ninputShape:{1}".format(inputName,inputShape))

    # FP16
    # migraphx.quantize_fp16(model)

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

    # 预处理并转换为NCHW
    pathOfImage ="../Resource/Images/ImageNet_01.jpg"
    image = Preprocessing(pathOfImage)

    # 推理
    results = model.run({inputName: image}) # 推理结果，list类型

    # 获取输出节点属性
    result=results[0] # 获取第一个输出节点的数据,migraphx.argument类型
    outputShape=result.get_shape() # 输出节点的shape,migraphx.shape类型
    outputSize=outputShape.lens() # 每一维大小，维度顺序为(N,C,H,W),list类型
    numberOfOutput=outputShape.elements() # 输出节点元素的个数

    # 获取分类结果
    result=results[0].tolist() # 将migraphx.argument转换为list
    result=np.array(result)

    # 打印1000个类别的输出
    print(result)