PP-HumanSeg模型导出为ONNX格式并分割人像

前言

在上一篇内容中已经对PP-HumanSeg模型进行了训练和微调，为了方便后续的部署，在这篇文章中，会将训练产生的模型导出ONNX格式并调用，对RGB图和深度图进行人像分割。

一、模型导出

1.1 导出预测模型

PaddleSeg训练好模型后，需要将模型导出为预测模型（这里还是利用的PaddleSeg项目中的脚本），才可以进行模型部署。

python ../../export.py \--config configs/human_pp_humansegv1_server.yml \--model_path output/human_pp_humansegv1_server/best_model/model.pdparams \--save_dir output/human_pp_humansegv1_server_export \--without_argmax \--with_softmax

导出预测模型后，会在保存文件夹下产生四个文件预测模型相关文件，下一步中会用到。

1.2 导出为ONNX格式

1.2.1 安装paddle2onnx

pip install paddle2onnx

1.2.2 格式转换

paddle2onnx --model_dir ./output/human_pp_humansegv1_server_export \--model_filename model.pdmodel \--params_filename model.pdiparams \--save_file ./output/human_pp_humansegv1_server_export/model.onnx \--enable_dev_version True \--enable_onnx_checker True \--opset_version 11

转换成功后就会在保存文件下产生一个model.onnx文件。

二、ONNX模型属性结构

2.1 模型属性

2.2 模型结构

使用Netron可视化模型结构。

三、ONNX模型使用

3.1 代码

import os
import numpy as np
import onnxruntime
import cv2
import time
import onnx
from PIL import Image as PILImage#onnx模型加载检测
onnx_path = "output/human_pp_humansegv1_server_export/model.onnx"
onnx_model = onnx.load(onnx_path)
onnx.checker.check_model(onnx_model)
print('The model is checked!')#路径
rgb_img_path = "data/images/2275122.png"
depth_img_path = "data/images/2275188.png"
save_img_path = "data/images_result"sess = onnxruntime.InferenceSession(onnx_path,providers=['CUDAExecutionProvider']) #使用GPU#读取图片
rgb_img = cv2.imread(rgb_img_path)
depth_img = cv2.imread(depth_img_path)#输入数据预处理
#input_img = cv2.resize(rgb_img,(512,512)) # resize到网络模型合适的输入，这里也可以不缩放，已经调整为随机大小
input_img = cv2.cvtColor(rgb_img, cv2.COLOR_BGR2RGB) #BGR->RGB
input_img = np.transpose(input_img, (2,0,1))  # 转为[c,h,w]
input_img = ((input_img.astype(np.float32)/255.0) - 0.5) / 0.5 #归一化
img_data = input_img[np.newaxis, :, :, :] # shape:[1,3,720,1080]#onnx模型推理
ort_inputs = {sess.get_inputs()[0].name: img_data}
results = sess.run(None, ort_inputs)#保存图片
if not os.path.exists(save_img_path):os.makedirs(save_img_path)for i, result in enumerate(results):result = np.argmax(result, axis=1) #shape:[1,720,1080]if result.ndim == 3:result = np.squeeze(result) #shape:[720,1080] value:0,1result *= 255 # value:0,255#创建mask PILmask = PILImage.fromarray(result.astype(np.uint8), mode='L')#构建保存路径rgbname = os.path.basename(rgb_img_path)rgbname, _ = os.path.splitext(rgbname)depthname = os.path.basename(depth_img_path)depthname, _ = os.path.splitext(depthname)maskname = f'{rgbname}_mask.png'seg_rgbname = f'{rgbname}_seg_rgb.png'seg_depthname = f'{depthname}_seg_depth.png'#保存mask图片mask_save_path = os.path.join(save_img_path, maskname)mask.save(mask_save_path)print('mask image is saved in {}'.format(mask_save_path))#保存分割后的图片mask_img = cv2.imread(mask_save_path, 0) #以灰度图的形式读取maskseg_rgb_img = cv2.bitwise_and(rgb_img, rgb_img, mask = mask_img)seg_depth_img = cv2.bitwise_and(depth_img, depth_img, mask = mask_img)seg_rgb_save_path = os.path.join(save_img_path, seg_rgbname)seg_depth_save_path = os.path.join(save_img_path, seg_depthname)cv2.imwrite(seg_rgb_save_path, seg_rgb_img)print('rgb_seg image is saved in {}'.format(seg_rgb_save_path))cv2.imwrite(seg_depth_save_path, seg_depth_img)print('depth_seg image is saved in {}'.format(seg_depth_save_path))

3.2 运行

四、分割效果展示

4.1 原图

4.2 分割

参考

1.https://github.com/PaddlePaddle/PaddleSeg
2.https://github.com/PaddlePaddle/Paddle2ONNX