yolov5 onnx的部署文件（主要是onnx文件的使用）

第一步。导出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 nn

CLASSES = ['cat', 'dog']
model_path = '../output_onnx_model'


def draw(image, box_data, CLASSES):
    # -------------------------------------------------------
    #	取整，方便画框
    # -------------------------------------------------------

    boxes = box_data[..., :4].astype(np.int32)  # x1 x2 y1 y2
    scores = box_data[..., 4]
    classes = box_data[..., 5].astype(np.int32)
    for box, score, cl in zip(boxes, scores, classes):
        top, left, right, bottom = box
        print('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 image


def 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_name


class Yolov5ONNX(nn.Module):
    def __init__(self, onnx_session, input_name):
        super(Yolov5ONNX,self).__init__()
        self.model = onnx_session
        self.input_name = input_name

    def inference(self, img):
        """ 1.cv2读取图像并resize
        2.图像转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转cwh
        img = img.astype(dtype=np.float32)  # onnx模型的类型是type: float32[ , , , ]
        img /= 255.0
        img = np.expand_dims(img, axis=0)  # [3, 640, 640]扩展为[1, 3, 640, 640]
        # img尺寸(1, 3, 640, 640)
        return img,resize_img

    def 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 keep

    def xywh2xyxy(self,x):
        # [x, y, w, h] to [x1, y1, x2, y2]
        y = np.copy(x)
        y[:, 0] = x[:, 0] - x[:, 2] / 2
        y[:, 1] = x[:, 1] - x[:, 3] / 2
        y[:, 2] = x[:, 0] + x[:, 2] / 2
        y[:, 3] = x[:, 1] + x[:, 3] / 2
        return y
    def 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_cls
                    curr_cls_box.append(box[j][:6])  # 左闭右开，0 1 2 3 4 5

            curr_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 class
            curr_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 output

    def 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_img




if __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)