在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