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

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)
相关推荐
Dymc26 分钟前
【目标检测之Ultralytics预测框颜色修改】
人工智能·yolo·目标检测·计算机视觉
lianyinghhh6 小时前
yolo8实现目标检测
yolo·目标检测·macos
灵智工坊LingzhiAI8 小时前
基于YOLO的足球检测Web应用:从训练到部署的完整实战
yolo
19891 天前
【零基础学AI】第31讲:目标检测 - YOLO算法
人工智能·rnn·yolo·目标检测·tensorflow·lstm
cver1231 天前
CSGO 训练数据集介绍-2,427 张图片 AI 游戏助手 游戏数据分析
人工智能·深度学习·yolo·目标检测·游戏·计算机视觉
Natsuagin1 天前
【保姆级目标检测教程】Ubuntu 20.04 部署 YOLOv13 全流程(附训练/推理代码)
yolo·目标检测·ubuntu·计算机视觉
要努力啊啊啊3 天前
YOLOv3-SPP Auto-Anchor 聚类调试指南!
人工智能·深度学习·yolo·目标检测·目标跟踪·数据挖掘
加油吧zkf3 天前
AI大模型如何重塑软件开发流程?——结合目标检测的深度实践与代码示例
开发语言·图像处理·人工智能·python·yolo
要努力啊啊啊5 天前
YOLOv2 正负样本分配机制详解
人工智能·深度学习·yolo·计算机视觉·目标跟踪