作者:Jiahui Yu 1 ^1 1 Zhe Lin 2 ^2 2 Jimei Yang 2 ^2 2 Xiaohui Shen 2 ^2 2 Xin Lu 2 ^2 2 Thomas S. Huang 1 ^1 1
单位: 1 ^1 1University of Illinois at Urbana-Champaign 2 ^2 2Adobe Research
会议:CVPR 2018
最近的基于深度学习的方法已经显示出了很有前景的结果,用于填补图像中的大缺失区域的挑战性任务。这些方法可以生成视觉上可信的图像结构和纹理,但往往会产生与周围区域不一致的扭曲结构或模糊纹理。这主要是由于卷积神经网络不能从遥远的空间位置显式地借用或复制信息。另一方面,当需要从周围区域借用纹理时,传统的纹理和patch合成方法特别适合。基于这些观察,我们提出了一种新的基于深度生成模型的方法,该方法不仅可以合成新的图像结构,而且在网络训练中会明确地利用周围的图像特征作为参考,从而做出更好的预测。该模型是一种前馈、全卷积神经网络,可以在测试时间内处理任意位置和不同尺寸的多hole图像。在人脸(CelebA, CelebA- hq)、纹理(DTD)和自然图像(ImageNet, Places2)等多个数据集上的实验表明,本文提出的方法比现有方法产生了更高质量的填充结果。代码,demo和模型在https://github.com/JiahuiYu/generative_inpainting。
不幸的是,这些基于CNN的方法经常会产生与周围区域不一致的边界伪影、扭曲的结构和模糊的纹理。我们发现,这可能是由于卷积神经网络在 建模(model) 远程上下文信息(远离缺失区域的上下文信息)和缺失区域之间的长期相关性(long-term correlations)时变得无效。例如,要允许一个像素受64个像素之外的内容的影响,它至少需要6层3 × 3的卷积,其膨胀系数为2或等效于[15,40]。然而,膨胀卷积从规则和对称网格(regular and symmetric grid)中采样特征,因此可能无法权衡感兴趣的特征。注意,最近的一个工作[38]试图通过优化 生成的patch和已知区域中匹配的patch 之间的纹理相似性来解决外观差异。虽然提高了视觉质量,但该方法需要数百次梯度下降迭代,在GPU上处理分辨率为512 × 512的图像需要花费几分钟时间。
我们为图像修复提出了一个统一的前馈生成网络与一个新颖的上下文注意层(a novel contextual attention layer)。我们提出的网络包括两个阶段。第一阶段是简单的膨胀卷积网络,用重构损失训练出缺失内容。第二阶段协调了上下文的注意(contextual attention)。contextual attention的核心思想是利用已知patch的特征作为卷积过滤器来处理生成的patch。通过卷积将生成的patch与已知的上下文patch进行匹配,通道级(channel-wise)的softmax对相关patch进行加权,依据上下文patch通过反卷积将生成的patch进行重构。contextual attention模块还具有空间传播层,以促进attention的空间一致性。为了让网络产生新内容,我们有与contextual attention平行的另一个卷积通道。这两条路径被聚合并馈入单个解码器以获得最终输出。对整个网络进行从端到端的训练,其中有重构损失和两个Wasserstein GAN损失[1,12],其中一个loss关注全局图像,而另一个则关注缺失区域的局部patch。
我们的贡献总结如下:
我们首先通过对最新的修复模型[15]进行复制和改进,构建了我们的基线生成图像修复网络,[15]在修复人脸、建筑和自然图像方面显示了很好的视觉效果。
Coarse-to-fine network architecture 改进后的模型的网络架构如图2所示。我们按照[15]中相同的输入输出配置进行训练和测试,即生成网络以一张带有白色像素hole的图像和一个表示hole区域的二进制掩码作为输入对,输出为修复完成的图像。我们将输入与相应的二进制掩模配对,以处理大小、形状和位置不同的hole。网络的输入是一幅256 × 256的图像,在训练过程中随机采样一个矩形缺失区域,然后训练好的模型的输入可以为带有多个hole的不同大小的图像。
在图像修复任务中,接受域的大小应该足够大,Iizuka et al.[15]为此采用了膨胀卷积。为了进一步扩大接受域以及稳定训练,我们引入了两阶段的粗到细网络结构,其中第一个网络进行初始粗预测,第二个网络以粗预测作为输入,预测精细化结果。粗网络采用显式的重构损失训练,精化网络采用重构损失和GAN损失进行训练。从直观上看,细化网络可以看到比缺少区域的原始图像更完整的场景,因此其编码器会学到比粗网络更好的特征表示。这种两阶段网络架构类似于残差学习[14]或深度监督[22]。
此外,为了提高效率,我们的修复网络采用了thin and deep(薄和深的)方案,比[15]的修复网络参数更少。在层实现方面,我们对所有卷积层使用镜像填充(mirror padding),并删除批处理归一化层[16] (我们发现它会破坏颜色一致性)。此外,我们使用ELUs[7]作为激活函数,而不是在[15]中使用的ReLU,并裁剪输出滤波器的值,而不是使用tanh或sigmoid函数。此外,我们发现在[15]中,分离全局和局部特征表示的GAN训练效果比特征拼接更好。更多细节可以在补充材料中找到。
**Global and local Wasserstein GANs ** 受[15]启发,我们将WGAN-GP损失附加到第二阶段精化网络的全局和局部输出,以加强全局和局部一致性。众所周知,WGAN-GP损失在图像生成任务中优于现有的GAN损失,当它与 ℓ1 重建损失结合使用时效果很好,因为它们都使用 ℓ1 距离度量。
具体来说,WGAN 使用 Earth-Mover 距离(又名 Wasserstein-1) W ( P r , P g ) W(P_r,P_g) W(Pr,Pg) 来比较生成的和真实的数据分布。 它的目标函数是通过应用 Kantorovich-Rubinstein 对偶来构建的:
其中,D为1-Lipschitz函数集, P g P_g Pg为属于 x ~ = G ( z ) \tilde{x} = G(z) x~=G(z)的模型分布。z为生成器的输入。
由于这两种损失(重构损失和对抗损失)都是衡量像素的ℓ1 distance,因此组合损失更容易训练,并使优化过程更稳定。
Spatially discounted reconstruction loss 修复问题涉及像素的生成,所以在任何上下文的情况下,它都可能有许多看似合理的解决方案。在这种具有挑战性的情况下,一个貌似可信的完整图像可能有与原始图像非常不同的patches或像素。由于我们使用原始图像作为唯一的ground truth来计算重构损失,在这些像素上进行很强的重构损失可能会误导卷积网络的训练过程。
直观地说,hole边界附近的缺失像素比那些接近hole中心的像素有更少的模糊性。这与在强化学习中观察到的问题相似。当长期(long-term)奖励在采样过程中有很大的变化时,人们会在使用采样轨迹上使用temporal discounted rewards[36]。受此启发,我们使用权重掩码M引入spatially discounted重构损失。掩模中每个像素的权重计算为 γ l γ^l γl,其中 l l l为像素到最近的已知像素的距离。 γ γ γ在所有实验中均设为0.99。
[30,39]也探讨了类似的权重思想。在[39]中提出了重要度加权的上下文损失,在空间上由固定窗口(例如7×7)内未缺损像素的比率加权。Pathak等人[30]预测在边界区域有一个略大且损失权重较高的patch(×10)。对于大hole修复,提出的discounted损失对提高视觉质量更有效。我们在实现中使用discounted的ℓ1 重构loss。
有了以上所有的改进,我们的baseline生成修复模型收敛得比[15]快得多,结果是更准确。对于Places2[41],我们将训练时间从[15]报告的11,520 gpu小时(K80)减少到120 gpu小时(GTX 1080),这几乎是100倍的加速。此外,后期处理步骤(图像混合)[15]也不再需要。
卷积神经网络利用局部卷积核逐层处理图像特征,因此对于从遥远的空间位置借用特征并不有效。为了克服这一局限性,我们考虑了注意机制,在深层生成网络中引入了新的上下文注意层。在本节中,我们首先讨论上下文注意层的细节,然后讨论如何将其集成到统一的修复网络中。
上下文注意层学习从已知的背景patch中借用或复制特征信息以生成缺失的patch。它是可微分的,因此可以在深度模型中训练,并允许在任意分辨率上进行测试。
Match and attend 我们考虑的问题是,我们想要匹配缺失像素(前景)的特征到环境(背景)。如图3所示,我们首先从背景中提取3 × 3的patch,并将其reshape为卷积滤波器。为了匹配前景patch { f x , y } \{f_{x,y}\} {fx,y}和背景patch { b x ′ , y ′ } \{b_{x',y'}\} {bx′,y′},我们使用归一化内积(余弦相似度)进行度量:
其中 s x , y , x ′ , y ′ s_{x,y,x',y'} sx,y,x′,y′表示以背景 ( x ′ , y ′ ) (x',y') (x′,y′)和前景 ( x , y ) (x, y) (x,y)为中心的patch的相似性。然后我们用缩放的softmax沿着 x ′ y ′ x'y' x′y′维度衡量相似性,以获得每个像素的注意力评分 s x , y , x ′ , y ′ ∗ = s o f t m a x x ′ , y ′ ( λ s x , y , x ′ , y ′ ) s^*_{x,y,x',y'}= softmax_{x',y'} (λs_{x,y,x',y'}) sx,y,x′,y′∗=softmaxx′,y′(λsx,y,x′,y′),其中λ是一个常量值。这被有效地实现为卷积和逐通道的softmax。最后,我们再利用提取的patch { b x ′ , y ′ } \{b_{x',y'}\} {bx′,y′}作为反卷积滤波器来重构前景。重叠像素的值被平均。
**Attention propagation ** 我们进一步通过传播(融合)来鼓励注意力的一致性。一致性的概念是,在注意力上一个前景patch的移动很可能对应于一个相等的背景patch的移动。例如, s x , y , x ′ , y ′ ∗ s^*_{x,y,x',y'} sx,y,x′,y′∗通常与 s x + 1 , y , x ′ + 1 , y ′ ∗ s^*_{x+1,y,x'+1,y'} sx+1,y,x′+1,y′∗具有近似值。为了建模和鼓励注意力映射的一致性,我们先进行从左至右传播(left-right propagation),然后进行核大小为k的自顶向下传播(top-down propagation)。以left-right传播为例,我们得到新的注意分数:
该方法以单位矩阵为核,利用卷积的方法有效地实现了传播。注意力传播能显著提高修复效果,丰富训练的梯度。
Memory efficiency 假设在一张128 × 128的feature map中缺失了一个64 × 64的region,那么从背景中提取卷积滤波器的个数为12288。这可能会导致gpu的内存不够用。为了克服这一问题,我们提出了两种方法:1)用strides提取背景patch以减少滤波器的数量;2)在卷积前降低前景输入的分辨率,在传播后提高注意力映射map的分辨率。
为了集成注意模块,我们在图2的基础上引入了如图4所示的两个并行编码器。底部的编码器通过一层一层(膨胀)的卷积专门关注生成的内容,而顶部编码器试图关注感兴趣的背景特征。两个编码器的输出特征被聚合到一个单解码器中以获得最终输出。为了解释上下文注意,我们将其可视化,如图4所示。我们使用颜色来表示每个前景像素最感兴趣的背景patch的相对位置。例如,白色(颜色编码地图的中心)表示自己的像素,左下角为粉色,右下角为绿色。对于不同的图像,offset值的比例是不同的,以最好地可视化最有趣的范围。
对于训练,给定一个原始图像x,我们在一个随机位置采样一个二值图像掩码m。输入图像z为z = x⊙m。修复网络G以z和m的拼接作为输入,输出与输入大小相同的预测图像 x ′ = G ( z , m ) x' = G(z,m) x′=G(z,m)。将x’的掩码区域粘贴到输入图像中,得到修复的输出 x ~ = z + x ′ ⊙ ( 1 − m ) \tilde{x} = z + x '⊙(1−m) x~=z+x′⊙(1−m)。在所有实验中,输入和输出的图像值均线性缩放为[−1,1]。训练过程如算法1所示。
我们对四个数据集进行了评估,包括Places2 [41], CelebA faces [26], celebaq faces [20], DTD textures[6]和ImageNet[32]。
Qualitative comparisons
在训练中,我们使用分辨率为256 × 256的图像,最大hole为128×128,见第4.2节。方法都基于全卷积神经网络,因此可以填充不同分辨率图像上的多个hole。图6显示了验证集中各种复杂场景的视觉比较。为了测试的一致性,这些测试图像的尺寸都是512 × 680。所有报告的结果都是训练模型的直接输出,没有使用任何后处理。对于每个示例,我们还在最后一列可视化了模型的潜在注意图(颜色编码在第4.2节中解释)。
如图所示,我们的完整模型具有上下文关注,可以利用周围的纹理和结构,从而生成比基线模型更真实的结果。注意力地图的可视化揭示了我们的方法理解到了图像的背景结构,并能自适应地从周围区域借用信息,以帮助合成和生成。
Quantitative comparisons 像其他图像生成任务一样,图像修复缺乏良好的定量评价指标。为评估GAN模型引入的Inception评分[33]并不是评估图像修复方法的一个很好的度量标准,因为修复主要关注背景填充(例如物体移除情况),而不是生成各种类别物体的能力。
在重建误差方面的评价指标也不是完美的,因为有许多不同于原始图像内容的可能的解决方案。然而,我们在Places2上的验证集上设置了平均ℓ1误差、平均ℓ2误差、峰值信噪比(peak signal-to-noise ratio, PSNR)和总变差(total variation, TV)损失,仅供表1中的参考。从表中可以看出,基于学习的方法在ℓ1,ℓ2errors和PSNR方面表现较好,而直接复制原始图像patch的方法总变差损失较小。
我们的完整模型共有2.9M参数,大约是[15]中提出的模型的一半。模型在TensorFlow v1.3, CUDNN v6.0, CUDA v8.0上实现,运行在CPU Intel® Xeon® CPU E52697 v3 (2.60GHz)和GPU GTX 1080 Ti的硬件上。对于分辨率为512 × 512的图像,我们的完整模型在GPU上每帧运行0.2秒,在CPU上每帧运行1.5秒。
Contextual attention vs. spatial transformer network and appearance flow
Choice of the GAN loss for image inpainting
Essential reconstruction loss 我们的结论是,像素级的重建损失,虽然倾向于使结果模糊,是一个基本成分的图像修复。
Perceptual loss, style loss and total variation loss 在我们的框架中,我们没有发现感知机损失(VGG特征上的重构损失)、风格损失(VGG特征上计算的Gram矩阵的平方Frobenius范数)[19]和总变异(TV)损失对图像的inpainting有明显的改善,因此没有使用。
我们提出了一个由粗到细的生成图像修复框架,并引入了我们的基线模型和带有新的上下文注意模块的完整模型。我们发现,背景注意模块通过学习特征表示来显式匹配和关注相关的背景patch,显著提高了图像修复结果。作为未来的工作,我们计划使用类似于progressive growing GANs[20]的理念,将该方法扩展到高分辨率的修复应用中。所提出的inpainting框架和上下文注意模块也可以应用于条件图像生成、图像编辑和计算摄影任务,包括基于图像的渲染、图像超分辨率、引导编辑等。
个人总结:基于从粗到细两阶段的框架,对第一阶段和第二阶段的记过使用 空间距离加权的重构 l 1 l_1 l1 Loss,第二阶段的全局和局部判别器改用WGAN Loss,并在第二阶段整合了上下文注意力层(以与原来的膨胀卷积层平行的方式);对整体结构以及上下文注意力层都使用了技巧来减少参数,从而提高了计算效率。