DiffBIR: Towards Blind Image Restoration with Generative Diffusion Prior CVPR 2023

GitHub - XPixelGroup/DiffBIR: Official codes of the repository for DiffBIR, aimed at blind image restoration using a generative diffusion prior.

目录

Abstract

Contribution

1. DiffBIR的设计

2. IRControlNet的提出

3. 可控模块的引入

总结

Method

Restoration Module

Generation Module

Preliminary: Stable Diffusion.

**IRControlNet-**条件编码 (Condition Encoding)

**IRControlNet-**条件网络 (Condition Network)

**IRControlNet-**特征调制 (Feature Modulation)

**IRControlNet-**优化目标

变体1（与ControlNet相同）

变体2（不使用 zt）

变体3（随机初始化条件网络）

变体4（控制中间块特征和解码器特征）

Restoration Guidance

1. 设计目标

2. 区域自适应恢复指导

3. 去噪过程

4. 区域自适应均方误差（MSE）损失函数

5. 权重映射和损失函数

6. 影响机制

7.优化过程

Experiment

1. 数据集

2. 实现细节

3. 评估指标

VS StableSR

Idea

Abstract

我们开发出了一种名为DiffBIR的通用恢复管道系统，旨在应对多种盲图像恢复挑战。该系统通过统一框架整合多任务处理能力。其核心创新在于将复杂的技术难题分解为两个关键环节：**1）去噪优化：**精准去除与图像无关的信息；**2）信息重建：**有效补充缺失的数据细节。这两个环节采用分层设计并实现了高效协同工作。在去噪优化环节中，我们引入了先进的去模糊模块技术，在此过程中充分考虑了各向异性效应的影响机制，并通过多层网络迭代提升了去模糊效率和效果。随后，在信息重建环节中，我们开发出了一个创新性的IRControlNet模型架构（图1）。该模型基于特意生成的理想条件图像进行训练（这些条件图像不含非理想噪声干扰），从而实现了稳定可靠的细节重建能力。值得注意的是，在整个推理流程中我们实现了模块化设计：通过简单的参数调节即可动态调节去噪步骤（无需额外计算开销），这显著提升了系统的灵活性和适用性

Contribution

1. DiffBIR的设计

• 分离BIR任务：DiffBIR将盲图像恢复（BIR）任务划分为两个独立的部分：
  • 去噪模块：专门负责去除图像中的退化细节。
  • 重建模块：专注于补充缺失的图像细节。
• 单一框架实现多项突破：这种双阶段的设计使DiffBIR能够在单一框架中首次实现了多个关键领域的最佳性能，在盲图像超分辨率（BSR）、盲面部恢复（BFR）以及盲图像去噪（BID）等多个领域都展现了卓越的效果。

2. IRControlNet的提出

• 基于扩散机制：研究提出了一种名为IRControlNet的新架构，在图像重建领域展现了显著优势。该网络通过将文本信息转化为图像特征图并结合空间信息完成重建过程。
  IRControlNet的主要创新点在于其独特的控制方法：首先，在网络输入端引入一个可编程参数以调节生成过程；其次，在解码器模块中设计了一种基于对抗学习的信息融合机制；最后，在损失函数设计中加入了多尺度特征匹配项以提升重建质量。

3. 可控模块的引入

• 无训练的可控模块 ：引入了一种区域自适应恢复指导模块，通过区域自适应恢复指导 ，利用设计的区域自适应MSE损失 ，在每个采样步骤 中将生成结果与高保真指导图像进行比较 ，实现保真度与质量 之间的权衡。低频区域 更受高保真指导 图像的影响，而高频区域 则保持更多的生成能力 。
• 用户偏好的灵活性 ：这种设计使得用户可以根据个人需求，在生成的图像质量和真实感之间找到平衡，提供了更高的控制能力。

总结

该段文字突出了DiffBIR在盲图像恢复领域的创新性突破，在框架架构设计、生成模块的高效性和用户干预能力等方面展现了显著贡献。

Method

Restoration Module

• 去除了降质影响：该算法的核心功能是通过深度学习技术去除了低质量（LQ）图像中残留的降质干扰，并成功地重建出高质量（HQ）图像，并未引入任何新信息。
  • 条件图重建：该算法不仅负责修复损坏的主体图像，并且能够自动生成可作为训练数据输入的重建图。整个重建过程旨在为后续的训练提供稳定可靠的输入数据。

基于均方差（MSE）的数值差异数值作为评价标准**：该模块遵循经典的降质模型，在学习过程中旨在最小化恢复图像与高质量图像之间的MSE差异数值。其具体形式如下：...

Generation Module

Preliminary: Stable Diffusion.

基于SD模型的Stable Diffusion是一种潜在扩散模型类别，并由一套高效的自动编码器架构构成。该架构系统地整合了编码器组件_E_与解码器组件_D_的工作原理，在此基础之上实现了对输入图像_x_的潜在表示提取以及相应重建过程的功能完整性。

优化目标 ：潜在扩散模型的优化过程定义为：

取自数据集的变量_x和_c被选取；z等于_x的期望值；时间变量t被等时间间隔地选取；随机噪声ε由标准正态分布生成。

**IRControlNet-**条件编码 (Condition Encoding)

VAE编码器 ：在IRControlNet框架中采用预训练的变分自编码器（VAE）编码器 E ，以实现可靠条件图像 IRM 在潜在空间中的表示转换。数学推导过程如下：

输出

**IRControlNet-**条件网络 (Condition Network)

复制的UNet结构：基于ControlNet的方法构建了一个可训练且预训练的UNet编码器以及其中间块的副本（标记为_Fcond_）。该副本接收条件信息并生成控制信号。

* **权重初始化** ：这种复制策略为条件网络提供了良好的权重初始化。

改写说明

• 通道数量提升：通过连接操作（cat），实现了网络的通道数量提升。具体而言，在 Fcond 网络架构的第一层中引入了一定数量的参数，并将这些参数初始化为零值。这种初始策略其作用机制与ControlNet中的零卷积相似，以防止初始阶段由于随机噪声带来的梯度不稳定问题。

**IRControlNet-**特征调制 (Feature Modulation)

• 多尺度特徵输出：條件網絡生成多种規模特徵並用於對凜化的U-Net去噪器中間特徵之調控。
  • 調控方式：僅針對中間分層與跳躍特徵實施調控並采用加法操作實現。
  • 連接穩定性：為提升模型訓練穩定性，在條件網絡與固定型U-Net去噪器之間引入零卷積進行連接。
  • 訓練流程：在整个訓練過程中僅優化條件網絡相關參數及相應特性調控 parameter。

**IRControlNet-**优化目标

• 目标函数 ：优化的结果在这一阶段被表示为 *IGM，*生成模块的优化目标定义为：

IRControlNet基于条件编码、条件网络和特征调制等关键模块构建了创新体系，在稳定扩散模型框架下实现了生成能力的显著提升。该系统中各组件均经过精心设计与优化，在提升训练稳定性的同时显著提升了生成质量，在图像恢复任务中展现出卓越的效果

变体1（与ControlNet相同）

• condition encoding：通过引入一种基于零初始化训练的小型网络_E_ 来替代IRControlNet中的condition encoder模块。这种网络架构由多层堆叠的convolutional layers 和一个zero-convolution layer 组成。
• input processing：将经过编码处理的条件特征整合到目标条件网络的第一层特征图中。
• performance analysis：通过对比实验发现，在处理输入低质量图像时（PSNR值），变体1的表现略逊于IRControlNet；具体而言，在均值意义下PSNR值降低约3dB左右。

变体2（不使用 zt ）

• 条件网络 ：去除噪声潜在表示 zt 后，默认条件下采用条件潜在表示 cRM 作为输入。
• 训练损失 ：该变体2在整个训练周期中的训练损失始终保持高于IRControlNet水平，并由此可观察到这种设计不仅提升了收敛效果还显著提高了预测精度。
• 性能对比 ：
  • 从保真度评估的角度来看，在各测试基准上均达到最优表现。
  • 而从视觉效果的角度来看，在各测试基准上的定量分析结果显示该方法能够有效降低生成图像的整体模糊程度。
  • 定性结果进一步验证了这种方法的优势在于其能够更精确地捕捉目标物体的关键特征并有效避免因过平滑化导致的信息丢失问题。

变体3（随机初始化条件网络）

• 条件网络未采用UNet去噪器的原始权重进行模仿，在训练过程中从随机初始化阶段开始建立。
• 该变体在收敛速度和稳定性方面表现不佳，并且在各项关键指标上均未达到最佳水平。
• 良好的权重初始化对于提升生成模块的效果起到关键作用。

变体4（控制中间块特征和解码器特征）

• 特性调节 ：本研究通过调节中间层特性和解码器特性（而非依赖于跳跃连接机制）实现了更好的效果。
• 实验结果表明 ，在收敛速率及定量评估指标方面 ，变体4与IRControlNet表现接近 ，证明了其对跳线性连接特性和解码器特性的调节效果不相上下 。
• 观察发现 ，值得注意的是，在通道数量上 ，解码器特性所具有的优势将导致额外的参数量与计算开销增加 。由此可见 ，仅依靠IRControlNet在跳线性连接上的特性调节即可达到理想效果 。

💡观察结论：条件编码在调控潜在扩散先验方面发挥着关键作用。生成过程位于潜在空间中时，必然导致条件信息也必须被投影至该空间内以保持一致性。

Restoration Guidance

1. 设计目标

• 控制组件：该模块赋予用户根据需求调整生成图像细节的程度的能力。高频区域如纹理和边缘处通常需要更多细节；而平缓区域如天空和墙面则要求较少内容。

2. 区域自适应恢复指导

• 指导过程 ：该模块基于区域自适应策略进行操作，在指导去噪过程中促使结果朝向高质量的条件图像 IRM 发展。
• 可调节的指导尺度 ：这一调节参数能够根据需要调整其大小，并适用于所有采样阶段；此外，在整个过程中无需任何预先训练内容。

3. 去噪过程

噪声预测 ：在时间步 t 时，UNet去噪器首先预测噪声 ϵt ：

降噪处理：然后我们从噪声潜在表示 zt 中提取并预测噪声特征并通过线性变换 W_h 对其进行降噪处理得到干净潜在表示 z~N(0 I)

4. 区域自适应均方误差（MSE）损失函数

• 损失函数目标：本阶段的目标旨在引导去噪器输出_D(z~0)朝向具有高保真度条件图像_IRM的方向发展。
• 计算梯度：采用Sobel算子进行梯度幅值计算。因为强对比边缘处仅有极少数像素具有显著的对比强度差异，在这种情况下我们选择将_IRM_分割为若干互不重叠的小块区域（即互补块），从而为每个区域块_G(IRn)_计算其对应的对比强度分布。

5. 权重映射和损失函数

权重映射 ：利用梯度幅值计算权重映射：

可调损失函数 ：after adjustment, the mean squared error (MSE) loss function is defined as: where H, W, and C respectively denote the spatial dimensions of the image.

6. 影响机制

• 低频区域 ：将更大权重应用于梯度较弱的区域，则会导致低频区域产生更大的损失。这从而对高保真度图像有更强的影响。
  • 高频区域 ：相比之下，在梯度较强的地方应用权重则显得意义不大。其影响相对较小，并且能够在采样过程中更好地保持生成内容。

7.优化过程

梯度下降方法

• 参数解释：
  • z~0′ z ~0′：指经过优化后的新潜在空间。
  • s s ：代表尺度因子，在指导图像 IRM 中调节保留的关键信息比例。

该系统采用区域自适应方法，在平衡图像质量与细节方面表现出色。通过优化损失函数参数设置，在生成过程中对低频与高频区域分别给予不同重视。这种设计则赋予了用户更高的自主权，在特定场景下可灵活调节算法性能指标。

Experiment

1. 数据集

• 训练数据集：DiffBIR在其筛选后的laion2b-en 数据集中进行训练；该集合共计包含约150万张高质量图片样本；在整个训练过程中；所有图像均被随机裁剪至512×512像素的尺寸以确保一致性和均匀性。
• 评估任务：
  • BSR任务：该算法将在以下多个基准测试用例上展开性能评估：
    • 包括三个合成性数据集（DIV2K-Validation、DRealSR和RealSR）以及两个真实世界的数据集合（RealSRSet和自收集的数据集合）
    • 此外；还包括两个混合型基准测试用例
  • BFR任务：针对此算法的实际应用效果；将重点考察其在以下基准测试中的表现：
    • 真实世界测试用例包括LFW-Test和WIDER-Test
    • 这些测试用例具有高度代表性和广泛适用性
  • BID任务：为了全面验证算法的鲁棒性；计划在其筛选后的混合的真实世界数据集合中展开评估

2. 实现细节

• 训练过程 ：
  • 恢复模块 ：训练150k次迭代（批量大小为96）。
  • IRControlNet微调 ：采用Stable Diffusion 2.1-base作为生成先验，微调IRControlNet 80k次迭代（批量大小为256）。
  • 优化器 ：使用Adam优化器，前30k次迭代的学习率设置为10⁻⁴，之后降低到10⁻⁵，持续50k次迭代。
  • 硬件 ：训练在8个NVIDIA A100 GPU上进行，分辨率为512 × 512。
• 推理过程 ：
  • 在推理时，使用特定任务的现成恢复模型替换训练好的恢复模块：
    • BSR ：使用BSRNet。
    • BFR ：使用在DifFace中使用的SwinIR。
    • BID ：使用SCUNet-PSNR。
  • IRControlNet 保持不变，负面提示使用“low quality”和“blurry”等文本，正面提示设置为空。
  • 恢复指导尺度 ：在合成数据集上设置为0、0.5和1进行比较，真实场景中设置为0以获得更高质量。
  • 采样过程 ：采用间隔DDPM采样调度，需要50个采样步骤。对于边长大于512的图像，直接输入DiffBIR；对于边长小于512的图像，先将短边放大到512，然后在恢复后调整回原始尺寸。

3. 评估指标

• 合成数据评估: 基于经典评价准则(如PSNR, SSIM和LPIPS)进行分析。
  • PSNR: 用于量化图像去噪后的质量。
  • SSIM: 衡量图像结构相似程度的指数。
  • LPIPS: 反映感知系统对图像差异的主观感受。
• 无参考图像质量评估: 主要包含以下三种方法:
  • MANIQA: 一种基于全息对比度分析的质量评价方法。
  • MUSIQ: 利用多谱系信息提升视觉质量评估的技术。
  • CLIP-IQA: 结合视觉-语言模型进行高质量图像感知的评价体系。
• BFR评估: 主要使用以下方法:
  • 基于Fréchet Inception Distance(FID)的方法来计算生成模型与真实数据之间的分布差异。

VS StableSR

在第一阶段中采用了SwinIR网络结构，在这一阶段中首先去除包含噪声和artifacts的影响的数据样本，在此处理后得到了一个较为干净的图片输出。随后通过扩散模型引入先验知识进行细节增强（超分）。

StableSR采用了自己设计的独特支路进行低分辨率信息注入；而diffBIR则基于控制网络构建；其核心组件模仿自扩散架构中的解码器模块——平行模块（Parallel Module）。为了提升性能，在解码器特征中加入了若干个零初始化（0×1×1卷积核）层，并将其整合到解码器特征中；其中平行模块采用了已预训练好的模型，在微调过程中能够迅速收敛并生成高质量的结果。

Parallel Module将z_t作为一个独立的输入加入到其处理流程中,这有助于使每个支路所获得的信息更加合理和具有代表性

保真度-真实性权衡的潜在图像引导

Idea

首先，在第一个阶段中, 我们会对图像进行退化处理, 去除一部分细节. 这种做法的好处在于, 它使得后续扩散网络接收的信息更为平滑, 整体一致性较高. 同时, 这一过程也降低了训练难度. 然而, 不管如何优化, 前面网络带来的计算开销也不可忽视. 实验结果表明, DiffBIR在保真度方面优于StableSR, 但仍存在部分生成质量不佳的现象.

**多尺度处理：**一种系统性策略旨在通过多层次特征提取来分别分析不同尺寸的信息，并通过这一过程确保细节与全局信息能够得到充分的保留。该方法的好处是可以有效去除噪声并尽可能减少信息损失。

**优化训练数据集：**通过应用数据增强技术——包括随机裁剪、旋转以及亮度调节等方法——能够有效提高模型的抗干扰能力。这种方法不仅有助于模型更加有效地提取具有代表性的图像特征，并且还能降低生成图像质量的下降现象。

3. 优化损失函数 ：构建融合感知损失与对抗损失的复杂损失函数以提升生成图像质量。促进模型聚焦于图像高层特征同时提升生成图像的真实感。

4. 引入注意力机制 ：通过在网络中引入注意力机制来提升模型性能，在此过程中使其能够聚焦于关键的特征区域，并有效地降低干扰因素。这不仅有助于生成更加清晰的画面数据还能显著提升生成图像细节的保留程度

5. 分阶段完成 ：将修复过程划分为若干阶段，并逐一修复细节。例如，在每个阶段中优先修复低频区域。从而有效地降低生成过程中的复杂度

完成图像生成之后，在后续阶段实施相应的优化措施。例如采用降噪技术或高分辨率建模。有助于减少伪影并保留图像细节。

7. 优化模型结构 ：为了优化系统性能并提高其处理复杂性的能力, 我们可以通过调整参数规模来实现这一目标. 这种方法可能会帮助我们更好地识别复杂的模式, 并降低潜在问题的影响程度.