显式框检测统一端到端多人姿态估计

这是我们ICLR 2023论文"显式框检测统一端到端多人姿态估计"的官方PyTorch实现。

⭐ ED-Pose

我们提出了ED-Pose，一个用于多人姿态估计的端到端框架，具有显式框检测功能。ED-Pose将这个任务重新考虑为两个具有统一表示和回归监督的显式框检测过程。总的来说，ED-Pose在概念上简单，无需后处理和密集热图监督。

ED-Pose作为一个完全端到端的框架，首次使用L1回归损失在相同骨干网络下超越了基于热图的自上而下方法，在COCO数据集上提高了1.2 AP。
ED-Pose在CrowdPose数据集上达到了76.6 AP的最先进性能，且无需测试时增强。

🔥 新闻

2023/08/08: 1. 我们在Human-Art数据集上支持ED-Pose。2. 我们上传了更快速可视化的推理脚本。

🐟 待办事项

该仓库包含以下进一步修改：

集成到detrex中。
使用Gradio集成到Huggingface Spaces 🤗。

🚀 模型库

我们已将模型检查点放在这里。

COCO val2017数据集上的结果

模型	骨干网络	学习率调度	mAP	AP<sup>50</sup>	AP<sup>75</sup>	AP<sup>M</sup>	AP<sup>L</sup>	时间 (ms)	下载
ED-Pose	R-50	60e	71.7	89.7	78.8	66.2	79.7	51	Google Drive
ED-Pose	Swin-L	60e	74.3	91.5	81.7	68.5	82.7	88	Google Drive
ED-Pose	Swin-L-5scale	60e	75.8	92.3	82.9	70.4	83.5	142	Google Drive

CrowdPose测试数据集上的结果

模型	骨干网络	学习率调度	mAP	AP<sup>50</sup>	AP<sup>75</sup>	AP<sup>E</sup>	AP<sup>M</sup>	AP<sup>H</sup>	下载
ED-Pose	R-50	80e	69.9	88.6	75.8	77.7	70.6	60.9	Google Drive
ED-Pose	Swin-L	80e	73.1	90.5	79.8	80.5	73.8	63.8	Google Drive
ED-Pose	Swin-L-5scale	80e	76.6	92.4	83.3	83.0	77.3	68.3	Google Drive

COCO测试开发数据集上的结果

模型	骨干网络	损失	mAP	AP<sup>50</sup>	AP<sup>75</sup>	AP<sup>M</sup>	AP<sup>L</sup>
DirectPose	R-50	回归	62.2	86.4	68.2	56.7	69.8
DirectPose	R-101	回归	63.3	86.7	69.4	57.8	71.2
FCPose	R-50	回归+热图	64.3	87.3	71.0	61.6	70.5
FCPose	R-101	回归+热图	65.6	87.9	72.6	62.1	72.3
InsPose	R-50	回归+热图	65.4	88.9	71.7	60.2	72.7
InsPose	R-101	回归+热图	66.3	89.2	73.0	61.2	73.9
PETR	R-50	回归+热图	67.6	89.8	75.3	61.6	76.0
PETR	Swin-L	回归+热图	70.5	91.5	78.7	65.2	78.0
ED-Pose	R-50	回归	69.8	90.2	77.2	64.3	77.4
ED-Pose	Swin-L	回归	72.7	92.3	80.9	67.6	80.0

COCO测试开发数据集上的结果

使用Human-Art和COCO数据集联合训练的结果

🥂 注意：在ED-Pose上使用Human-Art训练可以提升MSCOCO上的性能！

Human-Art验证集上的结果

架构	主干网络	mAP	AP<sup>50</sup>	AP<sup>75</sup>	AR	AR<sup>50</sup>	下载
ED-Pose	ResNet-50	0.723	0.861	0.774	0.808	0.921	Google Drive

COCO val2017 上的结果

架构	主干网络	AP	AP<sup>50</sup>	AP<sup>75</sup>	AR	AR<sup>50</sup>	下载
ED-Pose	ResNet-50	0.724	0.898	0.794	0.799	0.946	Google Drive

注意：

ED-Pose 没有使用任何测试时数据增强。
我们在 Swin-L-5scale 设置下使用 Object365 数据集来预训练 ED-Pose 的人体检测。

🚢 环境配置

我们使用 DN-Deformable-DETR 作为我们的代码基础。我们在 python=3.7.3,pytorch=1.9.0,cuda=11.1 环境下测试我们的模型。其他版本可能也适用。

克隆此仓库

git clone https://github.com/IDEA-Research/ED-Pose.git
cd ED-Pose

安装 Pytorch 和 torchvision

按照 https://pytorch.org/get-started/locally/ 的说明进行操作。

# 示例:
conda install -c pytorch pytorch torchvision

安装其他所需的包

pip install -r requirements.txt

编译 CUDA 算子

cd models/edpose/ops
python setup.py build install
# 单元测试（应该看到所有检查都为 True）
python test.py
cd ../../..

</details> <details> <summary>数据准备</summary>

对于 COCO 数据，请从 COCO 下载下载。 coco_dir 应该如下所示：

|-- EDPose
`-- |-- coco_dir
    `-- |-- annotations
        |   |-- person_keypoints_train2017.json
        |   `-- person_keypoints_val2017.json
        `-- images
            |-- train2017
            |   |-- 000000000009.jpg
            |   |-- 000000000025.jpg
            |   |-- 000000000030.jpg
            |   |-- ... 
            `-- val2017
                |-- 000000000139.jpg
                |-- 000000000285.jpg
                |-- 000000000632.jpg
                |-- ...

对于 CrowdPose 数据，请从 CrowdPose 下载下载， crowdpose_dir 应该如下所示：

|-- ED-Pose
`-- |-- crowdpose_dir
    `-- |-- json
        |   |-- crowdpose_train.json
        |   |-- crowdpose_val.json
        |   |-- crowdpose_trainval.json (由 util/crowdpose_concat_train_val.py 生成)
        |   `-- crowdpose_test.json
        `-- images
            |-- 100000.jpg
            |-- 100001.jpg
            |-- 100002.jpg
            |-- 100003.jpg
            |-- 100004.jpg
            |-- 100005.jpg
            |-- ...

</details>

🥳 运行

在 COCO 上训练：

#对于 ResNet-50:
export EDPOSE_COCO_PATH=/path/to/your/cocodir
  python main.py \
 --output_dir "logs/coco_r50" \
 -c config/edpose.cfg.py \
 --options batch_size=4 epochs=60 lr_drop=55 num_body_points=17 backbone='resnet50' \
 --dataset_file="coco"

#对于 Swin-L:
export EDPOSE_COCO_PATH=/path/to/your/cocodir
export pretrain_model_path=/path/to/your/swin_L_384_22k
  python main.py \
 --output_dir "logs/coco_swinl" \
 -c config/edpose.cfg.py \
 --options batch_size=4 epochs=60 lr_drop=55 num_body_points=17 backbone='swin_L_384_22k' \
 --dataset_file="coco"

</details> <details> <summary>分布式运行</summary>

#对于 ResNet-50:
export EDPOSE_COCO_PATH=/path/to/your/cocodir
  python -m torch.distributed.launch --nproc_per_node=4  main.py \
 --output_dir "logs/coco_r50" \
 -c config/edpose.cfg.py \
 --options batch_size=4 epochs=60 lr_drop=55 num_body_points=17 backbone='resnet50' \
 --dataset_file="coco"

#对于 Swin-L:
export EDPOSE_COCO_PATH=/path/to/your/cocodir
export pretrain_model_path=/path/to/your/swin_L_384_22k
  python -m torch.distributed.launch --nproc_per_node=4 main.py \
 --output_dir "logs/coco_swinl" \
 -c config/edpose.cfg.py \
 --options batch_size=4 epochs=60 lr_drop=55 num_body_points=17 backbone='swin_L_384_22k' \
 --dataset_file="coco"

</details>

在 CrowdPose 上训练：

#对于 ResNet-50:
export EDPOSE_CrowdPose_PATH=/path/to/your/crowdpose_dir
  python main.py \
 --output_dir "logs/crowdpose_r50" \
 -c config/edpose.cfg.py \
 --options batch_size=4 epochs=80 lr_drop=75 num_body_points=14 backbone='resnet50' \
 --dataset_file="crowdpose"

#对于 Swin-L:
export EDPOSE_CrowdPose_PATH=/path/to/your/crowdpose_dir
export pretrain_model_path=/path/to/your/swin_L_384_22k
  python main.py \
 --output_dir "logs/crowdpose_swinl" \
 -c config/edpose.cfg.py \
 --options batch_size=4 epochs=80 lr_drop=75 num_body_points=14 backbone='swin_L_384_22k' \
 --dataset_file="crowdpose"

</details> <details> <summary>分布式运行</summary>

#对于 ResNet-50:
export EDPOSE_CrowdPose_PATH=/path/to/your/crowdpose_dir
  python -m torch.distributed.launch --nproc_per_node=4  main.py \
 --output_dir "logs/crowdpose_r50" \
 -c config/edpose.cfg.py \
 --options batch_size=4 epochs=80 lr_drop=75 num_body_points=14 backbone='resnet50' \
 --dataset_file="crowdpose"

#对于 Swin-L:
export EDPOSE_CrowdPose_PATH=/path/to/your/crowdpose_dir
export pretrain_model_path=/path/to/your/swin_L_384_22k
  python -m torch.distributed.launch --nproc_per_node=4 main.py \
 --output_dir "logs/crowdpose_swinl" \
 -c config/edpose.cfg.py \
 --options batch_size=4 epochs=80 lr_drop=75 num_body_points=14 backbone='swin_L_384_22k' \
 --dataset_file="crowdpose"

</details>

我们已将在 ImageNet-22k 上预训练的 Swin-L 模型放在这里。

在 COCO 上评估：

<details> <summary>ResNet-50</summary>

export EDPOSE_COCO_PATH=/path/to/your/cocodir
  python -m torch.distributed.launch --nproc_per_node=4  main.py \
 --output_dir "logs/coco_r50" \
 -c config/edpose.cfg.py \
 --options batch_size=4 epochs=60 lr_drop=55 num_body_points=17 backbone='resnet50' \
 --dataset_file="coco" \
 --pretrain_model_path "./models/edpose_r50_coco.pth" \
 --eval

</details> <details> <summary>Swin-L</summary>

export EDPOSE_COCO_PATH=/path/to/your/cocodir
export pretrain_model_path=/path/to/your/swin_L_384_22k
  python -m torch.distributed.launch --nproc_per_node=4 main.py \
 --output_dir "logs/coco_swinl" \
 -c config/edpose.cfg.py \
 --options batch_size=4 epochs=60 lr_drop=55 num_body_points=17 backbone='swin_L_384_22k' \
 --dataset_file="coco" \
 --pretrain_model_path "./models/edpose_swinl_coco.pth" \
 --eval

</details> <details> <summary>Swin-L-5scale</summary> ``` export EDPOSE_COCO_PATH=/path/to/your/cocodir export pretrain_model_path=/path/to/your/swin_L_384_22k python -m torch.distributed.launch --nproc_per_node=4 main.py \ --output_dir "logs/coco_swinl" \ -c config/edpose.cfg.py \ --options batch_size=4 epochs=60 lr_drop=55 num_body_points=17 backbone='swin_L_384_22k' \ return_interm_indices=0,1,2,3 num_feature_levels=5 \ --dataset_file="coco" \ --pretrain_model_path "./models/edpose_swinl_5scale_coco.pth" \ --eval ``` </详情>

在CrowdPose上评估:

<详情> <概要>ResNet-50</概要>

export EDPOSE_CrowdPose_PATH=/path/to/your/crowdpose_dir
  python main.py \
 --output_dir "logs/crowdpose_r50" \
 -c config/edpose.cfg.py \
 --options batch_size=4 epochs=80 lr_drop=75 num_body_points=14 backbone='resnet50' \
 --dataset_file="crowdpose"\
 --pretrain_model_path "./models/edpose_r50_crowdpose.pth" \
 --eval

</详情>

<详情> <概要>Swin-L</概要>

export EDPOSE_CrowdPose_PATH=/path/to/your/crowdpose_dir
export pretrain_model_path=/path/to/your/swin_L_384_22k
  python main.py \
 --output_dir "logs/crowdpose_swinl" \
 -c config/edpose.cfg.py \
 --options batch_size=4 epochs=80 lr_drop=75 num_body_points=14 backbone='swin_L_384_22k' \
 --dataset_file="crowdpose" \
 --pretrain_model_path "./models/edpose_swinl_crowdpose.pth" \
 --eval

</详情>

<详情> <概要>Swin-L-5scale</概要>

export EDPOSE_CrowdPose_PATH=/path/to/your/crowdpose_dir
export pretrain_model_path=/path/to/your/swin_L_384_22k
  python -m torch.distributed.launch --nproc_per_node=4 main.py \
 --output_dir "logs/crowdpose_swinl" \
 -c config/edpose.cfg.py \
 --options batch_size=4 epochs=80 lr_drop=75 num_body_points=14 backbone='swin_L_384_22k' \
 return_interm_indices=0,1,2,3 num_feature_levels=5 \
 -- dataset_file="crowdpose" \
 --pretrain_model_path "./models/edpose_swinl_5scale_crowdpose.pth" \
 --eval

</详情>

通过COCO关键点格式进行可视化:

<详情> <概要>ResNet-50</概要>

export EDPOSE_COCO_PATH=/path/to/your/cocodir
export Inference_Path=/path/to/your/inference_dir
  python -m torch.distributed.launch --nproc_per_node=1  main.py \
 --output_dir "logs/coco_r50" \
 -c config/edpose.cfg.py \
 --options batch_size=1 epochs=60 lr_drop=55 num_body_points=17 backbone='resnet50' \
 --dataset_file="coco" \
 --pretrain_model_path "./models/edpose_r50_coco.pth" \
 --eval

</详情>

<详情> <概要>Swin-L</概要>

export EDPOSE_COCO_PATH=/path/to/your/cocodir
export Inference_Path=/path/to/your/inference_dir
  python -m torch.distributed.launch --nproc_per_node=1 main.py \
 --output_dir "logs/coco_swinl" \
 -c config/edpose.cfg.py \
 --options batch_size=1 epochs=60 lr_drop=55 num_body_points=17 backbone='swin_L_384_22k' \
 --dataset_file="coco" \
 --pretrain_model_path "./models/edpose_swinl_coco.pth" \
 --eval

</详情>

<详情> <概要>Swin-L-5scale</概要>

export EDPOSE_COCO_PATH=/path/to/your/cocodir
export Inference_Path=/path/to/your/inference_dir
  python -m torch.distributed.launch --nproc_per_node=1 main.py \
 --output_dir "logs/coco_swinl" \
 -c config/edpose.cfg.py \
 --options batch_size=1 epochs=60 lr_drop=55 num_body_points=17 backbone='swin_L_384_22k' \
  return_interm_indices=0,1,2,3 num_feature_levels=5 \
 --dataset_file="coco" \
 --pretrain_model_path "./models/edpose_swinl_5scale_coco.pth" \
 --eval

</详情>

💃🏻 引用ED-Pose

@inproceedings{
yang2023explicit,
title={Explicit Box Detection Unifies End-to-End Multi-Person Pose Estimation},
author={Jie Yang and Ailing Zeng and Shilong Liu and Feng Li and Ruimao Zhang and Lei Zhang},
booktitle={International Conference on Learning Representations},
year={2023},
url={https://openreview.net/forum?id=s4WVupnJjmX}
}