人体姿态估计与跟踪的简单基线
介绍
这个仓库是MSRA在**ECCV 2018上发表的简单基线用于单个RGB图像中二维多人人体姿态估计的TensorFlow**实现。
这个仓库提供:
- TensorFlow实现的人体姿态估计与跟踪的简单基线。
- 灵活且简单的代码。
- 兼容大多数公开可用的二维多人人体姿态估计数据集,包括**MPII、PoseTrack 2018、MS COCO 2017**。
- 人体姿态估计可视化代码(从Detectron修改)。
依赖项
此代码在Ubuntu 16.04、CUDA 9.0、cuDNN 7.1环境下用两块NVIDIA 1080Ti GPU进行测试。
开发中使用Python 3.6.5版本和Anaconda 3。
目录结构
根目录
${POSE_ROOT}
描述如下。
${POSE_ROOT}
|-- data
|-- lib
|-- main
|-- tool
`-- output
data
包含数据加载代码和指向图像和注释目录的软链接。lib
包含二维多人人体姿态估计系统的内核代码。main
包含用于训练或测试网络的高级代码。tool
包含数据集转换器。我设定了MS COCO为参考格式并提供了mpii2coco和posetrack2coco转换代码。output
包含日志、训练模型、可视化输出和测试结果。
数据
你需要遵循如下data
目录结构。
${POSE_ROOT}
|-- data
|-- |-- MPII
| `-- |-- dets
| | |-- human_detection.json
| |-- annotations
| | |-- train.json
| | `-- test.json
| `-- images
| |-- 000001163.jpg
| |-- 000003072.jpg
|-- |-- PoseTrack
| `-- |-- dets
| | |-- human_detection.json
| |-- annotations
| | |-- train2018.json
| | |-- val2018.json
| | `-- test2018.json
| |-- original_annotations
| | |-- train/
| | |-- val/
| | `-- test/
| `-- images
| |-- train/
| |-- val/
| `-- test/
|-- |-- COCO
| `-- |-- dets
| | |-- human_detection.json
| |-- annotations
| | |-- person_keypoints_train2017.json
| | |-- person_keypoints_val2017.json
| | `-- image_info_test-dev2017.json
| `-- images
| |-- train2017/
| |-- val2017/
| `-- test2017/
`-- |-- imagenet_weights
| |-- resnet_v1_50.ckpt
| |-- resnet_v1_101.ckpt
| `-- resnet_v1_152.ckpt
- 在
tool
中,运行python mpii2coco.py
将MPII注释文件转换为MS COCO格式(MPII/annotations
)。 - 在
tool
中,运行python posetrack2coco.py
将PoseTrack注释文件转换为MS COCO格式(PoseTrack/annotations
)。 - 在训练阶段,使用GT人类bbox,在测试阶段使用
human_detection.json
,该文件应在测试前准备好并遵循MS COCO格式。 - 从tf-slim下载imagenet预训练resnet模型并将其放在
data/imagenet_weights
中。 - 除了MPII和PoseTrack的
annotations
外,所有其他目录都是下载的原始版本。 - 如果你想添加自己的数据集,你必须将其转换为MS COCO格式。
- 你可以通过修改每个数据集文件夹的
dataset.py
来更改data
的默认目录结构。
输出
你需要遵循如下output
文件夹的目录结构。
${POSE_ROOT}
|-- output
|-- |-- log
|-- |-- model_dump
|-- |-- result
`-- |-- vis
- 建议以软链接形式创建
output
文件夹,而不是文件夹形式,因为它会占用大量存储空间。 log
文件夹包含训练日志文件。model_dump
文件夹包含每个时期保存的检查点。result
文件夹包含测试阶段生成的最终估计文件。vis
文件夹包含可视化结果。- 你可以通过修改
main/config.py
来更改output
的默认目录结构。
运行TF-SimpleHumanPose
开始
- 运行
pip install -r requirement.txt
安装所需模块。 - 运行
cd ${POSE_ROOT}/lib
并运行make
构建NMS模块。 - 在
main/config.py
中,你可以更改模型的设置,包括使用的数据集、网络骨干、输入大小等。
训练
在main
文件夹中,运行
python train.py --gpu 0-1
在GPU 0,1上训练网络。
如果你想继续实验,运行
python train.py --gpu 0-1 --continue
可以使用--gpu 0,1
代替--gpu 0-1
。
测试
将训练好的模型放在output/model_dump/$DATASET/
,将人类检测结果(human_detection.json
)放在data/$DATASET/dets/
。
在main
文件夹中,运行
python test.py --gpu 0-1 --test_epoch 140
以用140时期训练的模型在GPU 0,1上测试网络。可以使用--gpu 0,1
代替--gpu 0-1
。
结果
在这里,我报告了这个仓库与原始论文中的模型性能。此外,我还提供了预训练模型和人类检测结果。
由于此仓库输出与MS COCO和PoseTrack兼容的输出文件,你可以直接使用cocoapi或poseval评估MS COCO或PoseTrack数据集的结果。你必须将生成的mat
文件转换为MPII mat
格式,以评估MPII数据集的结果,按照此方法。
MSCOCO 2017数据集上的结果
对于所有方法,使用相同的人类检测结果(下载链接提供如下)。为了比较,我使用了原始仓库的预训练模型,以下表格是COCO val2017集上的AP。
方法 | AP | AP .5 | AP .75 | AP (M) | AP (L) | AR | AR .5 | AR .75 | AR (M) | AR (L) | 下载 |
---|---|---|---|---|---|---|---|---|---|---|---|
256x192_resnet50 (本仓库) | 70.4 | 88.6 | 77.8 | 67.0 | 76.9 | 76.2 | 93.0 | 83.0 | 71.9 | 82.4 | 模型 姿态 |
256x192_resnet50 (原始仓库) | 70.3 | 88.8 | 77.8 | 67.0 | 76.7 | 76.1 | 93.0 | 82.9 | 71.8 | 82.3 | - |
- val2017上人类检测结果(55.3 AP在人类类上)和test-dev2017(57.2 AP在人类类上)[bbox]
- val2017上的其他人类检测结果[Detectron_MODEL_ZOO]
PoseTrack 2018数据集上的结果
按照论文,在PoseTrack数据集上进行训练时,使用COCO数据集上的预训练模型。在COCO数据集上训练模型之后,我在config.py
中设置lr
、lr_dec_epoch
、end_epoch
为5e-5
、[150, 155]
、160
,然后运行python train.py --gpu $GPUS --continue
。下表是验证集上的AP。
方法 | Head | Shoulder | Elbow | Wrist | Hip | Knee | Ankle | Total | 下载 |
---|---|---|---|---|---|---|---|---|---|
256x192_resnet50 (检测器的bbox) | 74.4 | 76.9 | 72.2 | 65.2 | 69.2 | 70.0 | 62.9 | 70.4 | 模型 姿态 |
256x192_resnet50 (GT的bbox) | 87.9 | 86.7 | 80.2 | 72.5 | 77.0 | 77.8 | 74.6 | 80.1 | 模型 姿态 |
- 验证集上的人类检测结果[bbox]
故障排除
-
当你的机器在训练过程中占用越来越多的内存时,请添加graph.finalize。[问题]
-
对于那些在测试阶段遇到
FileNotFoundError: [Errno 2] No such file or directory: 'tmp_result_0.pkl'
的人,请正确准备人类检测结果。pkl文件会在测试阶段自动生成和删除,所以你不必准备它们。大多数这种错误来自不正确的人类检测文件。
致谢
这个仓库大部分修改自TensorFlow CPN仓库和PyTorch简单基线仓库。
参考文献
[1] Xiao, Bin, Haiping Wu, and Yichen Wei. "Simple Baselines for Human Pose Estimation and Tracking". ECCV 2018.