论文题目
:Resolution-robust Large Mask Inpainting with Fourier Convolutions

来源
:WACV 2022

作者
:Roman Suvorov, Elizaveta Logacheva 等

内容整理
:刘希贝

如何修复大的缺失区域、复杂的几何结构和高分辨率图像是图像修复领域的一个难点。本文作者提出,造成这种困难的主要原因之一是修复网络和损失函数都缺乏有效的感受野。为了解决这个问题,本文提出了一种称为大掩码修复 (LaMa) 的新方法,基于快速傅里叶卷积 (FFC) 的修复网络架构具有图像范围的感受野。该模的结果超过了以往的 SOTA 模型,鲁棒性更好(即使在比训练的分辨率 256x256 更高的图像上也有很好的效果),参数量和时间也更少。
目录
  • 引言
  • 网络结构
    • 快速傅里叶卷积(FFC)
    • 损失函数
    • 训练掩膜的生成
  • 实验
    • 实施细节
    • 对比实验
    • 消融性实验
    • 分辨率鲁棒性
  • 总结

引言

图像修复(Image Inpainting)问题即尽可能真实地填充图像的缺失部分。这个问题的解决分为两步,首先需要“理解”图像的大尺度结构,其次需要进行图像合成。通常的做法是在真实图像上随机添加掩膜(mask)来训练修复模型。如今效果较好的网络大多使用复杂的两阶段模型,在网络中间生成平滑图像、边缘、分割图等特征来辅助网络修复图像中缺失的部分。在这项工作中,我们仅通过一个简单的单阶段网络来完成修复过程。
一个大的有效感受野对于理解图像的全局结构并因此解决修复问题至关重要。我们注意到目前流行的卷积架构可能缺乏足够大的有效感受野,不足以覆盖生成高质量结构所需的信息。
文章的贡献主要分为三点:
  1. 提出了一个基于快速傅里叶卷积 (FFC) 的修复网络。FFC使网络即使在浅层也能拥有整个图像的感受野。研究表明,FFC 的这一特性提高了网络的感知质量和参数效率。并且,FFC 允许网络泛化到训练期间从未见过的图片分辨率,使网络具有分辨率鲁棒性。
  2. 提出使用基于具有高感受野的语义分割网络的感知损失。感受野不足不仅会损害修复网络,还会影响感知损失,因此使用高感受野的感知损失有利于促进全局结构和形状的一致性。
  3. 第三点属于数据增强的操作,使用更宽、更大的mask进行训练,迫使网络充分利用模型和损失函数的高感受野。
实验表明,LaMa 可以在仅对低分辨率数据进行训练后泛化到高分辨率图像,可以捕获和生成复杂的周期性结构,并且对大掩模具有鲁棒性。

网络结构

图像修复任务的目标是修复由二进制mask掩盖的彩色图像 中缺失的部分,在填充大块损伤或图片场景本身比较复杂时,生成适当的修复结果需要考虑全局上下文。因此,我们认为,一个好的架构应该尽早在网络结构中拥有尽可能大的感受野的单元。而传统的全卷积模型,例如 ResNet ,由于卷积核通常很小(例如 3 × 3),有效感受野增长缓慢,尤其在网络的早期层中可能存在感受野不足的问题,因此,网络中的许多层是缺乏全局上下文信息的。对于大的掩模,生成器在某些位置的感受野可能全部在掩模内部,只能观察到丢失的像素。对于高分辨率图像,这个问题更加明显。因此,一个好的架构最好能在网络中尽可能早地提供尽可能大的感受野。
LaMa网络结构
本文的网络结构如图1所示,总体来说是一个前馈的encoder-decoder网络,引入了快速傅里叶卷积(Fast Fourier convolution, FFC)算子在较早的层提取图像全局信息。掩码图像表示为 ,掩码 与掩码图像 堆叠,产生一个四通道输入张量 。我们使用前馈修复网络 ,取 ,修复网络以全卷积方式处理输入,并生成修复后的三通道彩色图像 。训练数据是由在干净的图片上添加mask合成的。

快速傅里叶卷积(FFC)

FFC是最近提出的一个能较早使用全局信息的一个算子,基于通道级快速傅里叶变换 (FFT),具有覆盖整个图像的感受野。计算过程中,FFC将输入的特征在通道上分成两个并行分支:local分支使用传统卷积;global分支使用实序列的 FFT (Real FFT)来联系全局上下文信息。
由于图像是实数序列,对实序列做FFT,得到的频谱是共轭对称的,因此计算过程中仅需使用一半的频谱信息。具体来说,FFC计算步骤如下:
  1. 将 Real FFT2d 应用于输入图像
  1. 将得到的复值的实部和虚部在通道处拼在一起
  1. 在频域中应用卷积块
  1. 应用逆变换恢复空间结构
最后,将local和global分支的输出融合在一起。FFC 的说明见图2所示。
快速傅里叶卷积运算示意图
FFC 是传统卷积的完全可微分且易于使用的替代品。由于其图像范围的感受野,FFC 允许生成器从早期层就开始考虑全局上下文,这对于高分辨率图像修复至关重要:可训练的参数可以用于推理和生成,而不是“等待”信息的传播。研究表明,FFC 非常适合捕获周期性结构,例如砖块、梯子、窗户等。

损失函数

高感受野感知损失

传统的的有监督损失要求生成器精确地重建原始图像,然而,图像的可见部分通常不包含足够的信息来精确重建被遮盖的部分。并且修复问题本质上是一个不适定问题,对于相同的缺失区域,尤其是当mask很大的时候,可能有许多看似合理的填充。因此,盲目引入L1、L2损失会导致生成结果的模糊。
相比之下,感知损失通过基础预训练网络 评估从预测图像和目标图像中提取的特征之间的距离。它不需要网络逐像素的精确的重建,允许重建图像在合理范围内进行一定的变化,使网络对大掩模修复的重点转移到对全局结构的理解上。因此,我们认为使用感受野快速增长的基础网络很重要。我们引入了高感受野感知损失(High Receptive Field Perceptual Loss, HRF PL),它使用了高感受野基础模型 :
其中 是逐元素运算, 是顺序两阶段均值运算(层内均值的层间均值)。可以用傅里叶卷积或空洞卷积来实现。
另外,我们发现,训练感知损失的基础网络的选择非常重要。使用分割模型作为感知损失的基础网络可能有助于使模型关注到图像的一些高级语义信息,比如某个物体,或是它们的一部分。相反,使用分类模型则会使网络更多地关注纹理信息,这可能会引入对高级信息有害的偏差。

对抗损失

我们使用对抗性损失来确保修复模型 生成看起来自然的局部细节,判别器定义为 。由于有监督的 HRF 感知损失的作用,生成器可以快速学习复制输入图像的已知部分,为了使模型更加专注于掩膜内的区域,对于生成器生成的结果,只有掩膜里面的patch会获得“假”标签,掩膜外的区域,也就是输入时候有效的patch得到的也是“真”标签。所以损失函数前两项计算的都是获得“真”标签的鉴别器的损失,最后一项才是获得“假:标签部分的鉴别器的损失。总的来说,也就是非饱和对抗性损失:

总的损失函数

在最终的损失函数中,我们也添加了梯度惩罚算子来稳定训练过程,以及一项特征匹配损失(feature matching loss)来约束判别器的中间层特征,记为。可以稳定训练,并且在某些情况下可以略微提高模型性能。
因此,我们的修复网络的总的损失函数为:
其中, 负责生成看起来自然的局部细节,而 负责监督信号和全局结构的一致性。

训练掩膜的生成

第三个贡献点在于设计了种类各样,更大更宽的mask用于训练。我们采用了一种激进的 large mask 生成策略,随机生成遮挡面积较大的 wide 或 box 遮挡,如图3所示。实验表明,使用更大更宽的mask训练,测试的时候无论小mask还是大mask都可以达到一个比较好的效果。
来自不同训练 mask 生成策略的样例

实验

实施细节

对于 LaMa 修复网络,我们使用具有 3 个下采样块、6-18 个残差块和 3 个上采样块的类 ResNet 架构。在我们的模型中,残差块使用 FFC。训练过程中使用 Adam 优化器,修复网络和鉴别器的学习率分别固定为 0.001 和 0.0001。模型经过 1M 次迭代训练,batch size大小为 30,超参数 。模型在 Places 和 CelebA-HQ 数据集上进行训练,使用 PyTorch、PyTorchLightning 和 Hydra 架构实现。

对比实验

我们将提出的方法与下表中列出的方法进行比较。对于每个数据集,我们分别验证了窄掩膜、宽掩膜和基于分割的掩膜的性能。可以看出,LaMa-Fourier 的结果始终优于大多数模型,而参数则比其他效果最好的模型小很多,在大掩膜的情况下,LaMa的优势更加明显。

消融性实验

该表展示了不同 LaMa 架构的性能,而其他组件保持不变。▲ 表示恶化,▼ 表示与 Base-Fourier 模型(在第一行中显示)相比的改进。基于 FFC 的模型可能会在窄掩膜上牺牲一点性能,但明显优于在宽掩膜上进行常规卷积的较大模型。在视觉上,基于 FFC 的模型可以更好地恢复复杂的视觉结构,如下图所示。可以看到 FFCs 能够很好的生成重复的结构,例如玻璃和铁丝网。
其中,LaMa-Regular表示使用常规卷积组成的网络,LaMa-Fourier表示使用了FFC block的网络。

分辨率鲁棒性

模型直接在高分辨率下训练速度慢且计算量大,尽管如此,大多数现实世界的图像编辑场景都需要修复高分辨率的图像。因此,我们在更大的图像上评估我们的模型,这些模型使用将 512 × 512 图像随机裁剪出 256 × 256 的图像进行训练。我们以完全卷积的方式应用模型,即图像是单次处理的,而不是分块处理的,结果如下所示。随着分辨率的提高,基于传统卷积的模型开始产生致命的伪影,而基于 FFC 的模型可以继续生成精细的语义一致的图像。
将修复模型迁移到更高分辨率的图像上
其中,Big LaMa-Fourier 在三个方面与 LaMa-Fourier 不同:生成器的深度;训练数据集;和批次的大小。它有 18 个残差块(LaMa-Fourier有 9 个残差块),全部基于 FFC。该模型是在来自 Places-Challenge 数据集 的 450 万张图像的子集上进行训练的,batch size设为120(而不是我们其他模型的 30)。

总结

在这项研究中,我们提出了LaMa这个单阶段的图像修复模型,该方法在生成重复的视觉结构方面效果很好。然而,当图像中涉及到强烈的透视扭曲时,LaMa 通常会遇到困难。FFCs 是否可以解释周期性信号的这些形变仍然是一个值得研究的问题。有趣的是,FFC 允许该方法推广到模型训练时从未见过的高分辨率,并且与最先进的基线相比,参数效率更高。傅立叶卷积或空洞卷积并不是接收高感受野的唯一选择,例如,使用Transformer可以也获得高感受野,这也是未来研究的一个热门主题。我们相信,具有大感受野的模型将为开发高效的高分辨率计算机视觉模型开辟新的机遇。
附上论文链接:https://arxiv.org/pdf/2109.07161.pdf
附上github链接:https://github.com/saic-mdal/lama
继续阅读
阅读原文