WaterMamba: Visual State Space Model for Underwater Image Enhancement

TL;DR

WaterMamba introduces a novel state space model for underwater image enhancement that efficiently models pixel and channel dependencies, outperforming existing methods with fewer resources.

Contribution

It proposes a new linear-complexity state space model with SCOSS blocks for underwater image enhancement, addressing limitations of CNN and Transformer approaches.

Findings

Outperforms state-of-the-art methods on multiple datasets

Reduces parameters and computational resources

Demonstrates strong generalizability

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.