最完整CodeFormer复现指南：从零实现NeurIPS 2022人脸修复核心算法

你还在为老照片模糊不清、AI生成人脸细节失真而困扰吗？作为NeurIPS 2022的明星算法，CodeFormer凭借其创新的Codebook Lookup Transformer架构，在盲人脸修复领域实现了质量与保真度的完美平衡。本文将带你从零开始复现这一算法，读完你将掌握：环境搭建全流程、三阶段训练策略、核心模块代码实现，以及老照片修复/AI人脸优化的实战技巧。

算法原理与核心优势

CodeFormer创新性地将Transformer与矢量量化（Vector Quantization）结合，通过代码本（Codebook）学习人脸的先验知识，实现对低质量人脸图像的精准修复。其网络结构主要包含三个部分：

• 编码器（Encoder）：将输入人脸图像压缩为潜在特征向量
• Transformer模块：通过代码本查找（Codebook Lookup）预测高质量人脸的特征编码
• 解码器（Decoder）：将预测的特征编码重建为高分辨率人脸图像

相比传统方法，CodeFormer具有三大优势：

1. 强鲁棒性：对严重模糊、压缩 artifacts和极端光照条件均表现稳定
2. 可控性：通过保真度权重（w）调节修复效果，兼顾质量（小w）与真实感（大w）
3. 多功能性：支持人脸修复、上色、补全三大核心功能

环境准备与依赖安装

基础环境配置

# 克隆项目仓库
git clone https://gitcode.com/gh_mirrors/co/CodeFormer
cd CodeFormer

# 创建并激活conda环境
conda create -n codeformer python=3.8 -y
conda activate codeformer

# 安装核心依赖
pip3 install -r requirements.txt
python basicsr/setup.py develop

# 安装可选依赖（dlib人脸检测器）
conda install -c conda-forge dlib

预训练模型下载

项目提供了便捷的模型下载脚本，支持CodeFormer主模型、人脸检测库和dlib模型的一键下载：

# 下载CodeFormer主模型
python scripts/download_pretrained_models.py CodeFormer

# 下载人脸检测和解析模型
python scripts/download_pretrained_models.py facelib

# 可选：下载dlib人脸检测器（用于更精准的人脸对齐）
python scripts/download_pretrained_models.py dlib

模型将自动保存至weights/目录下，完整的模型文件结构如下：

• weights/CodeFormer/codeformer.pth：主修复模型
• weights/facelib/detection_Resnet50_Final.pth：人脸检测模型
• weights/facelib/parsing_parsenet.pth：人脸解析模型

核心算法复现步骤

1. 数据准备

复现论文效果需使用FFHQ（Flickr-Faces-HQ）数据集，包含7万张高质量人脸图像。可通过官方渠道下载后，使用项目提供的脚本进行预处理：

# 人脸裁剪与对齐（如需使用自定义数据集）
python scripts/crop_align_face.py -i [原始图像文件夹] -o inputs/cropped_faces

项目已预置多种测试数据，存放于inputs/目录：

• inputs/whole_imgs/：完整人脸图像
• inputs/cropped_faces/：预裁剪对齐的512x512人脸
• inputs/gray_faces/：黑白人脸图像（用于上色任务）
• inputs/masked_faces/：带掩码人脸图像（用于补全任务）

2. 三阶段训练流程

完整的CodeFormer训练分为三个阶段，对应不同的网络组件和训练目标：

阶段一：VQGAN预训练

VQGAN（Vector-Quantized GAN）负责学习人脸图像的离散编码表示：

python -m torch.distributed.launch --nproc_per_node=8 --master_port=4321 basicsr/train.py \
  -opt options/VQGAN_512_ds32_nearest_stage1.yml --launcher pytorch

训练完成后，可使用以下脚本生成训练数据的特征编码，加速后续训练：

python scripts/generate_latent_gt.py

阶段二：CodeFormer基础模块训练（w=0）

训练Transformer的代码序列预测模块，此时保真度权重w设为0，专注于生成高质量人脸：

python -m torch.distributed.launch --nproc_per_node=8 --master_port=4322 basicsr/train.py \
  -opt options/CodeFormer_stage2.yml --launcher pytorch

阶段三：可控修复模块训练（w=1）

训练保真度控制模块，使模型能够通过调节w参数平衡质量与保真度：

python -m torch.distributed.launch --nproc_per_node=8 --master_port=4323 basicsr/train.py \
  -opt options/CodeFormer_stage3.yml --launcher pytorch

完整训练文档可参考docs/train.md（英文）和docs/train_CN.md（中文）。

3. 推理代码核心实现

推理脚本inference_codeformer.py实现了完整的人脸修复流程，核心步骤包括：

1. 模型初始化：加载预训练权重并设置推理设备

net = ARCH_REGISTRY.get('CodeFormer')(
    dim_embd=512, 
    codebook_size=1024, 
    n_head=8, 
    n_layers=9, 
    connect_list=['32', '64', '128', '256']
).to(device)

ckpt_path = load_file_from_url(
    url=pretrain_model_url['restoration'], 
    model_dir='weights/CodeFormer'
)
net.load_state_dict(torch.load(ckpt_path)['params_ema'])
net.eval()

2. 人脸检测与对齐：使用FaceRestoreHelper类实现自动化人脸处理

face_helper = FaceRestoreHelper(
    upscale=2,
    face_size=512,
    det_model='retinaface_resnet50',
    use_parse=True,
    device=device
)
face_helper.read_image(img)
face_helper.get_face_landmarks_5()  # 检测5点人脸 landmarks
face_helper.align_warp_face()       # 对齐并裁剪人脸区域

3. 修复推理：核心推理代码，支持通过w参数调节修复强度

# 图像预处理
cropped_face_t = img2tensor(cropped_face / 255., bgr2rgb=True, float32=True)
normalize(cropped_face_t, (0.5, 0.5, 0.5), (0.5, 0.5, 0.5), inplace=True)
cropped_face_t = cropped_face_t.unsqueeze(0).to(device)

# 模型推理
with torch.no_grad():
    output = net(cropped_face_t, w=0.5, adain=True)[0]  # w=0.5为默认平衡值
    restored_face = tensor2img(output, rgb2bgr=True, min_max=(-1, 1))

功能实战与效果展示

1. 人脸修复基础用法

完整图像修复

对包含人脸的完整图像进行修复，程序会自动检测人脸区域并进行针对性修复：

# 基础修复（默认w=0.5）
python inference_codeformer.py -w 0.5 --input_path inputs/whole_imgs

# 增强背景细节（结合Real-ESRGAN）
python inference_codeformer.py -w 0.7 --bg_upsampler realesrgan --input_path inputs/whole_imgs/01.jpg

# 超高分辨率输出（人脸单独上采样）
python inference_codeformer.py -w 1.0 --face_upsample --input_path inputs/whole_imgs/02.png

修复效果对比： 左：输入低质量人脸 | 右：CodeFormer修复结果（w=0.5）

对齐人脸修复

对于已裁剪对齐的512x512人脸图像，可使用--has_aligned参数获得更精确的修复结果：

python inference_codeformer.py -w 0.8 --has_aligned --input_path inputs/cropped_faces

上排：输入模糊人脸 | 下排：CodeFormer修复结果

2. 人脸上色功能

CodeFormer提供专业的人脸上色功能，可将黑白老照片或褪色照片恢复鲜艳色彩：

python inference_colorization.py --input_path inputs/gray_faces

上色效果展示： 左：黑白输入 | 中：传统上色方法 | 右：CodeFormer上色结果

3. 人脸补全功能

针对人脸区域有缺失或损坏的图像，CodeFormer可智能补全缺失部分，保持人脸结构的完整性：

python inference_inpainting.py --input_path inputs/masked_faces

补全效果展示： 左：带掩码输入（白色区域为待补全部分） | 右：CodeFormer补全结果

4. 视频人脸增强

CodeFormer支持视频文件输入，可批量处理视频中的所有人脸帧：

# 安装视频处理依赖
conda install -c conda-forge ffmpeg

# 视频增强（默认w=1.0，优先保证真实感）
python inference_codeformer.py --bg_upsampler realesrgan --face_upsample -w 1.0 --input_path inputs/test_video.mp4

进阶训练与参数调优

训练参数配置

项目使用YAML配置文件管理训练参数，主要配置文件位于options/目录：

• options/CodeFormer_stage2.yml：第二阶段训练配置
• options/CodeFormer_stage3.yml：第三阶段训练配置
• options/VQGAN_512_ds32_nearest_stage1.yml：VQGAN预训练配置

关键可调参数说明：

• datasets.train.dataset.root：训练数据集路径
• optimizer.lr：学习率（建议初始值1e-4）
• train.batch_size_per_gpu：单GPU批次大小（根据GPU显存调整）
• model.codebook_size：代码本大小（默认1024，增大可提升细节但增加训练难度）

保真度权重（w）调节指南

CodeFormer的核心创新在于通过w参数控制修复效果，实际应用中可根据需求灵活调节：

w值	效果特点	适用场景
0.0	最高修复质量，可能轻微改变身份特征	艺术化修复、动漫人脸增强
0.3-0.5	平衡质量与保真度，默认推荐值	日常照片修复、老照片翻新
1.0	最高保真度，保留更多原始特征	证件照修复、珍贵历史照片还原

常见问题解决

1. 显存不足：降低--bg_tile参数（默认400），或使用--face_upsample替代全图上采样
2. 人脸检测失败：尝试更换检测模型--detection_model YOLOv5l（适合小人脸）
3. 修复结果过度平滑：增大w值（如从0.3调整至0.7），或减小输入图像尺寸
4. 背景模糊：启用--bg_upsampler realesrgan进行背景增强

总结与扩展应用

通过本文的步骤指南，你已掌握CodeFormer的核心功能和实现原理。该项目不仅可用于个人照片修复，还可扩展到多个专业领域：

• 数字存档：老照片、历史影像的数字化修复与保存
• 影视后期：低成本提升人脸镜头质量，减少重拍需求
• 安防监控：增强低清监控画面中的人脸特征，辅助身份识别
• 游戏开发：快速生成和修复游戏角色人脸贴图

项目持续维护更新，更多功能可关注历史更新日志。如需深入研究，建议结合原始论文《Towards Robust Blind Face Restoration with Codebook Lookup Transformer》阅读源码。

---

项目资源汇总：

• 官方文档：README.md
• 训练指南：docs/train.md
• 核心源码：basicsr/models/codeformer_model.py
• 推理脚本：inference_codeformer.py

如有技术问题，可通过项目issue或邮件联系开发团队。