YOLOX的onnx图片推理方式

默认已经通过yolox中的代码将训练完的模型文件转成了onnx格式
1.引入必要的库

import cv2
import torch
import numpy as np
import onnxruntime
from yolox.data.data_augment import preproc
from yolox.data.datasets.voc_classes import VOC_CLASSES
from yolox.utils import mkdir, multiclass_nms, demo_postprocess, vis
from exps.example.yolox_voc.yolox_voc_s import Exp

2.初始化参数信息

class ONNX:
    def __init__(self):
        exp = Exp()
        self.test_size = exp.test_size
        self.class_names = VOC_CLASSES
        self.score_thr = 0.3
        self.nms_thr = 0.45
        self.score_thr = 0.1
        self.model ="" #你的onnx模型文件位置

3.推理代码

@torch.no_grad()
    def detect(self, origin_img):
        img, ratio = preproc(origin_img, self.test_size)
        session = onnxruntime.InferenceSession(self.model)
        ort_inputs = {session.get_inputs()[0].name: img[None, :, :, :]}
        output = session.run(None, ort_inputs)
        predictions = demo_postprocess(output[0], self.test_size)[0]
        boxes = predictions[:, :4]
        scores = predictions[:, 4:5] * predictions[:, 5:]
        boxes_xyxy = np.ones_like(boxes)
        boxes_xyxy[:, 0] = boxes[:, 0] - boxes[:, 2] / 2.
        boxes_xyxy[:, 1] = boxes[:, 1] - boxes[:, 3] / 2.
        boxes_xyxy[:, 2] = boxes[:, 0] + boxes[:, 2] / 2.
        boxes_xyxy[:, 3] = boxes[:, 1] + boxes[:, 3] / 2.
        boxes_xyxy /= ratio
        dets = multiclass_nms(boxes_xyxy, scores, nms_thr=self.nms_thr, score_thr=self.score_thr)
        if dets is not None:
            final_boxes, final_scores, final_cls_inds = dets[:, :4], dets[:, 4], dets[:, 5]
            origin_img = vis(origin_img, final_boxes, final_scores, final_cls_inds,
                             conf=self.score_thr, class_names=self.class_names)
            # print("zuibiao:{} label:{}".format((final_boxes).astype(np.int), (final_cls_inds).astype(np.int)))
        cv2.imwrite('./test.jpg', origin_img)