收藏    

目录

一、环境安装

1.1 基本环境介绍

1.2 pth模型序列化导出转pt

1.2 下载libtorch

1.3 安装OpenCV

1.4 创建win32 C++控制台工程

二、完整推理代码

三、测试

一、环境安装
1.1 基本环境介绍
本文实验环境为windows10，IDE选择VS 2019，最终开发win32 C++程序，可以批量对人像照片进行抠图处理，无需用户干预。本文重点在于部署Pytorch训练好的pth模型，能够脱离python平台实现C++调用，完成真正的生产级部署。如何通过Pytorch训练人像抠图模型可以参考我的其它博客教程1、博客教程2。

1.2 pth模型序列化导出转pt
训练完模型以后可以得到pth模型，然后使用下面的python代码将其序列化导出（即将动态图转静态图）：

#!/usr/bin/env python
# -*- encoding: utf-8 -*-
'''
@文件 :pth2pt.py
@说明 :模型序列化导出，pth文件转pt文件
@时间 :2020/05/08 14:18:55
@作者 :钱彬
@版本 :1.0
'''


import torch.backends.cudnn as cudnn
import torch
from torch import nn
from model.shm import SHM
from utils import *
import time
import cv2


# 测试图像
img_id='16'
imgPath = './results/'+img_id+'.png'

# 模型参数
#device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
device = torch.device("cpu")


if __name__ == '__main__':

# 预训练模型
checkpoint = "./results/shm_temp.pth"

# 加载模型
checkpoint = torch.load(checkpoint)
model = SHM()

model = model.to(device)
model.load_state_dict(checkpoint['model'])

model.eval()

# 加载图像
img_org = cv2.imread(imgPath, cv2.IMREAD_COLOR)
width = img_org.shape[1]
height = img_org.shape[0]

# 图像预处理
img = cv2.resize(img_org, (320,320), interpolation = cv2.INTER_CUBIC)

img = (img.astype(np.float32) - (114., 121., 134.,)) / 255.0
h, w, c = img.shape
img = torch.from_numpy(img.transpose((2, 0, 1))).view(c, h, w).float()
img= img.view(1, 3, h, w)

# 转移数据至设备
img = img.to(device)

# 模型推理并序列化导出
with torch.no_grad():
alpha = model(img)
alpha = alpha.squeeze(0).float().mul(255).add_(0.5).clamp_(0, 255).permute(1, 2, 0).to('cpu', torch.uint8).numpy()
traced_script_module = torch.jit.trace(model, img)
traced_script_module.save("./results/matting.pt")
1.2 下载libtorch
在官网下载libtorch，由于后面我们需要使用GPU进行推理，因此下载对应cuda版本的libtorch。下载时尤其要注意版本一致问题，即最终的推理平台用的是哪个版本的cuda，那么这里也需要下载对应版本的cuda，本文下载cuda10.1对应的版本。具体如下图所示：



1.3 安装OpenCV
由于需要进行图像加载等操作，这里选择开源框架OpenCV实现。具体安装过程可以参考我的另一篇博客。本文使用opencv4.2版本。

1.4 创建win32 C++控制台工程
本文使用VS2019来创建C++工程，具体如下：



创建完成后切换成64位平台：



然后进行项目配置如下：

（1）属性配置——VC++目录



这里我将前面下载的libtorch放置在了本工程目录下，读者按照上面的路径然后结合自己的路劲进行对应的修改即可。

（2）属性配置——库目录



（3）链接器——输入

添加lib文件如下：

opencv_world420.lib
asmjit.lib
c10.lib
c10_cuda.lib
caffe2_detectron_ops_gpu.lib
caffe2_module_test_dynamic.lib
caffe2_nvrtc.lib
clog.lib
cpuinfo.lib
dnnl.lib
fbgemm.lib
gloo.lib
gloo_cuda.lib
libprotobuf.lib
libprotobuf-lite.lib
libprotoc.lib
mkldnn.lib
torch.lib
torch_cpu.lib
torch_cuda.lib
这里我将libtorch中的所有lib文件全部包含了进入，实际使用时可以不用这么多，这里为了不出问题可以先全部包含进来，等后面程序调通了后再逐步删除。接下来将libtorch中的所有dll文件拷贝到项目根目录下即可。如下图所示：



最后很重要的一点，在高版本VS中有个bug，需要手动添加点东西我们的代码才能正常的调用GPU。找到链接器——命令行——其它选项，然后输入： /INCLUDE:?warp_size@cuda@at@@YAHXZ

具体如下图所示：



如果不实现上述操作我们的代码后面会无法调用GPU。

二、完整推理代码
完成前面的配置以后就可以进行代码编辑了，完整代码如下，主要实现端到端抠图：

/*********************************************************************************************************************************
* Copyright(c) 2020-2025
* All rights reserved.
*
* 文件名称:PortraitOpt
* 简要描述:基于C++实现证件人像抠图（GPU和CPU通用）
*
* 创建日期:2020-10-09
* 作者:钱彬
* 版本:V1.0.0
*********************************************************************************************************************************/

//导入深度学习torch库
#undef UNICODE
#include <torch/script.h> // 引入libtorch头文件
#include <torch/torch.h> // cuda相关函数头文件

//导入系统库
#include <iostream>
#include "time.h"

//导入opencv图像处理库
#include <opencv2/opencv.hpp>
#include <opencv2/core/core.hpp>
#include <opencv2/highgui/highgui.hpp>
#include <opencv2/imgproc.hpp>
#include <opencv2/imgproc/types_c.h>
#include <opencv2/videoio/videoio.hpp>

//定义命名空间
using namespace std;
using namespace cv;


int main()
{
//确定当前设备类型（CPU或GPU）
torch::DeviceType device_type = at::kCPU; // 定义设备类型
if (torch::cuda::is_available())
device_type = at::kCUDA;

// 加载抠图模型
torch::jit::script::Module model;
try {
model = torch::jit::load("matting.pt"); //加载模型
}
catch (const c10::Error& e) {
cout << "无法加载模型" << endl;
return 0;
}
model.eval();
model.to(device_type);

//查找文件夹中所有jpg格式图片
cv::String pattern = "imgs/*.jpg";
vector<cv::String> fn;
glob(pattern, fn, false);
vector<Mat> images;
int imgNum = fn.size();
std::cout << "当前需要处理的图像总数： " << imgNum << endl;

//----------------------------------------------循环处理图片-------------------------------------------------
for (int picIndex = 0; picIndex < imgNum; picIndex++)
{
//读取图像
Mat img = imread(fn[picIndex]);
clock_t start, finish;
double duration;
char imgpath[256];

start = clock();
Mat orgimg = img.clone();

// 压缩图像
int org_width = img.cols;
int org_height = img.rows;
resize(img, img, Size(320, 320), 0, 0, INTER_CUBIC);

//数据预处理
std::vector<int64_t> sizes = { 1, img.rows, img.cols,3 };
at::TensorOptions options(at::ScalarType::Byte);
at::Tensor tensor_image = torch::from_blob(img.data, at::IntList(sizes), options);//将opencv的图像数据转为Tensor张量数据
tensor_image = tensor_image.toType(at::kFloat);//转为浮点型张量数据

tensor_image = tensor_image.permute({ 0, 3, 1, 2 });
tensor_image[0][0] = tensor_image[0][0].sub(114).div(255.0);
tensor_image[0][1] = tensor_image[0][1].sub(121).div(255.0);
tensor_image[0][2] = tensor_image[0][2].sub(134).div(255.0);

tensor_image = tensor_image.to(device_type);

// 推理
std::vector<torch::jit::IValue> inputs;
inputs.push_back(tensor_image);
at::Tensor out_tensor = model.forward(inputs).toTensor();

out_tensor = out_tensor.squeeze().detach();
out_tensor = out_tensor.mul(255).add_(0.5).clamp_(0, 255).to(torch::kU8);
out_tensor = out_tensor.to(torch::kCPU);

cv::Mat alphaImg(img.rows, img.cols, CV_8UC1);
std::memcpy((void*)alphaImg.data, out_tensor.data_ptr(), sizeof(torch::kU8) * out_tensor.numel());

// 调整抠图结果大小
resize(alphaImg, alphaImg, Size(org_width, org_height), 0, 0, INTER_CUBIC);

// 与白背景合成
Mat bg(orgimg.size(), orgimg.type(), Scalar(255, 255, 255));//全白图
Mat alpha, comp;
cvtColor(alphaImg, alpha, COLOR_GRAY2BGR);
comp = orgimg.clone();

for (int i = 0; i < alpha.rows; i++)
for (int j = 0; j < alpha.cols; j++)
{
Vec3b alpha_p = alpha.at<Vec3b>(i, j);
Vec3b bg_p = bg.at<Vec3b>(i, j);
Vec3b img_p = orgimg.at<Vec3b>(i, j);
if (alpha_p[0] > 210)
{
alpha_p[0] = 255;
alpha_p[1] = 255;
alpha_p[2] = 255;
}
else
{
alpha_p[0] = 0;
alpha_p[1] = 0;
alpha_p[2] = 0;
}
comp.at<Vec3b>(i, j)[0] = int(img_p[0] * (alpha_p[0] / 255.0) + bg_p[0] * (1.0 - alpha_p[0] / 255.0));
comp.at<Vec3b>(i, j)[1] = int(img_p[1] * (alpha_p[1] / 255.0) + bg_p[1] * (1.0 - alpha_p[1] / 255.0));
comp.at<Vec3b>(i, j)[2] = int(img_p[2] * (alpha_p[2] / 255.0) + bg_p[2] * (1.0 - alpha_p[2] / 255.0));
}
finish = clock();
duration = (double)(finish - start) / CLOCKS_PER_SEC;
printf("抠图推理耗时 f 毫秒\n", duration * 1000);

//保存抠图结果
sprintf_s(imgpath, "results/d_matting.jpg", picIndex);
imwrite(imgpath, comp);

//手动释放内存资源
inputs.clear();
img.release();
orgimg.release();
alphaImg.release();
alpha.release();
comp.release();
}

return 0;
}

三、测试
测试效果如下：



第一两张图像速度会比较慢，猜测可能正在执行一些GPU调度之类的，后面的每张图像推理速度就正常了。抠图效果如下所示：

原图：



抠图结果：




————————————————
版权声明：本文为CSDN博主「钱彬 (Qian Bin)」的原创文章，遵循CC 4.0 BY-SA版权协议，转载请附上原文出处链接及本声明。
原文链接：https://blog.csdn.net/qianbin3200896/article/details/112772953