YOLO 深度神经网络

部署预训练模型是机器学习中的一项常见任务，尤其是在使用不支持 PyTorch 等某些框架的硬件时。本指南全面介绍了如何从 PyTorch 导出预训练的 YOLO 系列模型，并使用 OpenCV 的 DNN 框架进行部署。出于演示目的，我们将以 YOLOX 模型为例，但该方法同样适用于其他受支持的模型。

注意： 目前 OpenCV 支持以下 YOLO 模型：

• YOLOX
• YOLONas
• YOLOv10
• YOLOv9
• YOLOv8
• YOLOv7
• YOLOv6
• YOLOv5
• YOLOv4

这一支持包含这些模型特有的预处理与后处理例程。其他更老版本的 YOLO 也可以由 OpenCV 以 Darknet 格式支持，但不在本教程的讨论范围内。

假设我们已经成功训练了 YOLOX 模型，接下来的步骤就是导出该模型并用 OpenCV 运行它。在继续之前，有几个关键问题需要解决。让我们逐一深入。

YOLO 的预处理与输出

理解 YOLO 系列检测器输入与输出的性质至关重要。与大多数深度神经网络（DNN）一样，这些检测器的输入尺寸通常随模型规模而变化。

模型规模	输入尺寸
小型模型（1）	416x416
中型模型（2）	640x640
大型模型（3）	1280x1280

上表简要列出了各种 YOLO 模型常用的输入尺寸。这些是标准输入形状。如果你的模型训练时使用的输入尺寸与表中不同，请务必使用训练时的尺寸。

下一个关键环节是理解 YOLO 检测器图像预处理的细节。虽然整个 YOLO 家族的基本预处理方法保持一致，但仍有一些细微而关键的差异必须考虑，否则会导致性能下降。其中最重要的是 resize type（缩放方式）和缩放后应用的 padding value（填充值）。例如，YOLOX 模型 使用 LetterBox 缩放方法和 114.0 的填充值。必须确保这些参数以及归一化常数与要导出的模型正确匹配。

关于模型输出，它通常是一个形状为 [BxNxC+5] 或 [BxNxC+4] 的张量，其中 B 表示批大小（batch size），N 表示锚点（anchors）数量，C 表示类别数（例如，在 COCO 数据集上训练的模型为 80 类）。前一种结构中外加的 5 对应于目标性得分（objectness score，obj）、置信度得分（confidence score，conf）以及边界框坐标 (cx, cy, w, h)。值得注意的是，YOLOv8 模型的输出形状为 [BxNxC+4]，其中没有显式的目标性得分，目标得分直接从类别得分推断。具体到 YOLOX 模型，还需要结合锚点把预测结果重新缩放回图像域。这一步将被集成进 ONNX 图中，我们会在后面的章节详细介绍这一过程。

PyTorch 模型导出

知道了预处理的参数之后，我们就可以继续把模型从 PyTorch 导出为 ONNX 图。既然本教程以 YOLOX 作为示例模型，我们就用它的导出过程来演示（除 YOLOv10 模型外，其余 YOLO 检测器的流程完全相同，YOLOv10 的导出细节见后文）。导出 YOLOX 可以直接使用导出脚本。具体来说，我们需要以下命令：

git clone https://github.com/Megvii-BaseDetection/YOLOX.git
cd YOLOX
wget https://github.com/Megvii-BaseDetection/YOLOX/releases/download/0.1.1rc0/yolox_s.pth # download pre-trained weights
python3 -m tools.export_onnx --output-name yolox_s.onnx -n yolox-s -c yolox_s.pth --decode_in_inference

注意： 这里的 --decode_in_inference 用于把锚框（anchor box）的生成包含在 ONNX 图内部。它把这个值设为 True，从而将锚点生成函数包含进来。

下面演示了导出脚本的极简版本（可用于 YOLOX 以外的模型），以备不时之需。不过，通常每个 YOLO 仓库都自带预定义的导出脚本。

import onnx
import torch
from onnxsim import simplify

# load the model state dict
ckpt = torch.load(ckpt_file, map_location="cpu")
model.load_state_dict(ckpt)

# prepare dummy input
dummy_input = torch.randn(args.batch_size, 3, exp.test_size[0], exp.test_size[1])

#export the model
torch.onnx._export(
    model,
    dummy_input,
    "yolox.onnx",
    input_names=["input"],
    output_names=["output"],
    dynamic_axes={"input": {0: 'batch'},
                  "output": {0: 'batch'}})

# use onnx-simplifier to reduce reduent model.
onnx_model = onnx.load(args.output_name)
model_simp, check = simplify(onnx_model)
assert check, "Simplified ONNX model could not be validated"
onnx.save(model_simp, args.output_name)

导出 YOLOv10 模型

要运行 YOLOv10，需要从 torch 模型中裁掉带动态形状的后处理部分，然后再转换为 ONNX。如果你在寻找如何裁掉后处理的方法，这里有官方 YOLOv10 的一个 fork 分支。该 fork 分支通过在后处理过程之前返回模型输出来裁掉后处理。把 torch 模型转换为 ONNX 的步骤如下：

git clone git@github.com:Abdurrahheem/yolov10.git
conda create -n yolov10 python=3.9
conda activate yolov10
pip install -r requirements.txt
python export_opencv.py --model=<model-name> --imgsz=<input-img-size>

默认情况下 --model="yolov10s"、--imgsz=(480,640)。这会生成 yolov10s.onnx 文件，可用于在 OpenCV 中推理。

用 OpenCV 示例运行 YOLO ONNX 检测器

有了模型的 ONNX 图之后，我们就可以直接用 OpenCV 的示例来运行。为此需要确保：

1. OpenCV 编译时带有 -DBUILD_EXAMPLES=ON 标志。
2. 进入 OpenCV 的 build 目录。
3. 运行以下命令：

./bin/example_dnn_yolo_detector --input=<path_to_your_input_file> \
                                --classes=<path_to_class_names_file> \
                                --thr=<confidence_threshold> \
                                --nms=<non_maximum_suppression_threshold> \
                                --mean=<mean_normalization_value> \
                                --scale=<scale_factor> \
                                --yolo=<yolo_model_version> \
                                --padvalue=<padding_value> \
                                --paddingmode=<padding_mode> \
                                --backend=<computation_backend> \
                                --target=<target_computation_device> \
                                --width=<model_input_width> \
                                --height=<model_input_height>

• --input：输入图像或视频的文件路径。如果省略，将从相机捕获帧。
• --classes：包含目标检测类别名称的文本文件路径。
• --thr：检测的置信度阈值（例如 0.5）。
• --nms：非极大值抑制阈值（例如 0.4）。
• --mean：均值归一化值（例如 0.0 表示不做均值归一化）。
• --scale：输入归一化的缩放因子（例如 1.0、1/255.0 等）。
• --yolo：YOLO 模型版本（例如 YOLOv3、YOLOv4 等）。
• --padvalue：预处理中使用的填充值（例如 114.0）。
• --paddingmode：图像缩放与填充的处理方式。选项：0（直接缩放到所需输入尺寸，不做额外处理），1（缩放后裁剪），2（保持原始图像宽高比缩放）。
• --backend：计算后端选择（0 为自动，1 为 Halide，2 为 OpenVINO 等）。
• --target：目标计算设备选择（0 为 CPU，1 为 OpenCL 等）。
• --device：相机设备编号（0 为默认相机）。如果未提供 --input，默认使用编号为 0 的相机。
• --width：模型输入宽度。不要与图像宽度混淆（例如 416、480、640、1280 等）。
• --height：模型输入高度。不要与图像高度混淆（例如 416、480、640、1280 等）。

这里的 mean、scale、padvalue、paddingmode 必须与我们在预处理章节中讨论的参数完全一致，模型才能得到与 PyTorch 一致的结果。

若要演示如何在没有自己预训练模型的情况下运行 OpenCV YOLO 示例，请按以下说明操作：

1. 确保你的平台上安装了 Python。
2. 确认 OpenCV 编译时带有 -DBUILD_EXAMPLES=ON 标志。

运行 YOLOX 检测器（使用默认值）：

git clone https://github.com/opencv/opencv_extra.git
cd opencv_extra/testdata/dnn
python download_models.py yolox_s_inf_decoder
cd ..
export OPENCV_TEST_DATA_PATH=$(pwd)
cd <build directory of OpenCV>
./bin/example_dnn_yolo_detector

这将用你的相机运行 YOLOX 检测器。 对于 YOLOv8（举例而言），还需要执行以下额外步骤：

cd opencv_extra/testdata/dnn
python download_models.py yolov8
cd ..
export OPENCV_TEST_DATA_PATH=$(pwd)
cd <build directory of OpenCV>

./bin/example_dnn_yolo_detector --model=onnx/models/yolov8n.onnx --yolo=yolov8 --mean=0.0 --scale=0.003921568627 --paddingmode=2 --padvalue=144.0 --thr=0.5 --nms=0.4 --rgb=0

对于 YOLOv10，按以下步骤操作：

cd opencv_extra/testdata/dnn
python download_models.py yolov10
cd ..
export OPENCV_TEST_DATA_PATH=$(pwd)
cd <build directory of OpenCV>

./bin/example_dnn_yolo_detector --model=onnx/models/yolov10s.onnx --yolo=yolov10 --width=640 --height=480  --scale=0.003921568627 --padvalue=114

这将在系统找到的第一个相机上运行 YOLOv10 检测器。如果你想在图像/视频文件上运行，可以使用 --input 选项指定文件路径。

构建自定义流水线

有时需要对推理流水线做一些自定义调整。使用 OpenCV DNN 模块，这也相当容易实现。下面概述示例的实现细节：

• 导入所需的头文件：

#include <opencv2/dnn.hpp>
#include <opencv2/imgproc.hpp>
#include <opencv2/imgcodecs.hpp>
#include <fstream>
#include <sstream>
#include "iostream"
#include "common.hpp"
#include <opencv2/highgui.hpp>

• 读取 ONNX 图并创建神经网络模型：

Net net = readNet(weightPath);
int backend = parser.get<int>("backend");
net.setPreferableBackend(backend);
net.setPreferableTarget(parser.get<int>("target"));

• 读取图像并进行预处理：

float paddingValue = parser.get<float>("padvalue");
bool swapRB = parser.get<bool>("rgb");
int inpWidth = parser.get<int>("width");
int inpHeight = parser.get<int>("height");
Scalar scale = parser.get<Scalar>("scale");
Scalar mean = parser.get<Scalar>("mean");
ImagePaddingMode paddingMode = static_cast<ImagePaddingMode>(parser.get<int>("paddingmode"));

Size size(inpWidth, inpHeight);
Image2BlobParams imgParams(
    scale,
    size,
    mean,
    swapRB,
    CV_32F,
    DNN_LAYOUT_NCHW,
    paddingMode,
    paddingValue);

// rescale boxes back to original image
Image2BlobParams paramNet;
        paramNet.scalefactor = scale;
        paramNet.size = size;
        paramNet.mean = mean;
        paramNet.swapRB = swapRB;
        paramNet.paddingmode = paddingMode;

inp = blobFromImageWithParams(img, imgParams);

• 推理：

std::vector<Mat> outs;
std::vector<int> keep_classIds;
std::vector<float> keep_confidences;
std::vector<Rect2d> keep_boxes;
std::vector<Rect> boxes;

net.setInput(inp);
net.forward(outs, net.getUnconnectedOutLayersNames());

• 后处理：

所有后处理步骤都在函数 yoloPostProcess 中实现。请注意，NMS 步骤不包含在 ONNX 图中，示例使用 OpenCV 函数来完成。

yoloPostProcessing(
    outs, keep_classIds, keep_confidences, keep_boxes,
    confThreshold, nmsThreshold,
    yolo_model,
    nc);

• 绘制预测的边界框：

for (auto box : keep_boxes)
{
    boxes.push_back(Rect(cvFloor(box.x), cvFloor(box.y), cvFloor(box.width - box.x), cvFloor(box.height - box.y)));
}

paramNet.blobRectsToImageRects(boxes, boxes, img.size());

for (size_t idx = 0; idx < boxes.size(); ++idx)
{
    Rect box = boxes[idx];
    drawPrediction(keep_classIds[idx], keep_confidences[idx], box.x, box.y,
            box.width + box.x, box.height + box.y, img);
}

const std::string kWinName = "Yolo Object Detector";
namedWindow(kWinName, WINDOW_NORMAL);
imshow(kWinName, img);