如何用Matplotlib优雅地实现NTU-RGB D骨架可视化

骨架数据集由于其强鲁棒性而被广泛运用于动作识别和预测等领域，其中NTU RGB-D数据是最常用的骨架动作数据集。无论是在论文中描述自己的数据，亦或是分析特征变化过程，可视化都是必不可少的一步，本文将以NTU RGB-D数据集为基准使用Python的Matplotlib库对骨架数据进行可视化

NTU RGB-D的数据集非常规范，其文件后缀名为‘.skeleton’，用记事本直接打开结果如下：

从上到下依次为帧数、骨架数量（不超过2，即在该数据集上动作由1个或2个人完成）、骨架其他信息（原文指’bodyID’, ‘clipedEdges’, ‘handLeftConfidence’, ‘handLeftState’, ‘handRightConfidence’, ‘handRightState’,‘isResticted’, ‘leanX’, ‘leanY’, ‘trackingState’ ）、一个骨架中有几个关节点、每个关节点的信息（ ‘x’, ‘y’, ‘z’, ‘depthX’, ‘depthY’, ‘colorX’, ‘colorY’, ‘orientationW’, ‘orientationX’, ‘orientationY’, ‘orientationZ’, ‘trackingState’）

其中我们只需要其中的x，y，z即可（2D图只需要x和y），于是我们的程序大概为：

1. 导入数据文件
2. 创建一个0填充列表用于存储数据
3. 循环每一帧，每一个骨架，每一个关节点将数据赋值给数组
4. 绘制，用散点图绘制骨架节点，用直线图绘制骨骼连接
5. 逐帧展示

首先导入依赖包和数据文件路径，设定好参数

import numpy as np
import matplotlib.pyplot as plt

file_name = r'C:UserslenovoDesktopvisiual/S001C001P001R001A009.skeleton'
max_V = 25 #节点数
max_M = 2 #骨架数量

然后循环遍历获取关节点坐标信息

with open(file_name, 'r') as fr:
    frame_num = int(fr.readline())
    point = np.zeros((3, frame_num, 25, 2))
    for frame in range(frame_num):
        person_num = int(fr.readline())
        for person in range(person_num):
            fr.readline()
            joint_num = int(fr.readline())
            for joint in range(joint_num):
                v = fr.readline().split(' ')
                if joint < max_V and person < max_M:
                    point[0,frame,joint,person] = float(v[0])#一个关节的一个坐标
                    point[1,frame,joint,person] = float(v[1])
                    point[2,frame,joint,person] = float(v[2])

准备绘制，根据坐标值选取合适的坐标轴

xmax= np.max(point[0, :, :, :])+ 0.5
xmin= np.min(point[0, :, :, :])- 0.5
ymax= np.max(point[1, :, :, :])+ 0.3
ymin= np.min(point[1, :, :, :])- 0.3

通过观察数据集判断哪几各节点相连为骨骼

# 根据NTU骨架结构确定哪几个节点相连为骨骼
# 注意序号从0开始，需要减1
arms= np.array([24,12,11,10,9,21,5,6,7,8,20])-1 #胳膊
rightHand= np.array([12,25])-1 #右手
leftHand= np.array([8,23])-1 #左手
legs= np.array([20,19,18,17,1,13,14,15,16]) - 1 #腿
body= np.array([4,3,21,2,1]) -1  #身体

设定好绘图参数后，根据节点和连接图逐帧绘制图像：

n= 0     # 从第n帧开始展示
m= 2 # 到第m帧结束，n 
结果展示，可以根据上面的绘制参数调节展示帧数和输出颜色线条粗细等，这里只输出了两帧：
 
 
整体代码：
 
import numpy as np
import matplotlib.pyplot as plt

file_name = r'C:UserslenovoDesktopvisiual/S001C001P001R001A009.skeleton'
max_V = 25 #节点数
max_M = 2 #骨架数量
with open(file_name, 'r') as fr:
    frame_num = int(fr.readline())
    point = np.zeros((3, frame_num, 25, 2))
    for frame in range(frame_num):
        person_num = int(fr.readline())
        for person in range(person_num):
            fr.readline()
            joint_num = int(fr.readline())
            for joint in range(joint_num):
                v = fr.readline().split(' ')
                if joint < max_V and person < max_M:
                    point[0,frame,joint,person] = float(v[0])#一个关节的一个坐标
                    point[1,frame,joint,person] = float(v[1])
                    point[2,frame,joint,person] = float(v[2])

print('read data done!')
print(point.shape)

#通过最大最小的坐标值来选取合适的坐标轴 
xmax= np.max(point[0, :, :, :])+ 0.5
xmin= np.min(point[0, :, :, :])- 0.5
ymax= np.max(point[1, :, :, :])+ 0.3
ymin= np.min(point[1, :, :, :])- 0.3
zmax= np.max(point[2, :, :, :])
zmin= np.min(point[2, :, :, :])
 
row= point.shape[1] #有多少帧
 
 
# 根据NTU骨架结构确定哪几个节点相连为骨骼
# 注意序号从0开始，需要减1
arms= np.array([24,12,11,10,9,21,5,6,7,8,20])-1 #胳膊
rightHand= np.array([12,25])-1 #右手
leftHand= np.array([8,23])-1 #左手
legs= np.array([20,19,18,17,1,13,14,15,16]) - 1 #腿
body= np.array([4,3,21,2,1]) -1  #身体
 
n= 0     # 从第n帧开始展示
m= 2 # 到第m帧结束，n 
4.说明 
部分代码参考了博客园博主https://www.cnblogs.com/picassooo/p/14044566.html，稍加修改并添加注释
 
本代码的核心思想在于：使用散点图绘制关节点，使用直线图绘制骨骼，如果有不同的图像需求请参考Matplotlib库自行调节参数。