| # CodeFormer Face Restoration - Project Documentation |
|
|
| ## 1. Introduction |
|
|
| **CodeFormer** is a robust blind face restoration algorithm designed to restore old, degraded, or AI-generated face images. It utilizes a **Codebook Lookup Transformer** (VQGAN-based) to predict high-quality facial features even from severe degradation, ensuring that the restored faces look natural and faithful to the original identity. |
|
|
| This project wraps the core CodeFormer research code into a deployable, user-friendly **Flask Web Application**, containerized with **Docker** for easy deployment on platforms like Hugging Face Spaces. |
|
|
| ### Key Features |
| * **Blind Face Restoration:** Restores faces from low-quality inputs without knowing the specific degradation details. |
| * **Background Enhancement:** Uses **Real-ESRGAN** to upscale and enhance the non-face background regions of the image. |
| * **Face Alignment & Paste-back:** Automatically detects faces, aligns them for processing, and seamlessly blends them back into the original image. |
| * **Adjustable Fidelity:** Users can balance between restoration quality (hallucinating details) and identity fidelity (keeping the original look). |
|
|
| --- |
|
|
| ## 2. System Architecture |
|
|
| The application is built on a Python/PyTorch backend served via Flask. |
|
|
| ### 2.1 Technology Stack |
| * **Framework:** Flask (Python Web Server) |
| * **Deep Learning:** PyTorch, TorchVision |
| * **Image Processing:** OpenCV, NumPy, Pillow |
| * **Core Libraries:** `basicsr` (Basic Super-Restoration), `facelib` (Face detection/utils) |
| * **Frontend:** HTML5, Bootstrap 5, Jinja2 Templates |
| * **Containerization:** Docker (CUDA-enabled) |
|
|
| ### 2.2 Directory Structure |
| ``` |
| CodeFormer/ |
| ├── app.py # Main Flask application entry point |
| ├── Dockerfile # Container configuration |
| ├── requirements.txt # Python dependencies |
| ├── basicsr/ # Core AI framework (Super-Resolution tools) |
| ├── facelib/ # Face detection and alignment utilities |
| ├── templates/ # HTML Frontend |
| │ ├── index.html # Upload interface |
| │ └── result.html # Results display |
| ├── static/ # Static assets (css, js, uploads) |
| │ ├── uploads/ # Temporary storage for input images |
| │ └── results/ # Temporary storage for processed output |
| └── weights/ # Pre-trained model weights (downloaded on startup) |
| ├── CodeFormer/ # CodeFormer model (.pth) |
| ├── facelib/ # Detection (RetinaFace) and Parsing models |
| └── realesrgan/ # Background upscaler (Real-ESRGAN) |
| ``` |
|
|
| ### 2.3 Logic Flow |
| 1. **Input:** User uploads an image via the Web UI. |
| 2. **Pre-processing (`app.py`):** |
| * Image is saved to `static/uploads`. |
| * Parameters (fidelity, upscale factor) are parsed. |
| 3. **Inference Pipeline:** |
| * **Detection:** `facelib` detects faces in the image using RetinaFace. |
| * **Alignment:** Faces are cropped and aligned to a standard 512x512 resolution. |
| * **Restoration:** The **CodeFormer** model processes the aligned faces. |
| * **Upscaling (Optional):** The background is upscaled using **Real-ESRGAN**. |
| * **Paste-back:** Restored faces are warped back to their original positions and blended. |
| 4. **Output:** The final image is saved to `static/results` and displayed to the user. |
|
|
| --- |
|
|
| ## 3. Installation & Deployment |
|
|
| ### 3.1 Docker Deployment (Recommended) |
| The project is optimized for Docker. |
|
|
| **Prerequisites:** Docker, NVIDIA GPU (optional, but recommended). |
|
|
| 1. **Build the Image:** |
| ```bash |
| docker build -t codeformer-app . |
| ``` |
| |
| 2. **Run the Container:** |
| ```bash |
| # Run on port 7860 (Standard for HF Spaces) |
| docker run -it -p 7860:7860 codeformer-app |
| ``` |
| *Note: To use GPU, add the `--gpus all` flag to the run command.* |
| |
| ### 3.2 Hugging Face Spaces Deployment |
| This repository is configured for direct deployment to Hugging Face. |
|
|
| 1. Create a **Docker** Space on Hugging Face. |
| 2. Push this entire repository to the Space's Git remote. |
| ```bash |
| git remote add hf git@hf.co:spaces/USERNAME/SPACE_NAME |
| git push hf main |
| ``` |
| 3. The Space will build (approx. 5-10 mins) and launch automatically. |
| |
| ### 3.3 Local Development |
| 1. **Install Environment:** |
| ```bash |
| conda create -n codeformer python=3.8 |
| conda activate codeformer |
| pip install -r requirements.txt |
| ``` |
| 2. **Install Basicsr:** |
| ```bash |
| python basicsr/setup.py install |
| ``` |
| 3. **Run App:** |
| ```bash |
| python app.py |
| ``` |
| |
| --- |
|
|
| ## 4. User Guide (Web Interface) |
|
|
| ### 4.1 Interface Controls |
|
|
| * **Input Image:** Supports standard formats (JPG, PNG, WEBP). Drag and drop supported. |
| * **Fidelity Weight (w):** |
| * **Range:** 0.0 to 1.0. |
| * **0.0 (Better Quality):** The model "hallucinates" more details. Results look very sharp and high-quality but may slightly alter the person's identity (look less like the original). |
| * **1.0 (Better Identity):** The model sticks strictly to the original features. Results are faithful to the original photo but might be blurrier or contain more artifacts. |
| * **Recommended:** 0.5 is a balanced default. |
| * **Upscale Factor:** |
| * Scales the final output resolution (1x, 2x, or 4x). |
| * *Note: Higher scaling requires more VRAM.* |
| * **Enhance Background:** |
| * If checked, runs Real-ESRGAN on the non-face areas. |
| * *Recommendation:* Keep checked for full-photo restoration. Uncheck if you only care about the face or are running on limited hardware. |
| * **Upsample Face:** |
| * If checked, the restored face is also upsampled to match the background resolution. |
|
|
| ### 4.2 Viewing Results |
| The result page features an interactive **Before/After Slider**. Drag the handle left and right to compare the pixels of the original versus the restored image directly. |
|
|
| --- |
|
|
| ## 5. Technical Details |
|
|
| ### 5.1 Model Weights |
| The application automatically checks for and downloads the following weights to the `weights/` directory on startup: |
|
|
| | Model | Path | Description | |
| | :--- | :--- | :--- | |
| | **CodeFormer** | `weights/CodeFormer/codeformer.pth` | Main restoration model. | |
| | **RetinaFace** | `weights/facelib/detection_Resnet50_Final.pth` | Face detection. | |
| | **ParseNet** | `weights/facelib/parsing_parsenet.pth` | Face parsing (segmentation). | |
| | **Real-ESRGAN** | `weights/realesrgan/RealESRGAN_x2plus.pth` | Background upscaler (x2). | |
|
|
| ### 5.2 Performance Notes |
| * **Memory:** The full pipeline (CodeFormer + Real-ESRGAN) requires significant RAM/VRAM. On CPU-only environments (like basic HF Spaces), processing a single image may take 30-60 seconds. |
| * **Git LFS:** Image assets in this repository are tracked with Git LFS to keep the repo size manageable. |
|
|
| --- |
|
|
| ## 6. Credits & References |
|
|
| * **Original Paper:** [Towards Robust Blind Face Restoration with Codebook Lookup Transformer (NeurIPS 2022)](https://arxiv.org/abs/2206.11253) |
| * **Authors:** Shangchen Zhou, Kelvin C.K. Chan, Chongyi Li, Chen Change Loy (S-Lab, Nanyang Technological University). |
| * **Original Repository:** [sczhou/CodeFormer](https://github.com/sczhou/CodeFormer) |
|
|
