CAFE-GAN-PyTorch

This is An unofficial PyTorch implementation of CAFE-GAN - Arbitrary Face Attribute Editing with Complementary Attention Feature

Inverting 13 attributes respectively. From left to right: Input, Reconstruction, Bald, Bangs, Black_Hair, Blond_Hair, Brown_Hair, Bushy_Eyebrows, Eyeglasses, Male, Mouth_Slightly_Open, Mustache, No_Beard, Pale_Skin, Young

To-Do

• Upload experimental results
  • FID
  • Single Attribute Editing Accuracy
• Upload Pre-training model

Requirements

• Python 3.7
• PyTorch 1.7.0
• TensorboardX

pip3 install -r requirements.txt

• Dataset
  • CelebA dataset
    • Images should be placed in ./data/img_align_celeba/*.jpg
    • Attribute labels should be placed in ./data/list_attr_celeba.txt
  • HD-CelebA (optional)
    • Please see here.
  • CelebA-HQ dataset (optional)
    • Please see here.
    • Images should be placed in ./data/celeba-hq/celeba-*/*.jpg
    • Image list should be placed in ./data/image_list.txt
• Pretrained models: we will upload pretrained models to Google Cloud.

Usage

To train an CAFE-GAN on CelebA 128x128

CUDA_VISIBLE_DEVICES=0 \
python train.py \
--data CelebA \
--use_stu \
--shortcut_layers 3 \
--data_path data/img_align_celeba \
--attr_path data/list_attr_celeba.txt \
--img_size 128 \
--batch_size 32 \
--kernel_size 3 \
--num_workers 4 \
--save_interval 5000 \
--sample_interval 1500 \
--data_save_root output \
--experiment_name 128_ShortCut3_KernelSize3

To train an CAFE-GAN on CelebA-HQ 256x256

CUDA_VISIBLE_DEVICES=0 \
python train.py \
--data CelebA-HQ \
--n_samples 8 \
--gpu \
--use_stu \
--img_size 256 \
--batch_size 8 \
--epochs 100 \
--shortcut_layers 3 \
--kernel_size 3 \
--one_more_conv \
--data_path data/celeba-hq/celeba-256 \
--attr_path data/list_attr_celeba.txt \
--image_list_path data/image_list.txt \
--data_save_root output \
--experiment_name 256_ShortCut3_KernelSize3

To visualize training details

tensorboard \
--logdir path/to/experiment/summary

To test with single attribute editing

CUDA_VISIBLE_DEVICES=0 \
python test.py \
--use_model 'G' \
--experiment_name your_experiment \
--data_path data/img_align_celeba \
--attr_path data/list_attr_celeba.txt \
--data_save_root output \
--test_int 0.75 \
--load_epoch epoch \
--gpu

To test with attention

CUDA_VISIBLE_DEVICES=0 \
python test.py \
--use_model 'D' \
--data_path data/img_align_celeba \
--attr_path data/list_attr_celeba.txt \
--experiment_name your_experiment \
--data_save_root output \
--gpu

Your custom images are supposed to be in path/to/Image_data and you also need an attribute list of the images(.txt) path/to/label_Data. Please crop and resize them into square images in advance.

Acknowledgement

The code is built upon AttGAN, STGAN and ABN, thanks for their excellent works!