PaddleCV官方人体姿态估计预训练模型转ONNX
项目地址
https://github.com/PaddlePaddle/models/tree/develop/PaddleCV/human_pose_estimation
论文地址
Pose ResNet: https://arxiv.org/abs/1804.06208
网络结构
简单来说，这个网络把分类用的ResNet的最后几层全连接层替换为反卷积层输出heatmap，如果训练集有16个人体关键点，那么最后一层output的channel就是16，分别对应16个关键点，把置信度最高的一个像素采用为对应部位的关键点。
网络输入与输出
| Type | Size(batch_size、channel、height、width) |
|---|---|
| input | -1, 3, 384, 384 |
| output | -1, 16, 96, 96 |
预训练参数下载
不含网络结构，只有权重和偏置数据
使用mpii数据集训练的ResNet-50: https://paddlemodels.bj.bcebos.com/pose/pose-resnet50-mpii-384x384.tar.gz
下载后解压到项目根目录checkpoints文件夹下面
环境安装
# 使用的是paddle 2.1版本
pip install paddlepaddle-gpu -i https://mirror.baidu.com/pypi/simple
pip install onnx onnxruntime opencv-python
pip install paddle2onnx
导出网络结构和权重
import sys
import argparse
import functools
import paddle
import paddle.fluid as fluid
import paddle.fluid.layers as layers
from lib import pose_resnet
from utils.transforms import flip_back
from utils.utility import *
# 开启静态图模式，2.0版本默认关闭
paddle.enable_static()
# 预训练参数路径
weights_path = './checkpoints/pose-resnet50-mpii-384x384'
# 定义网络的输入
image = layers.data(name='image', shape=[3, 384, 384], dtype='float32')
# 定义使用的resnet版本，这里使用resnet50，kps_num是keypoints数量，对应mpii的16个关键点，test_mode开启，不需要返回loss
model = pose_resnet.ResNet(layers=50, kps_num=16, test_mode=True)
# 获得网络的output
output = model.net(input=image, target=None, target_weight=None)
# 选择使用CUDA执行，也可改为cpu
exe = fluid.Executor(fluid.CUDAPlace(0))
prog = fluid.default_main_program()
startup_prog = fluid.default_startup_program()
exe.run(startup_prog)
# 加载预训练参数
fluid.io.load_persistables(exe, weights_path, prog)
# 将网络结构与参数导出到__model__目录下
fluid.io.save_inference_model(dirname='./__model__', feeded_var_names=['image'], target_vars=[output], executor=exe)
print('模型导出完毕')
使用Paddle2ONNX将静态图导出为onnx模型
根据下面的教程把上一步的导出结果转换为.onnx文件
官方教程地址
Terminal输入以下命令
paddle2onnx --model_dir ./__model__ --save_file pdkp.onnx --opset_version 10 --enable_onnx_checker True
测试代码(CPU版本 - 尚未添加注释)
import onnxruntime
import cv2
import numpy as np
from scipy.ndimage import gaussian_filter, maximum_filter
import torch
from numba import jit
SIZE = 384
HEIGHT = 480
WIDTH = 640
@jit(nopython=True)
def load_image(image):
image = image / 255.0
image = image.transpose((2, 0, 1))
return image
def post_process_heatmap(heatMap):
kplst = list()
for i in range(heatMap.shape[0]):
_map = heatMap[i, :, :]
_map = gaussian_filter(_map, sigma=1)
_nmsPeaks = non_max_supression(_map, windowSize=3, threshold=1e-6)
y, x = np.where(_nmsPeaks == _nmsPeaks.max())
if len(x) > 0 and len(y) > 0:
kplst.append((int(x[0]), int(y[0]), _nmsPeaks[y[0], x[0]]))
else:
kplst.append((0, 0, 0))
kp = np.array(kplst)
return kp
def non_max_supression(plain, windowSize=3, threshold=1e-6):
under_th_indices = plain < threshold
plain[under_th_indices] = 0
return plain * (plain == maximum_filter(plain, footprint=np.ones((windowSize, windowSize))))
def render_kps(cvmat, kps, scale_x, scale_y):
for _kp in kps:
_x, _y, _conf = _kp
if _conf > 0.2:
cv2.circle(cvmat, center=(int(_x*4*scale_x), int(_y*4*scale_y)), color=(0, 0, 255), radius=5)
return cvmat
def create_ort_session(onnx_model):
global ort_session
ort_session = onnxruntime.InferenceSession(onnx_model)
def run_inference(image_path):
image = load_image(image_path)
image = image[np.newaxis, :, :, :]
ort_inputs = {ort_session.get_inputs()[0].name: image.astype('float32')}
ort_outs = ort_session.run(None, ort_inputs)
ort_kps = post_process_heatmap(ort_outs[0][0, :, :, :])
return ort_kps
if __name__ == '__main__':
x_scale = WIDTH/SIZE
y_scale = HEIGHT/SIZE
model_path = 'pdkp.onnx'
image_path = 'test.jpeg'
create_ort_session(model_path)
ori_image = cv2.imread(image_path)
ori_image = cv2.resize(ori_image, (SIZE, SIZE))
image = ori_image[:, :, ::-1]
ort_kps = run_inference(image)
ori_image = cv2.resize(ori_image, (HEIGHT, WIDTH))
img = render_kps(ori_image, ort_kps, y_scale, x_scale)
cv2.imshow('test', img)
cv2.waitKey(5000)
cv2.destroyAllWindows()
测试效果
