在 OpenCV 中运行经过预训练的 YOLO 模型

部署预训练模型是机器学习中的一项常见任务，尤其是在使用不支持某些框架（如 PyTorch）的硬件时。本指南全面概述了如何从 PyTorch 导出经过预训练的 YOLO 系列模型，以及如何使用 OpenCV 的 DNN 框架部署它们。出于演示目的，我们将重点关注 YOLOX 模型，但该方法适用于其他受支持的模型。

注意

目前，OpenCV 支持以下 YOLO 模型

此支持包括针对这些模型特有的预处理和后处理例程。虽然 OpenCV 也在 Darknet 格式中支持 YOLO 的其他旧版本，但它们超出了本教程的范围。

假设我们已成功训练了 YOLOX 模型，后续步骤包括使用 OpenCV 导出和运行此模型。在继续此过程之前，有几个关键的注意事项需要解决。让我们深入研究这些方面。

YOLO 的预处理和输出

了解与 YOLO 系列检测器关联的输入和输出的本质至关重要。这些检测器类似于大多数深度神经网络 (DNN)，其输入大小通常会因模型的规模而异。

模型比例	输入大小
小型模型 ¹	416x416
中型模型 ²	640x640
大型模型 ³	1280x1280

此表格提供了一个快速参考，以了解各种 YOLO 模型输入中常用的不同输入维度。这些是标准输入形状。请务必使用训练模型时的输入大小，如果它与表中提到的输入大小不同。

此过程中下一个关键因素，涉及 YOLO 检测器图像预处理的具体内容。尽管基本预处理方法在 YOLO 家族中保持一致，但有一些细微但至关重要的区别必须加以考虑，以避免性能下降。其中关键的要素包括 调整大小的类型 和调整大小后应用的 填充值。例如，YOLOX 模型使用 信箱 调整大小方法和 114.0 的填充值。必须确保这些参数，以及归一化常量，适当匹配要导出的模型。

关于模型的输出，它通常采用维度为 [BxNxC+5] 或 [BxNxC+4] 的张量的形式，其中，“B”表示批大小，“N”表示锚点数量，“C”表示类数（例如，如果在 COCO 数据集上训练模型，则为 80 类）。前一个张量结构中的其他 5 对应于物体得分（obj）、置信度得分（conf）和边界框坐标（cx、cy、w、h）。值得注意的是，YOLOv8 模型的输出形状为 [BxNxC+4]，其中没有明确的物体得分，并且物体得分直接从类得分推断出来。具体来说，对于 YOLOX 模型，还需要合并锚点以将预测重新调整回图像域。此步骤将集成到 ONNX 图中，我们会在后续部分进一步详细阐述这一过程。

PyTorch 模型导出

现在我们已经知道了预处理的参数，我们可以继续将模型从 PyTorch 导出到 ONNX 图。由于在本教程中我们使用 YOLOX 作为示例模型，因此让我们使用其导出进行演示（其余 YOLO 检测器的流程相同）。要导出 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 # 下载预先训练的权重
python3 -m tools.export_onnx --output-name yolox_s.onnx -n yolox-s -c yolox_s.pth --decode_in_inference

注意：此处 --decode_in_inference 用于在 ONNX 图中本身包括锚框创建。它将该值设为 True，这随后包括锚生成函数。

在下面，我们演示了导出脚本的最小版本（可用于非 YOLOX 的其他模型），以防需要。但是，通常每个 YOLO 储存库都有预先定义的导出脚本。

import onnx
import torch
from onnxsim import simplify
 
# 加载模型状态字典
ckpt = torch.load(ckpt_file, map_location="cpu")
model.load_state_dict(ckpt)
 
# 准备虚拟输入
dummy_input = torch.randn(args.batch_size, 3, exp.test_size[0], exp.test_size[1])
 
# 导出模型
torch.onnx._export(
model,
dummy_input,
    "yolox.onnx",
input_names=["input"],
output_names=["output"],
dynamic_axes={"input": {0: 'batch'},
                  "output": {0: 'batch'}})
 
# 使用 onnx-simplifier 简化冗余模型。
onnx_model = onnx.load(args.output_name)
model_simp, check = simplify(onnx_model)
assert check, "无法验证简化的 ONNX 模型"
onnx.save(model_simp, args.output_name)

使用 OpenCV 样本运行 Yolo ONNX 检测器

一旦有了该模型的 ONNX 图，我们可以直接使用 OpenCV 的样本运行。为此，我们需要确保

OpenCV 使用 -DBUILD_EXAMLES=ON 标记生成。
导航至 OpenCV 的build目录
运行以下命令

./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>

视频演示

–input：指向输入图像或视频的文件路径。如果省略，它将从摄像头捕获帧。
–classes：包含对象检测类别的文本文件的路径。
–thr：检测的置信阈值（例如 0.5）。
–nms：非最大抑制阈值（例如 0.4）。
–mean：平均归一化值（例如 0.0，表示无平均归一化）。
–scale：输入归一化的缩放比例（例如 1.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 的摄像头。

在此，mean、scale、padvalue、paddingmode 应与我们在预处理部分讨论的完全一致，以使该模型与 PyTorch 中的结果相匹配。

若要了解如何在没有自己预训练的模型的情况下运行 OpenCV YOLO 样本，请按照以下说明进行操作

确保您的平台上已安装 Python。
确认 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 <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 <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

构建自定义管道

有时需要对推理管道进行一些自定义调整。使用 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<float>("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);
 
    // 将边界框缩放到原始图像大小
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);

绘制预测框

        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 目标检测器";
namedWindow(kWinName, WINDOW_NORMAL);
imshow(kWinName, img);


原作者	Alessandro de Oliveira Faria
扩展者	Abduragim Shtanchaev
兼容性	OpenCV >= 4.9.0