YOLOv5转ONNX后推理Bounding Box异常问题求助

我训练的YOLOv5模型用yolo detect.py在新图像上推理完全正常，但导出为ONNX格式后，自行编写代码加载模型推理无法得到正确的Bounding Box。由于需要获取Bounding Box的面积等信息，无法直接使用detect.py。我尝试用YOLO的non_max_suppression来筛选有效框，但不确定这是否为合适的解决方案，恳请帮忙排查问题。

加载图像并预处理 #

import cv2
import numpy as np

img = cv2.imread("image.jpg", cv2.IMREAD_UNCHANGED)
resized = cv2.resize(img, (640,640), interpolation = cv2.INTER_AREA).astype(np.float32)
resized = resized.transpose((2, 0, 1))
resized = np.expand_dims(resized, axis=0)  # Add batch dimension

运行ONNX会话 #

import onnxruntime as ort
ort_session = ort.InferenceSession("yolov5.onnx", providers=["CUDAExecutionProvider"])
# compute ONNX Runtime output prediction
ort_inputs = {ort_session.get_inputs()[0].name: resized}
ort_outs = ort_session.run(None, ort_inputs)

此时已得到包含所有Bounding Box的张量

筛选有效Bounding Box #

import torch
from yolov5.utils.general import non_max_suppression, xyxy2xywh

output= torch.from_numpy(np.asarray(ort_outs))
out = non_max_suppression(output, conf_thres=0.2, iou_thres=0.5)[0]

# convert xyxy to xywh
xyxy = out[:,:4]
xywh = xyxy2xywh(xyxy)
out[:, :4] = xywh

可视化Bounding Box #

from PIL import Image, ImageDraw, ImageFont

tmp_image = Image.fromarray(img)
draw = ImageDraw.Draw(tmp_image)

for i,(x,y,w,h,score,class_id) in enumerate(out):
  real_x = x * w_ratio # resize from model size to image size
  real_y = y * h_ratio

  shape = (real_x, real_y, (x + w) * w_ratio, (y + h) * h_ratio) # shape of the bounding box to draw

  class_id = round(float(class_id))
  
  class_string = list(class_list.keys())[list(class_list.values()).index(class_id)]
  color = CLASS_RECTANGLE_COLORS[class_string]

  draw.rectangle(shape, outline=color, width=4)
  fnt = ImageFont.load_default()

  draw.multiline_text((real_x + 8, real_y + 8), f"{class_string} {score*100:.2f}%", font=fnt, fill=color)

tmp_image.show()

自行实现的算法生成的图像与detect.py生成的图像存在明显差异。

内容的提问来源于stack exchange，提问作者Timothee W