onnx模型测试代码

import io
import torch
import torch.onnx
from model import resnet34
import onnx
import onnxruntime
import numpy as np
from PIL import Image
from torchvision import transforms,datasets 
import os


class_names = ['empty-run', 'have','empty-stop'] #这个顺序很重要，要和训练时候的类名顺序一致
 
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
img_path = 'data/1/1-0086.jpg'

##载入模型并读取权重

os.environ['CUDA_DEVICE_ORDER'] = 'PCI_BUS_ID'
os.environ['CUDA_VISIBLE_DEVICES'] = '1' 
device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')

def to_numpy(tensor):
    return tensor.detach().cpu().numpy() if tensor.requires_grad else tensor.cpu().numpy()

def softmax(x):
    x = x.reshape(-1)
    e_x = np.exp(x - np.max(x))
    return e_x / e_x.sum(axis=0)

def postprocess(result):
    return softmax(np.array(result)).tolist()


tf = transforms.Compose([
    transforms.Resize([224,224]), #先放大再旋转后裁剪
    transforms.ToTensor(),
    transforms.Normalize(mean = [0.485,0.456,0.406], std = [0.229,0.224,0.225])
])


image = Image.open(img_path).convert('RGB')

# w, h = image.size
# scale = 224./max(w, h)
# img = image.resize([int(x) for x in [w*scale, h*scale]])
# img_size = img.size
# isize = 224
# imgx= Image.new('RGB', (isize,isize))
# imgx.paste(img,((isize - img_size[0]) // 2,
#                                     (isize - img_size[1]) //2 ))
# image = imgx


img_ = tf(image)#拓展维度
img_.unsqueeze_(0)


session = onnxruntime.InferenceSession("resnet18.onnx")
# input = {session.get_inputs()[0].name: to_numpy(img_)}
# output = session.run(None, input)
output = np.array(session.run(None, { "input": to_numpy(img_)}))
print(output)


res = postprocess(output)#后处理 softmax
idx = np.argmax(res)
result=class_names[idx]
print(result)