O-Mamba: O-shape State-Space Model for Underwater Image Enhancement

TL;DR

O-Mamba addresses underwater image enhancement by modeling spatial and cross-channel information with an O-shaped dual-branch architecture that combines Spatial Mamba and Channel Mamba blocks. It introduces the Multi-scale Bi-mutual Promotion module to fuse multi-scale features and enable inter-branch interaction through MS-MoE, Mutual Promotion, and a Cyclic Multi-scale Optimization strategy. The approach delivers state-of-the-art results across multiple datasets, validated by quantitative metrics (PSNR, SSIM, LPIPS, FID, UIQM, Uranker) and extensive ablations confirming the importance of cross-channel modeling and multi-scale fusion. This work offers a scalable, efficient framework for robust UIE with strong practical impact in underwater imaging applications.

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.

Figures (4)

• Figure 1: Visualization of high-resolution image. Compared to the latest SOTA methods, our method achieves the best color correction. Please zoom to view.
• Figure 2: Overall O-Mamba is an O-shaped dual-branch network consisting of a Spatial Mamba Branch and a Channel Mamba Branch. The Spatial Mamba Branch is designed to capture spatial dependencies in images, while the Channel Mamba Branch, utilizing the Channel Mamba Block, focuses on capturing cross-channel dependencies. The interaction between the two branches is achieved through the Multi-scale Bi-mutual Promotion Module, which consists of three parts: MS-MoE for integrating multi-scale information within branches, the Mutual Promotion module for integrating spatial and channel information between branches and the Cyclic Multi-scale (CMS) optimization strategy for optimize multi-scale losses.
• Figure 3: The architecture of the Feed Forward MoE and Multi-scale Mamba MoE.
• Figure 4: Comparison of inputs and results from different methods.