TokenHMR

TokenHMR：通过标记化姿势表示推进人体网格重建

计算机视觉与模式识别（CVPR 2024）

Sai Kumar Dwivedi^1,*, Yu Sun^2,*, Priyanka Patel¹, Yao Feng^1,2,3, Michael J. Black¹

（*贡献相同）

¹马克斯·普朗克智能系统研究所，图宾根，德国
²Meshcapade ³苏黎世联邦理工学院

核心思想

大多数3D人体姿势估计方法使用带有2D关键点和/或3D伪真实标注的真实图像进行训练，这有助于提高泛化能力。然而，在这种数据上训练的方法虽然在图像对齐方面表现良好，但3D精度较差。TokenHMR通过引入阈值自适应损失缩放（TALS）损失并将身体回归重新表述为标记预测来解决这个问题。我们的方法分为两个阶段：

标记化： 编码器将连续的姿势映射为离散的姿势标记。
TokenHMR： 在人体姿势估计训练过程中，预训练的解码器提供有效姿势的"词汇表"，而不引入偏见。

更新

2024年2月7日：发布改进的TokenHMR模型 - 更好的图像对齐
2024年5月6日：发布TokenHMR代码和模型 - 论文中使用的模型

设置和安装

克隆仓库将仓库克隆到本地机器：

git clone https://github.com/saidwivedi/TokenHMR.git

创建Conda环境

重要提示：请勿使用高于3.10的Python版本。
```
conda create -n tkhmr python=3.10
```

安装PyTorch

经测试使用PyTorch 2.1.0和CUDA 11.8，但也可以使用较低版本。

pip install torch==2.1.0 torchvision==0.16.0 torchaudio==2.1.0 --index-url https://download.pytorch.org/whl/cu118

安装其他依赖

使用提供的requirements.txt文件安装其他依赖项。
```
pip install -r requirements.txt
```
安装Detectron2用于图像演示

确保CUDA_HOME设置为与PyTorch一起安装的CUDA版本。
```
pip install git+https://github.com/facebookresearch/detectron2
```
安装人体跟踪器的分叉版本（修复了依赖项）（PHALP, CVPR 2022）用于视频演示

确保CUDA_HOME设置为与PyTorch一起安装的CUDA版本。
```
pip install git+https://github.com/saidwivedi/PHALP
```

代码结构

TokenHMR/
├── tokenization/               # 包含训练、评估和运行标记化演示的代码。[方法(a)]
│   └── tokenization_data/      # 标记化的训练/评估数据。[查看下方如何下载]
│   ... 
├── tokenhmr/                   # 包含训练、评估和运行TokenHMR演示的代码 [方法(b)]
│   └── dataset_dir/            # TokenHMR训练的根目录
│       └── training_data/      # TokenHMR的训练数据。[查看下方如何下载]
│       └── evaluation_data/    # TokenHMR的训练数据。[查看下方如何下载]
│   ... 
├── data/                       # 基本设置数据
│   └── body_models/            # SMPL和SMPLH身体模型
│   └── checkpoints/            # 带配置的标记化和TokenHMR检查点
├── requirements.txt            # 依赖列表
└── ...

准备基本设置数据 [演示所需]

所有文件都上传到项目网页。下载需要注册并同意许可协议。使用脚本fetch_demo_data.sh下载运行演示所需的文件。这包括SMPL和SMPLH身体模型、最新的TokenHMR和标记化检查点。有关训练和评估，请参阅相应章节。

bash ./fetch_demo_data.sh

PHALP需要SMPL中性模型来运行视频演示。将模型复制到适当的位置。

cp data/body_models/smpl/SMPL_NEUTRAL.pkl $HOME/.cache/phalp/3D/models/smpl/

在图像上运行TokenHMR演示

在运行图像演示之前，请确保安装了Detectron2。查看安装指南了解更多详情。

python tokenhmr/demo.py \
    --img_folder demo_sample/images/ \
    --batch_size=1 \
    --full_frame \
    --checkpoint data/checkpoints/tokenhmr_model_latest.ckpt \
    --model_config data/checkpoints/model_config.yaml

在视频上运行TokenHMR演示

确保安装了原始PHALP, CVPR 2022的分叉版本。查看安装指南了解更多详情。

python tokenhmr/track.py \
    video.source=demo_sample/video/gymnasts.mp4 \
    render.colors=slahmr \
    +checkpoint=data/checkpoints/tokenhmr_model_latest.ckpt \
    +model_config=data/checkpoints/model_config.yaml

演示GIF

标记化

数据准备

我们使用AMASS和MOYO的身体姿势数据来训练tokenizer。用于训练和评估的处理后文件可以从这里下载（下载前需要在项目页面注册）。下载后解压，文件夹结构应如下所示：

TokenHMR/
├── tokenization/ 
│   └── tokenization_data/          # 用于Tokenization的训练/评估数据
│       └── smplh/
│           └── train/              # 训练集
│               └── train_CMU.npz
│               ...
│           └── val/                # 验证集
│               └── val_MPI-Mosh.npz            
│               ...
└── ...

训练

我们在单个A100 Nvidia GPU上训练tokenizer，迭代150K次，大约需要2天时间。

cd tokenization
python train_poseVQ.py --cfg configs/tokenizer_amass_moyo.yaml

数据准备

我们使用BEDLAM和4DHumans的训练数据来训练TokenHMR。参考这里下载4D humans的训练tar文件。对于BEDLAM的tar文件，请从我们的项目页面这里下载。评估时，请直接从各自的网站下载图片：3DPW和EMDB。评估所需的元数据可以从这里下载。最终的文件夹结构应如下所示：

TokenHMR/
├── tokenhmr/
│   └── dataset_dir/
│       └── training_data/                      # 训练数据
│           └── dataset_tars/
│               └── coco-train-2014-pruned/
│               └── aic-train-vitpose/
│               └── bedlam/
|               ...                          
│           ...
│       └── evaluation_data/                    # 评估数据
│           └── 3DPW/
│           └── EMDB/
│           └── emdb.npz
│           └── 3dpw_test.npz
└── ...

训练

训练完tokenizer后，我们可以开始训练TokenHMR。如果你想跳过tokenization训练，可以直接使用检查点中提供的预训练模型。使用4DHumans预训练的backbone（从官方仓库下载模型），在4个A100 Nvidia GPU上训练大约需要4天时间。如果你想更改任何默认设置，请更新tokenhmr/lib/configs_hydra/experiment/tokenhmr_release.yaml。

python tokenhmr/train.py datasets=mix_all experiment=tokenhmr_release

评估

要评估论文中使用的原始模型在3DPW和EMDB上的表现，请从这里下载。然后运行以下命令：

python tokenhmr/eval.py  \
    --dataset EMDB,3DPW-TEST \
    --batch_size 32 --log_freq 50 \
    --dataset_dir tokenhmr/dataset_dir/evaluation_data \
    --checkpoint data/checkpoints/tokenhmr_model.ckpt \
    --model_config data/checkpoints/model_config.yaml

致谢

本代码基于以下两个优秀的仓库构建。我感谢作者们开源他们的代码。

部分代码来自或改编自以下仓库：

我们诚挚感谢Perceiving Systems部门和Meshcapade GmbH的ML团队提供的有见地的讨论和反馈。我们感谢国际马克斯普朗克智能系统研究院（IMPRS-IS）对Sai Kumar Dwivedi的支持。我们感谢Meshcapade GmbH对Yu Sun的支持并提供GPU资源。这项工作部分得到了德国联邦教育和研究部（BMBF）的支持：图宾根AI中心，FKZ：01IS18039B。

引用

如果你发现这份代码对你的研究有用，请考虑引用以下论文：

@inproceedings{dwivedi_cvpr2024_tokenhmr,
    title={{TokenHMR}: Advancing Human Mesh Recovery with a Tokenized Pose Representation},
    author={Dwivedi, Sai Kumar and Sun, Yu and Patel, Priyanka and Feng, Yao and Black, Michael J.},
    booktitle={IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR)},
    year={2024},
}