关于PP-Humanseg是在正点原子的3568开发板AI测试例子里看到的,目的也是想自己训练并部署,这里记录一下训练和在onnxruntime部署运行的全过程,会转成ONNX,至于部署到rk3568上,会在另一篇文章説明ONNX转成RKNN并部署到RK3568.
一、训练模型
一、介绍
本文将PaddleSeg的人像分割(PP-HumanSeg)模型导出为onnx,并使用onnxruntime部署,实现人像分割,效果如下图所示。
二、环境搭建
使用的是AutoDL服务器,配置如下:
创建虚拟环境
conda create -n ppseg_env python=3.8 -y
激活环境
conda activate ppseg_env
三、安装PaddlePaddle
直接安装
python -m pip install paddlepaddle-gpu==2.4.2 -i https://pypi.tuna.tsinghua.edu.cn/simple
也可以源码其他的,参考
开始使用_飞桨-源于产业实践的开源深度学习平台
四、安装PaddleSeg
git clone https://github.com/PaddlePaddle/PaddleSeg
cd PaddleSeg
pip install -r requirements.txt
pip install -v -e .
五、下载数据
PP-Humanseg位于contrib/PP-HumanSeg目录下,在目录下有个 readme.md文件,根据文件,下载数据及模型
进入PP-HumanSeg目录下
cd PaddleSeg/contrib/PP-HumanSeg
1、下载Inference Model
python src/download_inference_models.py
2、下载测试数据
python src/download_data.py
六、测试
测试的数据可以是视频也可以是图片,这里测试图片
执行下面指令:
# Use PP-HumanSegV2-Lite
python src/seg_demo.py \
--config inference_models/portrait_pp_humansegv2_lite_256x144_inference_model_with_softmax/deploy.yaml \
--img_path data/images/portrait_heng.jpg \
--save_dir data/images_result/portrait_heng_v2.jpg
结果会保存到data/images_result目录下,结果如下:
视频可以自我测试,不在演示,使用如下 命令:
python src/seg_demo.py \
--config inference_models/portrait_pp_humansegv2_lite_256x144_inference_model_with_softmax/deploy.yaml \
--video_path data/videos/video_heng.mp4 \
--save_dir data/videos_result/video_heng_v2.avi
七、训练
配置文件保存在./configs
目录下,如下。配置文件中,已经通过pretrained
设置好预训练权重的路径。
执行如下命令,进行模型微调,模型训练的详细文档,请参考链接
export CUDA_VISIBLE_DEVICES=0 # Linux下设置1张可用的卡
# set CUDA_VISIBLE_DEVICES=0 # Windows下设置1张可用的卡
python ../../tools/train.py --config configs/human_pp_humansegv2_lite.yml --save_dir output/human_pp_humansegv2_lite --save_interval 100 --do_eval --use_vdl
训练完成后,模型会保存在output目录下
八、评估模型
python ../../tools/val.py --config configs/human_pp_humansegv2_lite.yml --model_path output/human_pp_humansegv2_lite/best_model/model.pdparams
看起来评估的结果还是可以的。
九、预测
python ../../tools/predict.py --config configs/human_pp_humansegv2_lite.yml --model_path output/human_pp_humansegv2_lite/best_model/model.pdparams --image_path data/images/human.jpg --save_dir ./data/images_result
接下来是重头戏,需要导出模型
十、导出静态模型
输入图片大小是192x192
python ../../tools/export.py --config configs/human_pp_humansegv2_lite.yml --model_path output/human_pp_humansegv2_lite/best_model/model.pdparams --save_dir output/human_pp_humansegv2_lite --input_shape 1 3 192 192
十一、预测导出静态模型
参考deploy下的测试
python ../../deploy/python/infer.py --config output/human_pp_humansegv2_lite/deploy.yaml --image_path ./human.jpg --save_dir ./data/images_result3
至此,导出的静态模型测试完成
二、ONNX模型转换
一、安装paddle2onnx
pip install paddle2onnx
二、导出ONNX模型
paddle2onnx --model_dir output/human_pp_humansegv2_lite/ \
--model_filename model.pdmodel \
--params_filename model.pdiparams \
--opset_version 12 \
--save_file output.onnx
三、测试
测试使用的是onnxruntime环境测试,所以需要先安装onnxruntime
pip install onnxruntime
测试代码是自己编写的,文件为predict.py
import cv2
import numpy as np
import onnxruntime as rt
def normalize(im, mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]):
im = im.astype(np.float32, copy=False) / 255.0
im -= mean
im /= std
return im
def resize(im, target_size=608, interp=cv2.INTER_LINEAR):
if isinstance(target_size, list) or isinstance(target_size, tuple):
w = target_size[0]
h = target_size[1]
else:
w = target_size
h = target_size
im = cv2.resize(im, (w, h), interpolation=interp)
return im
def preprocess(image, target_size=(192, 192)):
image = normalize(image)
image = resize(image, target_size=target_size)
image = np.transpose(image, [2, 0, 1])
image = image[None, :, :, :]
return image
def display_masked_image(mask, image, color_map=[255, 0, 0], weight=0.6):
mask = mask > 0
c1 = np.zeros(shape=mask.shape, dtype='uint8')
c2 = np.zeros(shape=mask.shape, dtype='uint8')
c3 = np.zeros(shape=mask.shape, dtype='uint8')
pseudo_img = np.dstack((c1, c2, c3))
for i in range(3):
pseudo_img[:, :, i][mask] = color_map[i]
vis_result = cv2.addWeighted(image, weight, pseudo_img, 1 - weight, 0)
return vis_result
onnx_model_path = './output.onnx'
sess = rt.InferenceSession(onnx_model_path,providers=['AzureExecutionProvider', 'CPUExecutionProvider'])
input_name = sess.get_inputs()[0].name
label_name = sess.get_outputs()[0].name
target_size = (192, 192)
#cap_video = cv2.VideoCapture('./human.jpg')
#if not cap_video.isOpened():
# raise IOError("Error opening video stream or file.")
path = "./human.jpg"
while 1:
#ret, raw_frame = cap_video.read()
ret = 1
raw_frame = cv2.imread(path)
pre_shape = raw_frame.shape[0:2][::-1]
if ret:
frame = cv2.cvtColor(raw_frame, cv2.COLOR_BGRA2RGB)
frame = preprocess(frame, target_size)
pred = sess.run(
[label_name],
{input_name: frame.astype(np.float32)}
)[0]
pred = pred[0]
raw_frame = resize(raw_frame, target_size)
image = display_masked_image(pred, raw_frame)
image = resize(image, target_size=pre_shape)
#cv2.imshow('HumanSegmentation', image)
cv2.imwrite('result.jpg',image)
print("finish! result save result.jpg")
break
if cv2.waitKey(1) & 0xFF == ord('q'):
break
else:
break
#cap_video.release()
执行python predict.py
运行正常