YOLOv5 ONNX推理实战：CUDA加速+可视化

原创

于 2025-06-20 19:12:32 发布 · 595 阅读

7 ·

CC 4.0 BY-SA版权

文章标签：

#YOLO

本篇记录了如何使用 onnxruntime-gpu 和 OpenCV 对导出的 YOLOv5 ONNX 模型进行推理，支持：

✅ CUDA 加速
✅ 自动 Letterbox 缩放与填充
✅ 非极大值抑制（NMS）
✅ 输出坐标自动还原到原图尺寸
✅ 实时可视化

📦 依赖安装

pip install onnxruntime-gpu opencv-python numpy

🧠 推理代码（支持 YOLOv5 ONNX）

import onnxruntime as ort
import numpy as np
import cv2
import time

class YOLOv5ONNX:
    def __init__(self, model_path, input_size=(640, 640), providers=['CUDAExecutionProvider']):
        self.input_size = input_size  # (w, h)
        self.session = ort.InferenceSession(model_path, providers=providers)
        self.input_name = self.session.get_inputs()[0].name

    def letterbox(self, image, new_shape=(640, 640), color=(114, 114, 114)):
        shape = image.shape[:2]  # current shape [height, width]
        r = min(new_shape[1] / shape[0], new_shape[0] / shape[1])
        new_unpad = (int(round(shape[1] * r)), int(round(shape[0] * r)))
        dw = new_shape[0] - new_unpad[0]
        dh = new_shape[1] - new_unpad[1]
        dw /= 2  # divide padding into 2 sides
        dh /= 2

        resized = cv2.resize(image, new_unpad, interpolation=cv2.INTER_LINEAR)
        top, bottom = int(round(dh - 0.1)), int(round(dh + 0.1))
        left, right =