# -*- coding: utf-8 -*-
"""
offload为false时的分类器示例
"""

import argparse
import cv2
import numpy as np
import migraphx

def AllocateOutputMemory(model):
    outputData = {}
    for key in model.get_outputs().keys():
        outputData[key] = migraphx.allocate_gpu(s=model.get_outputs()[key])

    return outputData

def Preprocessing(pathOfImage):
    # 读取图像
    image = cv2.imread(pathOfImage, cv2.IMREAD_COLOR)             
    image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
    
    # 调整大小，使短边为256，保持长宽比
    ratio = float(256) / min(image.shape[0], image.shape[1])
    if image.shape[0] > image.shape[1]:
        new_size = [int(round(ratio * image.shape[0])), 256]
    else:
        new_size = [256, int(round(ratio * image.shape[1]))]
    image = np.array(cv2.resize(image, (new_size[1],new_size[0])))
    
    # 裁剪中心窗口为224*224
    h, w, c = image.shape
    start_x = w//2 - 224//2
    start_y = h//2 - 224//2
    image = image[start_y:start_y+224, start_x:start_x+224, :]
    
    # transpose
    image = image.transpose(2, 0, 1)
    
    # 将输入数据转换为float32
    img_data = image.astype('float32')
    
    # normalize
    mean_vec = np.array([123.675, 116.28, 103.53])
    stddev_vec = np.array([58.395, 57.12, 57.375])
    norm_img_data = np.zeros(img_data.shape).astype('float32')
    for i in range(img_data.shape[0]):
        norm_img_data[i,:,:] = (img_data[i,:,:] - mean_vec[i]) / stddev_vec[i]
    
    # 调整尺寸
    norm_img_data = norm_img_data.reshape(1, 3, 224, 224).astype('float32')
    return norm_img_data

if __name__ == '__main__':
    
    # 量化方法选项
    use_int8 = False
    use_fp16 = False

    # 加载模型
    model = migraphx.parse_onnx("../Resource/Models/resnet50-v2-7.onnx")

    # 获取模型输入/输出节点信息
    inputs = model.get_inputs()
    outputs = model.get_outputs()
    inputName = model.get_parameter_names()[0]
    inputShape = inputs[inputName].lens()

    # INT8量化
    if use_int8:
        dic = dict()
        testofImage = "../Resource/Images/ImageNet_test.jpg"
        testimage = Preprocessing(testofImage)
        dic[inputName] = migraphx.argument(testimage)
        calibration = [dic]
        migraphx.quantize_int8(model, migraphx.get_target("gpu"), calibration)
    
    if use_fp16:
        migraphx.quantize_fp16(model)

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

    # 为输出节点分配device内存，用于保存输出数据
    modelData = AllocateOutputMemory(model)

    # 预处理并转换为NCHW
    pathOfImage = "../Resource/Images/ImageNet_01.jpg"
    image = Preprocessing(pathOfImage)
    
    modelData[inputName] = migraphx.to_gpu(migraphx.argument(image))

    # 推理
    results = model.run(modelData)

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

    # 获取分类结果
    print(np.array(result))