3D姿态估计——ThreeDPose项目简单易用的模型解析

前言

之前写过tensorflow官方的posenet模型解析，用起来比较简单，但是缺点是只有2D关键点，本着易用性的原则，当然要再来个简单易用的3D姿态估计。偶然看见了ThreeDPose的项目，感觉很强大的，所以把模型扒下来记录一下调用方法。

参考博客：

ThreeDPose官方代码

微软的ONNX模型解析库

ONNX解析库的pythonAPI文档

3D姿态估计最大的好处就是卡通角色的肢体驱动了，其实就是单目摄像头的动捕方法。

理论和代码解析

可以从官方去下载模型,戳这里，或者在文末的百度网盘下载。

模型结构

模型已经被作者转换成ONNX的模型文件了，所以下载netron软件去可视化模型，打开以后可以发现模型的网络结构和输入输出

这里说明一下各部分含义：

有三个input，其实都是一样，都要输入(448,448,3)的图片

有四个output，解析模型即提取关键点的时候，只需要第3和4个输出，具体解析方法看下面代码解析。

代码解析

在windows上为了读取ONNX的模型文件，可以使用微软提供的onnxruntime这个库，直接pip安装即可，这个库的文档见上面参考博客的2和3。

先引入相关的库文件

import numpy as np
import onnxruntime as rt
import cv2

import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D

首先就是加载模型，输入图片，推断当前图片的四个输出

#加载模型
sess = rt.InferenceSession("Resnet34_3inputs_448x448_20200609.onnx")
inputs = sess.get_inputs()
#读取图片
img = cv2.imread("D:/photo/pose/5.jpg")
img = cv2.cvtColor(img,cv2.COLOR_BGR2RGB)
img = cv2.resize(img,(448,448))
img = img.astype(np.float32)/255.0
img = img.transpose(2,1,0)
img = img[np.newaxis,...]
img.shape
#输入到网络结构中
pred_onx = sess.run(None,{
    inputs[0].name:img,
    inputs[1].name:img,
    inputs[2].name:img
})

获取输出

offset3D = np.squeeze(pred_onx[2])
heatMap3D = np.squeeze(pred_onx[3])
print(offset3D.shape)#(2016, 28, 28)
print(heatMap3D.shape)#(672, 28, 28)
print(offset3D.shape[0]/heatMap3D.shape[0])#3.0

按照参考博客1的unity代码解析输出，可以发现这个命名和posenet的解析一模一样，使用heatmap粗略定位关节位置，然后使用offset在heatmap结果上精确调整关节位置。

heatmap的 $(672,28,28)$ 代表 24个关节的28个大小为$(28,28)$的特征图。而offset比heatmap的特征图多三倍，很明显就是刚才说的精确定位，只不过需要在offset中定位到x,y,z三个坐标，所以就是三倍关系了。

定位原理就是：heatmap相当于把原图划分为$(28,28)$的网格点，每个网格点代表附近有一个关节的概率，通过找到某个关节最可能在heatmap哪一副特征图的哪个网格，我们根据heatmap的索引在offset中找到对应的3个精确矫正xyz坐标的特征图，这三个特征图的值就是对应关节的xyz坐标相对于当前heatmap网格位置的精确偏移量。

在写代码之前，着重关注一下heatmap和offset的特征图相对于关节是一个怎样的顺序。

• heatmap的顺序是第1个关节的第1个特征图、第1个关节的第2个特征图、…、第2个关节的第1个特征图、第二个关节的第2个特征图、…、第24个关节的第28个特征图
• offsetmap的顺序是第1个关节的第1个特征图对应的x坐标偏移、第1个关节的第2个特征图对应的x坐标偏移、第1个关节的第3个特征图对应的x坐标偏移、…、第1个关节的第28个特征图对应的x坐标偏移、…、第2个关节的第1个特征图对应的x坐标偏移、…、第24个关节的第28个特征图对应的x坐标偏移、第1个关节的第1个特征图对应的y坐标偏移、第1个关节的第2个特征图对应的y坐标偏移、…、第24个关节的第28个特征图对应的y坐标偏移、第1个关节的第1个特征图对应的z坐标偏移、第1个关节的第2个特征图对应的z坐标偏移、…、第24个关节的第28个特征图对应的z坐标偏移。

说了一堆，不如看代码简单明了：

比如提取第j个关节的3D坐标位置，对应的特征图是$[j\ast 28,(j+1)\ast 28-1]$ ，然后从这里面找到最大值的位置就是当前关节最可能在哪个特征图的哪个网格位置上。

# 找到第j个关节的28个特征图，并找到最大值的索引
joint_heat = heatMap3D[j*28:(j+1)*28,...]
[x,y,z] = np.where(joint_heat==np.max(joint_heat))
# 避免有多个最大值，所以取最后一组
x=int(x[-1])
y=int(y[-1])
z=int(z[-1])

然后按照找到的heatmap索引去查询offset，找到精确的矫正值

pos_x = offset3D[j*28+x,y,z] + x
pos_y = offset3D[24*28+j*28+x,y,z] + y
pos_z = offset3D[24*28*2+j*28+x,y,z] + z

每个坐标都是在当前网格的位置上加上精确的偏移量，因为xyz分别在offset上分别隔了$24\ast 28$个特征图，所以出现了y和z的第一个索引有变化。

如果还有疑问，建议去看看前面的2D姿态估计的解析，然后自己手写一遍解析算法就理解了。

把所有关节的位置存到kps然后可视化看看

%matplotlib inline

fig = plt.figure()
ax = fig.gca(projection='3d')
ax.scatter3D(kps[:,0],-kps[:,1],-kps[:,2],'red')
parent = np.array([0,1,2,3,3,  1,6,7,8,8,   12,15,14,15,24,  24,16,17,18,  24,20,21,22, 0])-1;
for i in range(24):
    if(parent[i]!=-1):
        ax.plot3D(kps[[i,parent[i]],0], -kps[[i,parent[i]],1], -kps[[i,parent[i]],2], 'gray')
        
ax.xaxis.set_tick_params(labelsize=10)
ax.yaxis.set_tick_params(labelsize=10)
ax.zaxis.set_tick_params(labelsize=10)

ax.view_init(elev=10., azim=180)

其实把关键点保存下来用matlab画更好看

%test
clear;clc;close all
% a=dlmread("D:\code\python\ThreeDPose\unity_data\kps100.txt");
a=[0.292807,14.994286,6.560671;
-0.123055,15.480020,8.367174;
-2.666008,19.526012,9.689341;
-3.994198,19.631011,9.902868;
-3.537779,20.058028,9.978683;
1.467720,11.799177,6.442903;
-1.372830,10.170244,5.793396;
-2.116019,9.653587,3.315798;
-1.332626,9.906492,1.420080;
-2.152191,9.461523,2.985400;
0.378545,12.237494,4.861950;
-0.739344,12.452946,4.971550;
-0.678598,14.203241,5.388392;
-0.671332,13.130285,4.690326;
-1.256000,12.918788,5.623796;
0.216525,13.818989,12.736559;
0.793380,13.383041,17.217848;
-0.301103,13.578390,22.771406;
-0.406112,13.623824,23.280164;
0.380849,11.442492,12.643937;
-2.110893,11.706800,18.199155;
-1.205901,12.321800,22.871755;
-2.116101,12.384523,24.476315;
-0.500211,12.706436,11.490892];

parent =[15,1,2,3,3,  15,6,7,8,8,   12,15,14,15,24,  24,16,17,18,  24,20,21,22, 0];
plot3(a(:,1),-a(:,2),-a(:,3),'ro','MarkerSize',2,'MarkerFaceColor','r')
for i=1:24
    if(parent(i)~=0)
        line([a(i,1) a(parent(i),1)],[-a(i,2) -a(parent(i),2)],[-a(i,3) -a(parent(i),3)])
    end
end
axis equal
axis off

24个关节的名称分别标注一下说一下吧，从unity代码中的VNectModel.cs能找到，玩过骨骼动画的基本知道每个单词的代表的关节，这里就不画骨骼结构图了，自己对应到上面的关键点图中即可。

rShldrBend,    rForearmBend,    rHand,    rThumb2,    rMid1,
lShldrBend,    lForearmBend,    lHand,    lThumb2,    lMid1,
lEar,    lEye,    rEar,    rEye,    Nose,
rThighBend,    rShin,    rFoot,    rToe,
lThighBend,    lShin,    lFoot,    lToe,
abdomenUpper,

后记

其实如果要做肢体驱动，还需要

• 根据上述关节计算额外的一些关节坐标，这一块不在本博文的学习范围内，博文只关注怎么简单的使用这个模型去解析关节位置。
• 将关节坐标映射到原图，这个根据比例原图比例缩放一下即可，与posenet一样，也不做阐述。

模型文件网盘地址：
链接：https://pan.baidu.com/s/1TsvALWJRIoCAtQ9ffcno7w
提取码：rfow

本博文同步更新到微信公众号中，有兴趣可关注一波，代码在微信公众号简介的github找得到，有问题直接公众号私信。