基于傅里叶卷积的鲁棒分辨率大Mask修补
WACV 2022
论文链接
代码链接
摘要: 现代图像修复技术主要受阻于大缺失区域、复杂几何结构和高分辨率图像,主要受阻原因是修复网络和损失函数都缺乏有效的感受野。为缓解该问题,本文提出了一个新方法:large mask inpainting (LaMa)。LaMa基于:
- 一种新的使用快速傅里叶卷积的修补网络,具有图像范围感受野,傅里叶卷积可以视为self-attention的轻量级替代。
- 一个高感受野感知loss。
- large training masks,能发挥前两个组件的潜力。
我们的修补网络提高了一系列数据集的sota,在挑战性场景(如周期性结构)下也具有出色性能,并能以比baseline更低的参数量和时间成本泛化到比训练集分辨率更高的图像上。
文章目录
- 基于傅里叶卷积的鲁棒分辨率大Mask修补
- 1. Introduction
- 2. Method
- 2.1. Global context within early layers
- 2.2. Loss functions
- 2.2.1 High receptive field perceptual loss
- 2.2.2 Adversarial loss
- 2.2.3 The final loss function
- 2.3. Generation of masks during training
1. Introduction
图像修复需要真实填充缺失部分,既需要“理解”自然图像的大规模结构,也需要执行图像合成。基于神经网络的图像修复通常在一个大型随机mask图像的数据集上训练两阶段网络。本工作使用简单的one-stage 网络实现了sota。
大感受野对于理解图像的全局结构并解决修复问题至关重要。 在大mask的情况下,大但有限的感受野可能都不足以访问生成修复的图像。当前主流的架构缺乏大感受野,因此我们干预其每个组件来缓解问题,并发挥 one-satge 解决方案的潜力。具体而言:
- 我们基于快速傅里叶卷积(FFCs)提出了一个修复模型,FFCs甚至允许网络在early layers覆盖整幅图像的感受野,这一特性提高了网络的感知质量和参数效率。 有趣的是,FFC的 inductive bias使网络可以泛化到训练期间从未见过的更高分辨率图像上(图5,图6),有效减少了训练数据和计算。
- 我们提出基于高感受野语义分割网络的感知loss。感受野不足不仅会损害修复网络,而且也损害感知loss,我们的loss促进了全局结构和形状的一致性。
- 我们引入了一种训练mask生成的策略以发挥前两个组件高感受野的潜力。生成宽大的mask迫使网络充分利用模型和损失函数的高感受野。
因此,LaMa的主要组成部分是:1.高感受野结构、2. 高感受野损失和 3. 训练mask生成的算法。通过评估,我们发现 LaMa 仅在低分辨率数据上训练,就能推广到高分辨率图像上。LaMa可以捕获和生成复杂的周期结构,并对大mask具有鲁棒性。此外,LaMa具有比baseline更低的训练参数和推理时间。
2. Method
我们的目标是修复一张被一个 pixel 未知的二进制掩码 m 掩盖住的彩色图像 x,被mask的图像表示为:x⊙m 。m与x⊙m堆叠,因此输入的是4通道tensor:x′ = stack(x ⊙ m, m),我们使用 feed-forward 修补网络 fθ(·),也称之为生成器,以完全卷积的方式处理x′ ,生成修复的三通道彩色图像 x ^ = f θ ( x ′ ) \hat{x}=fθ(x^′) x^=fθ(x′)。在真实图像和合成生成的mask 组成的 (image, mask) 对上进行训练。
2.1. Global context within early layers
正确修复大mask需要考虑全局上下文。因此,一个好的修复架构在pipeline中应尽早具有尽可能宽的感受野。传统的完全卷积模型,如ResNet,有效感受野增长缓慢,由于卷积核通常较小(例如3×3),感受野可能不足,尤其在网络的 early layers,因此,网络中的许多层会缺乏全局上下文信息,需要浪费计算和参数进行创建。对于宽mask,在特定位置生成器的感受野可能在mask内,因此只能观察到缺失的像素。
Fast Fourier convolution (FFC)。FFC允许网络在 early layers 使用全局上下文,FFC基于一个channel-wise fast Fourier transform (FFT),并具有覆盖整幅图像的感受野。FFC将 channel 分成两个并行分支:i)使用传统卷积的局部分支,ii)使用 real FFT 来获取全局上下文的全局分支。 real FFT 只能应用于实信号,inverse real FFT 确保输出是实值,与FFT相比,real FFT只使用一半的频谱。具体而言,FFC采取以下步骤:
a) 对一个输入 tensor 应用Real FFT2d,并 concatenate 实部和虚部:
b) 在频域中应用一个卷积block:
c) 应用逆变换恢复空间结构:
最后,局部(i)和全局(ii)分支的输出融合在一起。FFC如图2所示。
The power of FFCs。 FFC完全可微且易于使用,可直接替换传统卷积。由于 image-wide 感受野,FFC允许生成器在 early layers 考虑全局上下文,这对高分辨率图像修复至关重要,也提高了效率:可训练的参数可以用于推理和生成,而不是“等待”信息的传播。
FFC非常适合捕捉周期性结构,这种周期性结构常见于人造环境,例如砖块、梯子、窗户等(图4)。有趣的是,在所有频率上共享相同卷积使模型尺度等变(图5、6)。
2.2. Loss functions
修复问题本身就模棱两可,同样的缺失区域有许多似乎可行的填充物,尤其当“洞”变宽时。
2.2.1 High receptive field perceptual loss
普通监督损失需要生成器准确重建GT,但图像的可见部分通常不包含足够的信息来精确重建被 mask的部分,因此,由于修复内容的多可能性模式的均值,使用普通监督会导致结果模糊。
相反,感知损失通过预训练网络 ϕ(·) 评估预测图像特征和目标图像特征间的距离,无需精确重建,允许重建图像的变化。大 Mask 修复的重点转向了理解全局结构,因此使用感受野快速增长的基础网络很重要。我们引入了高感受野感知损失:high receptive field perceptual loss (HRF PL),使用一个高感受野的基础模型ϕHRF(·) :
[
⋅
−
⋅
]
2
[· − ·]^2
[⋅−⋅]2 是一个元素级别操作,M是连续的两阶段平均操作(interlayer mean of intra-layer means)。可以使用傅里叶卷积或空洞卷积来实现ϕHRF(x),消融实验表明 HRF感知损失对我们的大mask 修复系统至关重要(表3)。
Pretext problem。 训练感知损失的基础网络的 Pretext problem 很重要,例如,使用分割模型作为感知损失的backbone有助于关注高级语义信息,例如对象及其部分,而分类模型更关注纹理信息,可能会引入不利于高级信息的偏差。
2.2.2 Adversarial loss
我们使用对抗损失来确保修复模型fθ(x′) 生成自然的外观局部细节,我们定义一个局部 patch-level 鉴别器 Dξ(·) 来区分“real” 和 “fake”patch,只有与 mask 区域相交的 patch 才得到“fake”标签。由于有监督的HRF感知损失,生成器很快学会复制输入图像的已知部分,因此我们将生成图像的已知部分标记为“real”。最后,我们使用 non-saturating adversarial loss:
x 来自数据集的样本,m是合成生成的mask,
x
^
=
f
θ
(
x
′
)
\hat{x}=f_θ(x^′)
x^=fθ(x′)是x′=stack(x⊙m, m) 的修补结果,
s
g
v
a
r
sg_{var}
sgvar停止关于 var 的梯度,LAdv是要优化的 joint loss。
2.2.3 The final loss function
final loss 还使用
R
1
=
E
x
∣
∣
∇
D
ξ
(
x
)
∣
∣
2
R_1=E_x||∇D_ξ(x)||^2
R1=Ex∣∣∇Dξ(x)∣∣2 梯度惩罚,以及一个discriminator-based perceptual loss,或所谓的特征匹配损失——鉴别器网络LDiscPL的特征感知损失,LDiscPL训练稳定,某些情况下可以稍微提高性能。最终的loss如下:
LAdv和LDiscPL负责生成自然的外观局部细节,而LHRFPL负责监督信号和全局结构的一致性。
2.3. Generation of masks during training
我们系统的最后一个组件是 mask 生成策略,每个训练样本 x′ 都是真是照片与合成生成的mask的叠加。与数据增强对最终性能有很大影响的 discriminative 模型类似,mask 生成策略修复系统的性能影响很大。
因此,我们选择了一种 aggressive large mask 生成策略,统一使用由一个随机宽度(wide masks)和任意纵横比的矩形(box masks)扩展的多边形链的样本。mask示例如图3所示。使用large mask生成策略进行训练可以提高窄mask和宽mask的性能(表4),表明增加mask的多样性对修复系统有利。
论文剩余部分略,不关注