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
的运行耗时差不多。
结果可视化如下,可以正确地对车辆目标的位置进行修正。