如何使用 C++ 创建 ONNX 图像分割应用程序？

Question

我有一个 Pytorch 图像分割模型，我测试了它以使用 onnx 运行时进行分割。 我想使用 C++ 创建 图像分割，但我不知道如何将图像提供给分割模型以及如何从模型中获取图像。

如果有人知道，可以详细解释一下吗？

Answer 1

推理代码将完全取决于您的模型，以下是您必须注意的一些常见因素，

1. 创建内存分配器，初始化会话和环境。
2. 制作输入和输出张量。
3. 获取输入输出名称。

参考这个： https://github.com/leimao/ONNX-Runtime-Inference/blob/main/src/inference.cpp

尝试进行示例图像分类模型推理，然后进行图像分割以更好地理解工作原理，因为分类非常容易完成。

参考这个： https://github.com/microsoft/onnxruntime/issues/13098

这里是初始化会话的示例代码：

size_t inputTensorSize = size_t(IMG_SIZE * IMG_SIZE * 3);
// Initializing the session
auto memory_info = Ort::MemoryInfo::CreateCpu(OrtDeviceAllocator, OrtMemTypeDefault);
Ort::Env env(ORT_LOGGING_LEVEL_WARNING, "test");
Ort::SessionOptions session_options;
OrtCUDAProviderOptions setGPUoption;
setGPUoption.device_id = 1;

 session_options.SetGraphOptimizationLevel(GraphOptimizationLevel::ORT_ENABLE_BASIC);

// uncomment below only if you have Gpu and also make sure with the gpu options
//session_options.AppendExecutionProvider_CUDA(setGPUoption);


Ort::Session session(env, wstrOnnxPath.c_str(), session_options);

// tensor dimension and other allocations
auto memory_info = Ort::MemoryInfo::CreateCpu(OrtDeviceAllocator, OrtMemTypeDefault);

std::vector<int64_t> input_node_dims = { 1,IMG_SIZE,IMG_SIZE,3 };
Ort::AllocatorWithDefaultOptions allocator;
std::string input_name = session.GetInputNameAllocated(0, allocator).get();
std::string output_name = session.GetOutputNameAllocated(0, allocator).get();
std::array<const char*, 1> input_node_names = { input_name.c_str() };
std::array<const char*, 1> output_node_names = { output_name.c_str() };

// THE BELOW SECTION CHANGES ACCORDING TO YOUR MODEL

std::vector<Ort::Value> inputTensors;
cv::Mat img_tensor;
std::vector<float> inputTensorValues(inputTensorSize);
cv::dnn::blobFromImage(ImageMat, img_tensor, true);
cv::Mat dst;
cv::transposeND(img_tensor, { 0,2,3,1 }, dst);
copy(dst.begin<float>(), dst.end<float>(), inputTensorValues.begin());

// input tensors
inputTensors.push_back(Ort::Value::CreateTensor<float>(memory_info, inputTensorValues.data(), inputTensorSize, input_node_dims.data(), 4));

// session run
auto output_tensors = session.Run(Ort::RunOptions{ nullptr }, input_node_names.data(), inputTensors.data(), 1, output_node_names.data(), 1);

float dOutput = output_tensors[0].GetTensorMutableData<float>()[0];

All the best.