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;
//获取输入name
const 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 tensor
auto 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)
{
// argmax
unsigned 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;
}

原文地址：https://blog.csdn.net/shuaijieer/article/details/142481480

免责声明：本站文章内容转载自网络资源，如本站内容侵犯了原著者的合法权益，可联系本站删除。更多内容请关注自学内容网（zxcms.com）！