YOLO 模型 ONNX 导出与跨平台部署流程

1. 模型导出（YOLO .pt → ONNX）

导出工具：使用 YOLO 框架（如 PyTorch YOLOv5/v7/v8）自带的导出脚本。

示例命令 ：

bash 复制代码

export model=best.pt format=onnx task=detect dynamic=True  opset=17

输出文件 ：生成 .onnx 格式的模型文件，支持跨平台部署（如 ONNX Runtime、TensorRT、OpenVINO 等）。

2. 预处理流程

python 复制代码

import numpy as np
import cv2

def preprocessing(image_path):
    input_height, input_width = 1280, 1280  # 模型输入尺寸
    original_image = cv2.imread(image_path)
    orig_height, orig_width = original_image.shape[:2]
    
    global scale  # 全局变量用于后处理
    scale = min(input_height / orig_height, input_width / orig_width)
    new_height = int(orig_height * scale)
    new_width = int(orig_width * scale)
    
    # 等比缩放并填充为正方形
    resized = cv2.resize(original_image, (new_width, new_height))
    padded_image = np.zeros((input_height, input_width, 3), dtype=np.uint8)
    top = (input_height - new_height) // 2
    bottom = input_height - new_height - top
    left = (input_width - new_width) // 2
    right = input_width - new_width - left
    padded_image[top:top+new_height, left:left+new_width] = resized
    
    # 归一化与通道转换
    image_data = padded_image.transpose(2, 0, 1)  # HWC → CHW
    image_data = np.expand_dims(image_data, axis=0).astype(np.float32) / 255.0  # 归一化到 [0,1]
    return image_data

关键步骤 ：
1. 等比缩放：保持宽高比缩放图像。
2. 填充为正方形 ：使用黑色填充（np.zeros）确保输入尺寸为 (1280, 1280)。
3. 通道转换 ：HWC → CHW（符合 ONNX 模型输入格式）。
4. 归一化 ：除以 255 将像素值归一化到 [0, 1]。

3. ONNX 推理

python 复制代码

import onnxruntime as ort

def predict(model_path, image_data):
    session = ort.InferenceSession(model_path, providers=['CPUExecutionProvider'])
    input_name = session.get_inputs()[0].name
    output_name = session.get_outputs()[0].name
    outputs = session.run([output_name], {input_name: image_data})
    return outputs[0]  # 输出形状: (1, 5, 33600)

输入：预处理后的 image_data（形状 [1, 3, 1280, 1280]）。
输出：检测结果 detections（形状 [1, 5, 33600]），每行包含 [x_center, y_center, width, height, confidence]。

4. 后处理与结果判断

python 复制代码

def isdetected(boxes):
    confidence_threshold = 0.1
    if np.max(boxes[:, 4]) < confidence_threshold:
        print("未检测到任何目标")
    else:
        max_conf_index = np.argmax(boxes[:, 4])
        best_box = boxes[max_conf_index]
        print(f"检测到目标，最高置信度: {np.max(boxes[:, 4]):.4f}")
        print(f"检测框数据: {best_box}")

逻辑：
- 过滤置信度低于阈值的检测框。
- 打印置信度最高的检测框信息。

5. 检测框绘制

python 复制代码

def draw_boxes(image_path, boxes, confidence_threshold=0.5):
    global scale  # 使用预处理阶段的缩放比例
    original_image = cv2.imread(image_path)
    image = original_image.copy()
    filtered_boxes = boxes[boxes[:, 4] >= confidence_threshold]
    
    for box in filtered_boxes:
        x_center, y_center, width, height, confidence = box
        
        # 坐标反向转换（从模型输出到原始图像坐标）
        x_center = x_center / scale
        y_center = (y_center - 280) / scale  # ← 注意：此处可能需要修正
        width = width / scale
        height = height / scale
        
        x_min = int(x_center - width / 2)
        y_min = int(y_center - height / 2)
        x_max = int(x_center + width / 2)
        y_max = int(y_center + height / 2)
        
        # 绘制矩形框与标签
        cv2.rectangle(image, (x_min, y_min), (x_max, y_max), (0, 255, 0), 2)
        cv2.putText(image, f"{confidence:.2f}", (x_min, y_min - 5), 
                    cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 1)
    
    cv2.imshow("Detection", image)
    cv2.waitKey(0)
    cv2.destroyAllWindows()

关键点 ：
- 坐标反向转换 ：根据 scale 和填充信息（top, left）将检测框坐标还原到原始图像。

6. 完整调用示例

python 复制代码

# 1. 预处理图像
image_data = preprocessing("example.jpg")

# 2. ONNX 推理
model_path = "best.onnx"
detections = predict(model_path, image_data)

# 3. 判断检测结果
boxes = detections[0].T  # 转置为 (33600, 5)
isdetected(boxes)

# 4. 绘制检测框
draw_boxes("example.jpg", boxes)