Home

Awesome

PSFR-GAN in PyTorch

Progressive Semantic-Aware Style Transformation for Blind Face Restoration
Chaofeng Chen, Xiaoming Li, Lingbo Yang, Xianhui Lin, Lei Zhang, Kwan-Yee K. Wong

Changelog

Prerequisites and Installation

Quick Test

Download Pretrain Models and Dataset

Download the pretrained models from the following link and put them to ./pretrain_models

Test single image

Run the following script to enhance face(s) in single input

python test_enhance_single_unalign.py --test_img_path ./test_dir/test_hzgg.jpg --results_dir test_hzgg_results --gpus 1

This script do the following things:

Test image folder

To test multiple images, we first crop out all the faces and align them use the following script.

python align_and_crop_dir.py --src_dir test_dir --results_dir test_dir_align_results

For images (e.g. multiface_test.jpg) contain multiple faces, the aligned faces will be stored as multiface_test_{face_index}.jpg
And then parse the aligned faces and enhance them with

python test_enhance_dir_align.py --src_dir test_dir_align_results --results_dir test_dir_enhance_results

Results will be saved to three folders respectively: results_dir/lq, results_dir/parse, results_dir/hq.

Additional test script

For your convenience, we also provide script to test multiple unaligned images and paste the enhance results back. Note the paste back operation could be quite slow for large size images containing many faces (dlib takes time to detect faces in large image).

python test_enhance_dir_unalign.py --src_dir test_dir --results_dir test_unalign_results

This script basically do the same thing as test_enhance_single_unalign.py for each image in src_dir

Train the Model

Data Preparation

Note: you may change ../datasets/FFHQ to your own path. But images and masks must be stored under your_own_path/imgs1024 and your_own_path/masks512 respectively.

Train Script for PSFRGAN

Here is an example train script for PSFRGAN:

python train.py --gpus 2 --model enhance --name PSFRGAN_v001 \
    --g_lr 0.0001 --d_lr 0.0004 --beta1 0.5 \
    --gan_mode 'hinge' --lambda_pix 10 --lambda_fm 10 --lambda_ss 1000 \
    --Dinput_nc 22 --D_num 3 --n_layers_D 4 \
    --batch_size 2 --dataset ffhq  --dataroot ../datasets/FFHQ \
    --visual_freq 100 --print_freq 10 #--continue_train

Scripts for FPN

You may also train your own FPN and generate masks for the HQ images by yourself with the following steps:

python train.py --gpus 1 --model parse --name FPN_v001 \
    --lr 0.0002 --batch_size 8 \
    --dataset celebahqmask --dataroot ../datasets/CelebAMask-HQ \
    --visual_freq 100 --print_freq 10 #--continue_train

python generate_masks.py --save_masks_dir ../datasets/FFHQ/masks512 --batch_size 8 --parse_net_weight path/to/your/own/FPN

Citation

@inproceedings{ChenPSFRGAN,
    author = {Chen, Chaofeng and Li, Xiaoming and Lingbo, Yang and Lin, Xianhui and Zhang, Lei and Wong, Kwan-Yee~K.},
    title = {Progressive Semantic-Aware Style Transformation for Blind Face Restoration},
    Journal = {IEEE Conference on Computer Vision and Pattern Recognition (CVPR)},
    year = {2021}
}

License

<a rel="license" href="http://creativecommons.org/licenses/by-nc-sa/4.0/"><img alt="Creative Commons License" style="border-width:0" src="https://i.creativecommons.org/l/by-nc-sa/4.0/88x31.png" /></a><br />This work is licensed under a <a rel="license" href="http://creativecommons.org/licenses/by-nc-sa/4.0/">Creative Commons Attribution-NonCommercial-ShareAlike 4.0 International License</a>.

Acknowledgement

This work is inspired by SPADE, and closed related to DFDNet and HiFaceGAN. Our codes largely benefit from CycleGAN.