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)
相关推荐
XIAO·宝1 天前
深度学习------YOLOV3
人工智能·深度学习·yolo
newxtc1 天前
【广州公共资源交易-注册安全分析报告-无验证方式导致安全隐患】
开发语言·selenium·安全·yolo
总有刁民想爱朕ha1 天前
YOLO目标检测:一种用于无人机的新型轻量级目标检测网络
yolo·目标检测·无人机
baole9631 天前
YOLOv4简单基础学习
学习·yolo·目标跟踪
我叫侯小科1 天前
YOLOv4:目标检测界的 “集大成者”
人工智能·yolo·目标检测
麒羽7601 天前
YOLOv4:目标检测领域的 “速度与精度平衡大师”
yolo·目标检测·目标跟踪
前网易架构师-高司机1 天前
鸡蛋质量识别数据集,可识别染血的鸡蛋,棕色鸡蛋,钙沉积鸡蛋,污垢染色的鸡蛋,白鸡蛋,平均正确识别率可达89%,支持yolo, json, xml格式的标注
yolo·分类·数据集·缺陷·鸡蛋
Python图像识别1 天前
73_基于深度学习的水面漂浮垃圾检测系统(yolo11、yolov8、yolov5+UI界面+Python项目源码+模型+标注好的数据集)
python·深度学习·yolo
JANGHIGH2 天前
YOLO系列——Ubuntu20.04下通过conda虚拟环境源码安装opencv-4.10与opencv_contrib-4.10.0
opencv·yolo·conda
放羊郎2 天前
面向不同障碍物的自主导航策略
yolo·slam·识别·平滑·语义地图