C++ Qt / VS2019 +opencv + onnxruntime 部署语义分割模型【经验2】

前序工作

C++ Qt / VS2019 +opencv + onnxruntime 部署语义分割模型【经验】

引言

前序工作中介绍了Pytorch模型如何转为ONNX格式，以及在Python中如何使用onnx模型
介绍了如何在VA或QT中配置Onnxruntime运行库

本文重点列出全部源代码及其使用

依赖库

onnxruntime: 1.8.1
opencv: 330

头文件

#pragma once
#include <string>
#include <onnxruntime_cxx_api.h>
#include <iostream>
#include <vector>
#include <opencv2/opencv.hpp>
#include <opencv2/dnn.hpp>
#include <time.h>class ObjectSeg
{
protected:Ort::Env env_;Ort::SessionOptions session_options_;Ort::Session session_{ nullptr };Ort::RunOptions run_options_{ nullptr };std::vector<Ort::Value> input_tensors_;std::vector<const char*> input_node_names_;std::vector<int64_t> input_node_dims_;size_t input_tensor_size_{ 1 };std::vector<const char*> out_node_names_;size_t out_tensor_size_{ 1 };int image_h;int image_w;cv::Mat normalize(cv::Mat& image);cv::Mat preprocess(cv::Mat image);public:ObjectSeg() = delete;ObjectSeg(std::wstring model_path, int num_threads, std::vector<int64_t> input_node_dims);cv::Mat predict_image(cv::Mat& src);void predict_image(const std::string& src_path, const std::string& dst_path);void predict_camera();};

源文件

#include "Seg.h"ObjectSeg::ObjectSeg(std::wstring model_path, int num_threads = 1, std::vector<int64_t> input_node_dims = { 1, 3, 64, 64 }) {input_node_dims_ = input_node_dims;for (int64_t i : input_node_dims_) {input_tensor_size_ *= i;out_tensor_size_ *= i;}//std::cout << input_tensor_size_ << std::endl;session_options_.SetIntraOpNumThreads(num_threads);session_options_.SetGraphOptimizationLevel(GraphOptimizationLevel::ORT_ENABLE_EXTENDED);OrtCUDAProviderOptions cuda_options{0,OrtCudnnConvAlgoSearch::EXHAUSTIVE,std::numeric_limits<size_t>::max(),0,true};session_options_.AppendExecutionProvider_CUDA(cuda_options);std::cout << "************* Infer model on GPU! *************" << std::endl;try {session_ = Ort::Session(env_, model_path.c_str(), session_options_);}catch (...) {}Ort::AllocatorWithDefaultOptions allocator;//获取输入nameconst char* input_name = session_.GetInputName(0, allocator);input_node_names_ = { input_name };//std::cout << "input name:" << input_name << std::endl;const char* output_name = session_.GetOutputName(0, allocator);out_node_names_ = { output_name };//std::cout << "output name:" << output_name << std::endl;
}cv::Mat ObjectSeg::normalize(cv::Mat& image) {std::vector<cv::Mat> channels, normalized_image;cv::split(image, channels);cv::Mat r, g, b;b = channels.at(0);g = channels.at(1);r = channels.at(2);b = (b / 255. - 0.485) / 0.229;g = (g / 255. - 0.456) / 0.224;r = (r / 255. - 0.406) / 0.225;normalized_image.push_back(r);normalized_image.push_back(r);normalized_image.push_back(g);normalized_image.push_back(b);cv::Mat out = cv::Mat(image.rows, image.cols, CV_32F);cv::merge(normalized_image, out);return out;
}/*
* preprocess: resize -> normalize
*/
cv::Mat ObjectSeg::preprocess(cv::Mat image) {image_h = image.rows;image_w = image.cols;cv::Mat dst, dst_float, normalized_image;cv::resize(image, dst, cv::Size(int(input_node_dims_[3]), int(input_node_dims_[2])), 0, 0);dst.convertTo(dst_float, CV_32F);normalized_image = normalize(dst_float);return normalized_image;
}/*
* postprocess: preprocessed image -> infer -> postprocess
*/
cv::Mat ObjectSeg::predict_image(cv::Mat& src) {cv::Mat preprocessed_image = preprocess(src);cv::Mat blob = cv::dnn::blobFromImage(preprocessed_image, 1, cv::Size(int(input_node_dims_[3]), int(input_node_dims_[2])), cv::Scalar(0, 0, 0), false);//std::cout << "load image success." << std::endl;// create input tensorauto memory_info = Ort::MemoryInfo::CreateCpu(OrtArenaAllocator, OrtMemTypeDefault);input_tensors_.emplace_back(Ort::Value::CreateTensor<float>(memory_info, blob.ptr<float>(), blob.total(), input_node_dims_.data(), input_node_dims_.size()));std::vector<Ort::Value> output_tensors_ = session_.Run(Ort::RunOptions{ nullptr },input_node_names_.data(),input_tensors_.data(),input_node_names_.size(),out_node_names_.data(),out_node_names_.size());// post progress// 3. post process.Ort::Value& scores = output_tensors_.at(0); // (1,21,h,w)auto scores_dims = scores.GetTypeInfo().GetTensorTypeAndShapeInfo().GetShape();const unsigned int output_classes = scores_dims.at(1);const unsigned int output_height = scores_dims.at(2);const unsigned int output_width = scores_dims.at(3);// time cost!cv::Mat class_mat = cv::Mat(output_height, output_width, CV_8UC3, cv::Scalar(0));cv::Mat color_mat = class_mat.clone();for (unsigned int i = 0; i < output_height; ++i){uchar* p_class = class_mat.ptr<uchar>(i);cv::Vec3b* p_color = color_mat.ptr<cv::Vec3b>(i);for (unsigned int j = 0; j < output_width; ++j){// argmaxunsigned int max_label = 0;float max_conf = scores.At<float>({ 0, 0, i, j });for (unsigned int l = 0; l < output_classes; ++l){float conf = scores.At<float>({ 0, l, i, j });if (conf > max_conf){max_conf = conf;max_label = l;}}if (max_label == 0) continue;// assign label for pixel(i,j)p_class[j] = cv::saturate_cast<uchar>(max_label);// assign color for detected class at pixel(i,j).p_color[j][0] = cv::saturate_cast<uchar>(255); // ((max_label % 10) * 20);p_color[j][1] = cv::saturate_cast<uchar>(255);// ((max_label % 5) * 40);p_color[j][2] = cv::saturate_cast<uchar>(255); // ((max_label % 10) * 20);// assign names map}}//cv::imwrite("1.png", color_mat);input_tensors_.clear();return color_mat;
}void ObjectSeg::predict_image(const std::string& src_path, const std::string& dst_path) {cv::Mat image = cv::imread(src_path);cv::Mat mask = predict_image(image);cv::imwrite(dst_path, mask);std::cout << "predict image over" << std::endl;}

主函数

#include <windows.h>
#include <vector>
#include <iostream>
#include <opencv2/opencv.hpp>
#include "Seg.h"
#include <string>int main()
{std::wstring model_path(L"model.onnx");std::cout << "infer...." << std::endl;ObjectSeg object_seg(model_path, 1, { 1, 3, 512, 512 });for (int i = 0; i < 20; i++){DWORD star_time = GetTickCount();cv::Mat src = cv::imread("(1).jpg");int height = src.rows;int width = src.cols;cv::Mat mask = object_seg.predict_image(src);cv::imwrite("09051.png", mask);DWORD end_time = GetTickCount();std::cout << "这个程序" << i << "运行时间为：" << (end_time - star_time) << "ms." << std::endl;}return 0;
}