c++调用训练好的分割模型主要步骤如下:
1、将训练好的.pth 模型转换为.pt模型
2、opencv读取要推理的图像
3、model前向传播,获得推理结果
4、将推理结果转为opencv,查看推理结果。
开发环境:
windows10
vs 2015
pytorch 1.2.0 cpu
opencv 4.3.0
libtorch 1.3.0 cpu
配置开发环境时遇到的坑:
1、vs2015不支持最新的libtorch
2、ptorch与liborch 使用的cpu或者gpu版本要统一,否则出现如下错误:
将训练好的.pth模型模型转为.pt模型
见上一篇文章:
pytorch训练好的.pth模型转换为.pt
2.2、c++代码如下:
#include "torch/script.h" // One-stop header.
#include <iostream>
#include <opencv2\opencv.hpp>
#include <opencv2\imgproc\types_c.h>
using namespace cv;
using namespace std;
int main()
{
Deserialize the ScriptModule from a file
torch::jit::script::Module module = torch::jit::load("caoxie_weight.pt");
//assert(module != nullptr);
auto image = imread("D:\\UNET\\data\\org\\2.bmp");
if (!image.data)
{
cout << "image imread failed" << endl;
}
cvtColor(image, image, CV_BGR2RGB);
Mat img_transfomed;
resize(image, img_transfomed,Size(480, 320));
/*cout << img_transfomed.data;*/
//img_transfomed.convertTo(img_transfomed, CV_16FC3, 1.0f / 255.0f);
//Mat to tensor,
torch::Tensor tensor_image = torch::from_blob(img_transfomed.data, {img_transfomed.rows, img_transfomed.cols, img_transfomed.channels()}, torch::kByte);
tensor_image = tensor_image.permute({2, 0, 1});
tensor_image = tensor_image.toType(torch::kFloat);
tensor_image = tensor_image.div(255);
tensor_image = tensor_image.unsqueeze(0);
std::vector<torch::jit::IValue> inputs;
inputs.push_back(tensor_image);
torch::Tensor output = module.forward(inputs).toTensor();
torch::Tensor output_max = output.argmax(1);
//tensor to Mat
output_max = output_max.squeeze();
output_max = output_max.mul(255).to(torch::kU8);
output_max = output_max.to(torch::kCPU);
Mat result_img(Size(480, 320), CV_8UC1);
memcpy((void*)result_img.data, output_max.data_ptr(), sizeof(torch::kU8)*output_max.numel());
imshow("result", result_img);
imwrite("result.bmp", result_img);
system("pause");
}
2.3 cmake 编译
cmakelist 如下:
cmake_minimum_required(VERSION 3.0 FATAL_ERROR)
project(example)
find_package(Torch REQUIRED)
add_executable(example example.cpp)
target_link_libraries(example ${TORCH_LIBRARIES})
set_property(TARGET example PROPERTY CXX_STANDARD 11)
命令行cd到当前目录下,传入libtorch路径:
cmake -DCMAKE_PREFIX_PATH=libtorh路径.. -DCMAKE_BUILD_TYPE=Release -G"Visual Studio 14 Win64" //release版
输出结果: