最接近《银翼杀手》追踪杀人犯的AI？DiffBIR：统一的盲人脸修复和盲图像超分SOTA框架
"这是我们到目前为止最接近《银翼杀手》中的"增强"功能的时刻。"近日，美国金融科技公司 Brex 技术负责人在推特（现更名为：X）上介绍近日火热的 DiffBIR 算法时，如是说道。
“增强”技术是科幻电影的经典之作《银翼杀手》中最为精彩的构思之一，主人公里克·德卡德（Rick Deckard）使用称为“增强”的高级图像处理技术，仔细检查一张照片以追踪潜在的杀人犯。
在电影中，“增强”是指一种高级的图像处理技术，可以通过对图像进行放大、清晰化和分析，放大照片中的特定部分，并从不同的角度“旋转”或“移动”以观察场景中的隐藏细节。
里克·德卡德在调查照片以追踪疑似谋杀嫌疑人的情节中，使用了这个方法。这部极具前瞻性的科幻电影发行于 1982 年，“增强”的这个场景在当时令人耳目一新，因为它展示了一种超乎当时实际技术能力的图像分析手段，为电影的科技设定增加了深度。如今，这项“增强”技术真正地成为了现实。
最近，中国科学院深圳先进技术研究院和上海人工智能实验室提出了一种基于扩散模型先验的通用图像修复算法：DiffBIR。
该项目开源一个月内已经在 GitHub 上取得了超过 2000 颗星标，除此之外，DiffBIR 在实际应用中也受到了广泛的好评，海外知名技术大 V——camenduru 在 DiffBIR 开源后不久为其搭建了 Colab 项目，目前已有超过 13 万次浏览。
相比于传统的图像超分算法，DiffBIR 通过利用预训练的生成模型的先验知识，可以实现真实的修复效果。这项工作很好地结合了强大的开源文生图模型 stable diffusion，并采用 LAControlnet 来捕获关键信息。经过优化后的模型在图像修复任务上表现出色，不仅功能强大，还能实现高质量的图像重建。经过实验发现，无论是多么模糊的图像，DiffBIR 处理起来效果也是十分的惊艳。
同时，DiffBIR 也曾以每小时近 5 颗 GitHub 星标涨幅的表现，连续一周占领 Papers With Code（汇集 AI 领域研究论文及其对应的代码实现的权威平台）热度榜榜首。
为方便 AI 使用者、开发者更加方便、快捷地感受 DiffBIR 的的能力，论文作者基于 Gradio 为用户搭建了一款无需部署、在线体验的应用，并上线在了 OpenXLab 平台，考虑到目前应用的热度，平台提供了云端 A100 显卡进行推理支持，方便用户更加快速地体验 DiffBIR！
https://openxlab.org.cn/apps/detail/linxinqi/DiffBIR-official
（建议电脑端体验，文末点击阅读原文可直达）
未来，作者也会在更多的平台提供在线应用。以下是 DifBIR 的效果展示：
DiffBIR：统一的盲人脸修复和
盲图像超分 SOTA 框架
在这个大语言模型和 AIGC 热潮的阶段，为何这项关于图像超分和修复这个传统方向的工作突然又有了如此高的热度？
首先，我们先对盲图像修复这个领域进行简单的回顾。盲图像超分（blind image super-resolution）是盲图像修复中的热门方向。代表性的工作如 Real-ESRGAN 和 BSRGAN 会选择去设计退化过程来拟合真实场景的退化，以此来达到鲁棒的修复效果。除此之外，他们还引入对抗损失来对模型进行端到端的训练，企图让修复的结果更加真实。虽然这些盲图像超分方法具备强大的去退化能力，但是他们经常输出不够真实的纹理，或者抹去过多的细节。这说明简单引入对抗损失来对模型进行端到端训练难以获得令人满意的生成效果。
另一方面，盲人脸修复（blind face restoration）任务在最近几年取得了突破性的进展，一系列高影响力的工作，比如 GFPGAN，GPEN，CodeFormer 相继被提出。它们基于强大的人脸生成先验（StyleGAN/VQGAN），微调特殊设计的模块来达到高质量的人脸修复效果。然而，受限于生成模型的生成能力，这种基于生成先验的图像修复方法的应用局限在人脸领域，没有被成功拓展到通用的自然场景中。
现在，得益于文生图领域的巨大成功，DiffBIR 使用了强大的开源文生图模型——Stable Diffusion，作为通用自然场景下的生成先验。DiffBIR 提出了 LAControlnet 来学习条件信息，并对预训练好的文生图模型做有效的控制，使其适配到图像修复任务。除此之外，我们还借鉴盲图像超分领域中的退化设计和底层视觉骨干模型来完成真实场景得退化去除。总的来说，DiffBIR 具有强大的修复功能，兼具高质量的图像重建能力。另外，DiffBIR 统一了盲人脸修复和盲图像超分方法的框架，并且在这两个任务上都达到了 SOTA。
DiffBIR 方法：技术解读
DiffBIR 将盲图像修复问题拆分成两个子问题，分别是退化去除与细节生成，并提出了一个两阶段（two-stage）的框架。如下图所示：
图 2：DiffBIR 框架图
在第一阶段中，DiffBIR 使用由 MSE 损失函数训练得到的修复模块（Restoration Module）去除图像中包含的大部分退化（如噪声、失真等），获得一张干净但缺少局部纹理细节的平滑图像；然后在第二阶段中通过 LAControlnet 模块来利用 Stable Diffusion 的生成能力弥补丢失的纹理细节或语义信息。此外，我们引入了一个无需额外训练的可控模块 - latent image guidance，用于在图像质量和保真度之间进行权衡。
阶段1 - Restoration Module
第一阶段的目标是去除低质量图像中的大部分退化，尤其是噪声和压缩伪影，这样能够让第二阶段的微调生成模型排除噪声的干扰，更加专注于纹理/细节的生成，从而获得更加稳定的生成结果。我们借鉴了 DifFace 采用的 SwinIR 作为修复模块。具体地，首先利用 PixelUnshuffle 操作对原始的低质量输入 进行 8 倍下采样，然后在低分辨率空间中完成所有原始 SwinIR 的计算，最后通过三次最近邻插值恢复图像的分辨率，每次插值都跟随着一个卷积层和一个 Leaky ReLU 激活层。由于修复模块不需要输出高质量图像，我们没有使用对抗损失，而只使用了最简单的 MSE loss 作为损失函数。
为了验证修复模块的重要性，我们做了消融实验：移除修复模块，并直接对生成先验进行微调。从表格我们能够看出，修复模块的移除导致了在真实世界数据集中 FID/MANIQA 指标的明显下降。从下图可以看出，单阶段模型（无修复模块）错误地将退化视为语义信息，或者不能完全去除真实世界的噪声。这表明修复模块和二阶段训练能够显著提升修复的稳定性和效果。
阶段2 - LAControlNet
尽管第一阶段可以消除大部分退化，但得到的 通常过于平滑，仍然和自然图像分布离得很远。接下来，我们利用 数据对来微调预训练好的Stable Diffusion。首先，我们使用 Stable Diffusion 预训练的 VAE 的编码器将 映射到潜空间，获得条件潜变量 。Stable Diffusion 使用 UNet 去噪器完成逆扩散过程，它包含一个编码器，一个中间块，和一个解码器。参考 ControlNet (ICCV best paper)，我们创建了一个并行模块（在图 2 中以橙色表示），该模块包含与 UNet 去噪器相同的编码器和中间块。然后，我们将条件潜变量 与逆扩散中间状态 拼接作为并行模块的输入。由于这个拼接操作会增加并行模块中第一卷积层的参数，我们将新增加的参数初始化为零，同时其它参数使用预训练的UNet去噪器中对应的参数进行初始化。并行模块的输出会加到原始的 UNet 解码器上，从而实现对 Stable Diffusion 的微调。我们的优化目标与 Stable Diffusion 相同：
与 ControlNet 相比，我们的模型对于图像修复任务来说更加有效。ControlNet 使用一个从头开始训练的 CNN 来编码条件信息，而我们使用经过良好训练的 VAE 的编码器，它能够将条件信息直接映射到与潜变量相同的空间，极大程度上减轻了 Stable Diffusion 的内部知识与外部条件信息之间的对齐负担。从消融实验也能够看出，直接使用 ControlNet 进行图像修复会导致严重的颜色偏移（如下图所示）：
阶段3 - latent image guidance
尽管上述两阶段的方法已经能够取得良好的修复结果，但由于用户的各种偏好，我们仍然需要在图像质量和保真度之间做出权衡。因此，我们提出了一个无需额外训练的可控模块：它通过使用第一阶段得到的平滑图像 引导去噪过程，从而在两个阶段得到的结果之间进行保真度和图像质量的权衡。
具体地，我们通过引导中间变量 来控制扩散模型的生成过程。具体来说，在采样过程中，我们通过估算 中的噪声来估算出一个清晰的图像 。
然后使用基于潜空间的损失 来衡量引导图像 与估算的清晰潜变量 之间的距离，并在采样过程中优化该损失：
我们可以通过调整这个优化过程来实现图像质量和保真度的权衡，下图展示了完整的算法流程：
下图展示了可控模块的效果：
实验
在盲超分任务中，我们在真实场景超过了 Real-ESRGAN+、BSRGAN、SwinIR-GAN、FeMaSR 和 StableSR 等方法；而对于盲人脸修复任务，DiffBIR 也取得了比 DMDNet、GFP-GAN、GPEN、GCFSR、VQFR、CodeFormer、RestoreFormer 和 DifFace 更优越的修复效果。值得注意的是，DiffBIR 不仅能够修复人脸区域，对于背景、服装、饰品等非人脸区域也能够取得很好的修复效果。
图3：盲图像超分任务数值指标对比
图4：盲人脸修复任务数值指标对比
图5：盲图像超分任务视觉效果对比
图6：盲人脸修复任务视觉效果对比
结语
DiffBIR 通过利用预训练的生成模型的先验知识，可以实现真实的修复效果。它包括两个阶段：修复阶段和生成阶段，确保了图像的真实性和高保真度。实验结果表明 DiffBIR 在 BSR 和 BFR 任务上同时超越了现有方法。考虑到 Stable Diffusion 是文生图模型，而 DiffBIR 并没有使用具体的 prompt，因此我们需要进一步探索由文本驱动的图像修复。另一方面，DiffBIR 需要 50 步采样完成图像修复，这导致了更高的计算资源消耗和更长的推理时间，需要进一步优化。
2023-10-18
2023-10-17