OpenCV 调用onnx模型

本文最后更新于：2024年7月1日上午

OpenCV 支持 Onnx 模型推断，本文记录相关内容。

简介

OpenCV是一个基于BSD许可发行的跨平台计算机视觉和机器学习软件库(开源)，可以运行在Linux、Windows、Android和Mac OS操作系统上。可以将pytorch中训练好的模型使用ONNX导出，再使用opencv中的dnn模块直接进行加载使用。

dnn

cv2.dnn 是 OpenCV 库中的一个模块，专门用于深度学习计算。这个模块提供了一套用于执行深度学习模型的基础操作，比如卷积、池化、归一化、激活函数和各种前向传播等。使用 cv2.dnn，可以加载预训练的深度学习模型，如卷积神经网络（CNN），进行图像分类、目标检测、图像分割等任务。

支持模型种类：

Caffe
Onnx
Tensorflow
Pytorch
DarkNet
TFLite
ModelOptimizer

blob

blob 代表着一个 N 维数组，它可以包含图像数据、矩阵或其他类型的数据。在深度学习框架中，blob 通常用于表示神经网络的输入或输出。

在 cv2.dnn 中，blob 主要用于以下几个目的：

数据预处理：在将图像数据输入到深度学习模型之前，通常需要对图像进行缩放、裁剪、归一化等操作。blobFromImage 函数就是用来将图像转换为一个 blob，同时进行这些预处理步骤。
模型输入：深度学习模型需要标准的输入格式，比如固定的尺寸、归一化的值等。blob 可以确保输入数据符合模型的要求。
内存管理：blob 对象在内存中占用空间，当不再需要时，可以通过 blob.empty() 方法释放内存。
前向传播：在执行模型的前向传播时，blob 用于存储输入数据，并从网络中获取输出数据。

调用 Onnx

需要用到 cv2.dnn.readNetFromONNX 方法，参数为 onnx 文件路径（当前不支持动态尺寸 Onnx 模型）

1	`cv2.dnn.readNetFromONNX(onnx_path)`

准备图像数据

需要将图像转换为 Onnx 相同位深度的数据

image = vv.cv_rgb_imread(img_path)
image = image - mean
image = image / std
image = image.astype('float32')

创建 blob

1	`blob = cv2.dnn.blobFromImage(image, scalefactor=1.0, size=(224, 224), swapRB=True, crop=False)`

设置 blob 为网络输入

1	`net.setInput(blob)`

模型前向传播

1	`output = net.forward()`

示例代码

Python

import cv2
import vvdutils as vv


if __name__ == '__main__':
    onnx_path = 'assets/model-fp32.onnx'
    net = cv2.dnn.readNetFromONNX(onnx_path)
    net.setPreferableBackend(cv2.dnn.DNN_BACKEND_OPENCV)

    output_layer_names = net.getUnconnectedOutLayersNames()

    mean=[
            123.675,
            116.28,
            103.53,
        ]

    std=[
        58.395,
        57.12,
        57.375,
    ]

    img_path = 'assets/test.png'

    image = vv.cv_rgb_imread(img_path)
    image = image - mean
    image = image / std
    image = image.astype('float32')

    blob = cv2.dnn.blobFromImage(image, swapRB=False, crop=False)
    net.setInput(blob)
    output_probs = net.forward()

    vv.PIS(output_probs[0][0])

C++

#include <iostream>
#include <string>
#include <vector>
#include<opencv2/opencv.hpp>
#include <opencv2/dnn.hpp>
using namespace std;
void readImagesInFolder(const std::string& folderPath, std::vector<cv::Mat>& images)
{
    cv::String path(folderPath + "/*.jpg"); // 这里假设你的图片格式是.jpg，如果是其他格式请相应修改
    std::vector<cv::String> fileNames;
    cv::glob(path, fileNames, true); // 通过glob函数获取文件夹内所有符合格式的文件名
    for (const auto& fileName : fileNames)
    {   // 使用imread函数读取图片
        cv::Mat bgrImage = cv::imread(fileName, cv::IMREAD_COLOR); 
         // 图片格式转化bgr-->rgb
        if (!bgrImage.empty())
        {	
        	cv::Mat rgbImage;
            cv::cvtColor(bgrImage, rgbImage, cv::COLOR_BGR2RGB);
            images.push_back(rgbImage);
        }
    }
}

cv::Mat transformation(const cv::Mat& image, const cv::Size & targetSize, const cv::Scalar& mean, const cv::Scalar& std) {

    cv::Mat resizedImage;
    //图片尺寸缩放
    cv::resize(image, resizedImage, targetSize, 0, 0, cv::INTER_AREA);
    cv::Mat normalized;
    resizedImage.convertTo(normalized, CV_32F);
    cv::subtract(normalized / 255.0, mean, normalized);
    cv::divide(normalized, std, normalized);
    return normalized;
}
cv::dnn::Net loadModel(const string& onnx_path) {
    cv::dnn::Net net = cv::dnn::readNetFromONNX(onnx_path);
    return net;
}
int main()
{   // 图片存放文件路径
    string folderPath = "D:/C++_demo/opencv_onnx_gpu/CAMO/c";
    std::vector<cv::Mat> rgbImages;
    readImagesInFolder(folderPath, rgbImages);

   // string image_path = "./animal-1.jpg";
    // 加载ONNX模型
    string onnx_path = "D:/C++_demo/opencv_onnx_gpu/PFNet.onnx";
    cv::dnn::Net net = loadModel(onnx_path);
    // 设置CUDA为后端
    net.setPreferableBackend(cv::dnn::DNN_BACKEND_CUDA);
    net.setPreferableTarget(cv::dnn::DNN_TARGET_CUDA);
    cv::Mat output_prob;
    std::vector<cv::Mat> output_probs;
    std::vector<cv::String> output_layer_names = net.getUnconnectedOutLayersNames();

    // 定义目标图像大小
    cv::Size targetSize(416, 416);
    // 定义每个通道的归一化参数
    cv::Scalar mean(0.485, 0.456, 0.406); // 均值
    cv::Scalar std(0.229, 0.224, 0.225);  // 标准差

    // 开始计时
    auto start = chrono::high_resolution_clock::now();
    for (const auto& rgbImage : rgbImages) {
        // 获取图像的大小
        cv::Size originalSize(rgbImage.cols, rgbImage.rows);
        //cv::imshow("输入窗口", rgbImage);
        //cv::waitKey(0);
        //cv::destroyAllWindows();
        // 图片归一化
        cv::Mat normalized = transformation(rgbImage, targetSize, mean, std);
        std::cout << normalized.size() << std::endl;
        cv::Mat blob = cv::dnn::blobFromImage(normalized);
        // 将Blob设置为模型的输入
        net.setInput(blob);
        // 运行前向传播
        net.forward(output_probs, output_layer_names);
        // 获取最完整的预测
        cv::Mat prediction = output_probs[3];
        // 预测图变mask
        cv::Mat mask;
        cv::resize(prediction.reshape(1, 416) * 255.0, mask, originalSize, 0, 0, cv::INTER_AREA);
    }
    auto end = std::chrono::high_resolution_clock::now();
    // 计算耗时
    std::chrono::duration<double> elapsed = end - start;
    double elapsedTime = elapsed.count();
    // 打印耗时
    std::cout << "Elapsed time: " << elapsedTime << " seconds" << std::endl;
    return 0;
}