Yolov5-v5.0 导出onnx推理结果完全错误

楼主研究生在读，在搞模型部署方面的研究，记录一些遇到的问题。

最近使用yolov5-v5.0时，用官方提供的yolov5s.pt权重文件并利用export.py文件导出onnx模型时，虽然能够导出onnx模型但其onnx模型经实测输出是错误的，网上没搜到解决方案，所以楼主发文记录并提供自己的解决方式，有帮助欢迎点赞收藏。

首先在yolov5-v5.0运行官方export.py导出onnx如下：

python export.py --weights yolov5s.pt --img 640 --batch 1

Starting ONNX export with onnx 1.13.1...
F:\RV1126\yolov5-5.0\models\yolo.py:106: TracerWarning: Converting a tensor to a Python boolean might cause the trace to be incorrect. We can't record the data flow of Python values, so this value will be treated as a constant in the future. This means that the trace might not generalize to other inputs!
  if augment:
F:\RV1126\yolov5-5.0\models\yolo.py:131: TracerWarning: Converting a tensor to a Python boolean might cause the trace to be incorrect. We can't record the data flow of Python values, so this value will be treated as a constant in the future. This means that the trace might not generalize to other inputs!
  if profile:
F:\RV1126\yolov5-5.0\models\yolo.py:142: TracerWarning: Converting a tensor to a Python boolean might cause the trace to be incorrect. We can't record the data flow of Python values, so this value will be treated as a constant in the future. This means that the trace might not generalize to other inputs!
  if profile:
ONNX export success, saved as yolov5s.onnx

可以看到能够导出yolov5s.onnx的模型。

但推理时会出现一些IndexError问题，导致根本无法推理或者是能够推理但推理结果完全错误。本人使用的推理程序代码如下：

参考博客https://blog.csdn.net/qq128252/article/details/127105463：

import os
import cv2
import numpy as np
import onnxruntime
import time

CLASSES=['person', 'bicycle', 'car', 'motorcycle', 'airplane', 'bus', 'train', 'truck', 'boat', 'traffic light',
        'fire hydrant', 'stop sign', 'parking meter', 'bench', 'bird', 'cat', 'dog', 'horse', 'sheep', 'cow',
        'elephant', 'bear', 'zebra', 'giraffe', 'backpack', 'umbrella', 'handbag', 'tie', 'suitcase', 'frisbee',
        'skis', 'snowboard', 'sports ball', 'kite', 'baseball bat', 'baseball glove', 'skateboard', 'surfboard',
        'tennis racket', 'bottle', 'wine glass', 'cup', 'fork', 'knife', 'spoon', 'bowl', 'banana', 'apple',
        'sandwich', 'orange', 'broccoli', 'carrot', 'hot dog', 'pizza', 'donut', 'cake', 'chair', 'couch',
        'potted plant', 'bed', 'dining table', 'toilet', 'tv', 'laptop', 'mouse', 'remote', 'keyboard', 'cell phone',
        'microwave', 'oven', 'toaster', 'sink', 'refrigerator', 'book', 'clock', 'vase', 'scissors', 'teddy bear',
        'hair drier', 'toothbrush'] #coco80类别

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):
        img=cv2.imread(img_path)
        or_img=cv2.resize(img,(640,640))
        img=or_img[:,:,::-1].transpose(2,0,1)  #BGR2RGB和HWC2CHW
        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]
        return pred,or_img

#dets:  array [x,6] 6个值分别为x1,y1,x2,y2,score,class 
#thresh: 阈值
def nms(dets, thresh):
    x1 = dets[:, 0]
    y1 = dets[:, 1]
    x2 = dets[:, 2]
    y2 = dets[:, 3]
    #-------------------------------------------------------
	#   计算框的面积
    #	置信度从大到小排序
	#-------------------------------------------------------
    areas = (y2 - y1 + 1) * (x2 - x1 + 1)
    scores = dets[:, 4]
    keep = []
    index = scores.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(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(org_box,conf_thres,iou_thres): #过滤掉无用的框
    #-------------------------------------------------------
	#   删除为1的维度
    #	删除置信度小于conf_thres的BOX
	#-------------------------------------------------------
    org_box=np.squeeze(org_box)
    conf = org_box[..., 4] > conf_thres
    box = org_box[conf == True]
    #-------------------------------------------------------
    #	通过argmax获取置信度最大的类别
	#-------------------------------------------------------
    cls_cinf = box[..., 5:]
    cls = []
    for i in range(len(cls_cinf)):
        cls.append(int(np.argmax(cls_cinf[i])))
    all_cls = list(set(cls))     
    #-------------------------------------------------------
	#   分别对每个类别进行过滤
	#	1.将第6列元素替换为类别下标
	#	2.xywh2xyxy 坐标转换
	#	3.经过非极大抑制后输出的BOX下标
	#	4.利用下标取出非极大抑制后的BOX
	#-------------------------------------------------------
    output = []
    for i in range(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])
        curr_cls_box = np.array(curr_cls_box)
        # curr_cls_box_old = np.copy(curr_cls_box)
        curr_cls_box = xywh2xyxy(curr_cls_box)
        curr_out_box = nms(curr_cls_box,iou_thres)
        for k in curr_out_box:
            output.append(curr_cls_box[k])
    output = np.array(output)
    return output

def draw(image,box_data):  
    #-------------------------------------------------------
    #	取整，方便画框
	#-------------------------------------------------------
    boxes=box_data[...,:4].astype(np.int32) 
    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)



if __name__=="__main__":
    onnx_path='yolov5s.onnx'
    model=YOLOV5(onnx_path)
    output,or_img=model.inference('bus.jpg')
    outbox=filter_box(output,0.5,0.5)
    draw(or_img,outbox)
    cv2.imwrite('onnx_run2.jpg',or_img)
    
    

推理程序经测试是没有问题的，但使用之前导出的onnx模型时报错，根本无法推理。所以必定是onnx模型本身的问题，运行时会报索引错误：

Traceback (most recent call last):
  File "f:/RV1126/yolov5-5.0/onnx_run2.py", line 181, in
    draw(or_img,outbox)
  File "f:/RV1126/yolov5-5.0/onnx_run2.py", line 160, in draw
    scores=box_data[...,4]
IndexError: index 4 is out of bounds for axis 0 with size 0。

错误如下所示：

(tf) PS F:\RV1126\yolov5-5.0> & D:/Anoconda3/envs/tf/python.exe f:/RV1126/yolov5-5.0/onnx_run2.py
Traceback (most recent call last):
  File "f:/RV1126/yolov5-5.0/onnx_run2.py", line 181, in 
    draw(or_img,outbox)
  File "f:/RV1126/yolov5-5.0/onnx_run2.py", line 160, in draw
    scores=box_data[...,4]
IndexError: index 4 is out of bounds for axis 0 with size 0

楼主先是尝试解决不能推理的问题，最终虽然使得模型能够推理，但结果也是完全错误的。如下：

可以看到模型疯狂输出图中不存在的目标，并且score明显错误，bed标签置信度居然超过了1。

最终解决方法换成yolov5-v6.0版用export.py文件导出onnx模型即可。如下可以看到成功运行，推理结果正确：

(tf) PS F:\RV1126\yolov5-6.0> & D:/Anoconda3/envs/tf/python.exe f:/RV1126/yolov5-6.0/onnx_run2.py
class: person, score: 0.8907531499862671
box coordinate left,top,right,down: [535, 234, 639, 517]
class: person, score: 0.8881538510322571
box coordinate left,top,right,down: [175, 240, 273, 509]
class: person, score: 0.8382899761199951
box coordinate left,top,right,down: [44, 234, 180, 528]
class: person, score: 0.6425769329071045
box coordinate left,top,right,down: [0, 326, 72, 526]
class: bus, score: 0.6496226787567139
box coordinate left,top,right,down: [21, 121, 637, 443]

 推测是yolov5-v5.0版本和6.0版本onnx导出存在不兼容情况。对大家有帮助的话欢迎点赞收藏。