Towards Robust Blind Face Restoration with Codebook Lookup Transformer (NeurIPS 2022)

Paper | Project Page | Video

Shangchen Zhou, Kelvin C.K. Chan, Chongyi Li, Chen Change Loy

S-Lab, Nanyang Technological University

⭐ If CodeFormer is helpful to your images or projects, please help star this repo. Thanks! 🤗

[News]: 🐳 Due to copyright issues, we have to delay the release of the training code (expected by the end of this year). Please star and stay tuned for our future updates!

Update

• 2022.10.05: Support video input --input_path [YOUR_VIDOE.mp4]. Try it to enhance your videos! 🎬
• 2022.09.14: Integrated to 🤗 Hugging Face. Try out online demo!
• 2022.09.09: Integrated to 🚀 Replicate. Try out online demo!
• 2022.09.04: Add face upsampling --face_upsample for high-resolution AI-created face enhancement.
• 2022.08.23: Some modifications on face detection and fusion for better AI-created face enhancement.
• 2022.08.07: Integrate Real-ESRGAN to support background image enhancement.
• 2022.07.29: Integrate new face detectors of ['RetinaFace'(default), 'YOLOv5'].
• 2022.07.17: Add Colab demo of CodeFormer.
• 2022.07.16: Release inference code for face restoration. 😊
• 2022.06.21: This repo is created.

TODO

• Add checkpoint for face inpainting
• Add checkpoint for face colorization
• Add training code and config files
• Add background image enhancement

🐼 Try Enhancing Old Photos / Fixing AI-arts

Face Restoration

Face Color Enhancement and Restoration

Face Inpainting

Dependencies and Installation

• Pytorch >= 1.7.1
• CUDA >= 10.1
• Other required packages in requirements.txt

# git clone this repository
git clone https://github.com/sczhou/CodeFormer
cd CodeFormer

# create new anaconda env
conda create -n codeformer python=3.8 -y
conda activate codeformer

# install python dependencies
pip3 install -r requirements.txt
python basicsr/setup.py develop

Quick Inference

Download Pre-trained Models:

Download the facelib pretrained models from [Google Drive | OneDrive] to the weights/facelib folder. You can manually download the pretrained models OR download by runing the following command.

python scripts/download_pretrained_models.py facelib

Download the CodeFormer pretrained models from [Google Drive | OneDrive] to the weights/CodeFormer folder. You can manually download the pretrained models OR download by runing the following command.

python scripts/download_pretrained_models.py CodeFormer

Prepare Testing Data:

You can put the testing images in the inputs/TestWhole folder. If you would like to test on cropped and aligned faces, you can put them in the inputs/cropped_faces folder.

Testing on Face Restoration:

[Note] If you want to compare CodeFormer in your paper, please run the following command indicating --has_aligned (for cropped and aligned face), as the command for the whole image will involve a process of face-background fusion that may damage hair texture on the boundary, which leads to unfair comparison.

🧑🏻 Face Restoration (cropped and aligned face)

# For cropped and aligned faces
python inference_codeformer.py -w 0.5 --has_aligned --input_path [input folder]

🖼️ Whole Image Enhancement

# For whole image
# Add '--bg_upsampler realesrgan' to enhance the background regions with Real-ESRGAN
# Add '--face_upsample' to further upsample restorated face with Real-ESRGAN
python inference_codeformer.py -w 0.7 --input_path [image folder/image path]

🎬 Video Enhancement

# For video clips
# Set frame rate of saved video via '--save_video_fps 24'
python inference_codeformer.py --bg_upsampler realesrgan --face_upsample -w 1.0 --input_path [video path] --save_video_fps 24

Fidelity weight w lays in [0, 1]. Generally, smaller w tends to produce a higher-quality result, while larger w yields a higher-fidelity result.

The results will be saved in the results folder.

Citation

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

@inproceedings{zhou2022codeformer,
    author = {Zhou, Shangchen and Chan, Kelvin C.K. and Li, Chongyi and Loy, Chen Change},
    title = {Towards Robust Blind Face Restoration with Codebook Lookup TransFormer},
    booktitle = {NeurIPS},
    year = {2022}
}

License

This project is licensed under NTU S-Lab License 1.0. Redistribution and use should follow this license.

Acknowledgement

This project is based on BasicSR. Some codes are brought from Unleashing Transformers, YOLOv5-face, and FaceXLib. We also adopt Real-ESRGAN to support background image enhancement. Thanks for their awesome works.

Contact

If you have any question, please feel free to reach me out at shangchenzhou@gmail.com.