PSFR-GAN in PyTorch

PyTorch codes for "Progressive Semantic-Aware Style Transformation for Blind Face Restoration"

Getting Started

Prerequisites and Installation

  • Ubuntu 18.04
  • CUDA 10.1
  • Clone this repository
    git clone
    cd PSFR-GAN
  • Python 3.7, install required packages by pip3 install -r requirements.txt

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_img_path ./test_dir/test_hzgg.jpg --results_dir test_hzgg_results --gpus 1

This script do the following things:

  • Crop and align all the faces from input image, stored at results_dir/LQ_faces
  • Parse these faces and then enhance them, results stored at results_dir/ParseMaps and results_dir/HQ
  • Paste then enhanced faces back to the original image results_dir/hq_final.jpg
  • You can use --gpus to specify how many GPUs to use, <=0 means running on CPU. The program will use GPU with the most available memory. Set CUDA_VISIBLE_DEVICE to specify the GPU if you do not want automatic GPU selection.

Test image folder

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

python --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 --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.

python --src_dir test_dir --results_dir test_unalign_results

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


    author = {Chen, Chaofeng and Li, Xiaoming and Lin, Xianhui and Lingbo, Yang and Zhang, Lei and Wong, KKY},
    title = {Progressive Semantic-Aware Style Transformation for Blind Face Restoration},
    Journal = {arXiv preprint arXiv:2009.08709},
    year = {2020}