南洋理工大学 S-Lab 打出一套「还你漂漂拳」CodeFormer开源：SOTA人脸复原能力

原文链接：https://arxiv.org/pdf/2206.11253.pdf

project：https://shangchenzhou.com/projects/CodeFormer/

code：https://github.com/sczhou/CodeFormer

本文研究表明：当把盲人脸复原视作码本预测任务时，可学习离散码本先验(Learned Discrete Codebook Prior)可以极大降低学习的不确定性与混淆，同时为生成高质量人脸提供丰富的视觉原子。基于该框架，本文提出了一种基于Transformer的预测网络CodeFormer，它可以对低质人脸的组成与上下文信息进行建模，促进自然人脸挖掘以生成与目标人脸更相近的结果。为增强不同退化的自适应性，本文还提出了一种可控特征变换模块以寻求真实(fidelity)与高质(quality)之间的均衡。受益于码本先验与全局建模能力，CodeFormer取得了SOTA人脸复原能力。

本文方法

本文主要聚焦于对离散表达空间进行探索以降低复原映射的不确定性、为退化输入提供质量细节补偿。由于局部纹理与细节的缺失、低质输入的干扰，作者采用Transformer模块进行全局上下文建模以促进高质量人脸重建。上图给出了本文方案整体架构示意图，很明显，它是在VQ-VAE、VQ-GAN基础上演变而来。

Stage I : Codebook Learning

为降低LQ-HQ映射之间的不确定性，同时提供高质量细节补偿，我们首先预训练一个量化自编码器以学习具有丰富上下文信息的码本，它同时有助于提升网络表达能力以及鲁棒性。

如前述图a所示，HQ人脸图像 首先通过编码器 被编码为特征 ，参考VQVAE与VQGAN机制：对 的每个特征以最近邻方式从码本 中选择最相似特征进行替换以得到量化特征 。该过程可以描述成如下公式：

所得量化特征将被送入到解码器 中进行高质量人脸图像 重建。该阶段的主要是进行码本与解码器的训练，涉及两个维度的损失函数：图像层面与码本层面。

图像层面的损失函数由L1、感知损失以及对抗损失构成，描述如下：

码本层面的损失函数描述如下（ ）：

整体损失定义如下( )：

Stage II : Codebook Lookup Transformer Learning

受限于LQ人脸图像的退化问题，常规的最近邻匹配通常无法精确的定位到用于复原的码本信息。为缓解该问题，本文采用Transformer进行全局相关性建模以达成更好的code预测。

如前面图b所示，本文在Encoder后插入一个包含9个自注意力模块的Transformer。此时，我们固定码本 与解码器 并对 进行微调，我们将微调后的编码器表示为 。我们首先通过编码器 对LQ输入图像进行处理得到LQ特征 ，然后将其reshape为 并送入到Transformer模块中：

其中， 。此外，我们还在Q与K上添加位置嵌入 以提升表达能力。

总而言之，以 作为输入，Transformer以N路分类方式预测每个位置特征对应的code，进而构成量化特征 ，最终通过固定解码器处理得到高质人脸图像。

在该阶段，我们需要对Transformer进行训练，同时对编码器

此外，我们还需要为 特征引入码本约束损失（其实就是StageI阶段码本损失的一部分）：

最终，整体损失定义如下( )：

该部分是全文最核心的内容，不同于已有方案采用最近邻匹配方案，本文采用了分类机制进行最相似code匹配。受限于输入图像的低质特征，此时通过最近邻匹配方式得到的特征误差较大，进而导致重建效果次优；而本文则利用利用Transformer的全局建模特性缓解特征误差，同时将最近邻匹配机制替换为分类机制，进一步离散化。

此外，需要注意的是：为了确保量化特征与低质特征之间的色彩与亮度一致性，作者还引入AdaIN对量化特征进行调制。参考代码如下：

def adaptive_instance_normalization(content, style):
    style_mean, style_std = calc_mean_std(style)
    content_mean, content_std = calc_mean_std(content)
    normlized_feat = (content - content_mean) / content_std
    return normalized_feat * style_std + style_mean

Stage III : Controllable Feature Transformation

尽管Stage II方案已可以取得非常好的人脸复原，我们同时还希望在人脸复原的高质与真实性方面达成更灵活的均衡。因此，本文提出一种可控特征变换模块以控制从编码器 到解码器 的信息流。具体来说，通过SFT模块的参数 ，LQ特征 可以用于对解码特征 进行微调，该调制过程描述如下：

我们在多个尺度 采用上述CFT模块对编码特征与解码特征进行调制。通过这种方式，我们可以使得模型对于轻度退化具有高真实复原，对于重度退化具有高质复原。具体来说，我们可以调小 w 以降低对输入的依赖性，进而具有高质输出；调大 w 以引入更多LQ特征进而增强轻度退化时的高真实性。

在训练阶段，我们设置w=1，这使得我们可以在推理阶段通过对其在 间调整达成连续过度。如无特殊说明，在推理阶段我们默认设置w=0以提升高质输出。

该部分的code实现就比较简单了，参考如下：

class SFTBlock(nn.Module):
    def __init__(self, channels):
        super().__init__()
        self.fuse = ResBlock(channels * 2, channels, kernel_size=3)
        self.scale = nn.Sequential(
          nn.Conv2d(channels, channels, kernel_size=3, padding=1),
          nn.LeakyReLU(0.2, True),
          nn.Conv2d(channels, channels, kernel_size=3, padding=1)
        )
        self.shift = nn.Sequential(
          nn.Conv2d(channels, channels, kernel_size=3, padding=1),
          nn.LeakyReLU(0.2, True),
          nn.Conv2d(channels, channels, kernel_size=3, padding=1)
        )
    def forward(self, enc, dec, w=1):
        fuse = self.fuse(torch.cat([enc, dec], dim=1))
        scale = self.scale(fuse)
        shift = self.shift(fuse)
        residual = w * (dec * scale + shift)
        return dec + residual

本文实验

CodeFormer在人脸复原、人脸颜色增强以及人脸补全三个任务上均表现出了最优的效果。同时该方法也可以用于老照片和老电影的人脸增强。

在线Demo

作者还提供了一个可在线测试的Colab Demo(https://colab.research.google.com/drive/1m52PNveE4PBhYrecj34cnpEeiHcC5LTb "Colab Demo")(当然，这个需要梯子才能体验)，大家可以进去测试自己的照片哦~

概要总结

为更好理解CodeFormer，笔者结合源码绘制了上图。结合前述介绍对全文汇总如下：

首先，我们采用HQ人脸图像作为输入，采用自监督方式对Encoder、Dictionary以及Decoder进行训练。此时Dictionary部分的训练与VQVAE、VQGAN相同，Encoder与Decoder之间不存在跳过连接；

然后，基于上述预训练模型，我们固化Decoder与Dictionary，在Encoder后插入用于分类的Transformer模块。以LQ-HQ数据对作为输入，对Encoder进行微调，Transformer部分从头开始训练；

最后，基于步骤二的预训练模型，在Encoder与Decoder之间引入跳过连接与CFT模块（虚线），达成可调制功效。此时，仍然保持Dictionary与Decoder固化，其他部分部分微调，新加模块重头开始训练。

公众号后台回复“ECCV2022”获取论文分类合集下载～