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

• 前言
• 一、模型导出
• • 1.1 导出预测模型
  • 1.2 导出为ONNX格式
  • • 1.2.1 安装paddle2onnx
    • 1.2.2 格式转换
• 二、ONNX模型属性结构
• • 2.1 模型属性
  • 2.2 模型结构
• 三、ONNX模型使用
• • 3.1 代码
  • 3.2 运行
• 四、分割效果展示
• • 4.1 原图
  • 4.2 分割
• 参考

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

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

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

pip install paddle2onnx

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文件。

使用Netron可视化模型结构。

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,1
    result *= 255 # value:0,255

    #创建mask PIL
    mask = 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) #以灰度图的形式读取mask
    seg_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))

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