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

部署预训练模型是机器学习中的常见任务，特别是在使用不支持PyTorch等特定框架的硬件时。本指南全面概述了如何从PyTorch导出预训练的YOLO系列模型，并使用OpenCV的DNN框架进行部署。为了演示目的，我们将重点关注YOLOX模型，但此方法适用于其他受支持的模型。

注意

目前，OpenCV支持以下YOLO模型：

这种支持包括针对这些模型的预处理和后处理例程。虽然OpenCV也支持Darknet格式的其他旧版本YOLO，但它们不在本教程的范围之内。

假设我们已经成功训练了YOLOX模型，下一步是使用OpenCV导出并运行此模型。在进行此过程之前，需要解决几个关键的考虑因素。让我们深入探讨这些方面。

YOLO的预处理和输出

理解YOLO系列检测器输入和输出的性质至关重要。这些检测器，类似于大多数深度神经网络（DNN），通常根据模型的规模表现出输入大小的变化。

模型规模	输入尺寸
小型模型 ¹	416x416
中型模型 ²	640x640
大型模型 ³	1280x1280

此表提供了快速参考，以帮助理解各种YOLO模型输入中常用的不同输入尺寸。这些是标准输入形状。如果训练模型的输入大小与表中提到的不同，请确保使用您训练模型时的输入大小。

过程中的下一个关键要素是理解YOLO检测器图像预处理的细节。尽管YOLO系列的基本预处理方法保持一致，但仍存在微妙但关键的差异，必须加以考虑以避免性能下降。其中最重要的是resize type（调整大小类型）和padding value（填充值）。例如，YOLOX模型使用LetterBox调整大小方法和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作为示例模型，因此让我们使用它的导出进行演示（除YOLOv10模型外，其他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 # 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是为了将锚框创建包含在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的此分支。该分支通过在后处理程序本身之前返回模型输出来切断后处理。要将torch模型转换为ONNX，请按照以下步骤操作。

git clone [email protected]: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的示例来运行它。为此，我们需要确保：

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> \
--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示例，请遵循以下说明：

确保您的平台上已安装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 <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);


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