NeurIPS2024｜ECMamba 结合Retinex Theory和Mamba实现高效曝光矫正！

本篇分享 NeurIPS2024 论文ECMamba: Consolidating Selective State Space Model with Retinex Guidance for Efficient Multiple Exposure Correction，ECMamba 结合 Retinex Theory 和 Mamba 实现高效曝光矫正！

• 作者：Wei Dong等
• 论文链接：https://arxiv.org/pdf/2410.21535
• 代码链接：https://github.com/LowlevelAI/ECMamba

亮点直击

• 首次采用Mamba去处理多重曝光校正问题。
• 提出一种名为ECMamba的新型曝光校正模型，该模型采用双支路结构并且引入Retinex 信息作为指导ECMamba去分离光照信息和物体的内在反射特性。
• 为了使Mamba有效处理图像数据，我们将图像数据扫描视为一个“特征敏感型”问题，并提出了一个具有Retinex信息指导的特征感知的2D扫描机制
• 大量实验和消融研究证明了ECMamba在多重曝光校正和低光照图像增强任务上有优异表现。

在当前基于深度学习的曝光校正研究中，很少有模型完全将Retinex理论嵌入到其架构中。此外，如何平衡高性能与高效率仍然值得探索。

本文提出一种基于Retinex theory 和 Mamba的双通道高效曝光校正模型ECMamba。首先通过理论分析，本文将输入分别映射到与目标reflectance和 illumination map接近的两个过渡空间。

此外在Mamba架构中，为实现精细曝光校正，本文开发了一种基于可变形特征聚合的新型2D扫描策略。大量实验结果表明ECMamba的优异性能。

方法

本文提出的方法的概览如图2所示，可见本文提出的ECMamba是基于Mamba和Retinex理论设计的。

首先，通过分析Retinex 理论，作者提出了双支路的曝光校正架构，每个支路分别旨在修复reflectance 和illumination map。

其次，在每个支路中，本文利用ECMM模块实现对reflectance 或者illumination map的精确恢复。

此外，为了进一步增强效率和效果，本文开发了一种新的2D SSM 层以及一种基于可变形特征聚合的新型2D扫描策略。

基于Retinex 理论的双支路曝光校正架构

Retinex 理论可以表示为 , 其中 表示 Hadamard 乘积, 是一个理想的无退化图像， 和 分别代表 reflectance 和 illumination map。

然而，一个在非理想照明条件下（过曝光或欠曝光场景）拍摄的低质量图像 不可避免地会受到噪声、颜色失真和对比度不足的影响。

因此, 可以分别对 和 引入了扰动（ ( 和 来模拟这些退化图像,公式（1）如下：

当前一些基于 Retinex 的方法 将反射分量 视为最终的增强结果，从而忽略了公式（1）中的最后三项，只专注于使用网络 去模拟此映射：

但是，这些模型在多重曝光校正任务中只能达到次优性能，因为我们很难获得准确映射的困难，特别是多个欠曝光 （UE）和过曝光（OE）输入对应一个正常曝光（NE）图像时。

因此, 为了获得令人满意的结果, 本文选择同时恢复 和 。

具体来说, 将公式 (1) 的两边分别与 和 进行逐元素相乘公式 (2):

其中 和 是满足 和 的矩阵, 并且本文使用 Retinex 预测器 去估计这两个矩阵。

和 分别表示 和 ‘中剩余的退化, 因此可以采用深度学习网络来获得良好曝光的图片公式 (3) :

其中 和 是用来预测 和 中负退化的网络， 是从 中得到的 Retinex 指导信息。

如图2 所示, Retinex 估计器 将 及其沿通道维度的均值矩阵作为输入 (为清晰起见，图2中没有显示该矩阵）。 首先使用 卷积和一个 的 depthwise 卷积来提取特征，然后分别通过一个 卷积生成 和 。

更为关键的是， 和 被输入到 和 进行进一步的复原。除了优化 使其接近 ，本文的训练目标还包括了对 和 的限制。

讨论：

（i）许多基于Retinex的方法 [37]旨在学习输入与反射图像和照明图之间的映射，然后通过 Hadamard 乘积获得最终结果。然而，这种策略不适用于多重曝光校正任务。图1(a)展示了过曝光（OE）和欠曝光（UE）图像的分布很复杂且距离正常曝光图像很远。这种复杂的分布使得准确建议从输入到输出的映射是极其困难的。然而，通过仔细分析Retinex理论，本文构建了一个中间空间如图1(b)所示。该空间显著减少了与优化目标的距离并有利于后续的微调恢复过程。

(ii) 一些方法 将 视为最终增强结果, 这与 Retinex 理论的原始解释偏离。因此, 本文采用了一个双分支框架, 使用独立的深度学习网络去分别重建 和 。本文在消融研究中讨论了该框架中各个部分的重要性。

ECMM 以及RMB模块

为了使提出的曝光校正架构可以应用在资源有限的设备上并且实现高性能，本文提出了一个新颖的由Retinex信息指导的ECMM模块，它成功继承了Mamba的强大建模能力。

如图2所示, ECMM采用了双尺度U-Net架构。在编码过程中, 输入 首先通过一个 的卷积和一个RMB模块以获得初始特征 。然后, 通过一个步长为 2 的 卷积实现下采样, 并且将下采样后的特征输入到另一个RMB模块中, 以获得中间特征 。在解码阶段， 首先被一个步长为 2 的 反卷积上采样得到 。为了降低下采样过程中造成的信息丢失, 本文采用了一种自适应混合特征融合策略[46] 将编码信息传递到解码阶段公式（4）:

其中 是一个可学习的参数, 代表 sigmoid 方程。融合后的特征 依次进入 和卷积层，按照公式（3）通过残差叠加的方式来获得修复后的反射分量 。

作为 ECMM 模块的核心, RMB 块 采用了与 Transformer 相似的结构。然而，需要大量计算的自注意力和交叉注意力机制影响了 Transformer 的效率，使其无法应用于实时或资源受限的环境中。为此，本文去除了注意力过程，并引入了一种新型的由 Retinex 信息指导的 2D SSM 层来捕捉远距离依赖并促进动态特征聚合。因此, RMB 块可以这样表示公式（5）：

其中, LN 表示 LayerNorm, 是由 Retinex 估计器 提取的 Retinex 指导信息。此外，受到 ConvNext 的启发，本文移除了门控机制和深度卷积，引入了一个效率更高的前向反馈层（EFF）。该层遵循 卷积 卷积的流程, 其操作类似于 Transformers 中的 MLPs 但只需要更少的参数。

Retinex-SS2D 层和基于可变形特征聚合的新型 2D 扫描策略

图 3 详细展示了 Retinex-SS2D 层的结构。首先通过线性运算、depth-wise 卷积、逐元素相乘及 SiLU操作对输入特征 和 进行融合。然后, 融合后的特征 被输入到本文提出的 FA-SS2D 机制中，该机制旨在捕捉动态的远距离依赖并实现自适应空间聚合。此外, 还运用了门控信号 和线性运算来得到最终的聚合特征 。

基于可变形特征聚合的新型2D扫描策略选择性状态空间模型（S6）在处理涉及时间序列的自然语言处理等任务时表现出色。但是，将其应用于二维图像时面临不小挑战。为了更好地处理二维图像中的空间信息，当前研究提出多种扫描策略将图像块转换为一维序列。举例来说，文献[27]提出了一种交叉扫描策略。该策略沿四个不同的扫描路径生成四个序列，每个序列都单独进行S6处理。然而，这种策略显著增加了计算负担，这与S6本身高效、低计算需求的特性相悖。此外，这些方法只是简单地在不同方向进行图像扫描，导致在某些序列中局部纹理和全局结构发生分离。这种分离在一定程度上影响了S6模型对图像的建模效果。

现行扫描方法的局限性促使作者重新考虑如何更有效地应用S6于二维图像。在S6 过程中， 对于一维序列中的每一个元素，其输出y(t)取决于它的输入x(t)和之前的输入{x(1), x(2), …, x(t − 1)}。这种机制要求从二维图像转化而来的一维序列必须满足以下两个条件才能确保优良的表现：(1) 序列应在起始部分优先处理最关键的特征区域，而将次要信息置于末尾。(2) 空间相邻的特征应当被紧密地排序，以避免序列中出现显著的间隙。然而，现有的二维扫描策略未能满足这两个要求，这促使作者提出新的解决方案来弥补这一差距。

基于这些观察，本文引入了一种高效的FA-SS2D机制如图3所示。本文首先开发了一种由Retinex信息指导的可变形特征聚合操作。具体而言，采用可变形卷积（DCN）[48, 9]来捕捉融合特征的动态远程依赖。

例如，当DCN应用于图3中内的红框时，此时激活区域是一个不规则的并以蓝色标记。更为重要的是，当红框在上滑动时，这个不规则的区域会发生变化。最终可以得到每个小框的平均激活频率，进而得到如图3中所示的激活响应图，其中激活频率高的区域代表重要的特征。

具体而言，对于欠曝光图像的相对较亮区域或过曝光图像的相对正常曝光区域，这些区域包含重要特征从而具有比较大的激活响应。

基于获得的激活响应图，本文提出了一种新的扫描策略，与“方向敏感”的扫描方法[27]不同，本文提出的策略根据不同区域的激活频率进行排序，并将频率高的标记放在序列的前端，从而使得生成的序列有效地满足Mamba的要求，显著地提升了其对图像数据建模的能力。

损失函数

在本文中, 本文选择了一种端对端训练策略去同时优化 和 。本文的最终训练目标是使 接近 ，同时我们也在 和 上施加了若干约束以实现稳定训练。因此, 完整的损失函数如下所示公式（6）：

其中 和 是应用于 和 上的约束, 它们本质上采用自监督策略来学 习 和 。此外, 考虑到这一优化问题本质上是不适定问题，本文引入 引导优化朝向合适的方向进行。 是训练过程中的主要损失函数, 可以通过以下方式计算公式（7）：

其中， 表示结构相似性损失； 表示由 VGG19［30］提取的特征之间的差异。在本工作中，对应的损失函数系数设置为：

实验

与现有方法的比较

• 比较的方法： 包括 ZeroDCE[16], URetinexNet[37], Retiformer[6], LLFlow[34], LLFlow-SKF[38], DRBN[40], FECNet[20], LACT[4]等
• 数据集： ME[1], SICE[5], LOL[36]，LOL-v2-real[41], LOL-v2-synthetic[41], 分别在这些数据集的训练数据上训练，在相应测试集上评估。

定量比较

如表1所示， ECMamba在ME数据集和SICE数据集中处理欠曝光和过曝光图像时，显著优于当前的最先进方法。特别是在PSNR和SSIM方面，ECMamba在ME数据集上超过了排名第二的方法0.19分贝和0.007。此外，相比于排名第二的方法，本文提出的ECMamba在SICE数据集上的改进提升了0.83分贝和0.051。

表2总结了ECMamba与当前最先进方法在欠曝光校正（低光照图像增强）上的定量比较。特别地，ECMamba在仅使用4.4%的参数量的情况下，在PSNR上超过了排名第二的LLFlow-SKF 1.10分贝，显示出ECMamba的优异效果和高效率。表1和表2表明验证了本文提出的ECMamba和基于Retinex的双分支框架的有效性。

定性比较

图4（ME）和图5（SICE）展示了各种方法处理后的图像。总体来说，ECMamba处理后的图像具有令人愉悦的光照效果、正确的颜色恢复和更多的纹理细节，这些结果既吸引人又逼真。

例如，在图4的第3行可以看到云的丰富结构细节，在第6行可以看到山体表面的细节；在图5的第1行中桥梁及其边缘轮廓被很好地保留，在第2行中展示板上的文字呈现生动。

相比之下，先前的方法在维持颜色真实性和光照一致性方面存在困难。

消融实验

为了验证所提出的ECMamba的有效性，本文在SICE数据集上进行了广泛的消融实验。

双分支Retinex框架的贡献：作者首先移除了用于精确恢复照明图 的分支 ），因此剩余的网络旨在将 优化至真实值，其性能如表3所示。尽管其性能与表 1中当前的最好的表现相比仍有竞争力, 但与我们完整的ECMamba相比, 仅优化 不可避免地导致次优表现。此外, 本文还增加了 的结构复杂度, 使其参数与原始的双分支框架相当。然而, 与双分支ECMamba相比, 这个网络的性能仍然较差。最后，与其他基于Retinex的方法[37]相似，作者随后移除了Retinex估计器 , 直接采用剩余网络进行曝光校正。然而, 如表3所示, 这种调整大大降低了 ECMamba的性能，突显了在双分支Retinex框架中引入过渡空间（ 和 ）的重要性。

ECMM模块的重要性在双分支框架中，作者将原有的ECMM模块替换为了ViT [11] 和Retiformer[6]架构。如表3所示，ECMM模块的表现优于ViT。更重要的是，ECMamba的效率与著名的高效率低光照增强方法Retiformer相当。另外，为验证所提出的Retinex-SS2D层和FA-SS2D的重要性，本文使用VMamba[27]提出的交叉扫描机制替换了Retinex-SS2D层。这一改变导致参数数量增加和性能降低，从而显示出本文提出的Retinex-SS2D层和FA-SS2D策略的显著优势。

结论

针对曝光校正问题，本文提出了一种新型的基于Retinex的双分支Mamba架构 (ECMamba) 。首先，通过详细推导Retinex理论，本文开发了一个由Retinex信息指导的双分支框架。

为了更有效地平衡性能与效率，本文设计了ECMM作为主修复模块，配备了高效的Retinex引导SS2D层和特征感知扫描策略。

大量实验显示，本文的ECMamba在多重曝光和低光照图像增强数据集上均显著优于现有最先进技术。此外，本文有指出了未来的研究方向。

例如，与其他技术类似，在极端曝光案例中，由于图像退化导致的信息大量丢失，ECMamba也难以达到令人满意的效果。但近期一些图像恢复方法通过使用生成式的先验知识来推断退化细节，已经取得了不错的成果。

未来，作者计划将Mamba与生成式先验模型结合，以减轻极端曝光案例中的性能降低问题。

公众号后台回复“数据集”获取100+深度学习各方向资源整理