ONNX 运行模型

import onnx
import onnxruntime as ort
import numpy as np
import cv2

import sys


def preprocess(img_data):
    mean_vec = np.array([0.485, 0.456, 0.406])
    stddev_vec = np.array([0.229, 0.224, 0.225])
    norm_img_data = np.zeros(img_data.shape).astype('float32')
    for i in range(img_data.shape[0]):
         # for each pixel in each channel, divide the value by 255 to get value between [0, 1] and then normalize
        norm_img_data[i,:,:] = (img_data[i,:,:]/255 - mean_vec[i]) / stddev_vec[i]
    return norm_img_data


if __name__ == "__main__":
	if len(sys.argv) < 3:
        print("usage program model image")
        exit()
        
    model_name = sys.argv[1]
    img_name = sys.argv[2]
		
    img = cv2.imread(img_name)
    img = cv2.resize(img, (512,512), interpolation=cv2.INTER_AREA)
    #img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
    input_data = np.transpose(img, (2, 0, 1))
    input_data = preprocess(input_data)
    input_data = input_data.reshape([1, 3, 512, 512])

    sess = ort.InferenceSession(model_name)
    input_name = sess.get_inputs()[0].name
    result = sess.run([], {input_name: input_data})
    result = np.reshape(result, [1, -1])
    index = np.argmax(result)
    print("max index:", index)