基于Diffusion的图像修复方法

基于Diffusion的图像修复方法

本文介绍基于 Diffusion 的几个图像修复的工作。图像修复任务有两种应用的场景，一是图片的某部分真的缺失了，需要修复处这部分内容；二是想要修改图片中的某个部分，更换/新增/删除物体，这种情况其实可以看作是一种图像编辑人物。根据是否有条件引导修复的过程，又可分为无条件修复和条件修复，本文主要讨论的是基于 Diffusion 模型的文本条件引导的图像修复。整体上相当于有两个条件：文本提示词、非掩码区域的图像特征。

任务设定：给定一张原图 $x$，一个掩码 $m$ 和一个文本描述 $d$。将 $x$ 中 $m$ 位置出替换为 $t$ 中描述的内容。在这个任务中，有两个要求：一是要保证生成出的掩码区域内的内容与文本描述相符，二是要保证掩码外的区域与原图尽可能一致。

Blended Diffusion

Blended Diffusion 的研究主要是结合预训练好的 CLIP 和 DDPM 分别作语义引导和图像生成，通过设计特定的损失函数，引导图像生成的过程满足上述要求。

基线方法：Local CLIP Guided Diffusion

作者首先介绍了一种基线方法：Local CLIP Guided Diffusion，通过类似 Classifier Guidance 的方式来引导 inpainting 的过程。具体来说，作者设计了两项损失分别对采样过程进行引导，来分别满足任务的两个要求。

首先是 $D_{CLIP}$，用于引导掩码区域的生成结果与文本描述保持一致。分别计算掩码区域内生成结果的图像 embedding 和文本描述的文本 embedding，然后计算两者的余弦距离：
$${\mathcal{D}}_{CLIP}(x,d,m)=D_c(CLI{P}_{img}(x\odot m),CLI{P}_{txt}(d))$$
注意由于引导过程中的图像是带噪声的，所以原本的 Classifier Guidance 是在噪声数据上重新训练了一个分类器。而 CLIP 重新训练的成本太高，几乎无法实现。因此作者这里用的是按照 DDPM 形式化从噪声图 $x_t$ 估计的原始图像 ${\hat{x}}_0$，即：
$${\hat{x}}_0=\frac{x_t}{\sqrt{{\bar{\alpha }}_t}}-\frac{\sqrt{1-{\bar{\alpha }}_t}{ϵ}_{\theta }(x_t,t)}{\sqrt{{\bar{\alpha }}_t}}$$
第二，我们还要保持掩码外区域与原始图像尽可能一致，这里使用 ${\mathcal{D}}_{bg}$，通过计算掩码外区域的生成结果与原图的 MSE 和 LPIPS 两种距离来引导生成：
$$\begin{gathered} {\mathcal{D}}_{bg}(x_1,x_2,m)=d(x_1\odot (1-m),x_2\odot (1-m)); \\ d(x_1,x_2)=\frac{1}{2}(MSE(x_1,x_2)+LPIPS(x_1,x_2)) \end{gathered}$$
综合以上两项引导损失，基线方法 Local CLIP Guided Diffusion 通过二者的加权和 ${\mathcal{D}}_{CLIP}+\lambda {\mathcal{D}}_{bg}$ 来引导图像修复。通过调整两项损失的权重，就可以权衡同时保持该任务的两个要求。

然而，作者发现，无论怎么调整权重，基线方法的效果都不够理想。如下图实验（prompt = “a dog”）结果所示，在 $\lambda$ 较低时，背景的一致性完全无法保持，当增大 $\lambda$ 时，掩码区域内又无法生成出 prompt 所描述的内容，并且背景的保持也不是很完美。

Blended Diffusion

为了提升 inpainting 的效果，更好的兼顾区域语义生成和原图背景保持的两个要求，作者提出了 Blended Diffusion。顾名思义，Blended Diffusion 就是将两种噪声进行 混合。

在每个时间步 $t$，对噪声图 $x_t$，首先根据 CLIP Guidance 进行一次常规的去噪生成，得到的结果记为 $x_{t-1,fg}$，另外我们按照 DDPM 的原始公式 $x_t=\sqrt{\bar{\alpha }}{x}_0+\sqrt{1-{\bar{\alpha }}_t}ϵ$ 对输入原图进行加噪，得到结果记为 $x_{t-1,bg}$。然后根据给定的掩码 $m$ 对二者进行混合：
$$x_{t-1}=x_{t-1,fg}\odot m+x_{t-1,bg}\odot (1-m)$$

此外，为了避免对抗样本对 CLIP Guidance 的影响，作者采用的是多种数据增强结果的 CLIP 梯度的平均，相当于是一种 TTA。

Blended Latent Diffusion

Blended Latent Diffusion，是 Blend Diffusion 做在 Latent Diffusion 上的版本，混合噪声图修复部分的做法完全一致。区别就是：1. diffusion 过程做在隐层空间；2. 使用 Classifier-free Guidance（而非 Classifier Guidance）。

为了更好地重构背景，作者参考 GAN latent space optimization 系列方法，采用了一种 optimization-based 的方法来找到最优的 $z$。

还采用了 progressive mask shrinking，来改善 mask 区域较小时的效果。

Repaint

掩码区域根据 DDPM 进行正常去噪，非掩码区域对输入图进行加噪，再按掩码位置混合起来得到图片。
$$\begin{gathered} x_{t-1}^{\text{known}}\sim \mathcal{N}(\sqrt{{\bar{\alpha }}_t}{x}_0,(1-{\bar{\alpha }}_t)\mathbf{I}) \\ {x}_{t-1}^{\text{unknown}}\sim \mathcal{N}({\mu }_{\theta }(x_t,t),{\Sigma }_{\theta }(x_t,t)) \\ {x}_{t-1}=m\odot x_{t-1}^{\text{known}}+(1-m)\odot x_{t-1}^{\text{unknown}} \end{gathered}$$
为了使得掩码生成的区域充分考虑到输入的未掩码区域的语义信息，Repaint 提出了一种 Resampling 策略，即在混合得到 $x_{t-1}$ 之后，再加噪 $j$ 步，然后再降噪回来，从而可以利用其全图的上下文语义进行 inpainting。

GLIDE Inpainting

上述方法都是在训练好的扩散模型进行加噪采样，然后与原图的噪声混合起来，执行 inpainting 不需要重新训练。效果尚可，但是 mask 边缘终究是不太协调。以下介绍几种需要训练的 inpainting 方法，效果更进一步。

GLIDE 中介绍 Inpainting 是对整个扩散模型进行微调。具体来说，在微调时，随机选择数据样本中的区域进行擦除，然后将剩余的区域输入模型，而被擦除区域掩码则会通过额外添加一个输入通道作为条件输入给模型，额外添加的通道采用零初始化。

最近的 Kolors Inpainting 也是采用的这种方法。

SmartBrush

现有的 inpainting 训练方法一般都是随机生成 mask 区域，这无疑极大概率会导致 mask 区域内包含多种物体的部分，但是在实际图像编辑/修复中，我们一般是想要在指定的 mask 区域内大致生成一个目标物体。

SmartBrush 中为了保证 mask 区域内大致是一个主体物体，使用了图像分割数据集进行训练。训练时先随机选择图像 $x$ 中一个标注的掩码 $m$ 以及对应的类别标注 $c$。加噪时仅在掩码区域内加噪：
$$\begin{gathered} {\tilde{x}}_t=\sqrt{{\bar{\alpha }}_t}{x}_0+\sqrt{1-{\bar{\alpha }}_t}ϵ \\ {x}_t={\tilde{x}}_t\odot m+x_0(1-m) \end{gathered}$$
从而保证生图时能够根据清晰干净的未掩码区域上下文进行去噪。将 $x_t$、$m$ 和 $c$ 输入到扩散模型中，按照标准的 DDPM 优化目标进行训练。

形状精度控制 考虑到在进行 inpainting 时人工给出的 mask 区域一般都是不那么精确的，SmartBrush 在训练时采用了多种形状精度的策略，通过对分割掩码进行不同程度的高斯模糊来控制形状的精度 $m_s=\text{GaussianBlur}(m,k_s,{\sigma }_s)$，最精细的掩码是分割标注本身，最粗略的掩码是矩形框。

背景保持 当我们在掩码区域内生成文本 prompt 所指定的物体时，可能有些区域内是目标物体本身，有些部分是其周遭的背景。但这些生成的背景，很可能跟掩码区域外的背景是不匹配的，这会导致生成结果整体不够协调。为了解决这一问题，作者提出，在训练时同时要求网络预测出掩码区域内生成的目标物体本身的掩码。这需要网络的输出之上添加一个通道，用于表示二进制掩码，并添加一个分割的 DICE loss：
$$\begin{gathered} {\mathcal{L}}_{\text{prediction}}=H({ϵ}_{\theta }(m_s),m) \\ H(X,Y)=1-\frac{2|X\cap Y|}{|X|+|Y|} \end{gathered}$$

在推理生成时，我们就可以预测出给定掩码区域内实际目标物体的生成区域了，在推理过程中我们会将预测出的掩码代替给定的粗略掩码，将预测掩码外的区域替换为原图。（注意这里 SmartBrush 是需要完整原图的，也就是做图像编辑，而非真正缺掉某个区域的图像修复）。

BrushNet

除了训练专用的 inpainting 扩散模型（下图 Dedicated Inpainting Models）之外，还有一类通过添加并训练一个额外的网络分支来训练 inpainting 能力，比如 ControlNet Inpainting。将带掩码的图片，提取特征后输入到 ControlNet 中，将 ControlNet 逐层的输出作为条件输入给扩散模型进行条件生图。然而，ControlNet 更多地还是用于提取稀疏的空间结构条件，如人体姿势、Canny 边缘等，而 inpainting 要求的是像素级别的严格控制，因此 ControlNet 提取的信息可能不足，因此其效果跟专用的 inpainting 模型相比效果稍差。

BrushNet，也是提出了添加一个额外的网络分支进行 inpainting 任务，但其更进一步提出了适配 inpainting 任务的几点网络设计：

1. 为了对齐 UNet 的输入分布，使用 VAE 来提取图像特征，而非随机初始化的卷积网络
2. 为了实现像素级别的控制，采用了一种层级的方案。提取出每一层的特征给到生图模型
3. 为了保证 BrushNet 中是纯图像信息，去除了原始 UNet 中的 text cross attention。这进一步使得 BrushNet 成为一个即插即用的网络分支，并使得 inpainting 过程中对非掩码区域的控制更加灵活。

此外，为了实现更好的连贯性和非掩码区域的可控性，BrushNet 还采用了一种模糊混合策略。

具体来说，BrushNet 的整体结构如下图所示。

整体来看，还是输入一张图像和掩码图，输出 inpainting 后的完整图像。首先，由于 BrushNet 是基于 LDM 的，因此要先对掩码图进行降采样，以匹配 latent 的尺寸，同时输入图像经过 VAE encoder 提取（与隐空间特征对齐的）latent。然后将噪声 latent，图像 latent 和降采样的 mask 拼接起来作为 BrushNet 的输入。经由 BrushNet 提取并输出的特征会经过一个零卷积之后逐层对应输入到生图 UNet 模型中。在生图结束后，将生成图像和输入图像进行模糊混合，得到最终的结果。

这里的 BrushNet 所采用的模型与生图 UNet 结构一致，只是 BrushNet 是用于处理纯图像信息，因此去除了 text cross attention 模块。此外，由于对 mask 进行了降采样，且 VAE encode-decode 的过程本身就存在一定的损失，因此最终生成的图一般无法保证非掩码区域与原图完全一致。很多工作会在最后直接将原图按照 mask 拼到生图结果上，但这通常会导致结果语义不连贯。BrushNet 中提出先对 mask 进行模糊操作，然后按照模糊后的 mask 进行混合，这样虽然边界处会有一些细节丢失，但几乎看不出来，最终结果的一致性大大改善。

PowerPaint

图像修复有许多各不相同的细分任务，包括文本引导的图像修复、形状引导的图像修复、物体移除、扩图等。现有的图像修复工作一般只专注于上述细分任务的某一种。而 PowerPaint，主打的就是多功能，通过引入了三个可学习的 prompts：${\mathbf{P}}_{\mathbf{obj}},{\mathbf{P}}_{\mathbf{ctxt}},{\mathbf{P}}_{\mathbf{s}}\mathbf{hape}$ 来指示不同的任务，PowerPaint 在上述这些细分任务上都达到了领先的性能。

Context-aware Image Inpainting 仅根据非掩码区域的图像信息补全出掩码区域。输入是 mask 和原图，文本 prompt 是一个可学习的 ${\mathbf{P}}_{\mathbf{ctxt}}$。

Text-guided Object Inpainting 根据非掩码区域的图像信息和文本提示词在掩码区域生成特定的物体，mask 是矩形框，文本 prompt 是物体提示词在拼接上可学习的 ${\mathbf{P}}_{\mathbf{obj}}$。

Object Removal 去除掩码区域内的特定物体，利用了 Classifier Free Guidance 可以设置 negative prompt 的特点。positive prompt 是 ${\mathbf{P}}_{\mathbf{ctxt}}$，negative prompt 是 ${\mathbf{P}}_{\mathbf{obj}}$ ，即根据掩码外图像信息补全掩码区域，但是不能生成特定物体。

Controllable Shape Guided Object Inpainting 根据形状来生成特定物体。精确的分割 mask，使用 SmartBrush 的方法进行训练，文本 prompt 需要拼接可学习的 ${\mathbf{P}}_{\mathbf{shape}}$。但作者发现直接这样做会导致生成的物体过于接近于 mask 的形状，基本把 mask 区域填充满了，而丢失了该物体原本的形状。PowerPaint 提出 task prompt interpolation 来解决该问题。首先使用卷积操作对 mask 进行随机扩张，会得到一系列与物体形状适配度 $\alpha$ 不同的 mask，$\alpha$ 越大表示与 mask 的适配程度越高。将 ${\mathbf{P}}_{\mathbf{shape}}$ 与 ${\mathbf{P}}_{\mathbf{ctxt}}$ 拼接到文本描述后，各自输入到 CLIP text encoder 之后得到两个 embedding $e_{\text{shape}},e_{\text{ctxt}}$，然后对二者进行插值 $e=(1-\alpha )e_{\text{shape}}+\alpha e_{\text{ctxt}}$。在推理时，用户可以调整参数 $\alpha$ 来控制生成物体形状与 mask 的适配度。

总结

整体来说，基于 Diffusion 的 inpating 方法可以分为无需训练和需要训练两个大类。其中无需训练的方法主要是根据 mask 来调整采样策略，将掩码区域的去噪生成结果与掩码外区域的加噪结果进行混合（blend）。而需要训练的方法又可分为微调模型本身和添加额外的网路分支两类。核心就是更精细的像素级控制。PowerPaint 则通过设置可学习的 prompt，进一步实现了多功能的 inpainting 模型。

• 无需训练
  • 调整采样策略
    • Blended Diffusion：DDPM + CLIP Classifier Guidance
    • Blended Latent Diffusion：LDM + Classifier-free Guidance
    • Repaint
• 需要训练
  • 微调
    • GLIDE/Kolors inpainting
    • Smart Brush
    • PowerPaint：多功能
  • 额外的网络分支
    • ControlNet inpainting
    • BrushNet
    • PowerPaint-v2：PowerPaint + BrushNet