onnx模型测试代码

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)