【深度学习】脸部修复，CodeFormer，论文，实战

代码： https://github.com/sczhou/CodeFormer

论文：https://arxiv.org/abs/2206.11253

Towards Robust Blind Face Restoration with Codebook Lookup Transformer

论文

摘要

Blind face restoration是一个高度 ill-posed 的问题，通常需要辅助引导来 1) 改善从降级输入到期望输出的映射，或者 2) 补充在输入中丢失的高质量细节。在本文中，我们展示了在一个小的代理空间中学到的离散码本先验可以通过将盲目的面部恢复建模为 a code prediction task来大大减少恢复映射的不确定性和模糊性，同时为生成高质量面部提供丰富的视觉元素。在这个范式下，我们提出了一个基于 Transformer 的预测网络，命名为 CodeFormer，用于建模低质量面部的全局构成和上下文，以进行代码预测，使其能够在输入严重降级的情况下发现自然面部，从而紧密逼近目标面部。为了增强对不同降级的适应性，我们还提出了一个可控特征转换模块，允许在保真度和质量之间进行灵活的权衡。由于表达力丰富的码本先验和全局建模，CodeFormer在质量和保真度方面均优于现有技术，展现出对降级的卓越鲁棒性。大量的合成和真实数据集上的实验证明了我们方法的有效性。

1 引言

在野外捕捉的人脸图像往往受到各种降级的影响，如压缩、模糊和噪声。由于降级引起的信息丢失导致在给定低质量（LQ）输入的情况下存在无穷多合理的高质量（HQ）输出，因此恢复这样的图像是一个高度不适定的问题。在盲目恢复中，由于特定的降级是未知的，这种不适定性进一步增强。尽管随着深度学习的出现取得了一些进展，但在巨大的图像空间中学习没有附加指导的 LQ-HQ 映射仍然是不可行的，导致早期方法的恢复质量不佳。为了提高输出质量，辅助信息对于 1) 减少 LQ-HQ 映射的不确定性和 2) 补充高质量细节是不可或缺的。
各种先验已被用于缓解这个问题的不适定性，包括几何先验 [5, 6, 31, 45]、参考先验 [24–26] 和生成先验 [2, 38, 44]。尽管观察到了改进的纹理和细节，但这些方法通常对降级的敏感性较高，或者先验表现受限。这些先验对于面部恢复提供的指导不足，因此它们的网络基本上倾向于使用通常受损害的 LQ 输入图像的信息。因此，LQ-HQ 映射的不确定性仍然存在，并且输入图像的降级使输出质量恶化。**最近，基于生成先验，一些方法通过迭代潜在优化 [27] 或直接潜在编码 [29] 将受损面部投影到一个连续的无限空间。尽管输出非常逼真，但在严重降级的情况下难以找到准确的潜在向量，导致低保真度的结果（图1(d)）。**为了增强保真度，通常需要在这类方法中的编码器和解码器之间引入跳跃连接 [38, 44, 2]，如图1(a)（顶部）所示，然而，当输入严重降级时，这样的设计同时会在结果中引入伪影，如图1(e)所示。

与上述方法不同，本文将盲目的面部恢复视为在学习的离散码本先验的一个小有限代理空间中的代码预测任务，这表现出对降级的卓越鲁棒性以及丰富的表达能力。码本是通过使用矢量量化自重构 HQ 面部来学习的，它与解码器一起存储了用于面部恢复的丰富 HQ 细节。与连续生成先验 [11, 38, 44] 不同，码本项的组合形成了一个只有有限基数的离散先验空间。通过将 LQ 图像映射到一个更小的代理空间（例如，1024个代码），LQ-HQ 映射的不确定性显著减小，提高了对多样降级的鲁棒性，如图1(d-g)中所比较的。此外，码本空间具有更大的表达能力，感知上近似于图像空间，如图1(h)所示。这种性质使得网络能够减少对输入的依赖，甚至可以摆脱跳跃连接的限制。
尽管基于码本的离散表示已经用于图像生成 [4, 11, 35]，但对于图像恢复而言，准确的码组合仍然是一个非常棘手的问题。现有的工作通过最近邻（NN）特征匹配查找码本，但这在图像恢复中较不可行，因为 LQ 输入的固有纹理通常受到破坏。LQ 图像中的信息丢失和多样的降级不可避免地扭曲了特征分布，阻止了准确的特征匹配。如图1(b)（右侧）所示，即使在对 LQ 图像进行微调后，LQ 特征也不能很好地聚类到确切的代码，而是扩散到其他附近的代码簇，因此在这种情况下，最近邻匹配是不可靠的。
针对恢复，我们提出了一种基于 Transformer 的代码预测网络，称为 CodeFormer，以利用 LQ 面部的全局组成和长距离依赖关系进行更好的代码预测。具体而言，将 LQ 特征作为输入，Transformer 模块预测代码令牌序列，该序列被视为码本空间中面部图像的离散表示。由于全局建模弥补了 LQ 图像中的局部信息丢失，所提出的 CodeFormer 对严重降级表现出鲁棒性并保持整体一致性。与图1(f-g)中呈现的结果相比，所提出的 CodeFormer 能够恢复更多的细节，如眼镜，提高了恢复的质量和保真度。
此外，我们提出了一个可控特征转换模块，具有可调节系数，用于控制从 LQ 编码器到解码器的信息流。这种设计允许在保真度和质量之间进行灵活的权衡，以便可以实现它们之间的连续图像过渡。这个模块增强了 CodeFormer 在不同降级下的适应性，例如，在严重降级的情况下，可以手动减少携带降级的 LQ 特征的信息流以产生高质量的结果。
配备上述组件，所提出的 CodeFormer 在现有数据集和我们新引入的 WIDER-Test 数据集中表现出卓越的性能，该数据集包含从 WIDER-Face 数据集 [43] 中收集的 970 张严重受损的面部。除了面部恢复，我们的方法还在其他具有挑战性的任务上展示了其有效性，例如需要来自其他区域的长距离线索的面部修补。系统性的研究和实验证明了我们方法相对于先前工作的优点。

2 相关工作

盲目面部恢复。由于人脸具有高度结构化的特点，人脸的几何先验被用于盲目的面部恢复。一些方法引入了面部标志点 [6]、面部解析图 [5, 31, 42]、面部组件热图 [45] 或 3D 形状 [16, 28, 49] 在其设计中。然而，这样的先验信息无法从受损的面部准确获取。而且，几何先验无法为高质量的面部恢复提供丰富的细节。
为了绕过上述限制，提出了一些基于参考的方法 [9, 24–26]。这些方法通常要求参考图像与输入受损图像具有相同的身份。例如，Li等人 [26] 提出了一个引导面部恢复网络，包括一个变形子网络和一个重构子网络，使用与输入相同身份的高质量引导图像来更好地恢复面部细节。然而，这样的参考图像并不总是可用。DFDNet [24] 预先构建了由高质量面部组件特征组成的字典。然而，特定于组件的字典特征对于高质量面部恢复仍然不足，特别是对于字典范围之外的区域（例如皮肤、头发）。为了缓解这个问题，最近的基于 VQGAN 的方法 [40, 47] 探索了一个学到的高质量字典，其中包含更通用和丰富的面部恢复细节。
最近，通过预训练生成器（例如 StyleGAN2 [21]）学到的生成式面部先验已广泛用于盲目面部恢复。通过迭代潜在优化进行有效的 GAN 反演 [12, 27] 或对受损面部进行直接潜在编码 [29] 是利用这种先验的策略。然而，当将受损面部投影到连续的无限潜在空间时，保持恢复面部的高保真度是具有挑战性的。为了缓解这个问题，GLEAN [2, 3]、GPEN [44] 和 GFPGAN [38] 将生成先验嵌入到编码器-解码器网络结构中，同时利用输入图像的附加结构信息作为引导。尽管这些方法在保真度方面有所改善，但它们高度依赖输入通过跳跃连接，这在输入严重损坏时可能引入伪影。

字典学习。在图像恢复任务中，具有学到字典的稀疏表示已经证明了其优越性，如超分辨率 [13, 33, 34, 41] 和去噪 [10]。然而，这些方法通常需要迭代优化来学习字典和稀疏编码，具有较高的计算成本。尽管效率低下，对 HQ 字典的高级洞察启发了基于参考的恢复网络，例如 LUT [18] 和自引用 [48]，以及合成方法 [11, 35]。Jo 和 Kim [18] 通过将网络输出值转移为 LUT 中的值来构建查找表（LUT），因此在推断过程中只需要进行简单的值检索。然而，在图像域中存储 HQ 纹理通常需要一个庞大的 LUT，限制了其实用性。VQVAE [35] 首次引入了通过向量量化自编码器模型学习的高度压缩的码本。VQGAN [11] 进一步采用对抗损失和感知损失来增强感知质量，以较高的压缩率显著减小码本大小而不损失其表达能力。与大型手工制作的字典 [18, 24] 不同，可学习的码本自动学习了用于 HQ 图像重建的最佳元素，提供了卓越的效率和表达能力，并避免了繁琐的字典设计。受到码本学习的启发，本文研究了盲目面部恢复的离散代理空间。与最近的基于 VQGAN 的方法 [40, 47] 不同，我们通过全局建模预测代码序列来利用离散码本先验，并通过固定编码器来确保先验的有效性。这样的设计使得我们的方法能够充分利用码本，使其不依赖于与 LQ 提示的特征融合，显著增强了面部恢复的鲁棒性。

图2：CodeFormer框架。 (a) 我们首先通过 self-reconstruction learning ** 自重构学习**学习一个离散码本和一个解码器，用于存储面部图像的高质量视觉部分。(b) 在固定的码本和解码器的基础上，我们引入一个Transformer模块进行代码序列预测，对低质量输入建模全局面部组成。此外，我们使用一个可控特征转换模块来控制从LQ编码器到解码器的信息流。请注意，这种连接是可选的，可以在输入严重受损时禁用，用户可以调整标量权重w在质量和保真度之间进行权衡。

4 实验

4.1 数据集

训练数据集： 我们在FFHQ数据集[21]上训练模型，该数据集包含70,000张高质量（HQ）图像，所有图像都被调整为512×512进行训练。为了形成训练对，我们使用以下退化模型[24, 38, 44]从HQ图像Ih合成LQ图像Il：

[ Il = {[(Ih \otimes k_{\sigma}) \downarrow r + n_{\delta}] \text{JPEG}_q} \uparrow r, ]

其中HQ图像Ih首先与高斯核(k_{\sigma})卷积，然后进行尺度为r的降采样。之后，向图像添加加性高斯噪声(n_{\delta})，然后应用JPEG压缩，质量因子为q。最后，将LQ图像调整回512×512。我们从区间[1, 15]、[1, 30]、[0, 20]、[30, 90]中随机采样σ、r、δ和q。

测试数据集： 我们在一个合成数据集CelebA-Test和三个真实世界数据集上评估我们的方法：LFW-Test，WebPhoto-Test和我们提出的WIDER-Test。CelebA-Test包含从CelebA-HQ数据集[20]中选择的3,000张图像，其中LQ图像在与我们的训练设置相同的退化范围内合成。这三个真实世界数据集分别包含三个不同程度的退化，即轻度（LFW-Test），中度（WebPhoto-Test）和重度（WIDER-Test）。LFW-Test包含LFW数据集[17]中每个人的第一张图像，共包含1,711张图像。WebPhoto-Test [38]包含从互联网收集的407张低质量人脸。我们的WIDER-Test包含来自WIDER Face数据集[43]的970张严重受损的人脸图像，为评估盲目人脸修复方法的泛化能力和鲁棒性提供了更具挑战性的数据集。

4.2 实验设置和指标

设置： 我们将大小为512×512×3的面部图像表示为16×16的代码序列。对于所有训练阶段，我们使用批量大小为16的Adam [23]优化器。我们将学习速率设置为8×10^{-5，用于第I和II阶段，并在第III阶段采用较小的学习速率2×10}-5。三个阶段分别进行了150万、20万和2万次迭代的训练。我们的方法使用PyTorch框架实现，并使用四个NVIDIA Tesla V100 GPU进行训练。

指标： 对于在具有地面真实标签的CelebA-Test上的评估，我们采用PSNR、SSIM和LPIPS [46]作为指标。我们还使用余弦相似性评估身份保存，其中身份保存得分（IDS）和可视结果指示如Table 1所示。

4.3 与最先进的方法的比较

我们将提出的CodeFormer与最先进的方法进行比较，包括PULSE [27]，DFD-Net [24]，PSFRGAN [5]，GLEAN [3]，GFP-GAN [38]和GPEN [44]。我们在合成和真实世界数据集上进行了广泛的比较。

在合成数据集上的评估： 我们首先在CelebA-Test上展示了定量比较，如Table 1所示。在LPIPS、FID和MUSIQ等图像质量指标方面，我们的CodeFormer相较于现有方法取得了最佳分数。此外，它还在IDS和PSNR方面保持了最高的身份保存。此外，我们在Fig. 3中展示了定性比较。相比之下，其他方法未能产生令人愉悦的恢复结果，例如DFDNet [24]，PSFRGAN [5]，GFP-GAN [38]和GPEN [44]引入了明显的伪影，而GLEAN [3]产生了过度平滑的结果，缺乏面部细节。此外，所有对比方法都无法保持身份。由于表达丰富的码本先验和全局建模，CodeFormer不仅产生了高质量的面部，而且在输入受到严重损害时也能很好地保持身份。

在真实世界数据集上的评估： 如Table 2所示，我们的CodeFormer在具有轻度和中度退化的真实世界测试数据集上取得了可比较的FID分数的感知质量，并在具有重度退化的测试数据集上取得了最佳分数。虽然PULSE [27]在感知MUSIQ分数上也表现不错，但它无法保持输入图像的身份，如Table 1和Fig. 4所示。从Fig. 4中的视觉比较中可以观察到，我们的方法对真实的重度退化表现出卓越的鲁棒性，并产生大多数视觉上令人满意的结果。值得注意的是，CodeFormer成功保持了身份，并产生了具有丰富细节的自然结果。

4.4 消融研究

码本空间的有效性： 我们首先研究码本空间

的有效性。如Table 3中的Exp。 (a)所示，去除码本（即直接将编码器特征Zl馈送到解码器）导致更差的LPIPS和IDS分数。结果表明，码本的离散空间是确保我们模型的鲁棒性和有效性的关键。

基于Transformer的码本查找预测的优越性： 为验证我们基于Transformer的码本查找预测的优越性，我们将其与两种不同的解决方案进行比较，即最近邻（NN）匹配，即Exp。 (b)，以及基于CNN的码本预测模块，即Exp。 ©，该模块使用Linear层进行预测，按照编码器EL的顺序。如Table 3所示，Exps。 (b)和©的比较表明，采用码本查找的码本预测比NN特征匹配更为有效。然而，CNN的卷积操作的局部性质限制了其对长代码序列预测的建模能力。与纯CNN-based方法（即Exp。 ©）相比，我们基于Transformer的解决方案在LPIPS和IDS分数方面产生更高保真度的结果，以及在所有退化程度下更高的代码预测准确度，如Fig. 6所示。此外，CodeFormer的优越性也在Fig. 5和Fig. 9中的视觉比较中得到证明。

固定解码器的重要性： 与DFDNet [24]中的大型字典（∼3.2G）旨在存储大量面部细节的目标不同，我们故意采用了一个紧凑的码本C ∈ RN×d，其中N = 1024和d = 256，仅保留用于面部恢复的基本代码，然后激活预训练解码器中存储的详细线索。因此，码本必须与解码器一起使用，以充分发挥其潜力。为了证明我们的设计，我们进行了两项研究：1）同时固定码本和解码器，即Exp。 (g)，和2）固定码本但微调解码器，即Exp。 (e)。Table 3显示，微调解码器会导致性能下降，验证了我们的说法。这是因为微调解码器破坏了由预训练码本和解码器持有的学到的先验，导致次优性能。因此，我们在我们的方法中保持解码器固定。

**4.5 运行时间

我们比较了最先进的方法[27, 24, 5, 2, 38, 44]和提出的CodeFormer的运行时间。所有现有方法都在512^2张面部图像上使用其公开可用的代码进行评估。如Table 5所示，提出的CodeFormer的运行时间与PSFRGAN [5]和GPEN [44]相似，可以在0.1秒内推断一张图像。与此同时，我们的方法在Celeb-Test数据集上在LPIPS方面取得了最佳性能。

4.6 扩展

面部颜色增强： 我们在面部颜色增强上微调我们的模型，使用与GFP-GAN（v1）[38]相同的颜色增强（随机颜色抖动和灰度转换）。我们在具有颜色丢失的真实世界老照片（来自CelebChild-Test数据集[38]）上将我们的方法与GFP-GAN（v1）[38]进行比较。提出的CodeFormer生成具有更自然颜色和细节的高质量人脸图像。

面部修复： 提出的Codeformer可以很容易地扩展到面部修复，即使在大面具比例下也表现出色。为了构建训练对，我们使用一个公开可用的脚本[44]随机绘制不规则折线掩码以生成有掩码的面部。我们将我们的方法与两种最先进的面部修复方法CTSDG [14]和GPEN [44]以及用于码本查找的最近邻匹配进行比较。如Fig. 9所示，CTSDG和GPEN在大面具情况下表现困难。在我们的框架中使用最近邻匹配粗略地重建面部结构，但在恢复眼镜和眼睛等完整视觉部分时也失败了。相反，我们的CodeFormer生成了高质量的自然面部，没有笔画和伪影。

4.7 限制

我们的方法建立在一个带有码本的预训练自编码器上。因此，自编码器的能力和表达性可能影响我们方法的性能。1) 尽管通过Transformer的全局建模显着缓解了身份不一致问题，但在一些罕见的视觉部分（例如配饰）中仍存在不一致，其中当前的码本空间无法无缝地表示图像空间。使用码本空间中的多个尺度来探索更精细的视觉量化可能是一个解决方案。2) 尽管CodeFormer在大多数情况下表现出色，但在侧脸情况下，CodeFormer对其他方法的优势有限，也无法产生良好的结果，如Fig. 10中所示的失败案例。这是预期的，因为在FFHQ训练数据集中只有很少的侧脸，因此，码本无法学到足够的代码以处理这种情况，导致在重建和修复方面的效果较差。

5 结论

本文旨在解决盲目面部修复中的基本挑战。通过学习一个小的离散但表达丰富的码本空间，我们将面部修复转变为代码标记预测，显着降低了修复映射的不确定性，并简化了修复网络的学习。为了弥补局部信息损失，我们通过一个富有表现力的Transformer模块，探索了从受损面部获取全局构图和依赖的可能性，以更好地预测代码。由于这些设计的好处，我们的方法显示出很强的表达性和对严重退化的强大鲁棒性。为了增强我们的方法对不同退化的适应性，我们还提出了一个可控的特征转换模块，允许在保真度和质量之间进行灵活的权衡。实验结果表明了我们方法的卓越性和有效性。

论文总结

能和codeformer打的就只有gfpgan，但似乎感觉上codeformer还是更胜一筹。

脸部修复，既要考虑到不要太依赖于原来的降级图，如果太依赖很容易导致生成的脸非常丑。也要考虑到需要和原来的降级图要相似。如何在这二者之中平衡，就是codeformer所做的事情——codebook。

人脸修复所面临的问题：

在人脸修复领域，由于图像质量降低或损坏，传统的修复方法往往难以有效地还原高质量的人脸图像。特别是在缺乏高质量参考图像的情况下，传统方法无法提供令人满意的结果。本文提到的问题包括局部纹理和细节的丢失、面部结构的损坏以及对于重建的不确定性。

论文中提到的方法的新点：

论文提出了一种名为CodeFormer的新方法，该方法通过学习一个小而表达丰富的离散码本空间，将面部修复问题转化为代码标记的预测问题。与传统的方法相比，这种方法减少了修复映射的不确定性，并通过引入Transformer模块来全局建模受损面部的整体构图和依赖关系，从而更好地预测代码。

保真度

保真度权重w位于[0, 1]之间。通常，较小的w倾向于产生更高质量的结果，而较大的w则产生更高保真度的结果。

Codebook是什么：

Codebook是指学到的小而表达丰富的离散码本空间。在这个空间中，高质量的面部图像被编码为一系列离散的代码标记，这些标记在训练过程中学到，以便在修复阶段使用。Codebook的使用使得修复网络能够更有效地还原受损面部的高质量细节。

CodeFormer模型架构和其他模型的比较：

CodeFormer模型采用了一个两阶段的训练方法。第一阶段通过自重构学习预训练了一个量化自编码器，学到了码本和相应的解码器。第二阶段引入了Transformer模块，通过学习全局面部构图，更准确地预测了从低质量输入到高质量输出的代码标记。论文还介绍了一个可控的特征转换模块，允许在还原质量和还原保真度之间进行灵活的权衡。

与其他方法（如PULSE、DFD-Net、PSFRGAN、GLEAN等）相比，CodeFormer在合成和真实世界的测试数据集上表现出更好的质量、保真度和身份保持性。CodeFormer通过引入学得的码本和全局建模的Transformer模块，相对于传统的基于先验信息的方法和其他基于生成模型的方法，显著提高了面部修复的效果。

代码其他介绍

功能：

️ Whole Image Enhancement
# 对于整个图像 **Whole Image Enhancement**
# 添加'--bg_upsampler realesrgan'以使用Real-ESRGAN增强背景区域
# 添加'--face_upsample'以使用Real-ESRGAN进一步提高修复后的人脸
python inference_codeformer.py -w 0.7 --input_path [图像文件夹]|[图像路径]


 Face Colorization (cropped and aligned face)
# For cropped and aligned faces (512x512)
# Colorize black and white or faded photo
python inference_colorization.py --input_path [image folder]|[image path]



 Face Inpainting (cropped and aligned face)
# For cropped and aligned faces (512x512)
# Inputs could be masked by white brush using an image editing app (e.g., Photoshop) 
# (check out the examples in inputs/masked_faces)
python inference_inpainting.py --input_path [image folder]|[image path]

如果想自己训练，下面是训练过程，看一下有助于理解论文：

# :milky_way: 训练文档
[English](train.md) **|** [简体中文](train_CN.md)

## 准备数据集
- 下载训练数据集: [FFHQ](https://github.com/NVlabs/ffhq-dataset)

---

## 训练

###  阶段 I - VQGAN
- 训练VQGAN:
  > 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

- 训练完VQGAN后，可以通过下面代码预先获得训练数据集的密码本序列，从而加速后面阶段的训练过程:
  > python scripts/generate_latent_gt.py

- 如果你不需要训练自己的VQGAN，可以在Release v0.1.0文档中找到预训练的VQGAN (`vqgan_code1024.pth`)和对应的密码本序列 (`latent_gt_code1024.pth`): https://github.com/sczhou/CodeFormer/releases/tag/v0.1.0

###  阶段 II - CodeFormer (w=0)
- 训练密码本训练预测模块:
  > python -m torch.distributed.launch --nproc_per_node=8 --master_port=4322 basicsr/train.py -opt options/CodeFormer_stage2.yml --launcher pytorch

- 预训练CodeFormer第二阶段模型 (`codeformer_stage2.pth`)可以在Releases v0.1.0文档里下载: https://github.com/sczhou/CodeFormer/releases/tag/v0.1.0

###  阶段 III - CodeFormer (w=1)
- 训练可调模块:
  > python -m torch.distributed.launch --nproc_per_node=8 --master_port=4323 basicsr/train.py -opt options/CodeFormer_stage3.yml --launcher pytorch

- 预训练CodeFormer模型 (`codeformer.pth`)可以在Releases v0.1.0文档里下载: https://github.com/sczhou/CodeFormer/releases/tag/v0.1.0

---

:whale: 该项目是基于[BasicSR](https://github.com/XPixelGroup/BasicSR)框架搭建，有关训练、Resume等详细介绍可以查看文档: https://github.com/XPixelGroup/BasicSR/blob/master/docs/TrainTest_CN.md

推理调用