Effective Whole-body Pose Estimation with Two-stages Distillation

Zhendong Yang, Ailing Zeng, Chun Yuan, Yu Li

          DWPose                      DWPose + ControlNet (prompt: Ironman)

💃🏻 DWPose 💃🏻

This repository is the official implementation of the Effective Whole-body Pose Estimation with Two-stages Distillation (ICCV 2023, CV4Metaverse Workshop). Our code is based on MMPose and ControlNet.

⚔️ We release a series of models named DWPose with different sizes, from tiny to large, for human whole-body pose estimation. Besides, we also replace Openpose with DWPose for ControlNet, obtaining better Generated Images.

🔥 News

• 2023/08/17: Our paper Effective Whole-body Pose Estimation with Two-stages Distillation is accepted by ICCV 2023, CV4Metaverse Workshop. 🎉 🎉 🎉

• 2023/08/09: You can try DWPose with sd-webui-controlnet now! Just update your sd-webui-controlnet >= v1.1237, then choose dw_openpose_full as preprocessor.

• 2023/08/09: We support to run onnx model with cv2. You can avoid installing onnxruntime. See branch opencv_onnx.

• 2023/08/07: We upload all DWPose models to huggingface. Now, you can download them from baidu drive, google drive and huggingface.
• 2023/08/07: We release a new DWPose with onnx. You can avoid installing mmcv through this. See branch onnx.
• 2023/08/01: Thanks to MMPose. You can try our DWPose with this demo by choosing wholebody!

🐟 Installation

See installation instructions.

🚀 Results and Models

😎 DWPose on COCO. We release a series of DWPose models.

Results on COCO-WholeBody v1.0 val with detector having human AP of 56.4 on COCO val2017 dataset

🦈 DWPose for ControlNet.

First, you need to download our SOTA model DWPose-l_384x288 and put it into ControlNet-v1-1-nightly/annotator/ckpts. Then you can use DWPose to generate the images you like.

cd ControlNet-v1-1-nightly
python gradio_dw_open_pose.py

Non-cherry-picked test with random seed 12345 ("spider man"):

Comparison with OpenPose

🚢 Datasets

Prepare COCO in mmpose/data/coco and UBody in mmpose/data/UBody.

UBody needs to be tarnsferred into images. Don't forget.

cd mmpose
python video2image.py

If you want to evaluate the models on UBody

# add category into UBody's annotation
cd mmpose
python add_cat.py

⭐Train a model

Train DWPose with the first stage distillation

cd mmpose
bash tools/dist_train.sh configs/distiller/ubody/s1_dis/rtmpose_x_dis_l__coco-ubody-256x192.py 8

Train DWPose with the second stage distillation

cd mmpose
bash tools/dist_train.sh configs/distiller/ubody/s2_dis/dwpose_l-ll__coco-ubody-256x192.py 8

Tansfer the distillation models into regular models

cd mmpose
# if first stage distillation
python pth_transfer.py $dis_ckpt $new_pose_ckpt
# if second stage distillation
python pth_transfer.py $dis_ckpt $new_pose_ckpt --two_dis

⭐Test a model

# test on UBody
bash tools/dist_test.sh configs/wholebody_2d_keypoint/rtmpose/ubody/rtmpose-l_8xb64-270e_ubody-wholebody-256x192.py $pose_ckpt 8

# test on COCO
bash tools/dist_test.sh configs/wholebody_2d_keypoint/rtmpose/ubody/rtmpose-l_8xb64-270e_coco-ubody-wholebody-256x192.py $pose_ckpt 8

🥳 Citation

@article{yang2023effective,
  title={Effective Whole-body Pose Estimation with Two-stages Distillation},
  author={Yang, Zhendong and Zeng, Ailing and Yuan, Chun and Li, Yu},
  journal={arXiv preprint arXiv:2307.15880},
  year={2023}
}

🥂 Acknowledgement

Our code is based on MMPose and ControlNet.