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

1.简介

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

2.数据分析

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. 逐帧展示

3.实现过程

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

import numpy as np
import matplotlib.pyplot as plt

file_name = r'C:\Users\lenovo\Desktop\visiual/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<m<row，这个m可以选取一个小于最大帧数的阈值，便于查看，若m=1则展示一帧
plt.figure()
plt.ion() #使用plt.ion()这个函数，使matplotlib的显示模式转换为交互（interactive）模式。即使在脚本中遇到plt.show()，代码还是会继续执行。
color_point = '#03ff' #关节点颜色,可输入16进制调色板
color_bone = 'red' #骨骼颜色
for i in range(n, m):
    plt.cla() ## Clear axis即清除当前图形中的当前活动轴。其他轴不受影响。
    plt.scatter(point[0, i, :, :], point[1, i, :, :], c=color_point, s=40.0) #通过散点图绘制关节点
    #通过直线图绘制两点间的连接线，即骨骼
    plt.plot(point[0, i, arms,0], point[1, i, arms,0], c=color_bone, lw=2.0) 
    plt.plot(point[0, i, rightHand,0], point[1, i, rightHand,0], c=color_bone, lw=2.0)
    plt.plot(point[0, i, leftHand,0], point[1, i, leftHand,0], c=color_bone, lw=2.0)
    plt.plot(point[0, i, legs,0], point[1, i, legs,0], c=color_bone, lw=2.0)
    plt.plot(point[0, i, body,0], point[1, i, body,0], c=color_bone, lw=2.0)
    
    #第二个骨架，如果有的话
    plt.plot(point[0, i, arms,1], point[1, i, arms,1], c=color_bone, lw=2.0)
    plt.plot(point[0, i, rightHand,1], point[1, i, rightHand,1], c=color_bone, lw=2.0)
    plt.plot(point[0, i, leftHand,1], point[1, i, leftHand,1], c=color_bone, lw=2.0)
    plt.plot(point[0, i, legs,1], point[1, i, legs,1], c=color_bone, lw=2.0)
    plt.plot(point[0, i, body,1], point[1, i, body,1], c=color_bone, lw=2.0)
     
    plt.text(xmax-0.5, ymax-0.1,'frame: {}/{}'.format(i, row-1)) #这是第几帧
    plt.xlim(xmin, xmax) #坐标轴
    plt.ylim(ymin, ymax)
    plt.pause(0.001)
 
plt.ioff()
plt.show()

结果展示，可以根据上面的绘制参数调节展示帧数和输出颜色线条粗细等，这里只输出了两帧：

4.整体代码实现：

import numpy as np
import matplotlib.pyplot as plt

file_name = r'C:\Users\lenovo\Desktop\visiual/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<m<row，这个m可以选取一个小于最大帧数的阈值，便于查看，若m=1则展示一帧
plt.figure()
plt.ion() #使用plt.ion()这个函数，使matplotlib的显示模式转换为交互（interactive）模式。即使在脚本中遇到plt.show()，代码还是会继续执行。
color_point = '#03ff' #关节点颜色,可输入16进制调色板
color_bone = 'red' #骨骼颜色
for i in range(n, m):
    plt.cla() ## Clear axis即清除当前图形中的当前活动轴。其他轴不受影响。
    plt.scatter(point[0, i, :, :], point[1, i, :, :], c=color_point, s=40.0) #通过散点图绘制关节点
    #通过直线图绘制两点间的连接线，即骨骼
    plt.plot(point[0, i, arms,0], point[1, i, arms,0], c=color_bone, lw=2.0) 
    plt.plot(point[0, i, rightHand,0], point[1, i, rightHand,0], c=color_bone, lw=2.0)
    plt.plot(point[0, i, leftHand,0], point[1, i, leftHand,0], c=color_bone, lw=2.0)
    plt.plot(point[0, i, legs,0], point[1, i, legs,0], c=color_bone, lw=2.0)
    plt.plot(point[0, i, body,0], point[1, i, body,0], c=color_bone, lw=2.0)
    
    #第二个骨架，如果有的话
    plt.plot(point[0, i, arms,1], point[1, i, arms,1], c=color_bone, lw=2.0)
    plt.plot(point[0, i, rightHand,1], point[1, i, rightHand,1], c=color_bone, lw=2.0)
    plt.plot(point[0, i, leftHand,1], point[1, i, leftHand,1], c=color_bone, lw=2.0)
    plt.plot(point[0, i, legs,1], point[1, i, legs,1], c=color_bone, lw=2.0)
    plt.plot(point[0, i, body,1], point[1, i, body,1], c=color_bone, lw=2.0)
     
    plt.text(xmax-0.5, ymax-0.1,'frame: {}/{}'.format(i, row-1)) #这是第几帧
    # plt.text(xmax-0.8, ymax-0.4, 'label: ' + str(label[i]))
    plt.xlim(xmin, xmax) #坐标轴
    plt.ylim(ymin, ymax)
    plt.pause(0.001)
 
plt.ioff()
plt.show()

5.说明

如果需要添加Z轴，实现3D可视化，请查看我最新的博客：
NTU RGB-D 60 骨架数据集3D可视化

本代码的核心思想在于：使用散点图绘制关节点，使用直线图绘制骨骼，如果有不同的图像需求请参考Matplotlib库自行调节参数。

部分代码参考了博客园博主https://www.cnblogs.com/picassooo/p/14044566.html，稍加修改并添加注释