ZZUTK / Face-Aging-CAAE

Age Progression/Regression by Conditional Adversarial Autoencoder
636 stars 233 forks source link

Age Progression/Regression by Conditional Adversarial Autoencoder (CAAE)

TensorFlow implementation of the algorithm in the paper Age Progression/Regression by Conditional Adversarial Autoencoder.

Thanks to the Pytorch implementation by Mattan Serry, Hila Balahsan, and Dor Alt.

Pre-requisites

Datasets

Prepare the training dataset

You may use any dataset with labels of age and gender. In this demo, we use the UTKFace dataset. It is better to use aligned and cropped faces. Please save and unzip UTKFace.tar.gz to the folder data.

Training

$ python main.py

The training process has been tested on NVIDIA TITAN X (12GB). The training time for 50 epochs on UTKFace (23,708 images in the size of 128x128x3) is about two and a half hours.

During training, a new folder named save will be created, including four sub-folders: summary, samples, test, and checkpoint.

After training, you can check the folders samples and test to visualize the reconstruction and testing performance, respectively. The following shows the reconstruction (left) and testing (right) results. The first row in the reconstruction results (left) are testing samples that yield the testing results (right) in the age ascending order from top to bottom.

The reconstruction loss vs. epoch is shown below, which was passed through a low-pass filter for visualization purpose. The original record is saved in folder summary.

Custom Training

$ python main.py
    --dataset       default 'UTKFace'. Please put your own dataset in ./data
    --savedir       default 'save'. Please use a meaningful name, e.g., save_init_model.
    --epoch     default 50.
    --use_trained_model default True. If use a trained model, savedir specifies the model name. 
    --use_init_model    default True. If load the trained model failed, use the init model save in ./init_model 

Testing

$ python main.py --is_train False --testdir your_image_dir --savedir save

Note: savedir specifies the model name saved in the training. By default, the trained model is saved in the folder save (i.e., the model name). Then, it is supposed to print out the following message.

    Building graph ...

    Testing Mode

    Loading pre-trained model ...
    SUCCESS ^_^

    Done! Results are saved as save/test/test_as_xxx.png

Specifically, the testing faces will be processed twice, being considered as male and female, respectively. Therefore, the saved files are named test_as_male.png and test_as_female.png, respectively. To achieve better results, it is necessary to train on a large and diverse dataset.

A demo of training process

The first row shows the input faces of different ages, and the other rows show the improvement of the output faces at every other epoch. From top to bottom, the output faces are in the age ascending order.

Files

Citation

Zhifei Zhang, Yang Song, and Hairong Qi. "Age Progression/Regression by Conditional Adversarial Autoencoder." IEEE Conference on Computer Vision and Pattern Recognition (CVPR), 2017.

@inproceedings{zhang2017age,
  title={Age Progression/Regression by Conditional Adversarial Autoencoder},
  author={Zhang, Zhifei and Song, Yang and Qi, Hairong},
  booktitle={IEEE Conference on Computer Vision and Pattern Recognition (CVPR)},
  year={2017}
}

Spotlight presentation