UIE-UnFold: Deep Unfolding Network with Color Priors and Vision Transformer for Underwater Image Enhancement

TL;DR

This work introduces UIE-UnFold, a deep unfolding network for underwater image enhancement that explicitly embeds color priors via a Color Prior Guidance Block, models degradation with a Nonlinear Activation Gradient Descent Module, and enables cross-stage feature refinement through an Inter Stage Feature Transformer. By unrolling an optimization-like process and leveraging a vision-transformer-based inter-stage module, the method achieves robust improvements across diverse underwater conditions and datasets, outperforming many state-of-the-art approaches in both objective metrics and perceptual quality. The framework combines implicit color mappings, learned degradation approximations, and global context modeling to deliver more accurate color restoration, contrast, and detail preservation, with publicly available code to facilitate reproducibility. Overall, UIE-UnFold demonstrates that integrating physical priors with learnable components in a structured unfolding architecture yields reliable UIE performance with potential benefits for marine science and related applications.

Abstract

Underwater image enhancement (UIE) plays a crucial role in various marine applications, but it remains challenging due to the complex underwater environment. Current learning-based approaches frequently lack explicit incorporation of prior knowledge about the physical processes involved in underwater image formation, resulting in limited optimization despite their impressive enhancement results. This paper proposes a novel deep unfolding network (DUN) for UIE that integrates color priors and inter-stage feature transformation to improve enhancement performance. The proposed DUN model combines the iterative optimization and reliability of model-based methods with the flexibility and representational power of deep learning, offering a more explainable and stable solution compared to existing learning-based UIE approaches. The proposed model consists of three key components: a Color Prior Guidance Block (CPGB) that establishes a mapping between color channels of degraded and original images, a Nonlinear Activation Gradient Descent Module (NAGDM) that simulates the underwater image degradation process, and an Inter Stage Feature Transformer (ISF-Former) that facilitates feature exchange between different network stages. By explicitly incorporating color priors and modeling the physical characteristics of underwater image formation, the proposed DUN model achieves more accurate and reliable enhancement results. Extensive experiments on multiple underwater image datasets demonstrate the superiority of the proposed model over state-of-the-art methods in both quantitative and qualitative evaluations. The proposed DUN-based approach offers a promising solution for UIE, enabling more accurate and reliable scientific analysis in marine research. The code is available at https://github.com/CXH-Research/UIE-UnFold.

Figures (8)

• Figure 1: The underwater image enhancement results of our proposed model compared to several state-of-the-art methods on a sample underwater image. The input image (a) suffers from color distortion, low contrast, and haziness, which are typical issues in underwater imaging. Traditional methods like HLRP (b) struggle to effectively remove the haze and restore natural colors. Learning-based approaches such as CLUIE-Net (c) and PUGAN (d) show improved enhancement results but still exhibit some color deviations and artifacts. In contrast, our proposed model (e) successfully removes the haze, enhances the contrast, and restores vivid and natural colors, producing a visually pleasing result that closely resembles the target image (f).
• Figure 2: The overall framework of the proposed deep unfolding network (DUN), namely UIE-UnFold, for underwater image enhancement (UIE). The framework consists of three main stages, each containing key components such as the Color Prior Guidance Block (CPGB), Nonlinear Activation Gradient Descent Module (NAGDM), FeatureNet, and Inter Stage Feature Transformer (ISF-Former).
• Figure 3: The architecture of the Color Prior Guidance Block (CPGB).
• Figure 4: The architecture of the ResBlock in ResBlock (RIR).
• Figure 5: The architecture of the Nonlinear Activation ResBlock (NARB).