OpenCV 加载 ONNX 模型

OpenCV 加载 ONNX 模型

• 引言
• 模型获取
• 加载并推理模型

引言

在工业视觉领域 OpenCV 使用较为广泛，其 DNN 模块支持深度学习模型的推理，如果在项目中使用了 OpenCV，那么很容易添加深度学习支持。

模型获取

OpenCV 不适合用于搭建模型，通常使用其他框架训练模型。ONNX 作为通用的模型描述格式被众多框架支持，这里推荐使用 ONNX 作为模型保存格式。以 PyTorch 为例，保存 ONNX 格式方法如下：

x = torch.randn(1, 3, 640, 640)
torch.onnx.export(network, x, "model.onnx")

加载并推理模型

Python代码如下：

# opencv 推理
net = cv2.dnn.readNetFromONNX("model.onnx")  # 加载训练好的识别模型
image = cv2.imread("numbers/1.bmp")  # 读取图片
blob = cv2.dnn.blobFromImage(image)  # 由图片加载数据 这里还可以进行缩放、归一化等预处理
net.setInput(blob)  # 设置模型输入
out = net.forward()  # 推理出结果

C++代码如下：

// opencv 推理
cv::dnn::Net net = cv::dnn::readNetFromONNX("save.onnx");  // 加载训练好的识别模型
cv::Mat image = cv::imread("numbers/1.bmp");  // 读取图片
cv::Mat blob = cv::dnn::blobFromImage(image);  // 由图片加载数据 这里还可以进行缩放、归一化等预处理
net.setInput(blob);  // 设置模型输入
cv::Mat predict = net.forward(); // 推理出结果