CVPR2023 | 用于统一的图像恢复和增强的生成扩散先验

2023-11-27 15:34:57 浏览数 (1)

来源：CVPR 2023 论文题目：Generative Diffusion Prior for Unified Image Restoration and Enhancement 论文链接：https://arxiv.org/pdf/2304.01247.pdf 论文作者：Ben Fei等 内容整理：阳浩宁 现有的图像修复方法大多利用自然图像的后验分布。然而，它们通常假定已知退化情况，并且需要受监督的训练，这限制了它们对于复杂实际应用的适应性。本文提出了生成扩散先验（GDP），有效地以非监督的采样方式建模后验分布。GDP利用预训练去噪扩散生成模型（DDPM）解决线性反演、非线性或盲问题。具体来说，GDP系统地探索了条件引导的方案，比常用的引导方式更为实用。此外，GDP在去噪过程中优化退化模型的参数，实现了盲图像恢复。本文还设计了分层引导和基于分块的方法，使得GDP能够生成任意分辨率的图像。

目录

  • 简介
  • 方法
    • GDP基本原理
    • 单图像引导
    • 扩展应用
    • 损失函数
  • 实验
    • 线性和多线性退化任务
    • 曝光校正任务
    • HDR图片修复任务
    • 消融实验
  • 总结

简介

在拍摄、存储、传输和渲染过程中，图像质量往往会降低。图像恢复和增强的目标是逆转这种退化并改善图像质量。通常，恢复和增强任务可以分为两大类：1）线性反演问题，例如图像超分辨率（SR）、去模糊、修补、彩色化等，在这些任务中，退化模型通常是线性的且已知；2）非线性或盲问题，例如低光增强和HDR图像恢复，其中退化模型是非线性的且未知。对于特定的线性退化模型，可以通过对神经网络进行端到端的监督训练来解决图像恢复问题。然而，在现实世界中，受损图像往往存在多个复杂的退化情况，全面监督的方法很难泛化应用。近年来，通过生成模型寻找更通用的图像先验并在无监督设置下处理图像恢复问题引起了广泛的兴趣。在推理过程中，可以处理不同退化模型的多个恢复任务而无需重新训练。例如，经过大量干净图像数据集训练的生成对抗网络（GAN）通过GAN反演，在各种线性反演问题上取得了成功，学习到了真实世界场景的丰富知识。与此同时，去噪扩散概率模型（DDPMs）在GAN的基础上展现了令人印象深刻的生成能力、细节水平和多样性。作为早期尝试，现有的工作——去噪扩散恢复模型（DDRM）使用预训练的DDPMs进行变分推断，并在多个恢复任务上取得了令人满意的结果，但其在已知线性退化矩阵上利用奇异值分解（SVD），因此仍然局限于线性反演问题。本文进一步提出了一种高效的方法，名为生成扩散先验（GDP）。它利用经过良好训练的DDPM作为通用图像恢复和增强的有效先验，并以退化图像作为引导。作为一个统一的框架，GDP不仅适用于各种线性反演问题，还首次推广到非线性和盲目图像恢复和增强任务。GDP采用了一种盲退化估计策略，在去噪过程中随机初始化并优化GDP的退化模型参数。此外，为了进一步提高光真实性和图像质量，本文系统地研究了一种有效的指导扩散模型的方法。另外，借助提出的分层指导和基于分块的生成策略，GDP能够恢复任意分辨率的图像，其中首先预测低分辨率图像和退化模型，以引导高分辨率图像的生成过程。

方法

GDP基本原理

本文的目标是利用经过良好训练的DDPM作为统一图像恢复和增强的有效先验，特别是处理各种不同类型的退化图像。具体而言，假设图像

mathcal {y}

是通过

mathcal {y=D(x)}

获得的，其中

x

是原始的自然图像，

mathcal D

是退化模型。利用存储在某个先验中的

x

的信息，在

x

的空间中搜索最匹配

y

的

x

，并将

y

作为

x

的受损观测值。由于之前基于GAN的工作表现不够好并且应用范围有限，本文采用了更通用的图像先验，即在大规模自然图像上训练的扩散模型，用于图像合成。DDPM的反向去噪过程可以以退化的图片

y

作为条件，具体方式可以表示为一个条件分布，如公式1所示。

begin{aligned} log p_{theta}(x_{t-1}|x_t, y) =& log (p_{theta}(x_{t-1}|x_t)p(y|x_t)) K_1 \ approx & log p(r) K_2 end{aligned} quad (1)

其中

r sim mathcal {N}(r; mu_{theta}(x_t,t) Sigma g, Sigma), ; g=nabla log p(y|x_t), ; Sigma = Sigma_{theta}(x_t)

。

K_1, K_2

是常数，

p_{theta}(x_{t-1}|x_t)

代表扩散模型的反向生成过程。

p(y|x_t)

可以被看作

x_t

经过去噪后生成与

y

一致的高清晰度图片的概率。本文使用公式2近似表示这个概率。

pleft(boldsymbol{y}midboldsymbol{x}_tright)=frac1Zexpleft(-left[smathcal{L}left(mathcal{D}(boldsymbol{x}_t),boldsymbol{y}right) lambdamathcal{Q}(boldsymbol{x}_t)right]right) quad (2)

其中

mathcal{L}

是某种图像差异指标，

Z

是归一化因子，

s

是控制引导程度的比例因子。直观来讲这一定义鼓励

x_t

与受损的图片

y

一致，以获得高的

pleft(boldsymbol{y}midboldsymbol{x}_tright)

。

mathcal{Q}

是可选的质量增强损失，增强GDP的灵活性，可以用来控制某些属性(如亮度)或增强去噪图像的质量。

lambda

是调整图像质量的比例因子。等式两边的梯度可由公式3计算。

begin{aligned} &log pleft(boldsymbol{y}midboldsymbol{x}_tright)=-log Z-smathcal{L}left(mathcal{D}(boldsymbol{x}_t),boldsymbol{y}right)-lambdamathcal{Q}left(boldsymbol{x}_tright)\&nabla_{boldsymbol{x}_t}log pleft(boldsymbol{y}midboldsymbol{x}_tright)=-snabla_{boldsymbol{x}_t}mathcal{L}left(mathcal{D}(boldsymbol{x}_t),boldsymbol{y}right)-lambdanabla_{boldsymbol{x}_t}mathcal{Q}left(boldsymbol{x}_tright) end{aligned} quad (3)

用这种方法便可以通过调整无条件分布的均值来估计条件分布。

单图像引导

图像的超分辨率、修复、着色、去模糊和光照增强任务使用单图像引导。在以往的条件扩散模型中，方差

Sigma

被用来计算采样过程中的均值偏移。但本文发现

Sigma

可能会对生成图片的质量产生负面影响。因此，本文在引导去噪过程中去除方差以提高性能。在不存在

Sigma

并固定引导尺度

s

的情况下，引导去噪过程可以通过可变尺度

hat {s}

来控制。本文提出的引导方式包括在

x_t

上引导和在

x_0

上引导，分别称为

{rm {GDP-}}x_t

和

{rm {GDP-}}x_0

。

对于

{rm {GDP-}}x_t

，如图1和图2所示，这里条件扩散模型的种类的常用的一种，条件应用在

x_t

上，但移除了

Sigma

。但是，在

x_t

上应用引导的方式仍然可能产生不太令人满意的图像质量。这是因为

x_t

是一个具有特定噪声大小的噪声图像，但

y

通常是一个没有噪声或不同大小的噪声的损坏图像。我们缺乏可靠的方法来定义

x_t

和

y

之间的距离。一个朴素的MSE损失或感知损失将使

x_t

偏离其原始噪声大小，导致低质量的图像生成。

图1 第一种引导方式示意图

图2 第一种引导方式算法

为了解决上述问题，本文系统地研究了应用在

tilde{x}_0

上的条件信号。具体而言，在采样过程中，预训练的DDPM通常首先通过估计

x_t

中的噪声来从噪声图像

x_t

中预测出一个干净的图像

tilde{x}_0

，当每个时间步长

t

中给定

x_t

时，可以直接推断出该噪声。然后利用预测的

tilde{x}_0

与

x_t

一起对下一步潜在的

x_{t-1}

进行采样。

{rm {GDP-}}x_0

就是在这个中间变量

tilde{x}_0

上添加引导来控制DDPM的生成过程，如图3和图4所示。

图3 第二种引导方式示意图

图4 第二种引导方式算法

一些任务可以被认为是已知退化变换的。例如，图像去模糊、超分辨率任务的退化变换可以写作

boldsymbol{y}=(boldsymbol{x}otimesmathbf{k})downarrow_mathbf{s}

。它假设低分辨率(LR)图像是通过首先将高分辨率(HR)图像与高斯核(或点扩散函数)

rm k

卷积得到模糊图像

boldsymbol{x}otimesmathbf{k}

来获得的。图像修复的目的是恢复图像中缺失的像素。相应的退化变换是将原始图像与二值掩模

mathbf{m}

相乘，即

psi(boldsymbol{x})=boldsymbol{x}odotmathbf{m}

，其中

odot

表示哈达玛积。此外，图像着色旨在将灰度图像

y in mathbb{R}^{Htimes W}

恢复成有RGB通道的彩色图像

x in mathbb{R}^{3 times H times W}

。为了从彩色图像

x

得到

y

，退化变换

psi

是只保留

x

亮度的灰度变换。

然而在现实世界中，许多图像经历了复杂的退化，其中退化模型或退化模型的参数是未知的。在这种情况下，需要同时估计原始图像和退化模型参数。例如，在本文中，弱光图像增强和HDR恢复可以视为具有未知退化模型的任务。对此，本文设计了一个简单而有效的退化模型来模拟复杂的退化，可以表示为公式4。

boldsymbol{y}=fboldsymbol{x} mathcal{M}quad (4)

其中，光照因子

f

为标量，光照掩码

mathcal{M}

为与

x

相同维数的向量，

f

和

mathcal{M}

为退化模型的未知参数。本文使用这个简单的退化模型的原因是，只要

f

和

mathcal{M}

大小相同，任意一对损坏图像和相应的高质量图像之间的变换都可以被它们捕获。如果它们没有相同的大小，我们可以首先将

x

的大小调整为与

y

相同的大小，然后应用这个变换。值得注意的是，这种退化模型通常是非线性的，因为

f

和

mathcal{M}

依赖于

x

和

y

，需要对每个单独的损坏图像估计

f

和

mathcal{M}

。本文通过随机初始化它们并在DDPM的反向过程中同步优化它们来实现这一点。

扩展应用

在某些情况下，可以利用多个图像来引导单个图像的生成，这一内容很少被研究过，而且比单图像引导更具挑战性。鉴于此，本文提出了HDR-GDP，用于HDR图片修复任务，使用多个图像作为引导。它包含3个输入图像，分别为短、中、长曝光图像。与弱光增强类似，退化模型也被视为公式4的类型，因为确定HDR恢复的参数仍然未知。由于有3个损坏图像来引导生成，所以会对三幅LDR图像的三对盲参数进行随机初始化和优化。此外，本文采用基于分块的方法来解决扩散模型只能生成固定分辨率的图片这一问题。通过这种基于分块的策略的，GDP可以扩展到恢复任意分辨率的图像，以促进GDP的通用性。

损失函数

GDP的损失函数可以分为两部分，即重建损失和质量提升损失。前者的目标是恢复条件信号中包含的信息，而后者用来提升最终输出的质量。重建损失可以采用MSE、SSIM、感知损失等函数，本文采用的是MSE。对于不同的任务，GDP会采用不同的质量提升损失。主要有以下三种。 （1）曝光控制损失。为了提高GDP的通用性，曝光控制损失用于控制弱光图像增强的曝光水平。如公式5所示。

L_text{exp}=dfrac{1}{U}sum_{k=1}^{U}|R_k-E|quad(5)

其中

U

表示大小为8×8的不重叠局部区域的个数，

R

表示重建图像中局部区域的平均强度值，

E

是RGB色彩空间中的灰度级，可以通过调节

E

来控制亮度。 （2）色恒性损失。根据Gray-World颜色恒常性假设，颜色恒常性损失

L_{mathrm{col}}

用来校正恢复图像中潜在的颜色偏差，并在着色任务中连接三个调整通道之间的关系。如公式6所示。

L_{mathrm{col}}=sum_{forall(m,n)invarepsilon}(Y^m-Y^n)^2,varepsilon={(R,G),(R,B),(G,B)}quad (6)

其中

Y^m

是恢复后的图像中第

m

个通道的平均强度值，

(m,n)

是一对通道。 (3)照明平滑度损失。为了保持优化光掩模

mathcal{M}

中相邻像素之间的单调关系，对每个光方差

mathcal{M}

使用光照平滑损失。如公式7所示。

L_{tv_{mathcal{M}}}=frac{1}{N}sum_{n=1}^{N}sum_{cinxi}left(|nabla_{h}mathcal{M}_{n}^{c}| nabla_{v}mathcal{M}_{n}^{c}midright)^{2},xi={R,G,B}quad (7)

其中

N

是迭代次数，

nabla_{h}

和

nabla_{v}

分别是水平和垂直梯度运算。具体而言，图像着色任务利用颜色恒常性损失来获得更自然的颜色。出于同样的原因，弱光增强也需要使用色恒性损失。此外，弱光增强任务使用照明平滑度损失来让预测的光掩膜

mathcal{M}

更加平滑。曝光控制损失使我们能够手动控制恢复图像的亮度。

实验

本文系统地对比了GDP和其他多种应用于不同的图像和增强任务的方法，并验证了设计的有效性。本文使用的GDP基于一个使用ImageNet数据集预训练的无条件扩散模型。

线性和多线性退化任务

为了量化GDP的表现，出于ImageNet数据集的多样性，这部分的测试结果主要关注该数据集。这部分测试了各模型在4倍超分辨率、图像去模糊、图像修复和图像着色任务上的表现，结果如表1所示。结果表明GDP在FID和Consistency方面表现优于其他方法。

表1 ImageNet 1k数据集上线性图像恢复任务的定量比较

曝光校正任务

这部分评估了GDP在弱光增强任务上的表现，结果如表2所示。结果表明

{rm {GDP-}}x_0

方法在零次学习的方法中表现最优。

表2 多个数据集上图像亮度增强任务的定量比较

HDR图片修复任务

表3展示了GDP和一些SOTA方法在HDR图片修复任务上表现的对比。结果表明

{rm {HDR-GDP-}}x_0

取得了最好的效果。在图5的效果展示中可以发现，

{rm {HDR-GDP-}}x_0

重建的图像中，弱光部分得到了增强，而过度曝光区域得到了调整。并且

{rm {HDR-GDP-}}x_0

恢复的HDR图像细节更加清晰。

表3 NITRE数据集上的定量比较

图5 NITRE数据集中的一个示例

消融实验

表4展示了方差

Sigma

和引导方式对结果的影响。

表4 方差和引导方式的消融实验

表5展示了可训练退化算法和基于分块的策略的有效性。表中模型A设计为从参数不相关的块和块中简单地恢复图像。模型B对图像中的所有块进行了固定参数的设计。

表5 可训练退化算法和基于分块的策略的消融实验

总结

本文提出了用于统一图像恢复的生成扩散先验算法，可以用来解决线性逆、非线性和盲问题。本文提出的GDP能够通过分层指导和基于分块的方法恢复任何大小的图像。本文系统地研究了引导方式，从而充分利用DDPM的能力。GDP在图像超分辨率、图像去模糊、图像修复、图像着色、弱光增强、HDR恢复等多项任务上得到综合运用，体现了GDP在图像统一恢复方面的能力。

函数 模型 设计 数据 优化

0 人点赞