[论文解读] WaterMamba: Visual State Space Model for Underwater Image Enhancement

Meisheng Guan, Haiyong Xu+5 more|arXiv (Cornell University)|May 14, 2024

Image Enhancement Techniques被引用 6

一句话总结

WaterMamba 引入了具有线性复杂度的视觉状态空间模型及 SCOSS 块，用于高效的水下图像增强，在多数据集上实现最优 PSNR/SSIM，且参数量更少、FLOPs 更低。

ABSTRACT

Underwater imaging often suffers from low quality due to factors affecting light propagation and absorption in water. To improve image quality, some underwater image enhancement (UIE) methods based on convolutional neural networks (CNN) and Transformer have been proposed. However, CNN-based UIE methods are limited in modeling long-range dependencies, and Transformer-based methods involve a large number of parameters and complex self-attention mechanisms, posing efficiency challenges. Considering computational complexity and severe underwater image degradation, a state space model (SSM) with linear computational complexity for UIE, named WaterMamba, is proposed. We propose spatial-channel omnidirectional selective scan (SCOSS) blocks comprising spatial-channel coordinate omnidirectional selective scan (SCCOSS) modules and a multi-scale feedforward network (MSFFN). The SCOSS block models pixel and channel information flow, addressing dependencies. The MSFFN facilitates information flow adjustment and promotes synchronized operations within SCCOSS modules. Extensive experiments showcase WaterMamba's cutting-edge performance with reduced parameters and computational resources, outperforming state-of-the-art methods on various datasets, validating its effectiveness and generalizability. The code will be released on GitHub after acceptance.

研究动机与目标

• 为具有非均匀退化和颜色通道损失的水下图像增强（UIE）提供动机。
• 开发一个轻量、高效的 UIE 模型，利用状态空间建模捕获长距离依赖。
• 设计 SCOSS 块以联合建模空间与通道信息，实现鲁棒增强。
• 将 SCOSS 块整合到 U-Net 骨干网络中，以重建高质量的水下图像。
• 在多个基准数据集上进行评估，以证明相对于 SOTA 方法的效率和有效性。

提出的方法

• 采取具有线性复杂度的结构化状态空间模型（SSM）用于图像增强（WaterMamba）。
• 引入由 SCCOSS（空间）和 COOSS（通道）模块加上 MSFFN 进行多尺度特征处理的 SCOSS 块。
• 在 SCCOSS 中，实施具有四个方向序列的 2D 选择性扫描，以建模长程的空间依赖。
• 在 COOSS 中，应用基于坐标的通道注意力，结合空间池化与基于 Mamba 的通道处理。
• 使用 MSFFN 融合多尺度特征并避免模式崩溃，最后进行残差学习。
• 将 WaterMamba 组装成类似 U-Net 的编码-解码器，带跳跃连接和残差输出 DR+I。

实验结果

研究问题

• RQ1 WaterMamba 是否能够在基于线性复杂度的 SSM 架构下实现 UIE 的最先进性能？
• RQ2 SCOSS 块在非均匀退化下是否能有效建模水下图像的空间与通道依赖？
• RQ3 在 PSNR、SSIM、UIQM、UCIQE 和计算效率方面，WaterMamba 与基于 CNN 和 Transformer 的 UIE 方法相比如何？

主要发现

• WaterMamba 在 R90 上实现 PSNR/SSIM 为 24.715/0.931，在 U100 上为 21.992/0.843（表 1 中的 UIEB 派生数据集）。
• 在类似 UIEB 的评估中，WaterMamba 获得具有竞争力的 UIQM 和 UCIQE 分数（UIQM 0.555；基于 SQUID 的度量为 UCIQE 2.767）。
• WaterMamba 使用 3.69M 参数和 7.53G FLOPs，与像 Restormer 这样的 Transformer 基线相比显著降低了复杂度。
• 定性结果显示在多样水下场景（珊瑚礁和地形）中细节更锐利，颜色恢复更自然。
• WaterMamba 在多个数据集（UIEB、UCIOD、UCCS、SQUID）的 PSNR/SSIM 指标上胜过 10 种 SOTA 方法，且计算成本更低。