WaterMamba: Visual State Space Model for Underwater Image Enhancement

📄 arXiv: 2405.08419v1 📥 PDF

作者: Meisheng Guan, Haiyong Xu, Gangyi Jiang, Mei Yu, Yeyao Chen, Ting Luo, Yang Song

分类: cs.CV

发布日期: 2024-05-14

备注: arXiv admin note: substantial text overlap with arXiv:2403.06098

💡 一句话要点

提出WaterMamba，一种用于水下图像增强的视觉状态空间模型，兼顾性能与效率。

🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture)

关键词: 水下图像增强 状态空间模型 视觉SSM 全向选择性扫描 深度学习

📋 核心要点

  1. 现有基于CNN的水下图像增强方法难以建模长程依赖，而Transformer方法参数量大，计算效率低。
  2. WaterMamba利用状态空间模型，通过空间-通道全向选择性扫描块建模像素和通道信息流，解决依赖关系。
  3. 实验结果表明，WaterMamba在多个数据集上优于现有方法，同时降低了参数量和计算资源消耗。

📝 摘要（中文）

水下成像常因光在水中的传播和吸收而质量下降。为改善图像质量，已提出一些基于卷积神经网络(CNN)和Transformer的水下图像增强(UIE)方法。然而，基于CNN的UIE方法在建模长程依赖方面受限，而基于Transformer的方法涉及大量参数和复杂的自注意力机制，带来效率挑战。考虑到计算复杂度和严重的水下图像退化，本文提出了一种用于UIE的、具有线性计算复杂度的状态空间模型(SSM)，名为WaterMamba。我们提出了空间-通道全向选择性扫描(SCOSS)块，该模块包含空间-通道坐标全向选择性扫描(SCCOSS)模块和多尺度前馈网络(MSFFN)。SCOSS块建模像素和通道信息流，解决依赖关系。MSFFN促进信息流调整并促进SCCOSS模块内的同步操作。大量实验表明，WaterMamba以更少的参数和计算资源实现了最先进的性能，在各种数据集上优于最先进的方法，验证了其有效性和泛化性。代码将在接收后在GitHub上发布。

🔬 方法详解

问题定义：水下图像由于光线在水中的吸收和散射，质量严重下降，导致对比度低、色彩失真等问题。现有的基于CNN的方法难以捕捉长距离依赖关系，而基于Transformer的方法计算复杂度高，难以部署。因此，需要一种既能有效建模长距离依赖，又能保持计算效率的水下图像增强方法。

核心思路：WaterMamba的核心思路是利用状态空间模型（SSM）的线性计算复杂度和建模长距离依赖的能力，设计一种高效的水下图像增强网络。通过空间-通道全向选择性扫描（SCOSS）块，同时考虑空间和通道维度的信息交互，增强网络对水下图像特征的提取能力。

技术框架：WaterMamba主要由一系列SCOSS块组成。每个SCOSS块包含一个空间-通道坐标全向选择性扫描（SCCOSS）模块和一个多尺度前馈网络（MSFFN）。SCCOSS模块负责建模像素和通道之间的依赖关系，MSFFN负责调整信息流并促进SCCOSS模块内的同步操作。整个网络采用编码器-解码器结构，逐步提取和恢复图像细节。

关键创新：WaterMamba的关键创新在于SCOSS块的设计，特别是SCCOSS模块。SCCOSS模块通过全向选择性扫描，能够更全面地捕捉空间和通道维度上的信息依赖关系，从而更好地处理水下图像的复杂退化。此外，MSFFN的设计有助于信息流的调整和模块间的协同工作。

关键设计：SCCOSS模块采用坐标注意力机制，在空间和通道维度上分别进行注意力计算，然后将结果融合。全向扫描是指从多个方向（例如，水平、垂直、对角线）扫描图像，以捕捉不同方向上的依赖关系。MSFFN采用多尺度卷积，以提取不同尺度的特征。损失函数方面，可能采用了L1损失、感知损失或对抗损失等，具体细节未知。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明，WaterMamba在多个水下图像增强数据集上取得了state-of-the-art的性能，并且参数量和计算复杂度显著低于基于Transformer的方法。具体性能数据和对比基线未知，但摘要强调了其在性能、参数量和计算资源方面的优势。

🎯 应用场景

WaterMamba在水下机器人视觉、海洋生物研究、水下考古、水下工程检测等领域具有广泛的应用前景。它可以提高水下图像的清晰度和可读性，帮助研究人员更好地分析水下环境，提高水下作业的效率和安全性。未来，该方法可以进一步扩展到水下视频增强、水下目标检测等任务。

📄 摘要（原文）

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.