libtorch的初步使用

由于毕业设计的设计，学习使用了一下libtorch的使用。

主要参考链接 https://zhuanlan.zhihu.com/p/92374964 。

在写好cmake文件以及基本的加载代码之后，调用模型对单张图片进行深度预测，并将结果保存下来。代码如下：

#include <torch/script.h>
#include <iostream>
#include <string>
#include <memory>
#include <opencv2/opencv.hpp>

using namespace std;

int main(int argc, const char* argv[]) {
    if(argc!=2)
    {
        std::cerr<<"no argument!!!\n";
        return -1;
    }

    string rootPath="/media/yaoliqun/disk/project/project-private/pytorchc-test";

    torch::jit::script::Module module = torch::jit::load(argv[1]);

    std::cout<<"ok"<<std::endl;

    //输入图像
    cv::Mat im=cv::imread(rootPath+"/data/test_rgb.png",CV_LOAD_IMAGE_UNCHANGED);
    cv::resize(im,im,cv::Size(304,228));
    cv::cvtColor(im,im,cv::COLOR_BGR2RGB);

    //图像转化为tensor
    torch::Tensor tensor_image=torch::from_blob(im.data,{im.rows,im.cols,3},torch::kByte);
    tensor_image=tensor_image.permute({2,0,1});
    tensor_image=tensor_image.toType(torch::kFloat);

    tensor_image=tensor_image.unsqueeze(0);

    //网络前向计算
    at::Tensor output=module.forward({tensor_image}).toTensor();
//    cout<<output.sizes()<<endl;
//    cout<<output.to(torch::kCPU).squeeze(0).sizes()<<endl;
    output=output.to(torch::kCPU).squeeze(0).detach().permute({1,2,0});

    //保存图像
//    cout<<output.min()<<endl;
//    output=output.div(output.max()).mul(255).clamp(0,255).to(torch::kU8);
//    cv::Mat result(228,304,CV_8UC1);

    cv::Mat result=cv::Mat(228,304,CV_32FC1,output.data_ptr());
    cv::resize(result,result,cv::Size(640,480));
    result*=5000;
    result.convertTo(result,CV_16UC1);
    
//    std::memcpy((void *) result.data,output.data_ptr(),sizeof(torch::kU8) * output.numel());
    cv::imwrite(rootPath+"/data/result.png",result);

}