This is the project page for "Unsupervised Shadow Removal and Relighting for Webcam Portrait Images with GANs" (paper submission in progress). We propose a framework for lighting correction on portrait images. This work is based on existing GAN-based approaches for portrait image generation and styletransfer.

The basic idea of this work is the following:

1. Project an image into the StyleGAN2 latent space to get the latent vector
2. Create (k) versions of the latent vector with random manipulations and generate the associated images with the generator. There will be many images with different lighting conditions
3. Automatically select the image with the most homogenous lighting conditions

• The generator is based on StyleGAN2-ada-PyTorch and uses a pretrained model based on the FFHQ-Dataset (note: checkpoints can be found in its associated repositories )
• The encoder is provided by Tov et al.: Encoder4Editing. It is used to find the latent vector of an input image
• Image preprocessing contains two steps: alignment and background removal. The former is done by FFHQ-Alignment and face masks for background removal were generated by Face-Parsing-PyTorch
• In order to calculate how much shadow an image contains a ST-CGAN is used to generate shadow masks

CUDA is mandatory and for execution a Nvidia GTX 3070 (8GB) was used. If you want to train your own models you will need a much faster GPU. I experimented with two RTX Titans (24GB) and even then training of the generator and encoder takes each over one week.

• Ubuntu 20.04
• Python 3.6
• PyTorch 1.8.1
• CUDA 10.1

There are pretrained checkpoints which you can use. You will need a directory "checkpoints" containing the following files:

• e4e_ffhq_encode.pt (e4e Encoder)
• moco_v2_800ep_pretrain.pt (used by e4e)
• model_ir_se50.pth (used by e4e)
• stylegan2-ffhq-config-f.pt (StyleGAN2-ada-PyTorch weights)

In addition download the pretrained face parsing model and place it under "models/face_parsing". You will also need checkpoints for the ST-CGAN. Put them under "models/st_cgan"

After you have set up your environment and downloaded the checkpoints to the correct directories run the following command:

python inference.py --images_dir=/images/inputs --save_dir=/images/outputs/inversions --edit_attribute='inversion' --ckpt=/checkpoints/e4e_ffhq_encode.pt

This will iterate through all the images inside the "/images/inputs" directory. Under "/images/outputs/inversions" you will find the inversion aswell as the final result (the image with the most homogenous lighting).

You can improve the quality by creating many different variations of an input image. To tune it you can set the paramter "n_random_imgs" in "scripts/edit_latents.py". By default it is set to 300. Runtime will also increase if you set it to a higher value. Depending on the dataset you can also try to choose a different value for "value" (also in "scripts/edit_latents.py"). This affects how big or small the manipulation of the latent vector is.

• Inversion quality can vary depending on you dataset. If you are using your own images you will notice that there is quite a big distortion. You can try to train your own generator or encoder model to better suit your own dataset.

This work was part of my Master thesis at TU Dortmund University. I would like to thank my mentor aswell as all of the people that were involved.