人像分割是图像分割领域的高频应用,PaddleSeg推出在大规模人像数据上训练的人像分割系列模型PP-HumanSeg,包括超轻量级模型PP-HumanSeg-Lite,满足在服务端、移动端、Web端多种使用场景的需求。本项目将PP-HumanSeg-Lite模型转为onnx并部署到树莓派,实现人像抠图效果。效果如下图所示(这里只露个半脸):
目录
一、介绍
硬件
软件
二、具体流程
1、克隆PaddleSeg仓库
2、安装paddleseg
3、下载预训练模型
4、导出静态图模型
5、ONNX转换
6、移动模型的路径
7、预测代码
三、树莓派运行
四、问题解答
五、参考链接
本项目将PaddleSeg的轻量级人像分割模型转换为onnx,将其部署到树莓派实现实时人像分割。树莓派环境如下:
树莓派4B/4G
摄像头一个
Ubuntu Desktop 21.10
onnxruntime
opencv-python
%cd ~/
!git clone https://gitee.com/PaddlePaddle/PaddleSeg.git
!pip install paddleseg
%cd ~/PaddleSeg/contrib/PP-HumanSeg
!python pretrained_model/download_pretrained_model.py
导出静态图模型的时候,需要加上input_shape参数,本项目部署的模型是PP-HumanSeg-Lite,该模型的输入图片大小是192x192,如果要部署其他模型,需要更改model_path和config路径。
(为什么导出静态图模型?答:paddlepaddle在训练时可以使用动态图模型很方便的完成模型组网,完成训练模型的任务,但是动态图预测速度比静态图差,为了加快推理速度,可以将动态图转换为静态图。)
%cd ~/PaddleSeg/contrib/PP-HumanSeg
!python ../../export.py \
--config configs/fcn_hrnetw18_small_v1_humanseg_192x192_mini_supervisely.yml \
--model_path pretrained_model/fcn_hrnetw18_small_v1_humanseg_192x192/model.pdparams \
--save_dir export_model/fcn_hrnetw18_small_v1_humanseg_192x192 \
--with_softmax --input_shape 1 3 192 192
将paddlepaddle静态图模型转换为onnx,方便后续的树莓派部署。
# ① 安装paddle2onnx
!pip install paddle2onnx
# ② 转换为onnx
%cd ~/PaddleSeg/contrib/PP-HumanSeg
! paddle2onnx --model_dir ./export_model/fcn_hrnetw18_small_v1_humanseg_192x192/ \
--model_filename model.pdmodel \
--params_filename model.pdiparams \
--save_file onnx_model/model.onnx \
--opset_version 12
为了部署到树莓派,仅保留必要的文件,创建pp_humanseg_deploy文件夹,模型路径是~/pp_humanseg_deploy,文件夹目录结构如下:
|-onnx_model
|---model.onnx
|-predict.py ---- 这个文件下一步生成
%cd ~/
!mkdir pp_humanseg_deploy
%cd ~/pp_humanseg_deploy/
!mkdir onnx_model
!cp ~/PaddleSeg/contrib/PP-HumanSeg/onnx_model/model.onnx ~/pp_humanseg_deploy/onnx_model
按照上一步的目录结构,创建predict.py文件,将以下代码放入其中。
下方代码包含预处理和模型预测2个部分,其中预处理部分需要和PaddleSeg的预处理相同。
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 = 'onnx_model/model.onnx'
sess = rt.InferenceSession(onnx_model_path)
input_name = sess.get_inputs()[0].name
label_name = sess.get_outputs()[0].name
target_size = (192, 192)
cap_video = cv2.VideoCapture(0)
if not cap_video.isOpened():
raise IOError("Error opening video stream or file.")
while cap_video.isOpened():
ret, raw_frame = cap_video.read()
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)
if cv2.waitKey(1) & 0xFF == ord('q'):
break
else:
break
cap_video.release()
将~/pp_humanseg_deploy文件夹打包,放入树莓派环境运行predict.py,效果如下:
1、为什么图片看起来很模糊?
这里使用的模型是pp-humanseg-lite,该模型的输入像素是192x192,在树莓派部署 时,先将480x640的图片resize到192x192送入模型预测,再将预测结果resize到480x640,所以图片看起来模糊。
2、为什么转换为onnx?
也可以在树莓派安装paddlepaddle环境,然后使用paddlepaddle来实现推理功能。此处转换为onnx是考虑到个人的树莓派内存有限,并未安装过多的包。
3、如何训练自己的人像分割模型?
参考PaddleSeg
onnx模型导出AiStudio参考链接