MambaUIE&SR: Unraveling the Ocean's Secrets with Only 2.8 GFLOPs

Zhihao Chen,Yiyuan Ge

2024-05-24

Abstract:Underwater Image Enhancement (UIE) techniques aim to address the problem of underwater image degradation due to light absorption and scattering. In recent years, both Convolution Neural Network (CNN)-based and Transformer-based methods have been widely explored. In addition, combining CNN and Transformer can effectively combine global and local information for enhancement. However, this approach is still affected by the secondary complexity of the Transformer and cannot maximize the performance. Recently, the state-space model (SSM) based architecture Mamba has been proposed, which excels in modeling long distances while maintaining linear complexity. This paper explores the potential of this SSM-based model for UIE from both efficiency and effectiveness perspectives. However, the performance of directly applying Mamba is poor because local fine-grained features, which are crucial for image enhancement, cannot be fully utilized. Specifically, we customize the MambaUIE architecture for efficient UIE. Specifically, we introduce visual state space (VSS) blocks to capture global contextual information at the macro level while mining local information at the micro level. Also, for these two kinds of information, we propose a Dynamic Interaction Block (DIB) and Spatial feed-forward Network (SGFN) for intra-block feature aggregation. MambaUIE is able to efficiently synthesize global and local information and maintains a very small number of parameters with high accuracy. Experiments on UIEB datasets show that our method reduces GFLOPs by 67.4% (2.715G) relative to the SOTA method. To the best of our knowledge, this is the first UIE model constructed based on SSM that breaks the limitation of FLOPs on accuracy in UIE. The official repository of MambaUIE at this https URL.

Image and Video Processing,Computer Vision and Pattern Recognition

What problem does this paper attempt to address?

The paper aims to address the problem of Underwater Image Enhancement (UIE). Specifically, the paper proposes a new architecture based on the State Space Model (SSM) called MambaUIE to improve the efficiency and accuracy of underwater image enhancement. Existing methods often perform poorly in dealing with underwater image degradation due to the effects of light absorption and scattering. Although Convolutional Neural Networks (CNN) and Transformers have shown excellent performance in UIE tasks, they usually suffer from insufficient local feature extraction or high computational complexity. MambaUIE introduces Visual State Space Blocks (VSS) and Dynamic Interaction Blocks (DIB), as well as a Spatial Feedforward Network (SGFN), which can effectively integrate global and local information while maintaining efficiency, thereby significantly improving the quality of underwater images. Experimental results show that MambaUIE outperforms existing state-of-the-art methods on the UIEB dataset and consumes only very low computational resources (2.715 GFLOPs).