YOLOv8转ONNX后输出异常：如何获取检测框与置信度？

一、输出形状差异原因 #

你看到的[1,5,5040]是YOLOv8 ONNX模型的默认通道优先格式，并非转换错误：

• 你的预期形状[1,25200,7]是[batch, num_boxes, num_params]的格式，而ONNX输出默认采用[batch, num_params, num_boxes]的通道优先排列
• 数值验证：5040 * 5 = 25200，刚好对应你预期的总检测框数；当前输出通道数为5，而你预期为7，可能是以下原因：
  • 模型训练时实际仅配置了1类（或类别配置未正确加载）
  • 转换ONNX时未指定类别相关参数，导致输出未包含类别概率

二、验证模型转换正确性 #

可以通过以下步骤确认转换是否成功：

1. 用Netron工具打开生成的ONNX文件，查看输入输出节点：
   • 输入维度应匹配[1,3,384,640]（对应你的设定尺寸）
   • 输出节点的结构是否符合模型的层级逻辑
2. 用ONNX Runtime做推理测试：

   import onnxruntime as ort
   import numpy as np
   
   session = ort.InferenceSession("best.onnx")
   input_name = session.get_inputs()[0].name
   output_name = session.get_outputs()[0].name
   
   dummy_input = np.random.randn(1,3,384,640).astype(np.float32)
   output = session.run([output_name], {input_name: dummy_input})[0]
   print("输出形状:", output.shape)
   

   若推理过程无报错，说明模型转换基本正确

三、获取Bounding Boxes和置信度 #

针对当前[1,5,5040]的输出格式，需先转置为常规格式，再提取目标信息：

# 假设output是ONNX推理得到的数组，形状为[1,5,5040]
output = output.transpose(0,2,1)  # 转换为[1,25200,5]
boxes = output[..., :4]  # 提取x,y,w,h相对坐标
confidences = output[..., 4]  # 提取置信度

# 若模型实际包含2类，输出通道数应为7，此时额外提取类别信息：
# boxes = output[..., :4]
# confidences = output[..., 4]
# class_probs = output[..., 5:]
# class_ids = np.argmax(class_probs, axis=-1)

注意：YOLOv8输出的是相对坐标（x/y为框中心坐标，w/h为宽高，均相对于输入图像尺寸），需转换为绝对坐标才能直接使用：

img_h, img_w = 384, 640
abs_boxes = boxes.copy()
# 转换中心坐标为绝对位置
abs_boxes[..., 0] = boxes[..., 0] * img_w
abs_boxes[..., 1] = boxes[..., 1] * img_h
# 转换宽高为绝对尺寸
abs_boxes[..., 2] = boxes[..., 2] * img_w
abs_boxes[..., 3] = boxes[..., 3] * img_h
# 可选：转换为左上角+右下角坐标格式
abs_boxes[..., 0] -= abs_boxes[..., 2] / 2
abs_boxes[..., 1] -= abs_boxes[..., 3] / 2
abs_boxes[..., 2] += abs_boxes[..., 0]
abs_boxes[..., 3] += abs_boxes[..., 1]

四、修正转换以匹配预期输出 #

若希望ONNX输出直接为[batch, num_boxes, 5+num_classes]格式，可按以下方式调整转换代码：

方法1：使用Ultralytics官方export工具 #

确保模型类别配置正确，转换时添加简化参数：

from ultralytics import YOLO

m = YOLO("best.pt")
# 先确认模型的类别配置
print("模型类别列表:", m.model.names)
# 执行转换
m.export(format="onnx", imgsz=[384,640], opset=12, simplify=True)

方法2：手动包装模型后用torch.onnx.export #

通过包装模型手动转置输出维度，避免ONNX自动优化为通道优先：

import torch
from ultralytics import YOLO

m = YOLO("best.pt")
m.model.eval()

class WrappedModel(torch.nn.Module):
    def __init__(self, model):
        super().__init__()
        self.model = model
    def forward(self, x):
        output = self.model(x)[0]  # 获取YOLOv8的输出张量
        return output.transpose(1,2)  # 转置为[batch, num_boxes, num_params]格式

wrapped_model = WrappedModel(m.model)
dummy_input = torch.randn(1,3,384,640)
torch.onnx.export(wrapped_model, dummy_input, "yolo8.onnx", opset_version=12, verbose=True)

内容的提问来源于stack exchange，提问作者Martin Sedláček