手把手教你用Python处理NTU-RGB+D骨骼数据(附完整代码与可视化)

当你第一次拿到NTU-RGB+D数据集时,面对数千个.skeleton文件和复杂的25个关节点数据结构,可能会感到无从下手。这份实战指南将带你从零开始,用Python一步步解析这些骨骼数据,并实现2D/3D动态可视化。我们会用Matplotlib处理2D可视化,用Open3D实现3D骨骼动画,过程中还会分享处理坐标归一化、帧对齐等常见问题的解决方案。

1. 理解NTU-RGB+D数据格式

1.1 文件命名规则解析

每个.skeleton文件都遵循严格的命名规则,理解这些规则对后续数据处理至关重要:

S001C001P001R001A001.skeleton
  • S001 :设置编号(1-17),对应不同的相机高度和距离配置
  • C001 :相机视角(1-3),分别对应-45°、0°和45°视角
  • P001 :受试者ID(1-40),共40位不同年龄的参与者
  • R001 :动作重复次数(1-2),每个动作执行两遍
  • A001 :动作类别(1-60),共60种不同动作类型

1.2 数据结构详解

用文本编辑器打开.skeleton文件,你会看到类似如下的内容:

100
1
0 0 0 0 0 0 0 0 0 0
25
0.12 -0.34 2.56 ... 1
0.15 -0.32 2.54 ... 1
...

文件结构解析:

  1. 第一行 :该样本包含的帧数(如100表示100帧动作序列)
  2. 第二行 :执行动作的人数(NTU-RGB+D中通常为1或2)
  3. 第三行 :10个元数据字段,包括身体ID、手部状态等
  4. 第四行 :关节点数量(固定为25)
  5. 后续行 :每帧25个关节点的12维数据,包括:
    • 3D坐标(x,y,z)
    • 深度图坐标(depthX,depthY)
    • RGB图像坐标(colorX,colorY)
    • 四元数旋转(orientationW,X,Y,Z)
    • 追踪状态(trackingState)

2. Python数据读取与解析

2.1 安装必要依赖库

首先确保安装了以下Python库:

pip install numpy matplotlib open3d tqdm

2.2 构建数据解析器

创建一个 NTUReader 类来统一处理.skeleton文件:

import numpy as np
from pathlib import Path

class NTUReader:
    def __init__(self, filepath):
        self.filepath = Path(filepath)
        self.metadata = {}
        self.joints_data = []
        self._parse_file()
    
    def _parse_file(self):
        with open(self.filepath, 'r') as f:
            lines = [line.strip() for line in f.readlines()]
            
            # 解析基础信息
            self.metadata['frame_count'] = int(lines[0])
            self.metadata['person_count'] = int(lines[1])
            self.metadata['body_info'] = list(map(float, lines[2].split()))
            
            # 解析关节数据
            joint_count = int(lines[3])
            frame_data = []
            
            for line in lines[4:]:
                if not line: continue
                values = list(map(float, line.split()))
                
                if len(values) == 12:  # 单个关节点数据
                    frame_data.append(values)
                
                if len(frame_data) == joint_count:  # 完成一帧读取
                    self.joints_data.append(frame_data)
                    frame_data = []
        
        # 转换为numpy数组方便处理
        self.joints_data = np.array(self.joints_data)
    
    def get_3d_coordinates(self):
        """提取所有帧的3D坐标(x,y,z)"""
        return self.joints_data[:, :, :3]
    
    def get_normalized_coordinates(self):
        """获取归一化后的3D坐标"""
        coords = self.get_3d_coordinates()
        # 减去骨盆关节(第一个关节)坐标进行中心化
        pelvis = coords[:, 0:1, :]
        centered = coords - pelvis
        # 按最大绝对值进行缩放
        max_val = np.max(np.abs(centered))
        return centered / (max_val + 1e-8)

3. 2D骨骼可视化

3.1 使用Matplotlib绘制单帧骨骼

定义骨骼连接关系,绘制2D投影:

import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation

# 定义关节点连接关系 (基于NTU-RGB+D的25个关节点)
BONES = [
    (1, 2), (2, 21), (3, 21), (4, 3), (5, 21), 
    (6, 5), (7, 6), (8, 7), (9, 21), (10, 9),
    (11, 10), (12, 11), (13, 1), (14, 13), (15, 14),
    (16, 15), (17, 1), (18, 17), (19, 18), (20, 19),
    (22, 23), (21, 22), (23, 24), (24, 25)
]

def plot_single_frame(coords_2d, bones=BONES, ax=None):
    """绘制单帧2D骨骼图"""
    if ax is None:
        fig, ax = plt.subplots(figsize=(10, 10))
    
    # 绘制关节点
    ax.scatter(coords_2d[:, 0], coords_2d[:, 1], c='blue', s=50)
    
    # 绘制骨骼连接
    for i, j in bones:
        if i <= len(coords_2d) and j <= len(coords_2d):
            ax.plot([coords_2d[i-1, 0], coords_2d[j-1, 0]],
                    [coords_2d[i-1, 1], coords_2d[j-1, 1]], 
                    'r-', linewidth=2)
    
    ax.set_xlim(-1, 1)
    ax.set_ylim(-1, 1)
    ax.set_aspect('equal')
    ax.invert_yaxis()  # 保持与图像坐标系一致
    return ax

3.2 创建2D骨骼动画

将多帧序列转换为动态可视化:

def create_2d_animation(normalized_coords, save_path=None):
    """创建并保存2D骨骼动画"""
    fig, ax = plt.subplots(figsize=(10, 10))
    scat = ax.scatter([], [], c='blue', s=50)
    lines = [ax.plot([], [], 'r-', linewidth=2)[0] for _ in BONES]
    
    def init():
        ax.set_xlim(-1, 1)
        ax.set_ylim(-1, 1)
        ax.set_aspect('equal')
        ax.invert_yaxis()
        return [scat] + lines
    
    def update(frame):
        # 更新关节点位置
        scat.set_offsets(frame[:, :2])
        
        # 更新骨骼连接
        for line, (i, j) in zip(lines, BONES):
            if i <= len(frame) and j <= len(frame):
                line.set_data([frame[i-1, 0], frame[j-1, 0]],
                             [frame[i-1, 1], frame[j-1, 1]])
            else:
                line.set_data([], [])
        
        return [scat] + lines
    
    anim = FuncAnimation(fig, update, frames=normalized_coords,
                         init_func=init, blit=True, interval=50)
    
    if save_path:
        anim.save(save_path, writer='ffmpeg', fps=20)
    
    plt.close()
    return anim

4. 3D骨骼可视化

4.1 使用Open3D创建3D可视化

安装Open3D并实现3D骨骼渲染:

import open3d as o3d
from IPython.display import HTML

def create_3d_skeleton_visualizer(coords_3d, bones=BONES):
    """创建交互式3D骨骼可视化"""
    vis = o3d.visualization.Visualizer()
    vis.create_window(width=800, height=600)
    
    # 创建关节点几何体
    joints = o3d.geometry.PointCloud()
    joints.points = o3d.utility.Vector3dVector(coords_3d[0])
    
    # 创建骨骼线几何体
    bone_lines = []
    for i, j in bones:
        if i <= len(coords_3d[0]) and j <= len(coords_3d[0]):
            line = o3d.geometry.LineSet()
            line.points = o3d.utility.Vector3dVector([coords_3d[0][i-1], coords_3d[0][j-1]])
            line.lines = o3d.utility.Vector2iVector([[0, 1]])
            bone_lines.append(line)
    
    # 添加到可视化器
    vis.add_geometry(joints)
    for line in bone_lines:
        vis.add_geometry(line)
    
    # 更新函数
    def update_geometry(vis, frame_idx):
        joints.points = o3d.utility.Vector3dVector(coords_3d[frame_idx])
        
        new_lines = []
        for idx, (i, j) in enumerate(bones):
            if i <= len(coords_3d[frame_idx]) and j <= len(coords_3d[frame_idx]):
                line = o3d.geometry.LineSet()
                line.points = o3d.utility.Vector3dVector([
                    coords_3d[frame_idx][i-1], 
                    coords_3d[frame_idx][j-1]
                ])
                line.lines = o3d.utility.Vector2iVector([[0, 1]])
                new_lines.append(line)
        
        # 更新骨骼线
        for old_line, new_line in zip(bone_lines, new_lines):
            old_line.points = new_line.points
            old_line.lines = new_line.lines
            vis.update_geometry(old_line)
        
        vis.update_geometry(joints)
        return False
    
    # 创建动画
    def animate():
        for i in range(len(coords_3d)):
            update_geometry(vis, i)
            vis.poll_events()
            vis.update_renderer()
            time.sleep(0.05)
    
    return vis, animate

4.2 3D动画保存与交互

将3D动画保存为视频或直接在Jupyter中交互:

def save_3d_animation(coords_3d, output_path, fps=20):
    """将3D骨骼动画保存为视频"""
    vis, animate = create_3d_skeleton_visualizer(coords_3d)
    
    # 配置视频录制
    vis.get_view_control().set_zoom(0.8)
    vis.capture_screen_image("temp.png", do_render=True)
    
    # 创建视频写入器
    video_writer = cv2.VideoWriter(
        output_path, 
        cv2.VideoWriter_fourcc(*'mp4v'), 
        fps, 
        (800, 600)
    )
    
    # 录制动画
    for i in range(len(coords_3d)):
        update_geometry(vis, i)
        vis.poll_events()
        vis.update_renderer()
        
        # 捕获当前帧
        img = np.asarray(vis.capture_screen_float_buffer(do_render=True))
        img = (img * 255).astype(np.uint8)
        video_writer.write(cv2.cvtColor(img, cv2.COLOR_RGB2BGR))
    
    video_writer.release()
    vis.destroy_window()

5. 实战技巧与常见问题

5.1 坐标归一化处理

原始骨骼坐标可能存在以下问题:

  • 不同样本的绝对坐标差异大
  • 受试者距离相机远近不同导致尺度不一致
  • 坐标原点不统一

解决方案:

def normalize_coordinates(coords_3d):
    """
    标准化3D坐标:
    1. 以骨盆关节(关节1)为中心
    2. 统一缩放至[-1,1]范围
    """
    # 中心化
    centered = coords_3d - coords_3d[:, 0:1, :]
    
    # 缩放
    max_abs = np.max(np.abs(centered), axis=(0,1,2), keepdims=True)
    return centered / (max_abs + 1e-8)  # 避免除以零

5.2 帧对齐与插值

当处理不同长度的动作序列时,可能需要统一帧数:

from scipy.interpolate import interp1d

def align_frames(coords_3d, target_length=100):
    """将骨骼序列插值到目标长度"""
    original_length = coords_3d.shape[0]
    x_original = np.linspace(0, 1, original_length)
    x_target = np.linspace(0, 1, target_length)
    
    aligned = np.zeros((target_length, coords_3d.shape[1], coords_3d.shape[2]))
    
    for joint in range(coords_3d.shape[1]):
        for dim in range(coords_3d.shape[2]):
            interp_func = interp1d(
                x_original, 
                coords_3d[:, joint, dim], 
                kind='linear',
                fill_value="extrapolate"
            )
            aligned[:, joint, dim] = interp_func(x_target)
    
    return aligned

5.3 处理多人交互数据

对于包含多人交互的样本(person_count=2),需要先分离数据:

def split_multi_person_data(reader):
    """分离多人骨骼数据"""
    if reader.metadata['person_count'] == 1:
        return [reader.joints_data]
    
    # 多人情况下,每隔25行是一个人的数据
    frames = []
    for frame in reader.joints_data:
        persons = np.split(frame, reader.metadata['person_count'])
        frames.append(persons)
    
    # 重组为[person1所有帧, person2所有帧,...]
    return [np.stack([frame[i] for frame in frames], axis=0) 
            for i in range(reader.metadata['person_count'])]

6. 完整处理流程示例

6.1 从文件到可视化的端到端流程

# 1. 读取并解析.skeleton文件
reader = NTUReader("S001C001P001R001A001.skeleton")

# 2. 获取归一化后的3D坐标
coords_3d = reader.get_normalized_coordinates()

# 3. 帧对齐处理(可选)
aligned_coords = align_frames(coords_3d, target_length=100)

# 4. 创建2D动画
create_2d_animation(aligned_coords, save_path="2d_animation.mp4")

# 5. 创建3D可视化
vis, animate = create_3d_skeleton_visualizer(aligned_coords)
animate()  # 在交互窗口中播放动画

6.2 批量处理数据集

处理整个数据集的实用函数:

from tqdm import tqdm
import os

def process_dataset(dataset_dir, output_dir, max_samples=None):
    """批量处理数据集中的.skeleton文件"""
    os.makedirs(output_dir, exist_ok=True)
    skeleton_files = [f for f in os.listdir(dataset_dir) if f.endswith('.skeleton')]
    
    if max_samples:
        skeleton_files = skeleton_files[:max_samples]
    
    for filename in tqdm(skeleton_files, desc="Processing samples"):
        try:
            # 1. 读取文件
            reader = NTUReader(os.path.join(dataset_dir, filename))
            
            # 2. 处理坐标
            coords_3d = reader.get_normalized_coordinates()
            aligned_coords = align_frames(coords_3d)
            
            # 3. 保存可视化结果
            base_name = os.path.splitext(filename)[0]
            create_2d_animation(
                aligned_coords,
                save_path=os.path.join(output_dir, f"{base_name}_2d.mp4")
            )
            
            # 4. 保存处理后的数据
            np.save(
                os.path.join(output_dir, f"{base_name}_coords.npy"),
                aligned_coords
            )
            
        except Exception as e:
            print(f"Error processing {filename}: {str(e)}")

在实际项目中处理NTU-RGB+D数据时,最耗时的部分往往是第一次解析所有.skeleton文件。建议先批量转换为更高效的.npy格式存储,后续直接从这些预处理文件加载会快很多。

Logo

这里是“一人公司”的成长家园。我们提供从产品曝光、技术变现到法律财税的全栈内容,并连接云服务、办公空间等稀缺资源,助你专注创造,无忧运营。

更多推荐