第一步。导出onnx文件，然后检查一下
pip install torch onnx onnxruntime #导入onnxruntime
pip install onnx --upgrade -i https://pypi.tuna.tsinghua.edu.cn/simple #镜像安装
pip install onnx-simplifier ##导入onnx的优化器
运行export,py文件，进行onnx模型的导出
python export.py --weights runs/train/exp3/weights/best.pt --img 640 --batch 1 --include onnx ##
对模型进行优化，检查模型
python -m onnxsim runs/train/exp3/weights/best.onnx output_onnx_model #对模型进行优化，优化后会在根目录下输出一个output_onnx_model文件
model = onnx.load(output_onnx_model) #通过路径加载模型
onnx.checker.check_model(model) #检查模型的正确性，正确则不反回任何消息
第二部写一个class来实现模型的加载，通过一个model实现传入图片，模型检测，检测结果的数据及它的检测图片的输出，方便后期脱离yolo架构，将onnx模型对外进行部署
声明这个代码参考了 qq_22487889 博主写的《详细介绍 Yolov5 转 ONNX模型 + 使用ONNX Runtime 的 Python 部署（包含官方文档的介绍）》，链接是https://blog.csdn.net/qq_22487889/article/details/128011883
我只是作了一定的改动，放到这里是方便自己后期落地使用。
代码分三部分 ——在一个py文件里，通过两个方法，一个类来实现
1，onnx模型的加载，及输入变量名的获取
2，输入图片，对图片进行计算，并能输出它的预测框参数
3，绘图+显示。
import onnx
import onnxruntime as ort
import cv2
import numpy as np
from onnxruntime.transformers.models.gpt2.parity_check_helper import inference
from torch import nnCLASSES = ['cat', 'dog']
model_path = '../output_onnx_model'def draw(image, box_data, CLASSES):# -------------------------------------------------------# 取整，方便画框# -------------------------------------------------------boxes = box_data[..., :4].astype(np.int32) # x1 x2 y1 y2scores = box_data[..., 4]classes = box_data[..., 5].astype(np.int32)for box, score, cl in zip(boxes, scores, classes):top, left, right, bottom = boxprint('class: {}, score: {}'.format(CLASSES[cl], score))print('box coordinate left,top,right,down: [{}, {}, {}, {}]'.format(top, left, right, bottom))cv2.rectangle(image, (top, left), (right, bottom), (255, 0, 0), 2)cv2.putText(image, '{0} {1:.2f}'.format(CLASSES[cl], score),(top, left),cv2.FONT_HERSHEY_SIMPLEX,0.6, (0, 0, 255), 2)return imagedef load_onnx(model_path):#在这里还是先加载模型，然后check一下看看正确否onnx_model = onnx.load(model_path)try:onnx.checker.check_model(onnx_model)except Exception:print("Model incorrect")else:print("Model correct")## 模型进入运行阶段，准备测试onnx_session = ort.InferenceSession(model_path)input_name = onnx_session.get_inputs()[0].name #明确是一个输入张量，就这样写吧，简简单单return onnx_session, input_nameclass Yolov5ONNX(nn.Module):def __init__(self, onnx_session, input_name):super(Yolov5ONNX,self).__init__()self.model = onnx_sessionself.input_name = input_namedef inference(self, img):""" 1.cv2读取图像并resize2.图像转BGR2RGB和HWC2CHW(因为yolov5的onnx模型输入为 RGB：1 × 3 × 640 × 640)3.图像归一化4.图像增加维度5.onnx_session 推理 """resize_img = cv2.resize(img, (640, 640)) # resize后的原图 (640, 640, 3)image = cv2.cvtColor(resize_img, cv2.COLOR_BGR2RGB) #cv2读进来是BGR的，一定转成Rgb的，上次研究了一下午就是这里出了问题img = image.transpose(2, 0, 1) # whc转cwhimg = img.astype(dtype=np.float32) # onnx模型的类型是type: float32[ , , , ]img /= 255.0img = np.expand_dims(img, axis=0) # [3, 640, 640]扩展为[1, 3, 640, 640]# img尺寸(1, 3, 640, 640)return img,resize_imgdef nms(self,boxs, thresh):# dets:x1 y1 x2 y2 score class# x[:,n]就是取所有集合的第n个数据x1 = boxs[:, 0]y1 = boxs[:, 1]x2 = boxs[:, 2]y2 = boxs[:, 3]# -------------------------------------------------------# 计算框的面积# 置信度从大到小排序# -------------------------------------------------------areas = (y2 - y1 + 1) * (x2 - x1 + 1)scores = boxs[:, 4]# print(scores)keep = []index = scores.argsort()[::-1] # np.argsort()对某维度从小到大排序# [::-1] 从最后一个元素到第一个元素复制一遍。倒序从而从大到小排序while index.size > 0:i = index[0]keep.append(i)# -------------------------------------------------------# 计算相交面积# 1.相交# 2.不相交# -------------------------------------------------------x11 = np.maximum(x1[i], x1[index[1:]])y11 = np.maximum(y1[i], y1[index[1:]])x22 = np.minimum(x2[i], x2[index[1:]])y22 = np.minimum(y2[i], y2[index[1:]])w = np.maximum(0, x22 - x11 + 1)h = np.maximum(0, y22 - y11 + 1)overlaps = w * h# -------------------------------------------------------# 计算该框与其它框的IOU，去除掉重复的框，即IOU值大的框# IOU小于thresh的框保留下来# -------------------------------------------------------ious = overlaps / (areas[i] + areas[index[1:]] - overlaps)idx = np.where(ious <= thresh)[0]index = index[idx + 1]return keepdef xywh2xyxy(self,x):# [x, y, w, h] to [x1, y1, x2, y2]y = np.copy(x)y[:, 0] = x[:, 0] - x[:, 2] / 2y[:, 1] = x[:, 1] - x[:, 3] / 2y[:, 2] = x[:, 0] + x[:, 2] / 2y[:, 3] = x[:, 1] + x[:, 3] / 2return ydef filter_box(self,pred, conf_thres, iou_thres): # 过滤掉无用的框# 先筛选置信度# […,4]：代表了取最里边一层的所有第4号元素，…代表了对:,:,:,等所有的的省略。此处生成：25200个第四号元素组成的数组conf = pred[..., 4] > conf_thres # 0 1 2 3 4 4是置信度，只要置信度 > conf_thres 的box = pred[conf == True] # 根据objectness score生成(n, 9)，只留下符合要求的框print('box:符合要求的框')print(box.shape)# 通过argmax获取置信度最大的类别cls_cinf = box[..., 5:] # 左闭右开（5 6），就只剩下了每个grid cell中各类别的概率cls = cls_cinf.argmax(axis=-1)all_cls = list(set(cls)) # 去重，找出图中都有哪些类别# set() 函数创建一个无序不重复元素集，可进行关系测试，删除重复数据，还可以计算交集、差集、并集等。output = []for i in np.arange(len(all_cls)):curr_cls = all_cls[i]curr_cls_box = []curr_out_box = []for j in range(len(cls)):if cls[j] == curr_cls:box[j][5] = curr_clscurr_cls_box.append(box[j][:6]) # 左闭右开，0 1 2 3 4 5curr_cls_box = np.array(curr_cls_box) # 0 1 2 3 4 5 分别是 x y w h score class# curr_cls_box_old = np.copy(curr_cls_box)curr_cls_box = self.xywh2xyxy(curr_cls_box) # 0 1 2 3 4 5 分别是 x1 y1 x2 y2 score classcurr_out_box = self.nms(curr_cls_box, iou_thres) # 获得nms后，剩下的类别在curr_cls_box中的下标for k in curr_out_box:output.append(curr_cls_box[k])output = np.array(output)return outputdef forward(self, image,conf_thres, iou_thres):image_np,resize_img = self.inference(image)print(image_np.shape) #这里应该是[1,3,640,640]pred =self.model.run(None, {self.input_name: image_np})print(pred.shape) #[1,1,25200,7]pred =pred[0][0] #[25200,7]#开始进行置信度cfg_thr，iou_thr，非极大值抑制out_put =self.filter_box(pred, conf_thres, iou_thres)return out_put,resize_imgif __name__ == '__main__':CLASSES = ['cat', 'dog']model_path = '../output_onnx_model'image = cv2.imread('../data/cat.jpg')onnx_session,input_name = load_onnx(model_path)yolov5_model=Yolov5ONNX(onnx_session,input_name)boxs,resize_image = yolov5_model(image, conf_thres=0.5, iou_thres=0.5)image =draw(resize_image,boxs,CLASSES)
