YOLO系列笔记（九）—— yolov5转onnx并推理

前言

该笔记记录了将训练后得出的yolov5s.pt模型转为.onnx模型并进行推理的过程。教程与大部分代码均来自这位作者的这篇博客。笔者仅对该代码进行修改和分析，在此十分感谢大佬们的帮助和分享！

代码

import onnxruntime
import numpy as np
import cv2
import os
import glob

CLASSES = ['smoke']

class YOLOV5():
    def __init__(self,onnxpath):
        self.onnx_session=onnxruntime.InferenceSession(onnxpath)
        self.input_name=self.get_input_name()
        self.output_name=self.get_output_name()
    #-------------------------------------------------------
	#   获取输入输出的名字
	#-------------------------------------------------------
    def get_input_name(self):
        input_name=[]
        for node in self.onnx_session.get_inputs():
            input_name.append(node.name)
        return input_name
    def get_output_name(self):
        output_name=[]
        for node in self.onnx_session.get_outputs():
            output_name.append(node.name)
        return output_name
    #-------------------------------------------------------
	#   输入图像
	#-------------------------------------------------------
    def get_input_feed(self,img_tensor):
        input_feed={
   
   }
        for name in self.input_name:
            input_feed[name]=img_tensor
        return input_feed
    #-------------------------------------------------------
	#   1.cv2读取图像并resize
	#	2.图像转BGR2RGB和HWC2CHW
	#	3.图像归一化
	#	4.图像增加维度
	#	5.onnx_session 推理
	#-------------------------------------------------------
    def inference(self,img_path):
        # cv2 读取图像
        img=cv2.imread(img_path)
        #进行resize
        or_img=cv2.resize(img,(640,640))
        #BGR2RGB和HWC2CHW（具体细节看代码解析）
        img=or_img[:,:,::-1].transpose(2,0,1)  
        # 将图像数据类型转换为 float32
        img=img.astype(dtype=np.float32)
        # 将像素值归一化
        img/=255.0
        # 增加批次纬度
        img=np.expand_dims(img,axis=0)
        # 准备模型输入和推理
        input_feed=self.get_input_feed(img)
        pred=self.onnx_session.run(None,input_feed)[0]
        # 返回模型的预测结果 pred 以及原始图像 or_img。
        return pred,or_img

# 坐标转换
def xywh2xyxy(x):
    # [x, y, w, h] to [x1, y1, x2, y2]
    y = np.copy(x)
    y[:, 0] = x[:, 0] - x[