论文:https://arxiv.org/abs/2203.00867
主页:https://dqiaole.github.io/ZITS_inpainting
代码(已开源):
https://github.com/DQiaole/ZITS_inpainting
导读
近年来,图像修复取得了重大进展。然而,恢复具有生动纹理和合理结构的损坏图像仍然具有挑战性。由于卷积神经网络 (CNN) 的感受野有限,一些特定的方法只能恢复局部纹理,且会丢失图像的整体纹理结构。而基于注意力的模型可以有效地建模图像局部之间的长期依赖,从而得到更好的全局修复能力,但它们大都受到计算量的限制。为了解决这些问题,本文提出一个增量结构增强型修复模型,ZITS。与以前的方法相比,它使用一个基于 Transformer 的结构恢复器来获得更好的整体结构。然后,利用一种新的 ZeroRA 策略,通过几个微调步骤将辅助结构纳入预训练的inpainting 模型。实验证明,所提出的 masking positional encoding 可以进一步提高修复性能。ZITS 可以在不同分辨率的实验中实现基于最先进模型的显著改进。
贡献
图像修复旨在填充图像中的缺失区域并保持填充区域的结构合理性以及纹理一致性。得益于最近卷积神经网络以积极生成对抗网络的发展,图像修复模型在某些场景上已经可以表现得非常好了。然而,现存的模型仍然面临如下困境:
- 有限的感受野:受限于卷积的局部归纳偏置以及有限的感受野,利用传统的CNN学习语义相合的纹理是非常困难的。面对大的掩蔽区域或者高分辨图像,即使是膨胀卷积也会失效。
- 缺少对整体结构的理解:如果没有对图像的整体结构的理解,在高分辨弱纹理场景中恢复关键的边缘和线框是非常困难的。
- 大的计算量:在高分辨图像上训练生成对抗网络需要非常多的技巧,并且遇到更高分辨率的图像,模型的修复性能也会退化的非常严重。
- 缺少对掩蔽区域的位置信息:如果缺少掩蔽区域的位置信息,模型在大的掩蔽区域内部往往会倾向于生成重复的无意义的纹理。
这启发了我们为图像修复网络增量式地注入整体结构信息和掩蔽区域位置信息。具体来说,我们利用了一个基于Transformer的模型推断整体结构(边缘+线框)。利用简单的CNN,我们可以把这种结构图非常轻松的上采样到任意分辨率。进一步地,我们提出了利用零初始化残差连接技术增量式地将结构注入到后续CNN纹理修复网络。相比于从头训一个基于结构辅助信息的图像修复模型,这种增量式的技巧只需要继续训练几千步就可以达到非常好的性能。最后,为了提高模型对大掩蔽区域的修复性能,我们还对掩蔽区域引入了位置编码。
方法
ZITS的模型整体架构如图1所示。首先利用一个基于普通注意力和轴向注意力机制的Transformer修复低分辩图像的结构信息(边缘+线框),修复的结构信息可以用一个简单的CNN网络进行上采样到与输入图片相同的分布率。利用零初始化残差连接技术,可以通过一个编码器抽取结构信息的特征并增量式的加到预训练好的CNN纹理修复模型上,最后再简单的微调模型即可。此外,本文提出的针对掩蔽区域的位置编码也可以通过零初始化残差连接技术增量式的加到预训练好的CNN网络中。在本文中,作者选取了基于快速傅里叶卷积的LaMa作为我们的预训练CNN纹理修复网络(FTR),得益于傅里叶变换的性质,FTR可以拥有全图的感受野,从而极大的提升了图像修复的效果。
Transformer Structure Restoration
在最近的研究中,Transformer结构展现出对全局结构的较强恢复能力,因此,本文选择利用Transformer以相对较低的分辨率恢复整体结构。给定尺寸为256*256的输入图像Im,边缘Ie,线条Il,以及遮挡掩模M,作者首先使用三个卷积层进行下采样,并将位置编码与序列化的下采样特征进行拼接。
如上图左上角所示,作者选择交替使用轴注意力机制和标准注意力机制进行结构修复。通过调整行和列的张量形状,然后分别用基于点积的自注意力处理它们,可以轻松实现轴向注意力模块。对特征X,作者将其i列或i行的平均值表示为Xri或Xci,并使用相对位置编码强化轴注意力机制的计算:
其中W代表可学习参数。通过迭代使用基于注意力机制的Transformer模块进行建模,ZITS利用如下函数约束输入-输出之间的结构相似度:
Simple Structure Upsampler
为了捕获可能的高分辨率图像的整体结构,作者提出将生成的边缘和线条上采样到任意比例,而不会出现明显的退化。 然而,基于传统插值的上采样算法会导致锯齿形,如图 3(f)-(i) 所示。 这种伪影对于大图像尺寸更为严重,这会恶化修复结果。因此,作者提出针对图像边界训练一个卷积神经网络,并利用该网络进行任意尺度的上采样。
Fourier CNN Texture Restoration (FTR)
本文采用抑郁研究中的Fourier卷积构建backbone。FTR 是一种自动编码器模型,具有多个卷积用于下采样和上采样图像特征。FTR 的关键模块是快速傅里叶卷积(FFC)层,它由两个分支组成:1)局部分支使用常规卷积;2)全局分支对快速傅里叶变换后的特征进行卷积。 然后在修复过程中将两个分支组合起来以获得更大的感受野和局部不变性。 然而,如此强大的模型无法学习合理的整体结构。 我们进一步提出了一系列新颖的组件来改进它。
Structure Feature Encoder (SFE)
为了恢复特征中的上下文信息,作者利用一个特定的全卷积模块进行特征建模。SFE 也是一个自动编码器模型,具有 3 层下采样卷积(编码器)、3 层带有扩张卷积的残差块 [56](中)和 3 层上采样卷积(解码器)。 对于 SFE 中的编码器和解码器,门控卷积 (GC) 被用于自适应地传输有用的特征。 GC 学习另一个具有相同通道的 sig moid 激活。 然后将 sigmoid 特征乘以卷积特征作为输出。因为灰度草图空间是稀疏的,并不是所有的特征都是修复所必需的,因此,作者提出对SFE中多层次的特征进行输出:
Masking Positional Encoding (MPE)
已有研究大多认为,在修复过程中,已知区域的位置信息是不必要的。 然而,作者认为位置信息对于遮罩区域仍然至关重要,尤其是当遮罩区域对于高分辨率图像来说很大时。 受限于 CNN 的感受野,具有大 mask 的模型可能会丢失方向和位置信息,从而导致无意义的伪影。 虽然 FFC 可以将特征学习扩展到频域,但它对区分屏蔽区域或未屏蔽区域并不敏感。 因此,本文建议在称为 MPE 的遮蔽区域中使用位置编码来进行图像修复,这与 FTR 中的 FFC 正交并对其进行了改进。
实验
下表对比了ZITS与其他模型在Indoor和Places2俩个数据集上的结果,图片分辨率是 256 x 256 。可以看到我们模型的效果显著优于其他模型。
作者还在分辨率为 512 x 512 以及 1024 x 1024 的数据集上做了测试,实验结果如下表所示:
在图5和图6中,作者提供了模型在Indoor和Places2上的定性结果(图片分辨率为 256 x 256),可以看到本文模型对整体结构的恢复更加合理,且面对大掩蔽面积的图像表现效果也显著优于其他模型。
作者还展示了模型在1K图片上的表现,面对高分辨的图像,模型的修复结果也非常不错。
内容中包含的图片若涉及版权问题,请及时与我们联系删除
评论
沙发等你来抢