[论文解读] O-Mamba: O-shape State-Space Model for Underwater Image Enhancement

Chenyu Dong, Chen Zhao+2 more|arXiv (Cornell University)|Aug 23, 2024

Image Enhancement Techniques被引用 5

一句话总结

O-Mamba 通过一个 O 形双分支的 Mamba 网络（包含 Spatial Mamba 和 Channel Mamba）来建模水下图像增强中的空间信息和跨通道信息，使用多尺度 Bi-mutual Promotion 模块来融合多尺度特征，在多个数据集上实现最先进的结果（SOTA）

ABSTRACT

Underwater image enhancement (UIE) face significant challenges due to complex underwater lighting conditions. Recently, mamba-based methods have achieved promising results in image enhancement tasks. However, these methods commonly rely on Vmamba, which focuses only on spatial information modeling and struggles to deal with the cross-color channel dependency problem in underwater images caused by the differential attenuation of light wavelengths, limiting the effective use of deep networks. In this paper, we propose a novel UIE framework called O-mamba. O-mamba employs an O-shaped dual-branch network to separately model spatial and cross-channel information, utilizing the efficient global receptive field of state-space models optimized for underwater images. To enhance information interaction between the two branches and effectively utilize multi-scale information, we design a Multi-scale Bi-mutual Promotion Module. This branch includes MS-MoE for fusing multi-scale information within branches, Mutual Promotion module for interaction between spatial and channel information across branches, and Cyclic Multi-scale optimization strategy to maximize the use of multi-scale information. Extensive experiments demonstrate that our method achieves state-of-the-art (SOTA) results.The code is available at https://github.com/chenydong/O-Mamba.

研究动机与目标

• 应对由复杂照明、散射和颜色失真导致的水下图像退化。
• 开发一个双分支架构，以分别建模空间信息和跨通道信息。
• 通过 MSBMP 实现跨分支交互和多尺度信息融合，以提升 UIE 性能。

提出的方法

• 提出一个具有 Spatial Mamba 分支（SM）和 Channel Mamba 分支（CM）的 O 形双分支网络。
• 使用 Spatial Mamba Block 和 Channel Mamba Block 通过状态空间建模（SSM）捕捉空间和跨通道依赖关系。
• 引入由 MS-MoE（Multi-scale MoE）、Mutual Promotion（MP）和 Cyclic Multi-scale optimization（CMS）组成的 Multi-scale Bi-mutual Promotion（MSBMP）模块。
• 扩展 SSM，使其沿四个空间方向进行 2D 选择性扫描（S-SSM）以及一个跨通道扫描模块（C-SSM）。
• 在 Mamba 块内将 FF-MoE 作为前馈网络，以应对多样化的水下条件。
• 实现 MS-MoE，在每个分支内使用当前分支的专家块来融合多尺度特征。
• 应用循环多尺度损失（CMS-OS）以在训练中减小尺度间的干扰。

实验结果

研究问题

• RQ1 一个双分支的基于 Mamba 的架构是否能有效建模 UIE 的空间信息和跨通道信息？
• RQ2 通过 MSBMP 的多尺度跨分支交互是否比单分支或非 MSBMP 变体提升重建质量？
• RQ3 循环多尺度优化对训练稳定性和各尺度性能有何影响？
• RQ4 与若干 SOTA UIE 方法相比，O-Mamba 在多个水下基准数据集上有何表现？

主要发现

• O-Mamba 在多个 UIE 数据集和指标上实现了最先进的性能。
• 消融实验表明双分支设计优于单分支变体（O-Mamba-S 与 O-Mamba-C）。
• MP 模块显著提升跨分支信息交互，使得消融中 PSNR 从 29.82 提升到 30.00。
• MS-MoE 对多尺度信息融合至关重要，并带来显著的性能提升。
• 循环多尺度优化有助于降低尺度间干扰并提升整体结果。
• 与若干 SOTA UIE 方法相比，O-Mamba 在 LSUI 和 UIEB 等数据集上展现出更优的定量结果（如研究中所述）。