RoHM

基于扩散模型的鲁棒人体运动重建

项目主页 | 论文

RoHM是一种新型的基于扩散模型的运动模型，它能够在有噪声和遮挡的输入数据的条件下，重建完整、合理的运动序列，并保持一致的全局坐标系。 -- 我们将其分解为两个子任务，并学习两个模型，一个用于全局轨迹，一个用于局部运动。为了捕捉两者之间的相关性，我们引入了一个新颖的条件模块，并将其与迭代推理方案相结合。

安装

创建一个干净的conda环境并安装所有依赖项：

conda env create -f environment.yml

安装完成后，激活conda环境：

conda activate rohm

数据准备

AMASS

从AMASS数据集下载SMPL-X中性注释，并解压文件。
要将原始AMASS数据预处理成RoHM所需的格式，请对每个子集运行以下脚本，其中dataset_name表示每个子集的名称。它会将处理后的AMASS数据保存到datasets/AMASS_smplx_preprocessed。

python preprocessing_amass.py --dataset_name=SUBSET_NAME --amass_root=PATH/TO/AMASS --save_root=datasets/AMASS_smplx_preprocessed

PROX

下载PROX数据集的以下内容：

从官方PROX数据集下载cam2world、calibration和recordings
从这里下载keypoints_openpose和mask_joint
按以下方式组织内容：

PROX
├── cam2world
├── calibration
├── recordings
├── keypoints_openpose
├── mask_joint

EgoBody

下载EgoBody数据集的以下内容：

从官方EgoBody数据集下载kinect_color、data_splits.csv、calibrations、kinect_cam_params、smplx_camera_wearer_*、smplx_interactee_*
从这里下载keypoints_cleaned、mask_joint和egobody_rohm_info.csv
按以下方式组织内容：

EgoBody
├── kinect_color
├── data_splits.csv
├── smplx_camera_wearer_train
├── smplx_camera_wearer_test
├── smplx_camera_wearer_val
├── smplx_interactee_train
├── smplx_interactee_test
├── smplx_interactee_val
├── calibrations
├── kinect_cam_params
├── keypoints_cleaned
├── mask_joint
├── egobody_rohm_info.csv

egobody_rohm_info.csv包含我们用于评估RoHM的EgoBody录像信息。

SMPL-X身体模型

从这里下载SMPL-X身体模型。注意最新版本是1.1，而我们在实现中使用1.0。

从这里下载smplx顶点分割smplx_vert_segmentation.json。

其他数据（检查点、结果等）

从这里下载模型检查点。从这里下载其他处理/保存的数据并解压，包括：

init_motions，PROX和EgoBody上的初始化运动序列（RoHM输入）
test_results_release，PROX和EgoBody上的重建运动序列（RoHM输出）
eval_noise_smplx，用于AMASS上RoHM评估的预计算运动噪声

按以下方式组织所有下载的数据：

RoHM
├── data
│   ├── body_models
│   │   ├── smplx_model
│   │   │   ├── smplx
│   ├── checkpoints
│   ├── eval_noise_smplx
│   ├── init_motions
│   ├── test_results_release
│   ├── smplx_vert_segmentation.json
├── datasets
│   ├── AMASS_smplx_preprocessed
│   ├── PROX
│   ├── EgoBody

训练

RoHM在AMASS数据集上进行训练。

TrajNet训练

使用课程学习方案对vanilla TrajNet进行三个阶段的训练，逐步增加噪声比例：

python train_trajnet.py --config=cfg_files/train_cfg/trajnet_train_vanilla_stage1.yaml 
python train_trajnet.py --config=cfg_files/train_cfg/trajnet_train_vanilla_stage2.yaml --pretrained_model_path=PATH/TO/MODEL
python train_trajnet.py --config=cfg_files/train_cfg/trajnet_train_vanilla_stage3.yaml --pretrained_model_path=PATH/TO/MODEL

对于第2和第3阶段，将pretrained_model_path设置为上一阶段训练的检查点。为了获得报告的检查点，我们分别对第1/2/3阶段训练800k/400k/450k步。

使用TrajControl对TrajNet进行微调：

python train_trajnet.py --config=cfg_files/train_cfg/trajnet_ft_trajcontrol.yaml --pretrained_backbone_path=PATH/TO/MODEL

将pretrained_backbone_path设置为预训练的vanilla TrajNet检查点，我们训练400k步以获得报告的检查点。

PoseNet训练

使用课程学习方案对PoseNet进行两个阶段的训练，逐步增加噪声比例：

python train_posenet.py --config=cfg_files/train_cfg/posenet_train_stage1.yaml
python train_posenet.py --config=cfg_files/train_cfg/posenet_train_stage2.yaml --pretrained_model_path=PATH/TO/MODEL

对于第2阶段，将pretrained_model_path设置为上一阶段训练的检查点。为了获得报告的检查点，我们分别对第1/2阶段训练300k/200k步。

在AMASS上测试和评估

在AMASS上测试

使用不同配置在AMASS上测试（对应论文中的表1），并将重建结果保存到test_results/results_amass_full：注意，在不同机器上使用相同的随机种子运行给定配置无法保证完全相同的数字，但随机性很小。

输入噪声级别3，遮挡10%的帧（同时遮挡轨迹和局部身体姿势）：

python test_amass_full.py --config=cfg_files/test_cfg/amass_occ_0.1_noise_3.yaml

输入噪声级别3，遮挡下半身关节：

python test_amass_full.py --config=cfg_files/test_cfg/amass_occ_leg_noise_3.yaml

输入噪声级别5，遮挡下半身关节：

python test_amass_full.py --config=cfg_files/test_cfg/amass_occ_leg_noise_5.yaml

输入噪声级别7，遮挡下半身关节：

python test_amass_full.py --config=cfg_files/test_cfg/amass_occ_leg_noise_7.yaml

在AMASS上评估

计算评估指标并可视化/渲染AMASS上的重建结果。

输入噪声级别3，遮挡10%的帧（同时遮挡轨迹和局部姿势）：

python eval_amass_full.py --config=cfg_files/eval_cfg/amass_occ_0.1_noise_3.yaml --saved_data_path=PATH/TO/TEST/RESULTS

输入噪声级别3，遮挡下半身关节：

python eval_amass_full.py --config=cfg_files/eval_cfg/amass_occ_leg_noise_3.yaml --saved_data_path=PATH/TO/TEST/RESULTS

输入噪声级别5，遮挡下半身关节：

python eval_amass_full.py --config=cfg_files/eval_cfg/amass_occ_leg_noise_5.yaml --saved_data_path=PATH/TO/TEST/RESULTS

输入噪声级别7，遮挡下半身关节：

python eval_amass_full.py --config=cfg_files/eval_cfg/amass_occ_leg_noise_7.yaml --saved_data_path=PATH/TO/TEST/RESULTS

用于可视化和渲染的其他标志：

--visualize=True：使用open3d可视化输入/输出/GT运动（包括骨骼和身体网格）
--render=True：使用pyrender渲染输入/输出/GT运动，并将渲染结果保存到--render_save_path

在PROX/EgoBody上测试和评估

对应论文中表2和表3的实验设置。

初始化

为了获得PROX上的初始（有噪声和部分可见）运动，我们使用以下选项：

对于PROX上基于RGB的重建，我们从CLIFF获得初始身体姿势，从PIXIE获得身体形状，从MeTRAbs获得全局平移/旋转。
对于PROX上基于RGBD的重建，我们使用改编自LEMO的代码进行逐帧优化获得初始运动。
对于EgoBody上基于RGB的重建，我们使用HuMoR中的VPoser-t代码获得初始运动。

我们在data/init_motions文件夹中提供了预处理的初始运动序列，并在data/test_results_release文件夹中提供了RoHM的最终输出运动序列供参考。

注意，对于以下脚本，初始运动在PROX中应以z轴向上，在EgoBody中应以y轴向上。

在PROX/EgoBody上测试

使用RGB-D输入在PROX上测试(初始化序列通过逐帧优化获得)，结果将保存到 test_results/results_prox_rgbd:

python test_prox_egobody.py --config=cfg_files/test_cfg/prox_rgbd.yaml --recording_name=RECORDING_NAME

使用RGB输入在PROX上测试(初始化序列通过回归器获得)，结果将保存到 test_results/results_prox_rgb:

python test_prox_egobody.py --config=cfg_files/test_cfg/prox_rgb.yaml --recording_name=RECORDING_NAME

使用RGB输入在EgoBody上测试(初始化序列通过VPoser-t获得，如HuMoR中所述)，结果将保存到 test_results/results_egobody_rgb:

python test_prox_egobody.py --config=cfg_files/test_cfg/egobody_rgb.yaml --recording_name=RECORDING_NAME

在PROX/EgoBody上评估

计算评估指标并可视化/渲染PROX/EgoBody上的重建结果。

在PROX上使用RGB-D输入评估:

python eval_prox_egobody.py --config=cfg_files/eval_cfg/prox_rgbd.yaml --saved_data_dir=PATH/TO/TEST/RESULTS --recording_name=RECORDING_NAME

在PROX上使用RGB输入评估:

python eval_prox_egobody.py --config=cfg_files/eval_cfg/prox_rgb.yaml --saved_data_dir=PATH/TO/TEST/RESULTS --recording_name=RECORDING_NAME

在EgoBody上使用RGB输入评估:

python eval_prox_egobody.py --config=cfg_files/eval_cfg/egobody_rgb.yaml --saved_data_dir=PATH/TO/TEST/RESULTS --recording_name=RECORDING_NAME

注意: recording_name 可以设置为:

序列记录名称: 然后评估在这个特定序列上进行。
'all': 评估在子集的所有序列上进行(用于报告论文中的数据)。

其他用于可视化和渲染的标志:

--visualize=True: 使用open3d可视化输入/输出/GT动作
- --vis_option=mesh: 可视化身体
- --vis_option=skeleton: 可视化骨架
--render=True: 使用pyrender渲染输入/输出/GT动作，并将渲染结果保存到 --render_save_path

自定义输入

如果你想在自定义输入上运行RoHM:

步骤1: 按照 data/init_motions 中的数据格式准备初始SMPL-X序列
步骤2: 按照 datasets/PROX/mask_joint 中的数据格式准备关节遮挡掩码
- 如果你有3D场景网格，从相机视角为3D场景渲染深度图，并通过比较深度值来识别3D关节是否被遮挡 (我们使用 utils/get_occlusion_mask.py 来获取PROX数据集上的遮挡掩码)
- 如果你没有3D场景网格，你可以使用OpenPose或其他2D人体检测方法的置信度分数，并将低置信度的关节设置为被遮挡
步骤3: 根据坐标系进行自定义规范化:
- 当前实现支持对y轴向上(EgoBody)或z轴向上(PROX/AMASS)的初始序列进行规范化，规范化后的序列始终是z轴向上
- 如果你的输入初始序列不符合这一点，你需要首先进行适当的变换以获得z/y轴向上的序列

许可证

RoHM的大部分内容采用CC-BY-NC许可(包括代码、发布的检查点、发布的初始化/最终运动序列数据集)，然而，项目的某些部分采用单独的许可条款:

Trimesh、Guided Diffusion和MDM采用MIT许可证;
konia采用Apache许可证。

引用

如果你发现我们的工作对你的研究有用，请考虑引用:

@inproceedings{zhang2024rohm,
   title={RoHM: Robust Human Motion Reconstruction via Diffusion},
   author={Zhang, Siwei and Bhatnagar, Bharat Lal and Xu, Yuanlu and Winkler, Alexander and Kadlecek, Petr and Tang, Siyu and Bogo, Federica},
   booktitle={CVPR},
   year={2024}
 }