chenxx89 / BFRffusion

Official codes of Towards Real-World Blind Face Restoration with Generative Diffusion Prior

Towards Real-World Blind Face Restoration with Generative Diffusion Prior

Xiaoxu Chen, Jingfan Tan, Tao Wang, Kaihao Zhang, Wenhan Luo, Xiaochun Cao

Keywords: blind face restoration, face dataset, diffusion model, transformer

Abstracts: Blind face restoration is an important task in computer vision and has gained significant attention due to its wide-range applications. Previous works mainly exploit facial priors to restore face images and have demonstrated high-quality results. However, generating faithful facial details remains a challenging problem due to the limited prior knowledge obtained from finite data. In this work, we delve into the potential of leveraging the pretrained Stable Diffusion for blind face restoration. We propose BFRffusion which is thoughtfully designed to effectively extract features from low-quality face images and could restore realistic and faithful facial details with the generative prior of the pretrained Stable Diffusion. In addition, we build a privacy-preserving face dataset called PFHQ with balanced attributes like race, gender, and age. This dataset can serve as a viable alternative for training blind face restoration networks, effectively addressing privacy and bias concerns usually associated with the real face datasets. Through an extensive series of experiments, we demonstrate that our BFRffusion achieves state-of-the-art performance on both synthetic and real-world public testing datasets for blind face restoration and our PFHQ dataset is an available resource for training blind face restoration networks.

📖 Table Of Contents

• 📖 Table Of Contents
• 📷 Demo on Real-World Blind Face Restoration
• 🏠 Installation
• ⚡ Quick Inference
• 💻 Training
• 📏 Evaluation
• 📈 Results
• 🖼️ PFHQ
• 📋 Citation
• ✨ Acknowledgement
• 📧 Contact

📷 Demo on Real-World Blind Face Restoration

🏠 Installation

A suitable conda environment named BFRffusion can be created and activated with:

# clone this repository
git clone https://github.com/chenxx89/BFRffusion
cd BFRffusion

# create new anaconda env
conda env create -f environment.yaml
conda activate BFRffusion

You can install xformers for more efficiency and speed on GPUs. For more details you can visit xformers and stablediffusion.

⚡ Quick Inference

Download pre-trained models from BaiduNetdisk or OneDrive to the experiments/weights folder.

Inference!

# for real-world image
python inferience.py --input inputs -output results --sr_scale 2

# for cropped face
python inferience.py --input inputs -output results --sr_scale 2 --aligned

Usage: python inferience.py --input inputs -output results --sr_scale 2 [options]...

  --input             Input image or folder. Default: inputs/whole_imgs
  --output            Output folder. Default: results
  --ckpt               The dir of ckpt to load
  --config             Path to config which constructs model
  --ddim_steps         Number of ddpm sampling steps
  --only_center_face    Only restore the center face
  --ext                 Image extension. Options: auto | jpg | png, auto means using the same extension as inputs. Default: auto
  --aligned             Input are aligned faces
  --sr_scale           The final upsampling scale of the image. Default: 2
  --bg_upsampler        background upsampler. Default: realesrgan
  --bg_tile             Tile size for background sampler, 0 for no tile during testing. Default: 400
  --suffix              Suffix of the restored faces

💻 Training

We provide the training codes used in our paper for BFRffusion.
You could improve it according to your own needs.

1. Dataset preparation: FFHQ
2. Download pre-trained Stable Diffusion 2.1-base from Hugging Face to the experiments/weights folder
3. Modify the configuration file options/train.yaml accordingly
4. Training

python train.py

📏 Evaluation

We evaluate BFRffusion on the synthetic dataset CelebA-Test, and three real-world datasets LFW-Test, CelebAdult-Test and Webphoto-Test. For reproduce our evaluation results, you need to perform the following steps:

1. Download testing datasets by the following links:

Dataset	Usage	Quantity	Type
CelebA-Test	testing	3,000	paired synthetic dataset
LFW-Test	testing	1,711	non-paired real-world dataset
CelebAdult-Test	testing	180	non-paired real-world dataset
WebPhoto-Test	testing	407	non-paired real-world dataset

2. Download pretrained models for different metrics, and put them to experiments/weights

Metrics	Pretrained Weights	Download
FID	inception_FFHQ_512.pth	BaiduNetdisk
Deg	resnet18_110.pth
LMD	alignment_WFLW_4HG.pth

3. Generate restoration results:

• Specify the dataset_lq/dataset_gt to the testing dataset root in options/test.yaml

• We use multiprocessing to test on multi-GPUs, please reset --gpu_ids

• Then run the following command:

    python test.py

You can download our testing results from BaiduNetdisk or OneDrive.

4. Run evaluation:

    # LPIPS|PSNR/SSIM|LMD|Deg.
    python metrics/[calculate_lpips.py|calculate_psnr_ssim.py|calculate_landmark_distance.py|calculate_cos_dist.py]
    -restored_folder folder_to_results -gt_folder folder_to_gt

    # FID
    python metric_paper/calculate_fid_folder.py -restored_folder folder_to_results

For more testing details, you can visit VQFR

📈 Results

Quantitative comparison on CelebA-Test for blind face restoration. Red and Blue indicate the best and the second-best performance.

🖼️ PFHQ

We build a privacy-preserving paired face dataset called PFHQ with balanced race, gender, and age. Extensive experiments show that our PFHQ dataset can serve as an alternative to real-face datasets for training blind face restoration methods. Following are some representative face images of our proposed PFHQ dataset. You can download our PFHQ from BaiduNetdisk.

📋 Citation

If our work is useful for your research, please consider citing:

@misc{chen2023realworld,
      title={Towards Real-World Blind Face Restoration with Generative Diffusion Prior}, 
      author={Xiaoxu Chen and Jingfan Tan and Tao Wang and Kaihao Zhang and Wenhan Luo and Xiaocun Cao},
      year={2023},
      eprint={2312.15736},
      archivePrefix={arXiv},
      primaryClass={cs.CV}
}

✨ Acknowledgement

• This project is based on Controlnet, and Restormer.

• Some codes are brought from BasicSR, and arcface-pytorch.

• This Readme is inspired by RobustSIRR, and VQFR.

📧 Contact

If you have any questions, please feel free to reach out at [e-mail].