基於OpenCV的DNN網絡推理——C++實現

OpenCV從3.3版本就開始引入DNN模塊，現在已經是4.5版本了，DNN模塊的支持度也更好了。目前OpenCV已經支持ONNX格式的模型加載和推理，后端的推理引擎也有了多種選擇。

而Pytorch作為目前易用性相對最好的深度學習訓練框架，使用非常廣泛。Pytorch的pth格式模型沒法直接用OpenCV加載，但可以轉換成ONNX格式使用。

1. OpenCV編譯

首先在Ubuntu下編譯OpenCV，默認配置選項已經可以支持加載ONNX模型進行推理，不需要contrib庫支持或其他的特殊配置。

我在編譯過程中，按照習慣僅指定了CMAKE_INSTALL_PREFIX參數，用來指定編譯后的安裝路徑，方便后期管理，不干擾系統目錄。

2. ONNX模型導出

Pytorch已經原生支持導出ONNX模型，具體可以參見官方教程和文檔。

一個最簡化的導出Demo如下，核心的只有torch.onnx.export一個函數。

import torch
import torchvision

model = torchvision.models.mobilenet_v2(pretrained=True)
x = torch.rand(1, 3, 224, 224)
save_path = "./model_file/mobilenetv2.onnx"
torch.onnx.export(model, x, save_path,
    input_names=["input"], output_names=["output"], opset_version=11)

3. OpenCV推理

利用OpenCV的DNN模塊進行網絡推理，可以參考官方教程，以及源碼包中的samples/dnn/classification.cpp。

這邊以我之前做的一個目標修正的小網絡為例，演示一個簡化的Demo。

3.1 模型結構

模型借鑒MobilenetV2中的InvertedResidual模塊，搭建特征提取部分，輸入64*64的圖像，最后通過pooling層和fc層，輸出一個10維的向量。

輸出向量的含義分別為：

• 0：目標Confidence，是否包含目標；
• 1-5：目標Class，分類得分；
• 6-9：2D-BBox，按xywh方式表述。

3.2 模型推理

一個簡化的推理Demo如下所示。

#include <iostream>
#include <opencv2/opencv.hpp>

using namespace cv;

int main(int argc, char** argv)
{
    // load model
    String model = "../model_file/mobilenet.onnx";
    dnn::Net net = dnn::readNetFromONNX(model);
    
    // load image data
    float scale = 0.0078125;
    Scalar mean = Scalar(128.0, 128.0, 128.0);
    Mat frame = imread("../images/car.png");
    Mat blob;
    dnn::blobFromImage(frame, blob, scale, Size(64, 64), mean, false, false);

    // inference and time
    double t = getTickCount(); 
    net.setInput(blob);
    Mat output = net.forward();
    t = (getTickCount() - t) / getTickFrequency();
    std::cout << "Output shape: " << output.size() << ", Time-cost: " << t << std::endl;

    // visualize
    visualize_result(frame, output);
    return 0;
}

其中結果可視化的代碼不再贅述，主要是把回歸出的2DBBox畫出來。

結果輸出如下：

Output shape: [10 x 1], Time-cost: 0.0013024

輸出結果維度符合預期，運行時間也與python調用onnxruntime的運行耗時差不多。

結果可視化如下，可以正確地對車輛目標的位置進行修正。