Classifier_io_binding.py

# -*- coding: utf-8 -*-
"""
分类器示例
"""
import cv2
import argparse
import numpy as np
import onnxruntime as ort

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

def postprocess(scores,pathOfImage):
    '''
    Postprocessing with mxnet gluon
    The function takes scores generated by the network and returns the class IDs in decreasing order
    of probability
    '''
    with open('../Resource/synset.txt', 'r') as f:
        labels = [l.rstrip() for l in f]
    preds = np.squeeze(scores)
    a = np.argsort(preds)[::-1]
    print('class=%s ; probability=%f' %(labels[a[0]],preds[a[0]]))

    text = 'class=%s ' % (labels[a[0]])
    saveimage(pathOfImage,text)

def ort_seg_dcu(model_path,image):
    
    provider_options=[{'device_id':'0','migraphx_fp16_enable':'true'}]
    dcu_session = ort.InferenceSession(model_path, providers=['MIGraphXExecutionProvider'], provider_options=provider_options)
    
    images = [image]
    input_nodes = dcu_session.get_inputs()
    input_names = [i_n.name for i_n in input_nodes]
    output_nodes = dcu_session.get_outputs()
    output_names = [o_n.name for o_n in output_nodes]

    input_dict = {}
    for i_d, i_n in zip(images, input_names):
        input_dict[i_n] = i_d

    io_binding = dcu_session.io_binding()

    inputData = {}
    for key in input_dict.keys():
        inputData[key] = ort.OrtValue.ortvalue_from_numpy(input_dict[key], device_type='cuda')
        io_binding.bind_input(
            name = key,
            device_type = inputData[key].device_name(),
            device_id = 0,
            element_type = np.float32,
            shape = inputData[key].shape(),
            buffer_ptr = inputData[key].data_ptr())
        
    outputData = {}
    output_data = {}
    for index, o_n in enumerate(output_names):
        output_data[o_n] = np.empty(dcu_session.get_outputs()[index].shape).astype(np.float32)
        outputData[o_n] = ort.OrtValue.ortvalue_from_numpy(output_data[o_n], device_type='cuda')
        io_binding.bind_output(
            name = o_n,
            device_type = outputData[o_n].device_name(),
            device_id = 0,
            element_type = np.float32,
            shape = outputData[o_n].shape(),
            buffer_ptr = outputData[o_n].data_ptr())

    dcu_session.run_with_iobinding(io_binding)
    result = io_binding.copy_outputs_to_cpu()[0]
    scores = np.array(result)
    print("ort result.shape:",scores.shape)

    return scores

def saveimage(pathOfImage,text):
    iimage = cv2.imread(pathOfImage, cv2.IMREAD_COLOR)
    font = cv2.FONT_HERSHEY_SIMPLEX
    font_scale = 0.5
    font_color = (0, 0, 255)  
    font_thickness = 1
    text_position = (5, 20)
    cv2.putText(iimage, text, text_position, font, font_scale, font_color, font_thickness)
    cv2.imwrite("./output_image.jpg", iimage)
    cv2.destroyAllWindows()

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--imgPath', type=str, default='../Resource/Images/ImageNet_01.jpg', help="image path")
    parser.add_argument('--staticModelPath', type=str, default='../Resource/Models/resnet50_static.onnx', help="static onnx filepath")
    args = parser.parse_args()

    # 数据预处理
    image = Preprocessing(args.imgPath)

    # 静态推理
    result = ort_seg_dcu(args.staticModelPath,image)
    
    # 解析分类结果
    postprocess(result,args.imgPath)