如何在Python中使用YOLOv8导出的ONNX模型并复现一致检测框

要让ONNX模型输出和YOLOv8原生完全一致的检测框，必须严格对齐导出参数、预处理、后处理三个核心环节，以下是具体实现步骤：

一、确认ONNX模型导出的正确性 #

导出时需保证参数匹配且算子兼容：

• 补充指定opset版本（推荐12+，避免低版本算子不兼容问题），执行导出命令：

  !yolo export model="/content/runs/detect/train/weights/best.pt" format=onnx imgsz=[640,480] opset=12
  

• 导出过程无报错或警告，可通过Netron工具查看输入节点形状是否为[1,3,480,640]（对应CHW格式的480高、640宽图像）

二、严格对齐YOLOv8预处理流程 #

YOLOv8的预处理逻辑是边界框一致的关键，必须复现letterbox缩放、通道转换、归一化步骤：

import cv2
import numpy as np

def yolov8_preprocess(img_path, target_size=(480, 640)):
    # 读取图像并转换为RGB（YOLOv8训练用RGB，OpenCV默认读BGR）
    img = cv2.imread(img_path)
    img_rgb = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
    orig_h, orig_w = img.shape[:2]
    
    # Letterbox缩放：保持原比例，黑边填充到目标尺寸
    scale = min(target_size[0]/orig_h, target_size[1]/orig_w)
    new_h, new_w = int(orig_h * scale), int(orig_w * scale)
    resized_img = cv2.resize(img_rgb, (new_w, new_h), interpolation=cv2.INTER_LINEAR)
    
    # 填充黑边
    padded_img = np.zeros((target_size[0], target_size[1], 3), dtype=np.uint8)
    pad_top = (target_size[0] - new_h) // 2
    pad_left = (target_size[1] - new_w) // 2
    padded_img[pad_top:pad_top+new_h, pad_left:pad_left+new_w] = resized_img
    
    # 转换为CHW格式 + 归一化（像素值除以255）
    img_tensor = padded_img.transpose((2, 0, 1))  # HWC转CHW
    img_tensor = np.ascontiguousarray(img_tensor).astype(np.float32) / 255.0
    # 添加batch维度
    img_tensor = np.expand_dims(img_tensor, axis=0)
    
    return img_tensor, scale, pad_top, pad_left, orig_h, orig_w

三、复现YOLOv8后处理逻辑 #

ONNX模型输出原始检测张量，需执行和YOLOv8一致的置信度筛选、坐标映射、非极大值抑制（NMS）：

import onnxruntime as ort

def yolov8_postprocess(output, scale, pad_top, pad_left, orig_h, orig_w, conf_thres=0.25, iou_thres=0.7):
    # ONNX输出形状：[1, num_boxes, 85]（x,y,w,h,conf + 80类概率）
    pred = output[0]
    # 计算每个框的最大类别置信度
    conf = pred[:, 4:5] * pred[:, 5:]
    max_conf = np.max(conf, axis=1)
    max_cls_idx = np.argmax(conf, axis=1)
    
    # 筛选置信度大于阈值的框
    mask = max_conf > conf_thres
    boxes = pred[mask, :4]
    max_conf = max_conf[mask]
    max_cls_idx = max_cls_idx[mask]
    
    # 将中心坐标(x,y,w,h)转换为原图像的左上角/右下角坐标
    x_center, y_center, w, h = boxes[:, 0], boxes[:, 1], boxes[:, 2], boxes[:, 3]
    x1 = (x_center - w/2 - pad_left) / scale
    y1 = (y_center - h/2 - pad_top) / scale
    x2 = (x_center + w/2 - pad_left) / scale
    y2 = (y_center + h/2 - pad_top) / scale
    
    # 限制坐标在原图像范围内
    x1 = np.clip(x1, 0, orig_w)
    y1 = np.clip(y1, 0, orig_h)
    x2 = np.clip(x2, 0, orig_w)
    y2 = np.clip(y2, 0, orig_h)
    
    # 执行NMS（与YOLOv8默认参数一致：conf_thres=0.25, iou_thres=0.7）
    indices = cv2.dnn.NMSBoxes(x1.tolist(), y1.tolist(), max_conf.tolist(), conf_thres, iou_thres)
    
    # 整理最终结果
    final_boxes = []
    final_confs = []
    final_cls = []
    if len(indices) > 0:
        indices = indices.flatten()
        for idx in indices:
            final_boxes.append([int(x1[idx]), int(y1[idx]), int(x2[idx]), int(y2[idx])])
            final_confs.append(float(max_conf[idx]))
            final_cls.append(int(max_cls_idx[idx]))
    
    return final_boxes, final_confs, final_cls

四、完整推理流程 #

# 加载ONNX模型（GPU推理需安装ONNXRuntime-GPU，切换providers为["CUDAExecutionProvider"]）
sess = ort.InferenceSession("best.onnx", providers=["CPUExecutionProvider"])
input_name = sess.get_inputs()[0].name
output_name = sess.get_outputs()[0].name

# 预处理
img_path = "test_image.jpg"
img_tensor, scale, pad_top, pad_left, orig_h, orig_w = yolov8_preprocess(img_path)

# 推理
pred_output = sess.run([output_name], {input_name: img_tensor})

# 后处理
final_boxes, final_confs, final_cls = yolov8_postprocess(pred_output, scale, pad_top, pad_left, orig_h, orig_w)

# 绘制检测结果（可选）
img = cv2.imread(img_path)
for box, conf, cls_idx in zip(final_boxes, final_confs, final_cls):
    x1, y1, x2, y2 = box
    cv2.rectangle(img, (x1, y1), (x2, y2), (0, 255, 0), 2)
    cv2.putText(img, f"Class {cls_idx}: {conf:.2f}", (x1, y1-10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0,255,0), 2)
cv2.imwrite("detection_result.jpg", img)

常见问题排查 #

• 确保推理时输入尺寸与导出时imgsz完全一致，不可随意更改
• 若使用GPU推理，需确认ONNXRuntime-GPU版本与CUDA版本兼容
• 检查NMS参数是否与YOLOv8默认值一致，避免因阈值差异导致框的数量或位置不同

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