Only a Matter of Style: Age Transformation Using a Style-Based Regression Model (SIGGRAPH 2021)

The task of age transformation illustrates the change of an individual's appearance over time. Accurately modeling this complex transformation over an input facial image is extremely challenging as it requires making convincing and possibly large changes to facial features and head shape, while still preserving the input identity. In this work, we present an image-to-image translation method that learns to directly encode real facial images into the latent space of a pre-trained unconditional GAN (e.g., StyleGAN) subject to a given aging shift. We employ a pre-trained age regression network used to explicitly guide the encoder to generate the latent codes corresponding to the desired age. In this formulation, our method approaches the continuous aging process as a regression task between the input age and desired target age, providing fine-grained control on the generated image. Moreover, unlike other approaches that operate solely in the latent space using a prior on the path controlling age, our method learns a more disentangled, non-linear path. We demonstrate that the end-to-end nature of our approach, coupled with the rich semantic latent space of StyleGAN, allows for further editing of the generated images. Qualitative and quantitative evaluations show the advantages of our method compared to state-of-the-art approaches.

<a href="https://arxiv.org/abs/2102.02754"><img src="https://img.shields.io/badge/arXiv-2008.00951-b31b1b.svg" height=22.5></a> <a href="https://opensource.org/licenses/MIT"><img src="https://img.shields.io/badge/License-MIT-yellow.svg" height=22.5></a>

<a href="https://www.youtube.com/watch?v=zDTUbtmUbG8"><img src="https://img.shields.io/static/v1?label=Two Minute Papers&message=SAM Video&color=red" height=22.5></a>
<a href="https://youtu.be/X_pYC_LtBFw"><img src="https://img.shields.io/static/v1?label=SIGGRAPH 2021 &message=5 Minute Video&color=red" height=22.5></a>
<a href="https://replicate.ai/yuval-alaluf/sam"><img src="https://img.shields.io/static/v1?label=Replicate&message=Demo and Docker Image&color=darkgreen" height=22.5></a>

Inference Notebook:  <a href="http://colab.research.google.com/github/yuval-alaluf/SAM/blob/master/notebooks/inference_playground.ipynb"><img src="https://colab.research.google.com/assets/colab-badge.svg" height=22.5></a>
Animation Notebook: <a href="http://colab.research.google.com/github/yuval-alaluf/SAM/blob/master/notebooks/animation_inference_playground.ipynb"><img src="https://colab.research.google.com/assets/colab-badge.svg" height=22.5></a>

<p align="center"> <img src="docs/teaser.jpeg" width="800px"/> </p>

Description

Official Implementation of our Style-based Age Manipulation (SAM) paper for both training and evaluation. SAM allows modeling fine-grained age transformation using a single input facial image

<p align="center"> <img src="docs/2195.jpg" width="800px"/> <img src="docs/1936.jpg" width="800px"/> </p>

Table of Contents

• Getting Started
  • Prerequisites
  • Installation
• Pretrained Models
• Training
  • Preparing your Data
  • Training SAM
  • Additional Notes
• Notebooks
  • Inference Notebook
  • MP4 Notebook
• Testing
  • Inference
  • Side-by-Side Inference
  • Reference-Guided Inference
  • Style Mixing
• Repository structure
• Credits
• Acknowledgments
• Citation

Getting Started

Prerequisites

• Linux or macOS
• NVIDIA GPU + CUDA CuDNN (CPU may be possible with some modifications, but is not inherently supported)
• Python 3

Installation

• Dependencies:
  We recommend running this repository using Anaconda. All dependencies for defining the environment are provided in environment/sam_env.yaml.

Pretrained Models

Please download the pretrained aging model from the following links.

| Path | Description | :--- | :---------- |SAM | SAM trained on the FFHQ dataset for age transformation.

You can run this to download it to the right place:

mkdir pretrained_models
pip install gdown
gdown "https://drive.google.com/u/0/uc?id=1XyumF6_fdAxFmxpFcmPf-q84LU_22EMC&export=download" -O pretrained_models/sam_ffhq_aging.pt
wget "https://github.com/italojs/facial-landmarks-recognition/raw/master/shape_predictor_68_face_landmarks.dat"

In addition, we provide various auxiliary models needed for training your own SAM model from scratch.
This includes the pretrained pSp encoder model for generating the encodings of the input image and the aging classifier used to compute the aging loss during training.

| Path | Description | :--- | :---------- |pSp Encoder | pSp taken from pixel2style2pixel trained on the FFHQ dataset for StyleGAN inversion. |FFHQ StyleGAN | StyleGAN model pretrained on FFHQ taken from rosinality with 1024x1024 output resolution. |IR-SE50 Model | Pretrained IR-SE50 model taken from TreB1eN for use in our ID loss during training. |VGG Age Classifier | VGG age classifier from DEX and fine-tuned on the FFHQ-Aging dataset for use in our aging loss

By default, we assume that all auxiliary models are downloaded and saved to the directory pretrained_models. However, you may use your own paths by changing the necessary values in configs/path_configs.py.

Training

Preparing your Data

Please refer to configs/paths_config.py to define the necessary data paths and model paths for training and inference.
Then, refer to configs/data_configs.py to define the source/target data paths for the train and test sets as well as the transforms to be used for training and inference.

As an example, we can first go to configs/paths_config.py and define:

dataset_paths = {
    'ffhq': '/path/to/ffhq/images256x256'
    'celeba_test': '/path/to/CelebAMask-HQ/test_img',
}

Then, in configs/data_configs.py, we define:

DATASETS = {
	'ffhq_aging': {
		'transforms': transforms_config.AgingTransforms,
		'train_source_root': dataset_paths['ffhq'],
		'train_target_root': dataset_paths['ffhq'],
		'test_source_root': dataset_paths['celeba_test'],
		'test_target_root': dataset_paths['celeba_test'],
	}
}

When defining the datasets for training and inference, we will use the values defined in the above dictionary.

Training SAM

The main training script can be found in scripts/train.py.
Intermediate training results are saved to opts.exp_dir. This includes checkpoints, train outputs, and test outputs.
Additionally, if you have tensorboard installed, you can visualize tensorboard logs in opts.exp_dir/logs.

Training SAM with the settings used in the paper can be done by running the following command:

python scripts/train.py \
--dataset_type=ffhq_aging \
--exp_dir=/path/to/experiment \
--workers=6 \
--batch_size=6 \
--test_batch_size=6 \
--test_workers=6 \
--val_interval=2500 \
--save_interval=10000 \
--start_from_encoded_w_plus \
--id_lambda=0.1 \
--lpips_lambda=0.1 \
--lpips_lambda_aging=0.1 \
--lpips_lambda_crop=0.6 \
--l2_lambda=0.25 \
--l2_lambda_aging=0.25 \
--l2_lambda_crop=1 \
--w_norm_lambda=0.005 \
--aging_lambda=5 \
--cycle_lambda=1 \
--input_nc=4 \
--target_age=uniform_random \
--use_weighted_id_loss

Additional Notes

• See options/train_options.py for all training-specific flags.
• Note that using the flag --start_from_encoded_w_plus requires you to specify the path to the pretrained pSp encoder.
  By default, this path is taken from configs.paths_config.model_paths['pretrained_psp'].
• If you wish to resume from a specific checkpoint (e.g. a pretrained SAM model), you may do so using --checkpoint_path.

Notebooks

Inference Notebook

To help visualize the results of SAM we provide a Jupyter notebook found in notebooks/inference_playground.ipynb.
The notebook will download the pretrained aging model and run inference on the images found in notebooks/images.

In addition, Replicate have created a demo for SAM where you can easily upload an image and run SAM on a desired set of ages! Check out the demo here.

MP4 Notebook

To show full lifespan results using SAM we provide an additional notebook notebooks/animation_inference_playground.ipynb that will run aging on multiple ages between 0 and 100 and interpolate between the results to display full aging. The results will be saved as an MP4 files in notebooks/animations showing the aging and de-aging results.

Testing

Inference

Having trained your model or if you're using a pretrained SAM model, you can use scripts/inference.py to run inference on a set of images.
For example,

python scripts/inference.py \
--exp_dir=/path/to/experiment \
--checkpoint_path=experiment/checkpoints/best_model.pt \
--data_path=/path/to/test_data \
--test_batch_size=4 \
--test_workers=4 \
--couple_outputs
--target_age=0,10,20,30,40,50,60,70,80

Additional notes