repository·master·Indexed 11 days ago
https://github.com/sczhou/codeformerA blind face restoration framework using a Codebook Lookup Transformer to enhance low-quality face images and videos. It supports face restoration, colorization, and inpainting, featuring a controllable fidelity weight to balance restoration quality and identity preservation. The framework includes a three-stage training process involving VQGAN and a codebook prediction module.
Use inference_inpainting.py to restore faces that have been masked. Inputs should be masked by a white brush using an image editing app (e.g., Photoshop). This mode is designed for cropped and aligned face images (512x512).
# For cropped and aligned faces (512x512)
python inference_inpainting.py --input_path [image folder]|[image path]CodeFormer requires pre-trained models for facelib, dlib (optional), and CodeFormer itself. You can download them manually from the Releases page or use the provided scripts to download them into the correct directories.
Model Locations:
facelib and dlib models go into weights/facelib.CodeFormer models go into weights/CodeFormer.# Download facelib and dlib models
python scripts/download_pretrained_models.py facelib
python scripts/download_pretrained_models.py dlib
# Download CodeFormer models
python scripts/download_pretrained_models.py CodeFormerTo train CodeFormer, you need to download the FFHQ training dataset. You can find it at the following link:
To install CodeFormer, clone the repository, create a new Conda environment with Python 3.8, and install the required dependencies. Note that dlib is required if you intend to use dlib for face detection or cropping.
Requirements:
dlib (via conda-forge) for specific face detection/cropping features.# 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
conda install -c conda-forge dlibUse inference_colorization.py to colorize black and white or faded photos. This mode is designed for cropped and aligned face images (512x512).
# For cropped and aligned faces (512x512)
python inference_colorization.py --input_path [image folder]|[image path]CodeFormer supports three main modes of face restoration: cropped/aligned faces, whole images, and video clips.
Fidelity Weight (-w):
Lays in [0, 1].
w $\rightarrow$ higher quality.w $\rightarrow$ higher fidelity.Modes:
--has_aligned for 512x512 images.--bg_upsampler realesrgan to enhance background regions and --face_upsample to further upsample the restored face.ffmpeg. Supports .mp4, .mov, and .avi formats.# For cropped and aligned faces (512x512)
python inference_codeformer.py -w 0.5 --has_aligned --input_path [image folder]|[image path]
# For whole image
python inference_codeformer.py -w 0.7 --input_path [image folder]|[image path]
# Recommended for better background/face quality:
python inference_codeformer.py -w 0.7 --bg_upsampler realesrgan --face_upsample --input_path [image folder]|[image path]
# For video clips (requires ffmpeg)
conda install -c conda-forge ffmpeg
python inference_codeformer.py --bg_upsampler realesrgan --face_upsample -w 1.0 --input_path [video path]Stage I involves training the VQGAN.
Note for PyTorch >= 1.10: Replace python -m torch.distributed.launch with torchrun in the command below.
After training, you can run scripts/generate_latent_gt.py to pre-calculate the code sequence for the training dataset, which speeds up subsequent training stages.
If you prefer not to train your own VQGAN, you can download the pre-trained vqgan_code1024.pth and the corresponding code sequence latent_gt_code1024.pth from the v0.1.0 Releases.
# Replace gpu_num with your actual number of GPUs
python -m torch.distributed.launch --nproc_per_node=gpu_num --master_port=4321 basicsr/train.py -opt options/VQGAN_512_ds32_nearest_stage1.yml --launcher pytorch
# Pre-calculate code sequence
python scripts/generate_latent_gt.pyStage III involves training the Controllable Module.
Note for PyTorch >= 1.10: Replace python -m torch.distributed.launch with torchrun in the command below.
You can use the pre-trained codeformer.pth available in the v0.1.0 Releases.
# Replace gpu_num with your actual number of GPUs
python -m torch.distributed.launch --nproc_per_node=gpu_num --master_port=4323 basicsr/train.py -opt options/CodeFormer_stage3.yml --launcher pytorchStage II focuses on training the Code Sequence Prediction Module.
Note for PyTorch >= 1.10: Replace python -m torch.distributed.launch with torchrun in the command below.
You can use the pre-trained codeformer_stage2.pth available in the v0.1.0 Releases.
# Replace gpu_num with your actual number of GPUs
python -m torch.distributed.launch --nproc_per_node=gpu_num --master_port=4322 basicsr/train.py -opt options/CodeFormer_stage2.yml --launcher pytorchStage I involves training the VQGAN.
Note for PyTorch >= 1.10: Replace python -m torch.distributed.launch with torchrun in the command below.
After training VQGAN, you can run python scripts/generate_latent_gt.py to pre-generate the codebook sequences for the training dataset, which accelerates subsequent training stages.
If you prefer not to train your own VQGAN, you can download the pre-trained weights (vqgan_code1024.pth) and the corresponding codebook sequences (latent_gt_code1024.pth) from the v0.1.0 Release.
# Replace gpu_num with your actual number of GPUs
python -m torch.distributed.launch --nproc_per_node=gpu_num --master_port=4321 basicsr/train.py -opt options/VQGAN_512_ds32_nearest_stage1.yml --launcher pytorch
# Generate latent ground truth to accelerate later stages
python scripts/generate_latent_gt.pyIf you want to perform restoration on cropped and aligned faces (512x512), you can use the crop_align_face.py script to process your input images first.
# you may need to install dlib via: conda install -c conda-forge dlib
python scripts/crop_align_face.py -i [input folder] -o [output folder]weights/ directory of the repository.