Lama:《Resolution-robust Large Mask Inpainting with Fourier Convolutions》解读及实操

最近任务需求跟图像inpainting相关,因此调研使用了多个较新的开源模型,例如crfill、RePaint、Lama等。综合比较之下Lama的速度、效果都是最佳的,并且支持自定义输入尺寸进行推理(而非必须固定输入尺寸)。因此结合Lama论文进行实操,记录如下。

一、总体方法 & 创新点

1、总体方法流程:

对于输入原图 x ,使用一个二进制掩膜 进行遮罩 x  m ,形成一个四通道的输入tensor :

x′ = stack(x  m, m)

再使用一个前向infer网络(·)(也是一个生成器),以全卷积方式修复获得一个三通道彩色图像。

训练过程也是基于“图像与掩膜”这样的pair数据进行的。

2、创新点:

1)旧有方法都不具备足够广泛的感受野,因此对于大分辨率的图像或是大范围的inpainting而言,很容易会被局部附近的细节或干脆是mask所影响导致效果差。因此Lama考虑让模型在网络初始阶段就拥有更大的感受野,提出了基于快速傅立叶卷积(FFC)构造的网络结构。

2)损失函数:利用预训练分割网络进行特征提取实现损失函数构建,服务于大感受野和大掩膜

3)mask生成方法:动态的生成大掩膜,实现类似数据增强的效果

二、具体方法

1、Baseline:

图像与掩膜组成pair对输入网络并经过下采样后,进入FFC残差块。

在FFC块中,输入tensor被划分为两个分支进行运算。Local分支使用常规卷积;Global分支使用Real FFT进行全局上下文关注。其中在Global分支中经历了Real FFT2d和Inverse Real FFT2d的操作,实现了图像重建,具体张量变化看见论文2.1中的a)、 b)、c)。在FFC的输出中两分支进行结果合并。

2、损失函数

首先需要明确的是,对一个被掩膜遮盖的区域其实可以有多种合理的填充结果,就像口罩下的长相谁也无法准确预测。所以一旦掩膜变大,那么loss就必须被更加合理地设置以避免不符合事物逻辑的生成。

1)高感受野知觉损失 HRFPL

HRFPL

HRFPL是通过预训练的基础网络对输入图像x和生成图像x尖进行特征相似的计算。具体来说是使用空洞卷积或傅立叶卷积构建的HRF进行逐像素计算,然后M代表两阶段均值操作,即先取同层均值再取层间均值。

作者认为选取一个合适的预训练网络对于 HRFPL的效果至关重要。例如分类网络更关注局部的细节纹理,而难以理解全局的结构信息,缺乏对整体的认知。而分割模型则相反,具有更好的效果。

2)对抗损失

作者定义了一个在局部补丁级别上工作的鉴别器Dξ(·),用于区分“真实”和“虚假”补丁。对于输入pair对中的原图和掩膜图片,将原图上被掩膜覆盖的区域标记为“real”标签,将生成图片上对应区域标记为“fake”标签。

对抗损失

3)总体损失

在总体损失中,作者还使用

进行梯度惩罚,和基于识别器的感知实际损失或所谓的特征匹配损失-感知鉴别器网络LDiscPL的特征损失。众所周知,LDiscPL可以稳定训练,在某些情况下性能略有提高。

总体损失

LAdv和LDiscPL负责生成自然外观当地细节,而LHRFPL负责监督全局结构的信号和一致性。

3、训练中的动态掩膜生成

作者认为掩膜的生成类似于数据增强,对模型的效果非常重要。作者采取了多种大掩膜生成方式,但也同时注意避免生成大于原始图像50%的掩膜。

掩膜生成方式

三、实验与数据

训练使用Places和CelebA进行。下图结果中与其他模型在不同大小的掩膜上进行了对比,红色箭头代表表现逊于Lama,绿色则表示优于Lama。可以看到Lama在少参数量的情况下效果基本达到最优。

消融实验也证明了FFC在大掩膜上的效果。

论文还提供了大量实验细节,感兴趣的请查看原文。

四、使用记录

官方开源的模型中提供了一个名为Big-Lama的模型权重,效果最优。因为相比普通Lama,其生成器结构更复杂、训练数据规模更大。该模型是根据来自Places Challenge数据集的4.5M张图像的子集进行训练的,在八台NVidia V100 GPU上接受了约240小时的train。

我在Big-Lama上对infer refine的参数、mask生成方式做了反复测试,直观感受是其修复效果在其较快的infer速度上的确已经相当不错。同时我也尝试在训练中修改了mask的动态生成策略,实现了自定义mask生成来贴合我的任务场景,但考虑到训练成本尚未进行大规模训练,欢迎大家交流。

发表评论

您的电子邮箱地址不会被公开。 必填项已用*标注