qt+libtorch+opencv调用u2net实现图像二值分割

内容介绍

本篇文章主要介绍使用libtorch调用u2net网络并使用qt作为前端展示工具，最后打包实现在其他电脑中能正常运行全过程以及其中踩过坑，第一次写如有不足，请见谅。

整体流程

1. 选择训练好的模型，并通过traced_script将.pth文件转为.pt文件；
2. 在qt中配置opencv和libtorch环境；
3. 使用opencv对图片进行修改，使用libtorch读取模型；
4. 在qt编写ui，实现前端相关操作，将结果展示；
5. 使用qt内置的打包工具打包文件，然后将模型和相关依赖导入qt打包的文件夹中；

模型转化

import torch
from model import U2NETP
#print(torch.version.cuda)
#print(torch.__version__)
model = U2NETP(3, 1)#自己定义的网络模型
model.load_state_dict(torch.load("xx.pth"))#保存的训练模型
model.eval()
example = torch.rand(1, 3, 320, 320).cuda()
#example = torch.rand(1, 3, 320, 320).cpu()
print(example.device)
traced_script_module = torch.jit.trace(model.cuda(), example)
#traced_script_module = torch.jit.trace(model.cpu(), example)
traced_script_module.save("xx.pt")

注释为转化为CPU版本代码。注意torch和cuda版本后面ibtorch版本要要一致

qt安装

由于官网速度下载较慢，建议使用清华源下载，这里教程很多不细说。

在qt中配置libtorch环境和opencv环境

1.下载opencv和libtorch相应的文件；
2.注意配置环境变量，如下图；

3.环境设置如下图；这里的./换成你自己文件是所在路径，不同版本的可能含有的lib文件有区别，大体一样；
libtorch最好选择release版本的，debug版本可能回报莫名其妙的错误，这是我的血泪教训！！！！！！！还有构建设置里面shadow build最好不勾（勾了的话只是做简单项目基本没有好处），不然有可能在改了qt中ui但是构建出来的项目没有变，并且方便后面的打包，这也是教训！！！

INCLUDEPATH+=./libtorch/include\
           ./libtorch/include/torch/csrc/api/include
   DEPENDPATH+=./libtorch/include\
           ./libtorch/include/torch/csrc/api/include
   LIBS += -L./libtorch/lib   -lc10_cuda -lc10 \#-lc10
   -lclog -lcpuinfo \
   -llibprotoc \
   -ltorch \#-ltorch -ltorch_cuda
INCLUDEPATH+=./opencv/build/include\
             ./opencv/build/include/opencv\
             ./opencv/build/include/opencv2
       DEPENDPATH+=./opencv/build/include\
                   ./opencv/build/include/opencv\
                   ./opencv/build/include/opencv2
       LIBS += -L./opencv/build/x64/vc15/lib \
           -lopencv_world3414d \
           -lopencv_world3414\

图片转为tensor格式

torch::Tensor process(cv::Mat& image, torch::Device device,int img_size)
{
    cv::resize(image, image, cv::Size(img_size,img_size));//此处将图像尺寸转化为你模型输出的尺寸，不然会报错
    cv::cvtColor(image, image, cv::COLOR_BGR2RGB);// bgr -> rgb
    std::vector<int64_t> dims = {1,img_size,img_size, 3 };
    torch::Tensor img_var = torch::from_blob(image.data, dims, torch::kByte).to(device);//将图像转化成张量
    img_var = img_var.permute({ 0,3,1,2 });//将张量的参数顺序转化为 torch输入的格式 1,3,224,224
    img_var = img_var.toType(torch::kFloat);
    img_var = img_var.div(255);

    return img_var;

}

QImage和Mat格式互相转化

QImage cvMat2QImage(const cv::Mat& mat)
{
    // 8-bits unsigned, NO. OF CHANNELS = 1
    if(mat.type() == CV_8UC1)
    {
        QImage image(mat.cols, mat.rows, QImage::Format_Indexed8);
        // Set the color table (used to translate colour indexes to qRgb values)
        image.setColorCount(256);
        for(int i = 0; i < 256; i++)
        {
            image.setColor(i, qRgb(i, i, i));
        }
        // Copy input Mat
        uchar *pSrc = mat.data;
        for(int row = 0; row < mat.rows; row ++)
        {
            uchar *pDest = image.scanLine(row);
            memcpy(pDest, pSrc, mat.cols);
            pSrc += mat.step;
        }
        return image;
    }
    // 8-bits unsigned, NO. OF CHANNELS = 3
    else if(mat.type() == CV_8UC3)
    {
        // Copy input Mat
        const uchar *pSrc = (const uchar*)mat.data;
        // Create QImage with same dimensions as input Mat
        QImage image(pSrc, mat.cols, mat.rows, mat.step, QImage::Format_RGB888);
        return image.rgbSwapped();
    }
    else if(mat.type() == CV_8UC4)
    {
        qDebug() << "CV_8UC4";
        // Copy input Mat
        const uchar *pSrc = (const uchar*)mat.data;
        // Create QImage with same dimensions as input Mat
        QImage image(pSrc, mat.cols, mat.rows, mat.step, QImage::Format_ARGB32);
        return image.copy();
    }
    else
    {
        qDebug() << "ERROR: Mat could not be converted to QImage.";
        return QImage();
    }
}
cv::Mat QImage2cvMat(QImage image)
{
    cv::Mat mat;
    qDebug() << image.format();
    switch(image.format())
    {
    case QImage::Format_ARGB32:
    case QImage::Format_RGB32:
    case QImage::Format_ARGB32_Premultiplied:
        mat = cv::Mat(image.height(), image.width(), CV_8UC4, (void*)image.constBits(), image.bytesPerLine());
        break;
    case QImage::Format_RGB888:
        mat = cv::Mat(image.height(), image.width(), CV_8UC3, (void*)image.constBits(), image.bytesPerLine());
        cv::cvtColor(mat, mat, CV_BGR2RGB);
        break;
    case QImage::Format_Indexed8:
        mat = cv::Mat(image.height(), image.width(), CV_8UC1, (void*)image.constBits(), image.bytesPerLine());
        break;
    }
    return mat;
}

模型调用

 using torch::jit::script::Module;
 QString path;
 QDir dir;
 path=QDir::currentPath();//这里是获取当前执行文件所在地址
 std::cout << 1 << std::endl;
 path.replace("\\","//");
 std::cout << path.toStdString()<< std::endl;
 String in = path.toStdString() + "/u2netpgpu2.pt";
 std::cout << in << std::endl;
  Module module = torch::jit::load(in);
  //加载pytorch模型
   module.to(torch::kCUDA);
   module.eval();
 // torch::Tensor result = module.forward({img_var}).toTensor();  //前向传播获取结果，还是tensor类型
  //模型返回多个结果，用toTuple,其中elements()[i-1]获取第i个返回值
   auto result = module.forward({ imgs }).toTuple()->elements[0].toTensor();

这一步的坑很多，注意libtorch版本要和你训练模型时的torch和CUDA版本相同，不然会有很多错，还有注意原模型训练后的输出为单个还是多个，单个.toTensor()，多个使用toTuple()。

将预测出来的数据转化为Mat

auto ten_wrp = result.squeeze(1);
torch::Tensor result1 = torch::max(ten_wrp);
torch::Tensor result2 = torch::min(ten_wrp);
ten_wrp = (ten_wrp - result2) / (result1 - result2);
ten_wrp = ten_wrp.permute({ 1, 2, 0 });
ten_wrp = ten_wrp.mul(255).clamp(0, 255).to(torch::kU8);
ten_wrp = ten_wrp.to(torch::kCPU);
cv::Mat resultImg(x, x, CV_8U);//xx为你输入图像的宽高，若图片出现问题则改为模型输出的宽高
std::memcpy((void*)resultImg.data, ten_wrp.data_ptr(), sizeof(torch::kU8) *ten_wrp.numel());

qt文件打包成exe文件移植到其他电脑

1.将release文件下的exe文件移出来放进一个空文件夹中；

2.使用qt自带的编译器，然后命令行进入上面的文件夹中，使用cd /d 上面exe所在文件夹；
3.使用windeployqt xx.exe（xx为你exe文件的名字）；

4.生成的文件夹中放入相关依赖，缺少的依赖.dll文件有些在你添加的例如opencv和libtorch中未找到去图片中的位置去找；

5.打包成功后可是将其移植到其他未安装相关环境中；
6.运行成功的图片.

打包好的文件地址.