DiffBIR：用生成式扩散先验实现盲图像恢复

来源: arxiv 论文题目：DiffBIR: Towards Blind Image Restoration with Generative Diffusion Prior 论文链接：https://arxiv.org/pdf/2308.15070.pdf 论文作者：Xinqi Lin等 内容整理：阳浩宁 本文提出了DiffBIR，它利用预训练的文本到图像扩散模型来解决盲图像恢复问题。它的框架包含两阶段。第一阶段在多种退化中预训练恢复模块，以提高现实场景中的泛化能力。第二阶段利用潜在扩散模型的生成能力，实现真实的图像恢复。具体来说，本文引入了一个注入式调制子网络-LAControlnet进行微调，而预训练的Stable Diffusion可以保持其生成能力。最后，本文引入了一个可控模块，通过在推理过程中引入潜在图像引导，使用户能够平衡质量和保真度。大量的实验已经证明了它在合成和现实世界数据集上的盲图像超分辨率和盲人脸恢复任务方面优于最先进的方法。

简介

图像恢复的目的是从低质量的观测中重建出高质量的图像。典型的图像恢复问题，如图像去噪、去模糊和超分辨率，通常是在受限的环境下定义的，其中退化过程是简单和已知的(例如，高斯噪声和双三次降采样)。为了处理现实世界中退化的图像，盲图像恢复(BIR)成为一个很有前途的方向。BIR的最终目标是在具有一般退化的一般图像上实现真实的图像重建。BIR不仅扩展了经典图像恢复任务的边界，而且具有广泛的实际应用领域。BIR的研究还处于初级阶段。根据问题设置的不同，现有的BIR方法大致可以分为三个研究方向，即盲图像超分辨率(BSR)、零次图像恢复(ZIR)和盲人脸恢复(BFR)。它们都取得了显著的进步，但也有明显的局限性。BSR最初是为了解决现实世界的超分辨率问题而提出的，其中低分辨率图像包含未知的退化。根据最近的BSR调查，最流行的解决方案可能是BSRGAN和Real-ESRGAN。它们将BSR表述为一个有监督的大规模退化过拟合问题。为了模拟真实的退化，分别提出了退化洗牌策略和高阶退化建模，并用对抗性损失来以端到端方式学习重建过程。它们确实消除了一般图像上的大多数退化，但不能生成真实的细节。此外，它们的退化设置仅限于×4或者×8超分辨率，这对于BIR问题来说是不完整的。第二组ZIR是一个新出现的方向。代表有DDRM、DDNM、GDP。它们将强大的扩散模型作为附加先验，因此比基于GAN的方法具有更大的生成能力。通过适当的退化假设，它们可以在经典图像恢复任务中实现令人印象深刻的零次恢复。但是，ZIR的问题设置与BIR不一致。他们的方法只能处理明确定义的退化(线性或非线性)，但不能很好地推广到未知的退化。第三类是BFR，主要研究人脸修复。最先进的方法可以参考CodeFormer和VQFR。它们具有与BSR方法相似的求解方法，但在退化模型和生成网络上有所不同。由于图像空间较小，这些方法可以利用VQGAN和Transformer在真实世界的人脸图像上取得令人惊讶的好结果。然而，BFR只是BIR的一个子域。它通常假设输入大小固定，图像空间有限，不能应用于一般图像。由以上分析可知，现有的BIR方法无法在一般图像上实现一般退化的同时实现真实图像的重建。因此需要一种新的BIR方法来克服这些限制。本文提出了DiffBIR，将以往工作的优点整合到一个统一的框架中。具体来说，DiffBIR(1)采用了一种扩展的退化模型，可以推广到现实世界的退化;(2)利用训练良好的Stable Diffusion作为先验来提高生成能力;(3)引入了一个两阶段的求解方法来保证真实性和保真度。本文也做了专门的设计来实现这些策略。首先，为了提高泛化能力，本文将BSR的多种退化类型和BFR的广泛退化范围结合起来，建立了一个更实用的退化模型。这有助于DiffBIR处理各种极端退化情况。其次，为了利用Stable Diffusion，本文引入了一个注入调制子网络-LAControlnet，可以针对特定任务进行优化。与ZIR类似，预训练的Stable Diffusion在微调期间是固定的，以保持其生成能力。第三，为了实现忠实和逼真的图像重建，本文首先应用恢复模块(即SwinIR)来减少大多数退化，然后微调生成模块(即LAControlnet)来生成新的纹理。如果没有这个部分，模型可能会产生过度平滑的结果(删除生成模块)或生成错误的细节(删除恢复模块)。此外，为了满足用户多样化的需求，本文进一步提出了一个可控模块，可以实现第一阶段的恢复结果和第二阶段的生成结果之间的连续过渡效果。这是通过在去噪过程中引入潜在图像引导而无需重新训练来实现的。适用于潜在图像距离的梯度尺度可以调整以权衡真实感和保真度。在使用了上述方法后，DiffBIR在合成和现实数据集上的BSR和BFR任务中都表现出优异的性能。值得注意的是，DiffBIR在一般图像恢复方面实现了很大的性能飞跃，优于现有的BSR和BFR方法(如BSRGAN、Real-ESRGAN、CodeFormer等)。可以观察到这些方法在某些方面的差异。对于复杂的纹理，BSR方法往往会产生不真实的细节，而DiffBIR方法可以产生视觉上令人愉悦的结果。对于语义区域，BSR方法倾向于实现过度平滑的效果，而DiffBIR可以重建语义细节。对于微小的条纹，BSR方法倾向于删除这些细节，而DiffBIR方法仍然可以增强它们的结构。此外，DiffBIR能够处理极端的退化并重新生成逼真而生动的语义内容。这些都表明DiffBIR成功地打破了现有BSR方法的瓶颈。对于盲人脸恢复，DiffBIR在处理一些困难的情况下表现出优势，例如在被其他物体遮挡的面部区域保持良好的保真度，在面部区域之外成功恢复。综上所述，DiffBIR首次能够在统一的框架内获得具有竞争力的BSR和BFR任务性能。广泛而深入的实验证明了DiffBIR优于现有的最先进的BSR和BFR方法。

方法

本文的目标是利用强大的生成先验-Stable Diffusion来解决普通图像和人脸图像的盲恢复问题。提出的框架采用有效、稳健和灵活的两阶段方法。首先，使用回归损失训练恢复模块来去除损坏，例如噪声或失真伪影。由于丢失的局部纹理和粗/细细节仍然不存在，之后利用Stable Diffusion来弥补信息丢失。整个框架如图1所示。具体来说，本文首先在大规模数据集上预训练一个SwinIR，以实现跨多种退化的初步退化去除，然后利用生成先验来产生逼真的恢复结果。此外，还引入了一种基于潜在图像引导的可控模块，用于在真实感和保真度之间进行权衡。

图1 DiffBIR整体框架

预训练去除退化的模型

BIR旨在从具有未知和复杂退化的低质量(LQ)图像中恢复干净图像。通常图像退化包括模糊，噪声，压缩伪影和低分辨率几种。为了更好地覆盖LQ图像的退化空间，本文采用了考虑多样化退化和高阶退化的综合退化模型。在所有退化中，模糊、放缩和噪声是现实场景中的三个关键因素。本文中的多样化退化涉及以下几点。（1）模糊: 各向同性高斯核和各向异性高斯核;（2）放缩: 区域放缩，双线性插值和双立方放缩;（3）加噪: 加性高斯噪声、泊松噪声和JPEG压缩噪声。对于高阶退化，本文使用二阶退化，将经典退化模型模糊-放缩-加噪的过程重复两次。经过退化过程后，图像将被重新调整到原始大小。为了构建一个鲁棒的生成图像恢复方法，本文采用了一种保守但可行的解决方案，首先去除LQ图像中的大部分退化(特别是噪声和压缩伪影)，然后使用生成模块来重现丢失的信息。这种设计将促进潜在扩散模型更多地关注纹理/细节的生成，而不受噪声损坏的干扰，并且在没有错误细节的情况下获得更真实/清晰的结果。本文修改了SwinIR作为恢复模块。具体来说，本文利用像素解混洗操作，对原始低质量输入

进行8倍下采样。然后，采用

卷积层进行浅层特征提取。与潜在扩散模型相似，所有后续的变换操作都在低分辨率空间中进行。深度特征提取采用多个残差Swin Transformer块(Residual Swin Transformer block, RSTB)，每个RSTB有多个Swin Transformer层(STL)。将增加浅层和深层特征，以同时维护低频和高频信息。为了将深度特征上采样回原始图像空间，模型执行了三次最邻近插值，每次插值之后是一个卷积层和一个Leaky ReLU激活层。本文通过最小化

像素损失来优化恢复模块的参数，如公式1所示:

其中

和

分别表示高质量图像和低质量对应图像。

是通过回归学习得到的，将用于对潜在扩散模型的微调。

利用生成式先验进行图像恢复

虽然第一阶段可以消除大多数退化，但获得的图像

通常过于平滑，距离高质量自然图像的分布还很远。因此需要利用预训练的Stable Diffusion与获得的

对来重建图像。首先，利用Stable Diffusion的预训练VAE编码器将

映射到潜在空间中，得到潜在表示

。UNet去噪器执行潜在扩散，它包含一个编码器、一个中间块和一个解码器。特别地，解码器从编码器接收特征并在不同尺度上进行融合。本文在这里创建了一个并行模块(在图1中用橙色表示)，它包含与UNet去噪器中相同的编码器和中间块。然后将潜在表示

与随机采样的噪声

连接起来作为并行模块的输入。由于此串联操作将增加并行模块中第一个卷积层的通道数，因此将新添加的参数初始化为零，而所有其他权重都从预训练的UNet去噪检查点初始化。并行模块的输出被添加到原始UNet解码器。此外，在每个尺度的加法运算之前都有一个

的卷积层。在微调过程中，并行模块和这些

卷积层同时优化，其中条件设置为空。训练目标是最小化公式2所示的潜在扩散目标

这一阶段得到的结果记为

。总而言之，只有UNet去噪器中的跳过连接特征才会针对特定任务进行调整。该策略缓解了小训练数据集的过拟合问题，并继承了Stable Diffusion的高质量生成。更重要的是，与ControlNet相比，本文的条件调节机制对于图像重建任务更直接有效，ControlNet使用从头开始训练的附加条件网络来编码条件信息。在本文的LAControlnet中，训练过的VAE编码器能够将条件图像投影到与潜在变量相同的表示空间中。该策略显著减轻了潜在扩散模型中内部知识与外部条件信息对齐的负担。在实践中，直接使用ControlNet进行图像重建会导致严重的颜色偏移，消融实验也展示了这一点。

实现在保真度和真实性之间权衡的潜在图像引导

虽然上述两阶段的方法已经可以获得很好的恢复效果，但由于用户的不同偏好，仍然需要在真实性和保真度之间进行权衡。因此，本文提出了一个可控的模块，可以将降噪过程向在第一阶段获得的

引导，从而在真实和平滑的结果之间进行调整。Dhariwal和Nichol提出了分类器引导，它利用在噪声图像上训练的分类器来引导生成目标类的图像。而在大多数情况下，作为引导的预训练模型通常是在干净的图像上训练的。为处理这种情况，一些研究工作通过引导中间变量

来控制扩散模型的生成过程。具体来说，在采样过程中，他们通过估计

中的噪声，从噪声图像

中估计出一个干净的图像

。而本文中，扩散和去噪过程是基于潜在空间的。因此，本文的目标是通过公式3获得一个干净的潜在表示

。

然后，定义基于潜在表示的损失

为潜在图像引导

和估计的干净潜在表示

之间的

距离。

上述引导可以用迭代的方式保证潜在特征之间的空间对齐和颜色一致性，并引导生成的潜在表示保留参考的潜在表示的内容。因此，可以控制从参考图像

中保留多少信息(如结构、布局和颜色)，从而实现从生成的输出到更平滑的结果的过渡。整体算法如图2所示。

图2 潜在引导扩散过程算法

实验

实验设置

对于BIR任务，本文采用ImageNet数据集训练DiffBIR，分辨率设置为

。对于BFR任务，采用FFHQ数据集，图片放缩为

大小。训练时，采用所提出的退化方法由HQ图像生成LQ图像。对于BSR任务，采用RealSRSet数据集，在现实世界的设定中进行比较。为了进行更完善的比较,本文从互联网上收集了47张图像,组成Real47数据集。它包含各种场景的一般图像,如自然的户外景观,旧照片,建筑,从肖像到人群的人类图像,植物,动物,等等。对于BFR任务进行测试时，本文采用了合成的数据集CelebA-Test和三个真实世界数据集：LFW-Test, CelebChild-Test和WIDER-Test。其中CelebA-Test数据集包含3000张CelebA-HQ数据集中的图像，LQ图片则通过与本文训练时相同的方法产生。模型的恢复模块采用了8个RSTB，每个RSTB包含6个Swin Transformer层。测试指标方面，在保真度层面，本文采用PSNR，SSIM，LPIPS等传统指标；在真实性层面，本文也采用了无参照的IQA方法，如MANIQA和NIQE。对于人脸图像恢复，本文还测试了特征保留程度IDS和被广泛使用的感知指标FID。本文也开展了用户调研，以进行更全面地比较。

与SOTA方法的比较

表1展示了多种BSR和ZIR方法在真实世界数据集4倍超分辨率任务中的结果。DiffBIR在两个数据集中都取得了最好的MANIQA结果。

表1 真实世界数据集上4倍超分辨率结果

图3展示了DiffBIR和其他一些方法在Real47数据集上进行用户调研的结果。对于每张图像，用户被要求对四种方法的结果进行排序，并按升序为不同的方法分配1-4点，结果更优则分数更高。可以看出，DiffBIR的中位数得分最高，且其上四分位数超过3。这表明用户倾向于将DiffBIR的结果排在首位。

图3 用户调研结果

表2展示了在合成数据集和真实世界数据集上进行BFR任务的结果。结果表明，DiffBIR能够生成高真实性和高保真度的结果。

表2 人脸图像恢复结果

消融实验

表3展示了添加和去除第一阶段图像恢复模块的结果。结果表明这一模块有助于保持保真度。

表3 恢复模块的作用展示

图4展示了恢复模块、Stable Diffusion微调和LAControlnet的作用。

图4 视觉效果展示

图5展示了可控模块的灵活性。实验表明更大的梯度尺度

倾向于让模型生成接近

的高保真平滑结果。

图5 通过调节梯度尺度实现在真实性和保真度之间权衡

总结

本文提出了一种统一的盲图像恢复框架DiffBIR，该框架利用预训练的Stable Diffusion的先验知识获得真实的图像恢复结果。它包括恢复和生成两个阶段，保证了画面的保真性和真实感。大量的实验已经验证了DiffBIR在BSR和BFR任务上优于现有的最先进的方法。虽然DiffBIR已经达到很好的效果，但文本驱动图像恢复的潜力并没有被探索，仍然需要在图像恢复任务中进一步开发Stable Diffusion。另一方面，DiffBIR需要50个采样步骤来恢复低质量的图像，与其他图像恢复方法相比需要更高的计算资源消耗和更多的推理时间。