AutoLink: Self-supervised Learning of Human Skeletons and Object Outlines by Linking Keypoints (NeurIPS 2022 Spotlight)

AutoLink: Self-supervised Learning of Human Skeletons and Object Outlines by Linking Keypoints
Xingzhe He, Bastian Wandt, and Helge Rhodin
Thirty-sixth Conference on Neural Information Processing Systems (NeurIPS 2022 Spotlight)

[Paper][Website][🤗 Hugging Face Demo ]

Setup

Setup environment

conda create -n autolink python=3.8
conda activate autolink
pip install -r requirements.txt

Download datasets

The CelebA-in-the-wild, Taichi, Human3.6m, DeepFashion, CUB, 11k Hands, AFHQ, Horse2Zebra and Flower can be found on their websites. We provide the pre-processing code for CelebA-in-the-wild, CUB and Flower to make them h5 files. Others can be used directly.

Download pre-trained models

The pre-trained models can be downloaded from Google Drive.

Testing

To qualitatively test the model, you can run our demo by

python app.py --log celeba_wild/celeba_wild_k8_m0.8_b16_t0.0025_sklr512

where,

• --log specifies the checkpoint folder under `checkpoints/

The default is our model on face:

You can also generate multiple images at the same time. Run

python gen_detection.py --log celeba_wild/celeba_wild_k8_m0.8_b16_t0.0025_sklr512 --folder_name celeba_wild_k8_detection --data_root data/celeba_wild

where,

• --data_root specifies the location of the dataset,
• --folder_name specifies the folder where you want to save the detection images.

To numerically test the model performance, run

python test.py --log celeba_wild/celeba_wild_k8_m0.8_b16_t0.0025_sklr512 --data_root data/celeba_wild

Therefore, the above command will give the performance metric on CelebA-in-the-wild, which is described in the paper.

Training

Note: We notice that, on h36m w/o background, training on A100 and A6000 are not as stable as training on V100, and there might be overfitting. We suggest to stop early or use larger masking ratio if readers want to train on h36m w/o background on A100 or A6000. We acknowledge Yuchen Yang for valuable discussion and experiments.

To train our model on CelebA-in-the-wild, run

python train.py --n_parts 8 --missing 0.8 --block 16 --thick 2.5e-3 --sklr 512 --data_root data/celeba_wild --dataset celeba_wild

where,

• --n_parts specifies the number of keypoints,
• --missing specifies the ratio of the image masking,
• --block specifies number of patches to divide the image in one dimension,
• --thick specifies thickness of the edges,
• --sklr specifies the learning rate of the edge weights,
• --data_root specifies the location of the dataset,
• --dataset specifies name of the dataset.

The trained model can be found in checkpoints/celeba_wild_k8_m0.8_b16_t0.0025_sklr512.

Citation

@inproceedings{he2022autolink,
    title={AutoLink: Self-supervised Learning of Human Skeletons and Object Outlines by Linking Keypoints},
    author={He, Xingzhe and Wandt, Bastian and Rhodin, Helge},
    booktitle={Advances in Neural Information Processing Systems},
    year={2022}
}