Diff-Retinex: Rethinking Low-light Image Enhancement with A Generative Diffusion Model 论文阅读

Diff-Retinex:用生成式扩散模型重新思考低光照图像增强

摘要

本文中,我们重新思考了低光照图像增强任务,并提出了一种物理可解释的生成式扩散模型,称为 Diff-Retinex。我们的目标是整合物理模型和生成网络的优点。此外,我们希望通过生成网络补充甚至推断低光照图像中缺失的信息。因此,Diff-Retinex 将低光照图像增强问题表述为 Retinex 分解和条件图像生成。在 Retinex 分解中,我们整合了 Transformer 中注意力的优势,并精心设计了一个 Retinex Transformer 分解网络(TDN)将图像分解为光照图(illumination map)和反射图(reflectance map)。然后,我们设计了多路径生成扩散网络来重建正常光照下的 Retinex 概率分布,并分别解决这些分量中的各种退化问题,包括光照不足、噪声、颜色偏差、场景内容丢失等。得益于生成式扩散模型,Diff-Retinex 将低光照下细微细节的恢复付诸实践。在真实世界低光照数据集上进行的大量实验,从定性和定量两方面证明了所提方法的有效性、优越性和泛化性。

1. 引言

在低光照场景下拍摄的图像通常会受到各种退化的影响,如不定性噪声、低对比度、多变的颜色偏差等。在这些退化中,场景结构的丢失是最棘手的问题。如图 1 所示,场景结构的丢失不仅影响视觉效果,还会减少信息量。图像增强是减少退化对人类感知和后续视觉任务干扰的有效方法,最终呈现高质量图像。为了处理这些退化,许多低光照图像增强(LLIE)方法被提出[20, 26]。此外,还开展了一系列关于对比度增强、噪声去除和纹理保持的研究。主流的 LLIE 方法大致可分为传统方法[35, 3]和基于学习的方法[22, 15, 25, 23, 18]。传统算法通常基于图像先验或简单的物理模型。例如,灰度变换[35, 14]和直方图均衡化[5, 34]通过线性或非线性手段调整强度分布。Retinex 模型[12, 16, 24]将图像分解为光照图和反射图,并使用传统优化方法解决问题。然而,这些方法也受到手工设计和优化驱动效率的限制。它们通常泛化性和鲁棒性较差,限制了这些方法的应用范围。

图 1. URetinex [40] 和 Diff-Retinex 用于低光照图像增强(LLIE)的示例。Diff-Retinex 通过生成式扩散模型重新思考 LLIE,能够修复一些缺失的场景内容。

为了解决这些缺点,深度学习被用来构建从低光照到正常光照图像的复杂映射[22, 38]。一些方法完全将低光照图像增强视为通过整体拟合的恢复任务,缺乏物理模型的理论支持和可解释性。与基于物理模型的方法相比,它们通常表现出针对性较弱的增强性能,表现为光照不均对噪声不鲁棒等。其主要原因是对某些退化的具体定义不足以及缺乏针对性的处理基于物理模型的方法将图像分解为具有物理意义的分量。然后,对分量进行特定处理以实现更具针对性的增强。

然而,现有方法很难摆脱拟合的本质。更具体地说,现有方法可以通过去噪更好地渲染失真的场景,但无法修复缺失的场景内容。以图 1 为例,最先进的方法(URetinex [40])无法恢复微弱和缺失的细节,甚至在一定程度上加剧了信息失真。为了解决这个缺点,并考虑到 LLIE 是一个在低光照图像引导下恢复正常光照图像的过程,我们使用生成式扩散模型重新思考 LLIE。我们的目标是恢复甚至推断出原始低光照图像中微弱甚至丢失的信息。因此,LLIE 不仅被视为一个恢复拟合函数,也被视为一个带条件的图像生成任务。对于生成模型,生成对抗网络(GAN)[36, 46] 通过对抗机制训练生成器和判别器。然而,它们存在训练不稳定的问题,导致模式崩溃(mode collapse)、不收敛、梯度爆炸或消失等问题。此外,基于 GAN 的 LLIE 方法也存在通过整体拟合直接生成正常光照图像的问题,如前所述,缺乏物理可解释性。

为此,我们提出了一种物理可解释的生成式模型用于低光照图像增强,称为 Diff-Retinex。我们的目标是整合物理模型和生成网络的优点。因此,Diff-Retinex 将低光照图像增强问题表述为 Retinex 分解和条件图像生成。在 Retinex 分解中,我们整合了 Transformer [21, 41] 的特性,并精心设计了一个 Retinex Transformer 分解网络(TDN)以提高分解的适用性。TDN 将图像分解为光照图和反射图。然后,我们设计了基于生成扩散的网络,分别解决这些分量中的各种退化问题,包括光照不足、噪声、颜色偏差、场景内容丢失等。主要贡献总结如下:

  • 我们从条件图像生成的角度重新思考低光照图像增强。我们不仅限于增强原始的低质量信息,还提出了一个生成式 Retinex 框架,以进一步补偿由低光照引起的内容丢失和颜色偏差。
  • 考虑到 Retinex 模型中分解的问题,我们提出了一种新颖的 Transformer 分解网络。它能够充分利用注意力和层间依赖性来高效地分解图像,即使是高分辨率图像。
  • 据我们所知,这是首次将扩散模型与 Retinex 模型结合应用于低光照图像增强的研究。扩散模型被应用于指导光照图和反射图的多路径调整,以获得更好的性能。

2. 相关工作

基于 Retinex 的低光照图像增强方法

视网膜-皮层(Retinex)理论基于颜色不变性模型和人类视觉系统(HVS)对颜色的主观感知[13]。它将图像分解为光照图和反射图。它已广泛应用于低光照图像增强,并被证明是有效且可靠的。

传统方法。在一些方法中,光照和反射模式的解析是通过高斯滤波器或一组滤波器组实现的,例如 SSR [10] 和 MSR [9]。LIME [4] 通过初始化三个通道的最大值并应用结构先验细化来估计光照图以形成最终的光照图。JED [32] 通过结合序列分解和伽马变换来增强图像并抑制噪声。传统方法主要表现出泛化性差和鲁棒性差,限制了它们的应用。

基于深度学习的方法。Retinex-Net [39] 将 Retinex 分解范式与深度学习相结合。它采用分阶段的分解和调整结构,并使用 BM3D [2] 进行图像去噪。类似地,KinD [44] 和 KinD++ [43] 采用分解和调整范式,并使用卷积神经网络(CNN)来学习分解和调整中的映射。Robust Retinex [45] 将图像分解为三个分量,即光照、反射和噪声。然后,它在损失函数的指导下通过迭代来估计噪声并恢复光照,以达到去噪和增强的目的。尽管这些方法表现出优异的性能,但由于卷积的局限性,基于 CNN 的分解无法充分利用全局信息。此外,它们也面临一些棘手的问题,例如损失函数设计的困难以及完成某些缺失场景内容的挑战。

生成式低光照图像增强方法

随着变分自编码器(VAE)[11]、GAN [36, 46] 和其他生成模型的发展,图像生成可以取得优异的结果。从一个新的角度来看,生成模型可以将低光照图像作为条件,生成相应的正常光照图像,从而客观地实现低光照图像增强的目标。EnlightenGAN [8] 设计了一个单一的生成器,直接将低光照图像映射到正常光照图像。它与全局和局部判别器结合以实现功能。CIGAN [28] 使用循环交互式 GAN 来完成正常光照和低光照图像之间光的循环生成和信息传递。这些方法取得了成果。然而,GAN 的训练过程困难,损失函数的收敛不稳定。最近,扩散模型 [7, 29, 31] 已成为强大的生成模型家族,在包括图像生成、修复等在内的许多领域都取得了破纪录的性能。它克服了 GAN 的一些缺点,打破了 GAN 在图像生成领域的长期主导地位。在本文中,我们探索了一种将 Retinex 模型与扩散模型相结合的新方法。

3. 方法

Diff-Retinex 的总体框架如图 2 所示。一个通用的基于 Retinex 的增强框架应该能够灵活地分解图像并自适应地去除各种退化。因此,Transformer 分解网络首先根据 Retinex 理论将图像分解为光照图和反射图。然后,通过多路径扩散生成调整网络(包括反射扩散调整和光照扩散调整)对光照图和反射图进行调整。增强结果是调整后的分量的乘积。
Diff-Retinex: Rethinking Low-light Image Enhancement with A Generative Diffusion Model 论文阅读_第1张图片

图 2:Diff-Retinex 的总体框架。它包含三个可分离模块:Transformer 分解网络(TDN)、反射扩散调整(RDA)和光照扩散调整(IDA)。

3.1 Transformer 分解网络

经典的 Retinex 理论假设图像可以分解为反射图和光照图,如下所示:
I = R ⋅ L , ( 1 ) I=R\cdot L,(1) I=RL,(1)

其中 I I I 是输入图像。 R R R L L L 分别表示反射图和光照图。这本质上是一个不适定问题(ill-posed problem)。反射图反映了场景内容,因此在不同光照条件下往往保持不变光照图与光照条件相关,应呈现局部平滑性

特别地,一些退化图像可能还携带不同程度的复杂噪声。在这种情况下,我们倾向于遵循光照图局部平滑的分解特性。因此,噪声被分解到反射图中。在我们的方法中实现 Retinex 分解的优化目标通常通过公式 (2) 表示:
min ⁡ R , L τ ( R ⋅ L ) + α ϕ ( R ) + β ψ ( L ) , ( 2 ) \min_{R,L}\tau(R\cdot L)+\alpha\phi(R)+\beta\psi(L),(2) R,Lminτ(RL)+αϕ(R)+βψ(L),(2)

其中 τ ( R ⋅ L ) \tau(R\cdot L) τ(RL) 确保可以从分解的光照图和反射图重建图像。 ϕ ( R ) \phi(R) ϕ(R) 约束反射图的一致性。 ψ ( L ) \psi(L) ψ(L) 使光照图结构简单且分段平滑。 α \alpha α β \beta β 是超参数。损失函数的详细设计如下。

3.1.1 损失函数

基于公式 (2),我们设计了以下损失函数,包括重建损失、反射一致性损失和光照平滑度损失,以优化 Transformer 分解网络。考虑到不同光照条件下的反射一致性,我们使用成对的低光照和正常光照图像进行训练,分别表示为 I l I_{l} Il I n I_{n} In。从中分解出的反射图分别表示为 R l R_{l} Rl R n R_{n} Rn。对应的光照图由 L l L_{l} Ll L n L_{n} Ln 表示。

重建损失 τ ( R ⋅ L ) \tau(R\cdot L) τ(RL)。它保证分解出的 R R R L L L 能够重建原始图像。因此,该损失通过考虑图像保真度来表示:
L r e c = ∥ R n ⋅ L n − I n ∥ 1 + α r e c ∥ R l ⋅ L l − I l ∥ 1 + ξ ( L e r s ) , ( 3 ) L_{rec}=\|R_{n}\cdot L_{n}-I_{n}\|_{1}+\alpha_{rec}\|R_{l}\cdot L_{l}-I_{l}\|_{1 }+\xi(L_{ers}),(3) Lrec=RnLnIn1+αrecRlLlIl1+ξ(Lers),(3)

其中 α r e c \alpha_{rec} αrec 是超参数,用于调整不同光照的贡献。 ξ ( L e r s ) \xi(L_{ers}) ξ(Lers) 是一个用于低光和正常光下光照图和反射图交叉相乘的小型辅助函数。

反射一致性损失 ϕ ( R ) \phi(R) ϕ(R)。考虑到物体的反射在各种光照条件下是不变的,我们约束不同光照条件下反射图的一致性。具体来说,可以描述为:
L r c = ∥ R n − R l ∥ 1 , ( 4 ) L_{rc}=\|R_{n}-R_{l}\|_{1},(4) Lrc=RnRl1,(4)

光照平滑度损失 ψ ( L ) \psi(L) ψ(L)。考虑到光照应该是分段平滑的(piece-wise smooth),我们通过以下方式约束它:
L s m o o t h = ∥ W T l ⋅ ∇ L l ∥ + ∥ W T n ⋅ ∇ L n ∥ , ( 5 ) L_{smooth}=\|W^{l}_{T}\cdot \nabla L_{l}\|+\|W^{n}_{T}\cdot \nabla L_{n}\|,(5) Lsmooth=WTlLl+WTnLn,(5)

其中 W T l W^{l}_{T} WTl W T n W^{n}_{T} WTn 是加权因子。它可以表示为分数形式或指数形式。为了简化过程,我们设置 W T l ← e − c ⋅ ∇ I l W^{l}_{T} \leftarrow e^{-c\cdot\nabla I_{l}} WTlecIl W T n ← e − c ⋅ ∇ I n W^{n}_{T} \gets e^{-c\cdot\nabla I_{n}} WTnecIn c c c 是约束因子。 ∇ \nabla 表示导数滤波器。该损失确保在图像平滑的区域施加较大的惩罚,而在图像光照突变的区域放松约束。

最终,总体分解损失表示为:
L = L r e c + γ r c L r c + γ s m L s m o o t h , ( 6 ) L=L_{rec}+\gamma_{rc}L_{rc}+\gamma_{sm}L_{smooth},(6) L=Lrec+γrcLrc+γsmLsmooth,(6)

其中 γ r c \gamma_{rc} γrc γ s m \gamma_{sm} γsm 是超参数。

3.1.2 网络架构

如图 2 所示,Transformer 分解网络(TDN)由两个分支组成,即反射分解分支和光照分解分支。
Diff-Retinex: Rethinking Low-light Image Enhancement with A Generative Diffusion Model 论文阅读_第2张图片
给定待分解的图像 I ∈ R H × W × 3 I\in\mathbb{R}^{H\times W\times 3} IRH×W×3,TDN 首先通过卷积投影获得其嵌入特征 F i n i t ∈ R H × W × C F_{init}\in\mathbb{R}^{H\times W\times C} FinitRH×W×C。在光照分解分支中,它由几个卷积层组成,在保证分解效果的前提下减少计算量。为了确保光照图和反射图的内在特性,并提高反射图中的恢复性能和信息保留,反射分解分支由多阶段 Transformer 编码器和解码器组成。具体来说,Transformer 编码器和解码器由注意力( A t t e n Atten Atten)模块和前馈网络( F F N FFN FFN)模块组成。通常,我们将 TDN 块中的计算表示为:
F ^ i = A t t e n ( N o r m ( F i − 1 ) ) + F i − 1 , ( 7 ) \hat{F}_{i}=Atten(Norm(F_{i-1}))+F_{i-1},(7) F^i=Atten(Norm(Fi1))+Fi1,(7)

F i = F F N ( N o r m ( F ^ i ) ) + F ^ i , ( 8 ) F_{i}=FFN(Norm(\hat{F}_{i}))+\hat{F}_{i},(8) Fi=FFN(Norm(F^i))+F^i,(8)

其中 N o r m Norm Norm 表示归一化。 F i − 1 F_{i-1} Fi1 表示当前 TDN 块的输入特征图。

注意力模块
考虑到 Transformer 中较高的注意力计算开销,时间复杂度与图像大小的平方成正比。因此,它不适合高分辨率图像分解。为了解决这个问题,我们在 TDN 中设计了一种新颖的多头深度卷积层注意力(MDLA)来计算注意力形式,如图 3 所示。在保持分解性能的前提下,它在很大程度上降低了注意力计算复杂度。
Diff-Retinex: Rethinking Low-light Image Enhancement with A Generative Diffusion Model 论文阅读_第3张图片

图 3:MDLA 的详细网络架构。注意力在跨通道方向上计算,以实现高分辨率图像的高效分解。

在 MDLA 中,对于从层归一化(Layer-Norm)获得的特征 X ∈ R h × w × c X\in\mathbb{R}^{h\times w\times c} XRh×w×c,我们首先使用 1 × 1 1\times 1 1×1 卷积聚合其通道方向的信息。随后, 3 × 3 3\times 3 3×3 5 × 5 5\times 5 5×5 7 × 7 7\times 7 7×7 卷积聚合信息。多个卷积的输出是查询 Q = W p c w t W d c q X Q=W^{wt}_{pc}W^{q}_{dc}X Q=WpcwtWdcqX、键 K = W p c k i W d c k X K=W^{ki}_{pc}W^{k}_{dc}X K=WpckiWdckX 和值 V = W p c v i W d c v X V=W^{vi}_{pc}W^{v}_{dc}X V=WpcviWdcvX。我们通过 1 × 1 1\times 1 1×1 卷积降低特征维度,同时重塑特征并在层的方向上计算注意力。具体来说,可以表示为公式 (9):
X ^ = s o f t m a x ( Q R K R / d ) ⋅ V R + X , ( 9 ) \hat{X}=softmax(Q_{R}K_{R}/d)\cdot V_{R}+X,(9) X^=softmax(QRKR/d)VR+X,(9)

其中 Q R , V R ∈ R h × w × c Q_{R},V_{R}\in\mathbb{R}^{h\times w\times c} QR,VRRh×w×c K R ∈ R c × h × w K_{R}\in\mathbb{R}^{c\times h\times w} KRRc×h×w 是重塑后的 Q Q Q V V V K K K d d d 是一个比例因子。

前馈网络模块
我们采用一个简单但有效的深度可分离前馈网络。它主要由可分离的点卷积(point-wise convolution)和深度卷积(depth-wise convolution)组成,以最小化计算量。给定层归一化后的特征 X ∈ R h × w × c X\in\mathbb{R}^{h\times w\times c} XRh×w×c,输出特征可表示为:
X ^ = W d c ( ϕ ( W p c W d c ( X ) ) ) + X , ( 10 ) \hat{X}=W_{dc}(\phi(W_{pc}W_{dc}(X)))+X,(10) X^=Wdc(ϕ(WpcWdc(X)))+X,(10)

其中 W p c W_{pc} Wpc W d c W_{dc} Wdc 分别是点卷积和深度卷积。 ϕ \phi ϕ 是激活函数。

扩散生成调整

扩散生成调整旨在构建恢复多通道的 Retinex 模型的原始数据分布。通常,它可以分为两条路径,即反射扩散调整(RDA)和光照扩散调整(IDA)。

正常光照图像分量表示为 I 0 ∈ R H × W × C I_{0}\in\mathbb{R}^{H\times W\times C} I0RH×W×C(在 RDA 中 C = 3 C=3 C=3,在 IDA 中 C = 1 C=1 C=1)用于扩散。条件图像分别与带噪图像连接以形成引导。我们采用去噪扩散概率模型(DDPM)[7] 中提出的扩散过程来为每个通道构建 Retinex 数据的分布。更具体地说,它可以描述为前向扩散过程和反向扩散过程,如图 4 所示。
Diff-Retinex: Rethinking Low-light Image Enhancement with A Generative Diffusion Model 论文阅读_第4张图片

图 4:扩散生成调整的前向和反向扩散过程示例。 I 0 I_{0} I0 是获得的结果。

前向扩散过程
前向扩散过程可以看作是一个马尔可夫链,逐步向数据添加高斯噪声。步骤 t t t 的数据仅依赖于步骤 t − 1 t-1 t1 的数据。因此,在任何 t ∈ [ 0 , T ] t\in[0,T] t[0,T] 时,我们可以获得带噪图像 I t I_{t} It 的数据分布为:
q ( I t ∣ I t − 1 ) = N ( I t ; 1 − β t I t − 1 , β t Z ) , ( 11 ) q(I_{t}|I_{t-1})=\mathcal{N}(I_{t};\sqrt{1-\beta_{t}}I_{t-1},\beta_{t}\mathcal{Z }),(11) q(ItIt1)=N(It;1βt It1,βtZ),(11)

其中 β t \beta_{t} βt 是控制添加到数据的噪声方差的变量。当 β t \beta_{t} βt 足够小时,从 I t − 1 I_{t-1} It1 I t I_{t} It 是一个添加少量噪声的恒定过程,即步骤 t t t 的分布等于前一步骤的分布加上高斯噪声。通过引入一个新变量 α t = 1 − β t \alpha_{t}=1-\beta_{t} αt=1βt,这个过程可以描述为:
I t = α t I t − 1 + 1 − α t ϵ t − 1 , ϵ t − 1 ∼ N ( 0 , Z ) , ( 12 ) I_{t}=\sqrt{\alpha_{t}}I_{t-1}+\sqrt{1-\alpha_{t}}\epsilon_{t-1},\quad\epsilon_ {t-1}\sim\mathcal{N}(0,\mathcal{Z}),(12) It=αt It1+1αt ϵt1,ϵt1N(0,Z),(12)

通过参数重整化(renormalization),合并并简化了多个高斯分布。我们可以获得第 t t t 步的分布 q ( I t ∣ I 0 ) q(I_{t}|I_{0}) q(ItI0)。更具体地说,它可以表示为:
q ( I t ∣ I 0 ) = N ( I t ; α ‾ t I 0 , ( 1 − α ‾ t ) Z ) , ( 13 ) q(I_{t}|I_{0})=\mathcal{N}(I_{t};\sqrt{\overline{\alpha}_{t}}I_{0},(1-\overline {\alpha}_{t})\mathcal{Z}),(13) q(ItI0)=N(It;αt I0,(1αt)Z),(13)

其中 α ‾ t = ∏ i = 0 t α i \overline{\alpha}_{t}=\prod_{i=0}^{t}\alpha_{i} αt=i=0tαi。当分布 q ( I t ∣ I 0 ) q(I_{t}|I_{0}) q(ItI0) 接近 N ( 0 , Z ) \mathcal{N}(0,\mathcal{Z}) N(0,Z) 时,可以认为模型完成了扩散的前向过程。

反向扩散过程

Diff-Retinex: Rethinking Low-light Image Enhancement with A Generative Diffusion Model 论文阅读_第5张图片

反向扩散过程是从纯噪声的高斯分布恢复原始分布的过程。与前向扩散过程类似,去噪扩散过程也是分步进行的。在步骤 t t t,对数据 I t I_{t} It 应用去噪操作,在条件图像 I c I_{c} Ic 的引导下获得 I t − 1 I_{t-1} It1 的概率分布。因此,给定 I t I_{t} It,我们可以将条件概率分布 I t − 1 I_{t-1} It1 表示为:
p θ ( I t − 1 ∣ I t , I c ) = N ( I t − 1 ; μ θ ( I t , I c , t ) , σ t 2 Z ) , ( 14 ) p_{\theta}(I_{t-1}|I_{t},I_{c})=\mathcal{N}(I_{t-1};\mu_{\theta}(I_{t},I_{c},t ),\sigma^{2}_{t}\mathcal{Z}),(14) pθ(It1It,Ic)=N(It1;μθ(It,Ic,t),σt2Z),(14)

其中 μ θ ( I t , I c , t ) \mu_{\theta}(I_{t},I_{c},t) μθ(It,Ic,t) 是平均值,来自步骤 t t t 的估计。 σ t 2 \sigma^{2}_{t} σt2 是方差。在 RDA 和 IDA 中,我们遵循 DDPM 的设置,将其设为固定值。更详细地说,它们可以进一步表示为:
μ θ ( I t , I c , t ) = 1 α t ( I t − β t ( 1 − α ‾ t ) ϵ θ ( I t , I c , t ) ) , ( 15 ) \mu_{\theta}(I_{t},I_{c},t)=\frac{1}{\sqrt{\alpha}_{t}}(I_{t}-\frac{\beta_{t}}{ (1-\overline{\alpha}_{t})}\epsilon_{\theta}(I_{t},I_{c},t)),(15) μθ(It,Ic,t)=α t1(It(1αt)βtϵθ(It,Ic,t)),(15)

σ t 2 = 1 − α ‾ t − 1 1 − α ‾ t β t , ( 16 ) \sigma^{2}_{t}=\frac{1-\overline{\alpha}_{t-1}}{1-\overline{\alpha}_{t}}\beta _{t},(16) σt2=1αt1αt1βt,(16)

其中 ϵ θ ( I t , I c , t ) \epsilon_{\theta}(I_{t},I_{c},t) ϵθ(It,Ic,t) 是深度神经网络的估计值,给定输入 I t I_{t} It I c I_{c} Ic 和时间步 t t t

在反向扩散过程的每一步 t ∈ [ 0 , T ] t\in[0,T] t[0,T] 中,我们优化一个关于网络估计的噪声与实际添加的噪声 ϵ \epsilon ϵ 的目标函数。因此,反向扩散过程的损失函数为:
L d i f f ( θ ) = ∥ ϵ − ϵ θ ( α ‾ t I 0 + 1 − α ‾ t ϵ , I c , t ) ∥ , ( 17 ) L_{diff}(\theta)=\|\epsilon-\epsilon_{\theta}(\sqrt{\overline{\alpha}_{t}}I_{0 }+\sqrt{1-\overline{\alpha}_{t}}\epsilon,I_{c},t)\|,(17) Ldiff(θ)=ϵϵθ(αt I0+1αt ϵ,Ic,t),(17)

反向扩散过程中的去噪网络通常结合了 UNet 和注意力的特性。在 RDA 和 IDA 中,我们采用 SR3 [33] 的主干网络,并遵循扩散去噪网络的设计,由多个堆叠的残差块结合注意力组成。从网络预测的噪声中,我们可以估计近似的 I ~ 0 \widetilde{I}_{0} I 0。保持逼近的 I ~ 0 \widetilde{I}_{0} I 0 与正常光照图像具有一致的内容信息是有意义的。我们采用一个一致性网络(consistent network)来实现这个过程:
I ~ 0 = 1 α ‾ t ( I t − 1 − α ‾ t ϵ θ ( I t , I c , t ) ) , ( 18 ) \widetilde{I}_{0}=\frac{1}{\sqrt{\overline{\alpha}_{t}}}(I_{t}-\sqrt{1- \overline{\alpha}_{t}}\epsilon_{\theta}(I_{t},I_{c},t)),(18) I 0=αt 1(It1αt ϵθ(It,Ic,t)),(18)

L c o n t e n t = ∥ I 0 − ϵ c ( I ~ 0 , t ) ∥ 1 , ( 19 ) L_{content}=\|I_{0}-\epsilon_{c}(\widetilde{I}_{0},t)\|_{1},(19) Lcontent=I0ϵc(I 0,t)1,(19)

在一致性网络 ϵ c \epsilon_{c} ϵc 中,RDA 部分采用 Restormer [41] 的主干网络,并添加了与时间嵌入(time embedding)进行特征仿射(feature affine)。IDA 部分采用与去噪网络相同的结构。整个扩散模型网络的损失函数由下式给出:
L = L d i f f + γ e l L c o n t e n t , ( 20 ) L=L_{diff}+\gamma_{el}L_{content}, (20) L=Ldiff+γelLcontent,(20)

总的来说,整个扩散生成调整过程是从低光照 Retinex 分解分布恢复原始 Retinex 分解分布。我们可以将整个扩散过程表述为:
R ^ D G A = F R D A ( ϵ s ( r ) , R T D N ) , ( 21 ) \hat{R}_{DGA}=\mathcal{F}_{RDA}(\epsilon_{s}^{(r)},R_{TDN}),(21) R^DGA=FRDA(ϵs(r),RTDN),(21)

L ^ D G A = F I D A ( ϵ s ( i ) , L T D N ) , ( 22 ) \hat{L}_{DGA}=\mathcal{F}_{IDA}(\epsilon_{s}^{(i)},L_{TDN}),(22) L^DGA=FIDA(ϵs(i),LTDN),(22)

其中 ϵ s ( r ) ∈ R H × W × 3 \epsilon_{s}^{(r)}\in\mathbb{R}^{H\times W\times 3} ϵs(r)RH×W×3 ϵ s ( i ) ∈ R H × W × 1 \epsilon_{s}^{(i)}\in\mathbb{R}^{H\times W\times 1} ϵs(i)RH×W×1 是通过初始化生成的高斯噪声。 R T D N R_{TDN} RTDN L T D N L_{TDN} LTDN 是由 TDN 获得的反射图和光照图。

最终,增强图像作为扩散生成调整后的光照图和反射图的乘积获得,即 I ^ = R ^ D G A ⋅ L ^ D G A \hat{I}=\hat{R}_{DGA}\cdot\hat{L}_{DGA} I^=R^DGAL^DGA

4. 实验

实现细节与数据集

实现细节。所提出的 Diff-Retinex 是分开训练的。首先训练 TDN。根据经验,我们设置 γ r c = 0.1 \gamma_{rc}=0.1 γrc=0.1 γ s m = 0.1 \gamma_{sm}=0.1 γsm=0.1 α r e c = 0.3 \alpha_{rec}=0.3 αrec=0.3。学习率为 l r = 0.0001 lr=0.0001 lr=0.0001,批量大小为 16,使用 Adam 优化器。然后,我们训练与扩散生成调整相关的网络。IDA 和 RDA 的步数设置为 t = 1000 t=1000 t=1000 γ c l = 1 \gamma_{cl}=1 γcl=1。输入图像大小为 160 × 160 160\times 160 160×160,批量大小为 16。使用学习率为 0.0001 的 Adam 优化器在网络上训练 800K 次迭代。所有实验均在配备 PyTorch [30] 框架的 NVIDIA GeForce RTX 3090 GPU 上进行。

数据集。为了验证泛化性,我们在 LOL [39] 和 VE-LOL-L [17] 数据集上进行了实验。LOL 数据集中的所有图像都是在现实生活中拍摄的。我们使用 485 对图像进行训练,15 张低光照图像进行测试。VE-LOL 数据集包含用于高级和低级视觉任务的数据,分别称为 VE-LOL-H 和 VE-LOL-L。VE-LOL-L 也被用来评估我们方法的有效性。DICM 被用作通过交叉测试进行泛化性评估的数据集。

结果与分析

我们与最先进的方法进行了定量和定性比较,包括传统方法和基于深度学习的方法。传统方法包括基于光照估计的 LIME [4] 和基于 Retinex 分解与联合去噪的 JED [32]。基于学习的方法包括 RetinexNet [39]、KinD [44]、KinD++ [43]、RUAS [19]、EnlightenGAN [8]、URetinex [40] 和 LLFormer [38]。

定性比较。定性结果如图 5 和图 6 所示。我们的方法显示出三个明显的优势。首先也是最重要的,Diff-Retinex 具有对缺失场景进行纹理补全和推理生成的能力。 这是我们生成式扩散模型的一个显著特征,是现有方法所不具备的。如图 5 所示,右侧高亮区域是具有粗粒纹理瓷砖的地面(见真实值 Ground Truth)。所有竞争对手都无法恢复粗粒纹理瓷砖,而我们的方法可以生成与真实值相似的缺失纹理。类似地,图 6 中的跳水台和扶手在低光照图像中严重缺失和损坏。大多数方法无法补全清晰的纹理,而 Diff-Retinex 可以。其次,我们的方法表现出更好的光照和颜色保真度。 在图 5 中,低光照图像存在相当大的颜色偏差。在整个视图上,Diff-Retinex 的颜色最接近真实值。KinD、KinD++、RetinexNet 和 URetinex 出现不同程度的颜色偏差,例如 URetinex 和 KinD++ 倾向于偏黄。在图 6 中,Diff-Retinex 在场馆颜色方面也优于其他 SOTA 方法。最后,我们的结果展现出比其他方法更生动、噪声更少的纹理。 LIME 和 RetinexNet 在整个图像中残留大量噪声,影响了场景表达。EnlightenGAN 和 LLFormer 在平坦区域(例如图 5 中电脑桌下方和墙面的黑色区域)的去噪性能不尽如人意。总的来说,Diff-Retinex 在这些方面表现出明显的优势。

图 5:在 LOL 数据集上与最先进的低光照图像增强方法的定性比较。

===== 第 7 页 =====

定量比较。采用包括 FID [6]、LPIPS [42]、BIQI [27]、LOE [37] 和 PI [1] 在内的指标进行评估。FID 是用于评估图像相似度的机器特征相似度。LPIPS 是学习感知图像块相似度(learned perceptual image patch similarity),用于衡量图像差异。BIQI 是一种图像盲质量评估指数。LOE 是图像亮度的序列误差,反映了图像的自然度保留能力。PI 代表图像的主观感知质量。FID、LPIPS、BIQI、LOE 和 PI 越低,图像质量越好。在 LOL 和 VE-LOL-L 数据集上的定量结果报告在表 1 中。对于 LOL,我们的方法在生成指标 FID 和 LPIPS 上显示出比其他方法更大的优势。这表明我们的结果具有更好的机器视觉生成相似性。在亮度序列误差方面,我们的方法略低于 LLFormer。然而,得益于生成式扩散模型和 TDN,它在所有基于 Retinex 的方法中表现出最佳性能,包括 RetinexNet、KinD、KinD++ 和 URetinex。对于 VE-LOL-L,从指标角度来看,我们的方法也取得了全面的最佳性能。这表明我们的方法在各种场景下具有很强的泛化性和先进的生成增强性能。对于 DICM,如图 7 所示,我们的方法也展示了竞争力。此外,我们还在表 2 中提供了 PSNR 和 SSIM 的定量比较。

表 1:在 LOL 和 VE-LOL-L 数据集上的低光照图像增强方法的定量结果。

方法 EnlightenGAN JED Robust Retinex RUAS KinD KinD++ LIME RetinexNet Zero-DCE URetinex LLFormer Diff-Retinex
LOL
FID↓ 105.59 105.86 92.32 95.59 78.59 110.68 114.00 150.50 106.63 59.00 76.96 47.85
LPIPS↓ 0.129 0.190 0.157 0.167 0.083 0.095 0.211 0.183 0.133 0.050 0.067 0.048
BIQI↓ 30.95 32.33 42.29 43.32 26.70 26.81 37.83 29.33 34.80 23.05 28.81 19.97
LOE↓ 395.52 306.90 202.31 195.36 758.56 700.79 547.54 395.33 232.97 197.02 176.61 191.56
VELOL-L
FID↓ 92.58 110.46 79.64 100.07 65.56 98.10 98.90 158.99 93.81 48.36 79.83 47.75
LPIPS↓ 0.124 0.158 0.106 0.144 0.070 0.114 0.248 0.283 0.123 0.091 0.110 0.050
BIQI↓ 32.77 27.29 39.75 32.51 28.23 32.33 47.09 45.59 35.06 35.39 32.47 26.54
LOE↓ 422.77 330.25 128.73 168.99 239.33 623.63 554.69 531.92 228.88 166.02 177.87 149.60

图 6:在 VE-LOL-L 数据集上与最先进的低光照图像增强方法的定性比较。

图 7:在 DICM 上与 SOTA 低光照图像增强方法的 PI 定性和定量比较。

表 2:在 LOL 数据集上 PSNR 和 SSIM 的定量比较。

方法 主要类型 PSNR↑ SSIM↑
RetinexNet CNN 17.56 0.698
KinD CNN 17.64 0.829
KinD++ CNN 17.75 0.816
EnlightenGAN GAN 17.48 0.716
URetinex Unfolding 21.32 0.836
LLFormer Transformer 23.66 0.873
Diff-Retinex Diffusion 21.98 0.863

===== 第 8 页 =====

消融研究

Transformer 分解网络。为了验证 Transformer 分解网络(TDN)的有效性,我们对分解结果进行了可视化。Retinex 分解是一个不适定问题,没有确切的最优解。一个核心点是反射信息在不同光照水平下应严格一致。典型且有效的表示方法采用 CNN 进行分解,例如 RetinexNet 和 KinD++。反射图分解结果如图 8 所示。

图 8:分解网络反射图的定性比较。结构细节和噪声被假定被分解到反射图中。

生成式扩散模型。为了验证扩散模型的有效性,一方面,我们对 RDA 和 IDA 的生成过程进行了可视化,如图 9 所示。另一方面,我们比较了通过我们的扩散模型和其他一些一步式基于 Retinex 的 LLIE 方法对反射图的恢复结果。我们使用反射图进行比较,因为它包含大量对视觉感知更敏感的颜色和纹理信息。典型的基于 Retinex 的 LLIE 方法包括 RetinexNet 和 KinD++。对于反射图恢复,RetinexNet 采用 BM3D,KinD++ 采用 CNN。结果如图 10 所示。由于不同方法的 Retinex 分解结果差异很大,我们展示它们各自从正常光照图像分解出的反射图作为真实值进行比较。可以看出,我们的方法能更好地处理颜色偏差,并在纹理恢复方面表现出更好的性能。我们还计算了 FID、

表 3:使用扩散模型与其他低光照增强方法在反射图恢复上的定量比较。

方法 FID↓ LPIPS↓ BIQI↓
RetinexNet (BM3D) 111.29 0.225 24.80
KinD++ (CNN) 171.45 0.110 36.99
Diff-Retinex (Diff) 61.33 0.059 18.98

图 9:RDA 和 IDA 的示例。为了更好地展示扩散效果,采用了恒定周期采样的输出格式。从左到右:从纯噪声逐步恢复的迭代过程。

图 10:使用我们的扩散模型与其他 SOTA 方法恢复反射图的定性比较。

===== 第 9 页 =====

你可能感兴趣的:(论文阅读,人工智能,深度学习,学习,图像处理,计算机视觉)