RestoreFormer++: 基于未退化键值对的真实场景盲人脸修复

这个仓库是"RestoreFormer++: 基于未退化键值对的真实场景盲人脸修复"的官方实现。

RestoreFormer++是我们RestoreFormer的扩展版本。它提出使用强大的全空间注意力机制来建模人脸中丰富的上下文信息及其与我们面向重建的高质量先验之间的相互作用，从而在保真度和真实性方面修复退化的人脸图像。此外，它引入了一个扩展退化模型(EDM)，包含更多真实的退化场景用于训练数据合成，这有助于增强其在真实场景中的鲁棒性和泛化能力。我们的结果与最先进方法的比较以及有/无EDM的性能如下所示：

images/fig1.png

images/fig3.png

待办事项

20240122 提供本地gradio演示。
20230915 更新在线演示
20230915 提供用户友好的推理方法。
- 可用于带RealESRGAN的背景超分辨率。
- basicsr应升级到1.4.2。
20230914 上传模型
20230914 发布代码
20221120 介绍项目。

环境

python>=3.7
pytorch>=1.7.1
pytorch-lightning==1.0.8
omegaconf==2.0.0
~~basicsr==1.3.3.4~~ basicsr>=1.4.2
realesrgan==0.3.0

pip install -r RF_requirements.txt

❗❗❗ 警告不同版本的pytorch-lightning和omegaconf可能导致错误或不同的结果。

数据集和模型准备

数据集：

训练数据：我们工作中的ROHQD和RestoreFormer++都使用FFHQ进行训练，该数据集从FFHQ仓库获得。FFHQ中图像的原始尺寸为1024x1024。在我们的工作中，我们使用双线性插值将它们调整为512x512。将此数据集链接到./data/FFHQ/image512x512。
测试数据：CelebA-Test，LFW-Test，WebPhoto-Test和CelebChild-Test

模型：用于训练的预训练模型以及我们的RestoreFormer和RestoreFormer++的训练模型可以从Google Drive获得。将这些模型链接到./experiments。

Gradio演示

python gradio_demo/app.py

快速推理

python inference.py -i data/aligned -o results/RF++/aligned -v RestoreFormer++ -s 2 --aligned --save python inference.py -i data/raw -o results/RF++/raw -v RestoreFormer++ -s 2 --save python inference.py -i data/aligned -o results/RF/aligned -v RestoreFormer -s 2 --aligned --save python inference.py -i data/raw -o results/RF/raw -v RestoreFormer -s 2 --save

注意：相关代码借鉴自GFPGAN。

测试

sh scripts/test.sh

scripts/test.sh

exp_name='RestoreFormer'
exp_name='RestoreFormerPlusPlus'

root_path='experiments'
out_root_path='results'
align_test_path='data/aligned'
# unalign_test_path='data/raw'
tag='test'

outdir=$out_root_path'/'$exp_name'_'$tag

if [ ! -d $outdir ];then
    mkdir -m 777 $outdir
fi

CUDA_VISIBLE_DEVICES=0 python -u scripts/test.py \
--outdir $outdir \
-r $root_path'/'$exp_name'/last.ckpt' \
-c 'configs/'$exp_name'.yaml' \
--test_path $align_test_path \
--aligned

此代码库适用于RestoreFormer和RestoreFormerPlusPlus。通过exp_name确定具体模型。
通过root_path设置模型路径
恢复结果保存在out_root_path
将退化的人脸图像放在test_path中
如果退化的人脸图像已对齐，设置*--aligned*，否则从脚本中删除。data/aligned中提供的测试图像是对齐的，而data/raw中的图像未对齐且包含多个人脸。

训练

sh scripts/run.sh

scripts/run.sh

export BASICSR_JIT=True

# 对于RestoreFormer
# conf_name='HQ_Dictionary'
# conf_name='RestoreFormer'

# 对于RestoreFormer++
conf_name='ROHQD'
conf_name='RestoreFormerPlusPlus'

# gpus='0,1,2,3,4,5,6,7'
# node_n=1
# ntasks_per_node=8

root_path='PATH_TO_CHECKPOINTS'

gpus='0,'
node_n=1
ntasks_per_node=1

gpu_n=$(expr $node_n \* $ntasks_per_node)

python -u main.py \
--root-path $root_path \
--base 'configs/'$conf_name'.yaml' \
-t True \
--postfix $conf_name'_gpus'$gpu_n \
--gpus $gpus \
--num-nodes $node_n \
--random-seed True \

此代码库适用于RestoreFormer和RestoreFormerPlusPlus。通过conf_name确定训练模型。'HQ_Dictionary'和'RestoreFormer'用于RestoreFormer，而'ROHQD'和'RestoreFormerPlusPlus'用于RestoreFormerPlusPlus。
在训练'RestoreFormer'或'RestoreFormerPlusPlus'时，configs/中相应配置文件的*'ckpt_path'*应更新为'HQ_Dictionary'或'ROHQD'的训练模型路径。

评估指标

sh scripts/metrics/run.sh

注意。

如果您想获得IDD、PSRN、SSIM、LIPIS，需要在脚本中添加CelebA-Test数据集的路径。
相关的评估模型在*./experiments/pretrained_models/*中。

引用

@article{wang2023restoreformer++,
  title={RestoreFormer++: Towards Real-World Blind Face Restoration from Undegraded Key-Value Paris},
  author={Wang, Zhouxia and Zhang, Jiawei and Chen, Tianshui and Wang, Wenping and Luo, Ping},
  booktitle={IEEE Transactions on Pattern Analysis and Machine Intelligence (T-PAMI)},
  year={2023}
}

@article{wang2022restoreformer,
  title={RestoreFormer: High-Quality Blind Face Restoration from Undegraded Key-Value Pairs},
  author={Wang, Zhouxia and Zhang, Jiawei and Chen, Runjian and Wang, Wenping and Luo, Ping},
  booktitle={The IEEE Conference on Computer Vision and Pattern Recognition (CVPR)},
  year={2022}
}

联系方式

如有任何问题，请随时发送电子邮件至wzhoux@connect.hku.hk或zhouzi1212@gmail.com。