验证onnx格式的yolox

• 转onnx时decode_in_inference是False，所以需要decode_outputs函数
• 需要注意的是，yolox的Exp中input_size和test_size是先高后宽，即(h,w)

import onnxruntime   
import cv2, sys, torch, torchvision
import numpy as np

cls_num=3
conf_th=0.35
nms_th=0.45
model_w=640
model_h=352
ow=model_w/32
oh=model_h/32

def decode_outputs(outputs):
        grids = []
        strides = []
        dtype=outputs.type()
        for (hsize, wsize), stride in zip([(oh*4,ow*4),(oh*2,ow*2),(oh,ow)], [8, 16, 32]):
            yv, xv = torch.meshgrid([torch.arange(hsize), torch.arange(wsize)])
            grid = torch.stack((xv, yv), 2).view(1, -1, 2)
            grids.append(grid)
            shape = grid.shape[:2]
            strides.append(torch.full((*shape, 1), stride))

        grids = torch.cat(grids, dim=1).type(dtype)
        strides = torch.cat(strides, dim=1).type(dtype)

        outputs[..., :2] = (outputs[..., :2] + grids) * strides
        outputs[..., 2:4] = torch.exp(outputs[..., 2:4]) * strides
        return outputs

session = onnxruntime.InferenceSession('demo.onnx', None)
input_name = session.get_inputs()[0].name

path='/home/lwd/00000.jpg'
raw=cv2.imread(path)
print(raw.shape)
ih, iw, _ = raw.shape
im=cv2.resize(raw, (model_w, model_h))
print(im.shape)
im=im[..., ::-1]
im=im.astype('float32')/255.0
im-=(0.485, 0.456, 0.406)
im/=(0.229, 0.224, 0.225)
im=im.transpose((2,0,1))
im=np.ascontiguousarray(im, dtype=np.float32)
im=np.expand_dims(im, 0)
print(im.shape)
outputs = session.run(None, {input_name: im})
for i in outputs: print(i.shape)
# post
prediction=torch.from_numpy(outputs[0])
prediction=decode_outputs(prediction)
box_corner = prediction.new(prediction.shape)
box_corner[:, :, 0] = prediction[:, :, 0] - prediction[:, :, 2] / 2
box_corner[:, :, 1] = prediction[:, :, 1] - prediction[:, :, 3] / 2
box_corner[:, :, 2] = prediction[:, :, 0] + prediction[:, :, 2] / 2
box_corner[:, :, 3] = prediction[:, :, 1] + prediction[:, :, 3] / 2
prediction[:, :, :4] = box_corner[:, :, :4]
image_pred=prediction[0]
class_conf, class_pred = torch.max(image_pred[:, 5: 5 + cls_num], 1, keepdim=True)
conf_mask = (image_pred[:, 4] * class_conf.squeeze() >= conf_th).squeeze()
print(conf_mask.sum())
detections = torch.cat((image_pred[:, :5], class_conf, class_pred.float()), 1)
detections = detections[conf_mask]
if not detections.size(0): sys.exit()
nms_out_index = torchvision.ops.batched_nms(
            detections[:, :4],
            detections[:, 4] * detections[:, 5],
            detections[:, 6],
            nms_th,
)
detections = detections[nms_out_index]
print(detections.shape)
res=detections.numpy()
print(res)
res[:,0]=res[:,0]*iw/model_w
res[:,1]=res[:,1]*ih/model_h
res[:,2]=res[:,2]*iw/model_w
res[:,3]=res[:,3]*ih/model_h
for i in res:
	x1=int(i[0])
	y1=int(i[1])
	x2=int(i[2])
	y2=int(i[3])
	cv2.rectangle(raw, (x1,y1), (x2,y2), (0,0,255))
cv2.imshow('ss', raw)
cv2.waitKey()