A Graph-Aware Color Correction and Texture Restoration Framework for Underwater Image Enhancement

Abstract

Underwater imagery exhibits markedly more severe visual degradation than their terrestrial counterparts, manifesting as pronounced color aberration, diminished contrast and luminosity, and spatially non-uniform haze. To surmount these challenges, we propose the graph-aware framework for underwater image enhancement (GA-UIE), integrating specialized modules for color correction and texture restoration, a unified framework that explicitly utilizes the intrinsic graph information of underwater images to achieve high-fidelity color restoration and texture enhancement. The proposed algorithm is architected in three synergistic stages: (1) graph feature generation, which distills color and texture graph feature priors from the underwater image; (2) graph-aware enhancement, performing joint color restoration and texture sharpening under explicit graph priors; and (3) graph-aware fusion, harmoniously aggregating the graph-aware color and texture joint representations to yield the final visually coherent output. Comprehensive quantitative evaluations reveal that the output from our novel framework achieves the significant scores across a broad spectrum of metrics, including PSNR, SSIM, LPIPS, UCIQE, and UIQM on the UIEB and U45 datasets. These results decisively exceed those of all existing benchmark techniques, thereby validating the method’s exceptional efficacy in the enhancement of underwater imagery.

1. Introduction

Underwater images constitute a critical conduit of environmental information for oceanic exploration and scientific discovery, yet the pronounced degradation phenomena outlined above severely impede the advancement in underwater vision. Underwater image enhancement (UIE) [1] has become a critical pre-processing stage for reliable visual perception in aquatic environments. Because light propagating through water undergoes wavelength-dependent attenuation and scattering by suspended particles, raw images typically exhibit diminished contrast, and severe color casts, as well as haze-like artifacts which rapidly intensify with depth and turbidity [2,3]. These degradations not only hinder human inspection of benthic habitats, marine infrastructure, or archaeological sites but also undermine the accuracy of downstream computer-vision modules on autonomous underwater vehicles as well as remotely operated vehicles. Moreover, the variability of optical properties across different water bodies, ranging from clear oceanic waters to turbid coastal harbors, renders the design of a universal enhancement solution exceptionally challenging. Consequently, developing deep underwater image enhancement algorithms has become essential for advancing both scientific exploration and industrial applications beneath the surface.

Current methodologies for enhancing underwater imagery are generally divided into two main paradigms: conventional techniques relying on physical priors [4,5] and modern strategies utilizing data-driven learning [6,7]. Traditional approaches formulate physically grounded mathematical models that represent the process of underwater image formation, aiming to explicitly reconstruct the original clear scene, or alternatively exploit color-related priors to manipulate perceptual attributes such as luminance and contrast for visual enhancement. Although traditional approaches effectively enhance underwater imagery, their hand-crafted priors are tailored to specific acquisition conditions and therefore exhibit limited adaptability to the complex, highly variable scenes encountered in large-scale underwater datasets. In contrast, data-driven strategies utilize deep convolutional neural networks (CNNs) to directly approximate an enhancement function. These end-to-end systems learn the intrinsic mapping from a deteriorated input to its restored, pristine version, accomplishing image restoration. Although deep learning-driven algorithms demonstrate superior robustness, they lack sufficient discriminative capacity to identify distinct degradation types, thereby failing to deliver tailored, input-adaptive enhancement for images exhibiting different types of degradation.

To transcend these entrenched limitations, we propose the graph-aware underwater image enhancement (GA-UIE), a unified framework that explicitly excavates and exploits the graph prior of underwater scenes to deliver high-fidelity color restoration and texture revelation. Figure 1 offers a comparison of three learning-based image processing operations. It is unequivocally evident that graph-prior operation confers unprecedented flexibility while endowing the network with a comprehensive, dynamically adjustable receptive field. The proposed pipeline unfolds in three mutually reinforcing phases: (1) graph-aware feature generation, wherein the degraded image is first distilled by a hierarchical convolutional encoder into an intermediate feature map, and explicitly decomposed into color-centric and texture-centric graph features; (2) graph-aware enhancement, wherein the disentangled embeddings are independently propagated through dedicated color-restoration and texture-sharpening networks, each endowed with adaptive receptive fields that allocate higher nodal degrees to severely corrupted regions and sparser connectivity to mildly degraded areas; and (3) graph-aware fusion, wherein a global graph prior harmoniously amalgamates the refined color and texture streams into a coherent, photo-realistic output. The following is a brief summary of this work’s main contributions:

1.

We propose a method called graph-aware color correction and texture restoration for underwater image enhancement, a holistic framework that explicitly mines and leverages the intrinsic graph priors of underwater scenes to achieve high-fidelity color restoration and fine-grained texture recovery.

2.

First, we devise a graph-aware feature generation process in which the hierarchical convolutional encoder initially condenses the degraded input towards an element mapping, and is subsequently parsed into color-centric and texture-centric graph features through explicit disentanglement.

3.

Second, we formulate a graph-aware enhancement engine in which the disentangled embeddings are separately routed through color-restoration and texture-sharpening networks. Each network instantiates an adaptive receptive field: regions exhibiting severe degradation receive denser node connectivity, whereas mildly corrupted areas are assigned sparser links, enabling content-adaptive propagation.

4.

Finally, we introduce a graph-aware fusion stage that employs a unified global graph prior to seamlessly integrate the refined color and texture streams, yielding a coherent and photo-realistic output.

Extensive quantitative evaluation on multiple datasets demonstrates that GAUIE achieves record PSNR, SSIM, UCIQE, UIQM, and LPIPS scores, decisively surpassing state-of-the-art alternatives and corroborating its superiority in underwater image enhancement.

2. Related Work

2.1. UIE Methods via Traditional Priors

Employing physical and statistical priors, model-based techniques [4] seek to reverse the equation governing underwater image formation. Early efforts can be traced to single-image dehazing techniques transplanted from atmospheric vision. The dark channel prior (DCP), which was first developed in hazy scenes, was adapted for underwater imaging as the underwater dark channel prior by He et al. [8] using underwater scene imaging parameters. UDCP relaxes the white-atmospheric-light assumption and re-estimates the transmission map by accounting for wavelength-selective attenuation; however, its reliance on the assumption that at least one color channel is dark often fails in scenes dominated by sand or artificial illumination. Bright channel prior (BCP) [9] and maximum-intensity prior [10] are subsequently introduced to cope with bright seabed or coral regions, yet they remain fragile when the scene violates the prescribed statistical regularities. CLAHE (Contrast-Limited AHE) [11] adapts histogram equalization locally and restricts contrast to avoid noise amplification, thereby aiding image enhancement.

Beyond intensity priors, several studies incorporated optical parameters directly. Schechner et al. [12] exploited polarization cues to estimate backscatter and recover object radiance; although effective in controlled lighting, the method requires multiple acquisitions and polarized hardware. Li et al. [13] fused histogram-distribution priors with minimal-information-loss constraints to compensate for non-uniform illumination, while Peng et al. [14] coupled blur kernel sparsity with normalized total variation to jointly dehaze and deblur. Color-line or haze-line priors further refined transmission estimation by assuming pixels of the same object form tight clusters in RGB space; nevertheless, these priors degrade in the presence of strong artificial lighting or wavelength-dependent scattering that violates the linear color constancy assumption.

Collectively, traditional prior-based approaches offer interpretable, explainable solutions, yet their performance hinges on accurate prior validity and parameter calibration. In complex, unconstrained underwater environments, which are characterized by depth-varying turbidity, artificial light sources, and dynamic color spectra, these rigid priors frequently misfit, yielding over-enhanced colors or residual haze.

2.2. Learning-Based UIE Approaches

Deep learning’s introduction has shifted research emphasis from hand-crafted priors to data-driven representations. Early CNN-based pipelines adopted encoder–decoder architectures for acquiring end-to-end transformations from corrupted underwater images to their enhanced versions. Li et al. [7] released the first large-scale paired underwater dataset (UIEB) and trained a residual CNN with composite perceptual and edge losses, establishing a benchmark for subsequent learning methods. Building on this, multi-scale dense networks [15,16] and attention-based U-Nets [17] are proposed to enlarge receptive fields and suppress artifacts, yet they still struggle to generalize across unseen degraded types.

Generative adversarial networks (GANs) [18] further enriched the learning landscape. CycleGAN-based frameworks [19] leveraged unpaired data to transfer style distributions from clear to degraded domains, while Funie-GAN [20] introduced lightweight generators for real-time enhancement. Despite impressive perceptual quality, GAN-based methods remain susceptible to hallucinated textures and unstable training dynamics when domain gaps widen.

Recent transformer architectures have attempted to reconcile global context modeling with locality. U-shape transformers [21] embedded self-attention inside a hierarchical encoder–decoder to capture long-range dependencies; however, quadratic complexity impedes deployment on high-resolution imagery. More recent diffusion-based paradigms recast underwater image restoration as a denoising process conditioned on degradation levels. DACA-Net [22] introduced degradation-aware conditional diffusion, utilizing lightweight CNNs to predict continuous degradation scores that steer Swin-UNet denoisers. These models achieve state-of-the-art fidelity by integrating physical priors (e.g., histogram matching and contrast consistency) into diffusion training objectives, yet they incur substantial computational overhead and require the careful scheduling of noise levels.

In summary, learning-based approaches excel at capturing complex, high-order statistics, and exhibit superior robustness to unseen scenes. Nevertheless, they demand copious annotated data, are computationally intensive, and can overfit to dataset-specific artifacts. The emerging consensus advocates synergistic frameworks that infuse physics-informed priors (e.g., dynamic kernels and graphs) into expressive deep architectures, aiming to retain interpretability while harnessing the representational power of modern neural networks.

2.3. Graph-Aware Image Enhancement Algorithms

Recently, graph neural networks (GNNs) [23,24,25] have surfaced as a potent alternative to CNNs, applicable to a broad spectrum of computer vision applications. In high-level scene understanding, graph-based representations have proven effective in point-cloud segmentation tasks, skeleton-based action recognition, and image classification via region graphs [26,27,28]. These methods typically cast diverse visual entities such as points, joints, or object proposals into graph nodes and distill discriminative embeddings through iterative message passing.

For low-level vision, researchers have devoted increasing attention to exploiting graph structures to model remote relations, which are challenging to capture with conventional small-kernel CNNs [29]. A dominant line of work focuses on intra-scale non-local aggregation, where every node exchanges information with spatially distant but appearance-similar neighbors inside the same feature level. To enrich contextual cues, several recent papers further introduce cross-scale edges that link nodes across different resolutions, allowing fine details to be enhanced by coarse semantics and vice versa. Beyond single-image contexts, some studies [30] extend the idea to cross-layer message passing inside a CNN or even inter-image graphs that connect patches from different images for joint restoration. In all these scenarios, edge-conditioned aggregation where the weight of a message is modulated by learned edge features such as appearance affinity or geometric distance has become the de facto choice for stable training and expressive power.

Despite these advances, two key design choices remain unsettled. First, graph construction from images is predominantly performed with a fixed k-nearest-neighbor (KNN) rule [31]. While simple, KNN assigns the same out-degree to every node, ignoring the inherently unbalanced relevance of image regions, especially in low-level tasks where textures, edges, and flat areas contribute unequally to the final quality. Second, existing works usually resample images into regular patches and treat each patch as a graph node. This patch-wise abstraction is vulnerable to misalignment errors caused by sub-pixel shifts or object motion, leading to artifacts in tasks such as super-resolution or deblurring. In this work, each pixel is elevated to an individual node in a learnable graph, from which we derive two dedicated graph feature priors: one is tailored to texture improvement, while the other is geared toward correcting color. The former encodes long-range chromatic dependencies to counteract wavelength-selective absorption, while the latter captures fine-grained structural correlations to restore high-frequency details degraded by scattering. These priors are propagated through graph convolutions to yield separate color-corrected and texture-enhanced representations. Ultimately, a unified graph feature prior that seamlessly integrates both color and texture cues fuses the two streams, producing the final high-fidelity underwater image with faithful hues and vivid details.

3. Methodology

As shown in Figure 2 and Algorithm 1, the proposed GA-UIE framework operates through a three-stage process: graph feature generation, graph-aware enhancement, and graph-aware fusion. First, the input underwater-degraded image is processed by a feature extraction module to generate features, which are then fed into a graph feature generation module to produce global graph prior features. Subsequently, a color–texture disentanglement module separates graph prior features suitable for color enhancement and those suitable for texture enhancement. These two types of features are then employed to enhance color and texture, respectively. Finally, in the graph-aware fusion stage, the enhanced texture and color features are fused under the guidance of the global graph prior, yielding the final enhanced image.

Algorithm 1:Graph-aware underwater image enhancement.

Input:  Degraded image I; Reference image $I_{gt}$. Output:  Enhanced image R; Loss function $\mathcal{L}$.   1: Extract features: $f_1=F_{enc}\left(I\right)$;   2: Generate graph: $f_g\leftarrow \mathrm{Graph}\left(f_1\right)$;   3: Decompose: $f_c=F_{3\times 3}\left(f_g\right)$, $f_t=f_g-f_c$;   4: Graph-aware color restoration: $I_c\leftarrow \mathrm{Color}(f_c,f_1)$;   5: Graph-aware texture enhancement: $I_t\leftarrow \mathrm{Texture}(f_t,f_1)$;   6: Graph-aware color-texture fusion: $R\leftarrow \mathrm{Fusion}(f_c,f_t,f_g,I_c,I_t)$.   7: if $I_{gt}$ is available then   8:      ${\mathcal{L}}_{\mathrm{Total}}\leftarrow {\lambda }_1·{\mathcal{L}}_1(R,I_{gt})+{\lambda }_2·{\mathcal{L}}_{\mathrm{VGG}}(R,I_{gt})$   9:      return $R,{\mathcal{L}}_{\mathrm{Total}}$ 10: else 11:       return R 12: end if

3.1. Graph Feature Generation Stage

A long-overlooked premise in image restoration is that assigning uniform importance to every pixel or region proves both inappropriate and inefficient. Consequently, we present an adaptive graph prior refer to [32]. We use a feature extractor module to encode the degraded underwater image I,

$$f_1=F_{enc}\left(I\right),$$

(1)

and this module, denoted as $F_{enc}$, comprises three convolutional layers utilizing the ReLU activation function. Then, we treat pixels as nodes to generate the graph prior.

We introduce a detail-richness indicator to quantify the local complexity of an intermediate feature map $F\in {\mathbb{R}}^{H\times W\times C}$. Formally, for a given downsampling ratio s, each pixel indicator $D_F\in {\mathbb{R}}^{H\times W}$ is characterized with ${\mathsf{\ell }}_1$ norm (across channel dimension) of the residual between the original feature map and its bilinearly downsampled-and-upsampled counterpart:

$$D_F=\sum _{c=1}^C\left|F_c-F_c^{\downarrow s\uparrow s}\right|,$$

(2)

where $F_c$ denotes the c-th channel slice. Following common practice to mitigate excessive information loss, we set $s=2$ throughout our experiments.

The global degree budget is distributed across all pixel nodes according to the magnitude of $D_F$. Specifically, the degree assigned to a pixel node v in characteristic mapping is proportional to the value of its matching entry in the detail-richness map:

$$deg\left(v\right)\propto D_F\left(v\right).$$

(3)

This strategy ensures that regions with richer structural details receive proportionally larger degrees, thereby guiding the subsequent sparse message-passing process to focus on informative locations. For each pixel node v, we adopt a local window strategy to balance local detail preservation and global context. The detail operation in the $deg\left(v\right)\propto D_F\left(v\right)$ can be defined as follows:

$$deg\left(v\right)=\sigma \left(\alpha ·D_F\left(v\right)+\beta \right),$$

(4)

where $\sigma$ is the sigmoid function (bounds [0, 1]), $\alpha$ and $\beta$ are learned scaling/offset parameters. We use approximate nearest neighbor (ANN) search to reduce computation time, and apply edge sparsification (removing edges with affinity < 0.1) to lower memory usage.

Hence, we obtain the global graph feature $f_g$. To more precisely restore color and repair texture for the obtained features, we decompose the global graph prior feature into a color graph feature and a texture graph feature, each of which stands for its low- and high-frequency components. We use $3\times 3$ convolutional layer $F_{3\times 3}$ to extract the low-frequency components as the color graph features $f_c$,

$$f_c=F_{3\times 3}\left(f_g\right).$$

(5)

then, the remaining features can be treated as the high-frequency component, which is the texture graph features $f_t$ in this manner:

$$f_t=f_g-f_c.$$

(6)

In the graph feature generation stage, we obtain the global graph features $f_g$, color graph features $f_c$, and texture graph features $f_t$. We employ these three graph prior features for fusion, color correction, and texture restoration, respectively.

3.2. Graph-Aware Enhancement Stage

To address the color and texture degradation in the input underwater image, we designed a graph-aware module for color correction and texture enhancement as illustrated in Figure 3. Given the encoded features $f_1$ and the color graph features $f_c$, we first use the global average pooling operation and $1\times 1$ convolution to generate the color-adjustment factors $m_c\in {\mathbb{R}}^{1\times 1\times C}$ across all pixels,

$$m_c^1=F_{1\times 1}\left(F_{avg}\left(f_1\right)\right).$$

(7)

Then, we use the color-adjusted factors $m_c$ to enhance the color of all pixels,

$$m_c^2=m_c^1{f}_1.$$

(8)

However, color degradation in underwater images is not uniformly distributed across pixels, as some pixels suffer severe degradation with strong yellow or green casts, while others are only mildly affected. We leverage the obtained color graph features to generate pixel-adaptive adjustment factors.

Specifically, the graph aggregation module processes the input image features and can be formulated as follows:

$$h_c={\displaystyle \frac{1}{C^c}}\sum _{u\in N\left(v\right)}exp\left(f^c(u,v)\right)f_1,$$

(9)

where $f^c:{\mathbb{R}}^d\times {\mathbb{R}}^d\to \mathbb{R}$ denotes the learnable function designed to quantify the relationship among any pair of nodes $(u,v)$. Formally, the correlation score is normalized by a constant $C^c$ to yield a valid probability distribution. $C^c$ denotes the color graph affinities. In this work, we instantiate $f^c$ as the cosine similarity,

$$f^c({\mathbf{h}}_u,{\mathbf{h}}_v)={\displaystyle \frac{{\mathbf{h}}_u^{\top }{\mathbf{h}}_v}{\parallel {\mathbf{h}}_u\parallel \parallel {\mathbf{h}}_v\parallel }},C^c=\sum _{(u^{\prime },v^{\prime })\in \mathcal{E}}{f}^c({\mathbf{h}}_{u^{\prime }},{\mathbf{h}}_{v^{\prime }}),$$

(10)

where $\mathcal{E}$ indicates the group of candidate edges and ${\mathbf{h}}_u,{\mathbf{h}}_v\in {\mathbb{R}}^d$ are the feature vectors for nodes u and v, respectively. This choice endows the model with scale-invariance and empirically yields stable gradients during training.

After obtaining the color graph aggregated features, we use a $3\times 3$ convolution for generating the residual color restored features,

$$k_c=F_{3\times 3}\left(h_c\right),$$

(11)

and the final color-corrected features can be defined as the addition of adjusted features and residual color-restored features,

$$l_c=k_c+m_c^2.$$

(12)

Next, we introduce the graph-aware texture enhancement module. Given the encoded features $f_1$ and the texture graph features $f_t$, we first use the $1\times 1$ convolution to generate the texture-adjustment factors $m_t\in {\mathbb{R}}^{H\times W\times 1}$ across all channels,

$$m_t^1=F_{1\times 1}\left(f_1\right),$$

(13)

then, we use the texture-adjustment factors $m_t$ for texture enhancement for all channels,

$$m_t^2=m_t^1{f}_1.$$

(14)

However, texture degradation in underwater images is not uniformly distributed across channels. We leverage the obtained texture graph features to generate channel-adaptive adjustment factors. Specifically, the graph aggregation module processes the input image features. And it can be mathematically formulated as follows:

$$h_t={\displaystyle \frac{1}{C}}\sum _{u\in N\left(v\right)}exp\left(f^t(u,v)\right)f_1,$$

(15)

where $f^t:{\mathbb{R}}^d\times {\mathbb{R}}^d\to \mathbb{R}$ denote the learnable function designed to quantify the relationship among any pair of nodes $(u,v)$. Formally, the correlation score is normalized by a constant $C^t$ to yield a valid probability distribution. $C^t$ represents the texture graph affinities. All affinities use softmax normalization with temperature scaling ($T=0.5$) to emphasize strong edges. We also instantiate $f^t$ as the cosine similarity,

$$f^t({\mathbf{h}}_u,{\mathbf{h}}_v)={\displaystyle \frac{{\mathbf{h}}_u^{\top }{\mathbf{h}}_v}{\parallel {\mathbf{h}}_u\parallel \parallel {\mathbf{h}}_v\parallel }},C^t=\sum _{(u^{\prime },v^{\prime })\in \mathcal{E}}{f}^t({\mathbf{h}}_{u^{\prime }},{\mathbf{h}}_{v^{\prime }}).$$

(16)

After obtaining the texture graph-aggregated features, we use a $3\times 3$ convolution for generating the residual texture-enhanced features,

$$k_t=F_{3\times 3}\left(h_t\right),$$

(17)

and the final texture-enhanced features can be defined as the addition of adjusted features and residual texture-restored features,

$$l_t=k_t+m_t^2.$$

(18)

During the graph-aware color- and texture-enhancement process, we use the ReLU function and batch normalization for the generation process of $m_c$, $k_c$, $m_t$, and $k_t$.

3.3. Graph-Aware Fusion Stage

After obtaining the color-corrected features and texture-enhancement features, the graph-aware fusion stage aims to combine these two features for generating the final output clear underwater images via global graph features, as shown in Figure 4.

Considering the texture, color, and global graph aspects available, the color and texture graph features are each concatenated with the global graph features. Adaptive fusion weights are then generated by applying a $1\times 1$ convolution followed by a sigmoid activation function:

$$W_c=Sig(F_{1\times 1}\left(Cat(f_c,f_g)\right),$$

(19)

$$W_t=Sig(F_{1\times 1}\left(Cat(f_t,f_g)\right).$$

(20)

Then, we use these two adaptive weights for obtaining the final clear underwater image R:

$$R=W_c{l}_c+W_t{l}_t.$$

(21)

3.4. Loss Function

We supervise the training of GA-UIE network with a composite objective consisting of two complementary terms: L1 loss and perceptual loss.

3.4.1. L1 Loss

The first term is the standard pixel-wise error,

$${\mathcal{L}}_1={\displaystyle \frac{1}{N}}\sum _{i=1}^N\left(R_i-U_i\right),$$

(22)

where the batch size is N, and the restored and reference pictures are indicated by ${\widehat{R}}_i$ and $U_i$, respectively.

3.4.2. Perceptual Loss

To counteract the oversmoothing tendency of pure ${\mathcal{L}}_1$, a perception loss calculated in characteristic space of a trained VGG-19 network is incorporated. Let ${\varphi }_j(·)$ denote a feature graph extracted from the j-th ReLU layer within the VGG-19 architecture. The perceptual loss function is formulated as follows:

$${\mathcal{L}}_{\mathrm{VGG}}={\displaystyle \frac{1}{C_j{H}_j{W}_j}}\sum _{i=1}^N{∥{\varphi }_j\left(R_i\right)-{\varphi }_j\left(U_i\right)∥}_2,$$

(23)

where $C_j\times H_j\times W_j$ is the dimensionality of the chosen feature map.

3.4.3. Total Loss

The weighted sum for two losses serves as the ultimate training goal:

$${\mathcal{L}}_{\mathrm{Total}}={\lambda }_1{\mathcal{L}}_1+{\lambda }_2{\mathcal{L}}_{\mathrm{VGG}},$$

(24)

with empirical weights ${\lambda }_1=1$ and ${\lambda }_2=0.01$.

4. Experiment

4.1. Datasets

For both training and evaluation, the publicly available Underwater Image Enhancement Benchmark (UIEB) is employed [7]. The UIEB dataset comprises 950 real-world underwater image pairs captured under diverse optical conditions, varying turbidity, color casts, lighting, and depth. It also provides high-quality reference images obtained through user studies. It has a balanced distribution of challenges: greenish, bluish, low-contrast, haze-laden, etc. A subset of 800 images is utilized for training, while the remaining 150 are reserved for testing. We use the script provided at https://github.com/Lo-ongWay/division-of-data (accessed on 2 June 2025) to perform the data split. Furthermore, to assess cross-dataset generalization, performance is also reported on publicly available U45 test set [33]. Extensive experiments demonstrate that the model, trained exclusively on the UIEB dataset, attains state-of-the-art results on the U45 test set, evidencing strong generalizability.

4.2. Experimental Settings

All experiments are conducted on the NVIDIA RTX 3090 GPU. The network is trained using the Adam optimizer [34] (${\beta }_1=0.9$, ${\beta }_2=0.999$) over 200 epochs with a batch size of 32. A starting learning rate of $4\times {10}^{-4}$ is employed, which is halved every 20 epochs for a total of 200 epochs. We use the same data augmentation methods refer to [6]. The software environment is built on PyTorch 2.1.0, CUDA 11.8, and cuDNN 8.7. During both training and testing, images are resized into patches measuring $256\times 256$. We adopt both full-reference and no-reference measures to comprehensively assess underwater image enhancement performance: for full-reference metrics, the peak signal-to-noise ratio (PSNR [35]) quantifies pixel-wise fidelity between reference and enhanced images; the structural similarity index (SSIM [36]) evaluates luminance, contrast, and structural preservation; and learned perceptual image patch similarity (LPIPS [37]) measures the distance of image patches in deep feature space. For no-reference metrics, the underwater image quality measure (UIQM [38]) combines sharpness, colorfulness, and contrast to gauge overall underwater visual quality, and underwater color image quality evaluation (UCIQE [39]) focuses on chromatic fidelity and clarity in underwater scenes. Citations of these metrics are provided, and the detailed definitions can be found in the literature. We confirm that the judgment criteria align with your observation—higher values of PSNR, SSIM, UIQM, and UCIQE indicate better image quality, while a lower LPIPS value reflects stronger perceptual fidelity. The training process requires 5233 MB of memory and 78 h. During the testing process, the model uses 4786 MB of memory and processes a single 256 × 256 image in 10.2 ms.

4.3. Compared Methods

In this section, we benchmark GA-UIE against eight leading underwater image enhancement techniques that span two principal research lines. Classic UIE methods: CBF, ERH, IBLA [14], and MMLE are first examined for their handcrafted priors and physical interpretability. We use published weights for learning-based baselines. All pretrained weights can be found at https://github.com/CXH-Research/Underwater-Image-Enhancement (accessed on 22 June 2025). Subsequently, we turn to recent data-driven approaches: Water-Net [7] and Ucolor [40], both CNN-based pipelines; SGUIE [16], which couples semantic priors with convolutional architectures; and PUGAN [18], a physics-based GAN framework. Together, these eight baselines provide a diverse and stringent testbed for validating the advantages of the suggested approach.

4.4. Qualitative Evaluation

UIEB dataset provides 150 challenging underwater images for head-to-head evaluation between our proposed approach and recent traditional and deep techniques. Of them, 70 are selected for qualitative visual comparisons (UIEB-70 test set), whereas the quantitative test bed is the remaining 80 (UIEB-80 test set). The same 80 images are also reused in the ablation study for objective measurements and subjective examination. To further corroborate robustness in our approach, the U45 dataset is additionally employed for supplementary qualitative and quantitative assessments.

Figure 5 and Figure 6 provide a comprehensive visual comparison, revealing that our method substantially outperforms both prior-driven and learning-based counterparts. In the first column, CBF, MMLE, Water-Net, and PUGAN yield desaturated palettes, whereas IBLA suffers from pronounced under-exposure. In contrast, ERH, Ucolor, SGUIE, and our GA-UIE restore vivid chrominance, with GA-UIE delivering the most vibrant and natural rendition. The second column further exposes residual greenish artifacts in the results of CBF, IBLA, and Water-Net. The last three columns, designed to stress-test dense haze, show that CBF, ERH, and MMLE introduce milky veils, with MMLE in particular producing a washed-out, low-contrast appearance, while SGUIE and PUGAN exhibit distant white streaks that degrade perceptual quality. Taken together, these results confirm the superior visual fidelity of GA-UIE on the UIEB-70 subset.

Figure 7 and Figure 8 corroborate that the proposed GA-UIE preserves its decisive edge on the real-world U45 collection, consistently surpassing all competing methods in chromatic fidelity, detail retention, and haze suppression. The first scene is dominated by stubborn green casts in IBLA and SGUIE, whereas CBF, ERH, and Water-Net merely thin the haze without dispelling it; MMLE, Ucolor, PUGAN, and our GA-UIE jointly raise luminance and contrast, yet GA-UIE alone injects the vivid yet natural coloration that most faithfully evokes clear-water appearance. In the second case, all competitors except IBLA recover plausible chrominance, but GA-UIE surpasses them in both saturation and brightness, lending coral and sediment a lifelike pop. Columns three and five expose the brittleness of prior-based methods: IBLA is effectively null, CBF, ERH, Ucolor, and SGUIE plunge the scene into gloom, and Water-Net sprinkles conspicuous red noise across the frame; GA-UIE counters with balanced exposure, clean dehazing, and fine texture that invite closer inspection. The fourth column, however, documents a collective breakdown under extreme green turbidity. IBLA and SGUIE remain pallid, Water-Net mis-paints rock faces with lurid red blotches, Ucolor desaturates the scene, and even GA-UIE succumbs to residual color shift, an edge-case failure we ascribe to the large domain gap between severely degraded inputs and their scarce counterparts in the UIEB reference set. Future work will therefore incorporate an unsupervised color-consistency loss to tame such distribution outliers without sacrificing the method’s prevailing strengths.

4.5. Quantitative Evaluation

Table 1. Various approaches are evaluated at UIEB-80 test set using performance metrics including PSNR (±std), SSIM, UCIQE, UIQM, and LPIPS (↑: higher is better; ↓: lower is better. This convention applies to the following tables).

Method	Metrics
	PSNR ↑	SSIM ↑	UCIQE ↑	UIQM ↑	LPIPS ↓
CBF	23.42 (±4.69)	0.9192	0.5891	1.202	0.1532
ERH	22.53 (±4.13)	0.9321	0.6023	1.189	0.1825
IBLA	21.05 (±6.06)	0.7832	0.5618	1.125	0.2994
MMLE	20.15 (±4.08)	0.8176	0.5782	1.353	0.2706
Water-Net	24.79 (±5.03)	0.9275	0.5679	1.136	0.1435
Ucolor	24.23 (±4.94)	0.9254	0.5593	1.142	0.1263
SGUIE	24.94 (±4.27)	0.9348	0.5352	1.148	0.1137
PUGAN	24.62 (±5.10)	0.9274	0.5710	1.158	0.1249
GA-UIE	25.12(±4.36)	0.9391	0.5732	1.171	0.1099

and

Table 2. Evaluations of various techniques at U45 test set are performed using the UCIQE, UIQM, and parameter metrics.

Method	Metrics
	UCIQE ↑	UIQM ↑	Parameters ↓
CBF	0.5841	1.192	-
ERH	0.5813	1.180	-
IBLA	0.5630	1.125	-
MMLE	0.5797	1.198	-
Water-Net	0.5679	1.162	1.10 M
Ucolor	0.5714	1.155	145.65 M
SGUIE	0.5733	1.179	45.69 M
PUGAN	0.5702	1.171	160.36 M
GA-UIE	0.5762	1.184	10.33 M

report quantitative comparisons on UIEB-80 as well as U45 datasets, respectively. Through paired t-tests on the data in Table 1, we found that there exists an extremely significant difference in PSNR between GA-UIE and other methods, a marginally significant difference in SSIM between GA-UIE and other methods, an extremely significant difference in UIQM between GA-UIE and other methods, and a significant difference in LPIPS between GA-UIE and other methods. It is evident that the proposed GA-UIE attains the highest PSNR and SSIM scores, alongside the lowest LPIPS values on both benchmarks, decisively outperforming all traditional-prior and deep learning counterparts. Although traditional methods slightly eclipse their deep counterparts in UCIQE and UIQM, their outputs are typically washed out, blurry, and noisy, deficiencies that are faithfully reflected in markedly lower PSNR, SSIM, and higher LPIPS scores. Among the learning-based baselines, GA-UIE attains the highest UCIQE and UIQM scores, corroborating its superior restoration of luminance, contrast, and chrominance. Collectively, these metrics demonstrate that our method secures a comprehensive lead across the full objective spectrum. As shown in Table 2, while our algorithm is not as lightweight as Water-Net in terms of parameter quantity, it achieves the best enhancement performance among all compared baseline methods—reflecting the trade-off between model complexity and performance.

4.6. Ablation Study

To analyze the distinct contributions of every stage and sub-module within proposed GA-UIE framework, we conducted a systematic ablation study. Both qualitative visual assessments and quantitative evaluations were carried out on the UIEB-80 test set.

Ablation experiments were structured around three axes: (1) component ablation study, (2) graph-prior ablation study, and (3) loss-function ablation study. Quantitative metrics and corresponding visual comparisons for component ablations are consolidated in

Table 3. Evaluations of component ablation study on UIEB-80 test set indicated by the PSNR, SSIM, UCIQE, UIQM, and LPIPS metrics. w/o denotes without.

Method	Metrics
	PSNR ↑	SSIM ↑	UCIQE ↑	UIQM ↑	LPIPS ↓
w/o Color	23.87	0.9251	0.5521	1.097	0.1315
w/o Texture	24.63	0.9177	0.5663	1.119	0.1207
w/o Fusion	24.99	0.9306	0.5699	1.135	0.1134
GA-UIE	25.12	0.9391	0.5732	1.171	0.1099

and

Table 4. Evaluations of component ablation study on UIEB-80 test set indicated by the parameters, FLOPs, running time, and memory. w/o denotes without.

Method	Metrics
	Parameters ↓	FLOPs ↓	Running Time ↓	Memory ↓
w/o Color	9.57 M	30.64 G	8.2 ms	3211 M
w/o Texture	9.83 M	32.87 G	8.9 ms	3798 M
w/o Fusion	10.12 M	33.15 G	9.7 ms	4233 M
GA-UIE	10.33 M	33.80 G	10.2 ms	4786 M

, and Figure 9; the graph-prior elimination results are provided in

Table 5. Evaluations of graph-prior ablation study on the UIEB-80 test set indicated by the PSNR, SSIM, UCIQE, UIQM, and LPIPS metrics.

Method	Metrics
	PSNR ↑	SSIM ↑	UCIQE ↑	UIQM ↑	LPIPS ↓
KNN-Based Graph Prior	24.88	0.9311	0.5609	1.120	0.1171
GA-UIE	25.12	0.9391	0.5732	1.171	0.1099

and Figure 10, and the loss-function elimination analysis is detailed in

Table 6. Evaluations of loss function ablation study on UIEB-80 test set indicated by the PSNR, SSIM, UCIQE, UIQM, and LPIPS metrics. w/o denotes without.

Method	Metrics
	PSNR ↑	SSIM ↑	UCIQE ↑	UIQM ↑	LPIPS ↓
w/o VGG Loss	23.99	0.9177	0.5540	1.110	0.1198
GA-UIE	25.12	0.9391	0.5732	1.171	0.1099

and Figure 11.

As illustrated in Figure 9 and corroborated by Table 3 and Table 4, removing the color-graph-aware branch yields palettes that remain pallid and chromatically impoverished; excising the texture-graph-aware module erases fine surface structures and leaves conspicuous detail loss; and ablating the global graph-aware fusion network prevents the coherent coupling of color and texture cues, resulting in residual chromatic aberrations and fragmented textural patterns. Figure 10 and Table 5 jointly reveal that substituting our learnable graph prior with a KNN-based surrogate undermines local chromatic and textural fidelity, yielding markedly less vivid reconstructions. Further, Figure 11, together with Table 6, demonstrates that the introduction of the perceptual loss galvanizes the generator, endowing the enhanced images with livelier coloration and more nuanced detail. Collectively, these ablations substantiate the indispensability of every architectural component, the proposed graph prior, and the composite loss design within GA-UIE.

As shown in

Table 7. Evaluations of complexity ablation study on UIEB-80 test set indicated by the running time, memory, PSNR, SSIM, and UCIQE metrics.

Method	Metrics
	Running Time (ms) ↓	Memory (MB) ↓	PSNR ↑	SSIM ↑	UCIQE ↑
Patchsize ($64\times 64$)	6.8 ms	1280 M	24.35	0.9133	0.5512
Patchsize ($128\times 128$)	9.1 ms	2680 M	24.78	0.9219	0.5575
$k=3\times 3$	8.5 ms	3850 M	24.60	0.9182	0.5557
$k=5\times 5$	11.4 ms	4320 M	24.95	0.9255	0.5604
Degree budget (1/16 pixels)	9.8 ms	2100 M	24.50	0.9150	0.5532
Degree budget (1/4 pixels)	10.6 ms	3550 M	24.85	0.9231	0.5580
GA-UIE	10.2 ms	4786 M	25.12	0.9391	0.5732

, we add the evaluation of a complexity ablation study to demonstrate the effectiveness of the chosen hyperparameters. It is apparent that reducing the sizes of patchsize, local window, and degree budget all affect the final enhancement performance, although this reduces memory usage and runtime. Therefore, the algorithm proposed in this paper adopts a patchsize of 256 × 256, a local window of $k=7\times 7$, and all pixels as the degree budget.

4.7. Further Study

To test our proposed method in real scenarios involving manipulators and real-time systems, we note that our initial experiments were primarily conducted on simulated datasets and offline image sequences—and we fully acknowledge the necessity of extending these experiments to real-world environments. To this end, we have collaborated with the robotics laboratory to establish a dedicated real testbed, which comprises a 6-degree-of-freedom (DoF) robotic manipulator, a real-time vision sensor, and a computing unit integrated with real-time software. We designed two experimental tasks, which were carried out using the aforementioned equipment in both a water tank and a real-world underwater environment. The comparative results of underwater grasping experiments are presented in

Table 8. Comparative test of underwater grasping success rates: Task 1 is conducted in a tank, whereas Task 2 is performed in a real-world underwater environment.

Method	Metrics
	Task 1 ↑	Task 2 ↑
Degraded image	0.52	0.34
Image enhanced by GA-UIE	0.71	0.63

. It can be clearly observed that images enhanced by our algorithm not only improve visual perception but also increase the grasping success rate.

Although our method is comprehensive and achieves good performance, there are still some shortcomings that we need to improve in future work. Figure 12 shows the performance of the proposed algorithm in complex scenarios with strong degradation. It can be observed that although the proposed algorithm provides some improvements, its overall enhancement remains limited, such as color restoration and defogging effects, remains unsatisfactory. Specifically, in the field of underwater image enhancement, existing datasets mostly focus on mild degradation phenomena (e.g., slight color cast or weak haze), while paired datasets that cover severe degradation—such as thick haze, non-uniform haze, and significant color distortion—remain scarce. In our future work, we will establish a dedicated paired dataset for underwater image enhancement that focuses on severe degradation scenarios; this dataset will include both synthetic and real-world samples, where synthetic data will be generated by simulating physical processes of severe underwater degradation (e.g., modeling light scattering for thick fog and color absorption for extreme color distortion) and real-world data will be collected via underwater robots. Additionally, we intend to integrate semi-supervised learning and self-supervised learning strategies into our framework. For semi-supervised learning, we will use a small subset of labeled data to initialize the model and leverage a large amount of unlabeled data for consistency regularization. For self-supervised learning, we will consider adding SSIM/MS-SSIM or a Lab-space color-consistency loss to stabilize colors, and pretraining the backbone on unlabeled data using pretext tasks.

Moreover, we fully recognize that high computational demands could limit the deployment of our method on mid-range or low-end equipment, and we have already initiated targeted optimizations to address this issue. Specifically, we plan to introduce lightweight network architectures (e.g., MobileNetV3 or EfficientNet-Lite) to replace the current heavy backbone in the future work, which can significantly reduce parameter counts and floating-point operations (FLOPs) without compromising performance.

Furthermore, our current implementation is overly complex. We will adopt a “simplicity-first” design principle for the future framework. First, we will use off-the-shelf, well-optimized libraries (e.g., TensorRT for inference acceleration and Hugging Face Datasets for data loading) to reduce engineering complexity. Second, we will provide a complete, user-friendly implementation pipeline, including pre-configured training/inference scripts, detailed documentation on environment setup (with Docker containers to avoid dependency conflicts), and step-by-step tutorials for deploying the model on common hardware (e.g., PCs, edge devices, and manipulators).

5. Conclusions

This paper presents the Graph-Aware Underwater Image Enhancement framework, designed for comprehensive color correction and texture restoration, a principled solution that leverages the intrinsic graph structure of underwater imagery to counteract the severe color aberrations, low contrast, and spatially non-uniform haze that plague sub-aquatic scenes. The proposed method is carefully architected into three synergistic stages that progressively refine visual quality. First, the graph feature generation stage distills robust graph priors from an input image via a cascade of convolutional encoders and a dedicated graph-generation module, followed by a color–texture graph decomposition that explicitly disentangles color-specific and texture-specific graph features. Second, the graph-aware enhancement stage independently propagates these disentangled priors through specialized graph-aware color and texture enhancement modules, yielding corrected color representations and high-frequency texture representations with restored structural fidelity. Third, the graph-aware fusion stage harmoniously aggregates the enhanced color and texture features under the guidance of global graph priors, producing a final output that is both visually coherent and photometrically accurate. Comprehensive quantitative evaluations corroborate the framework’s superior capacity in underwater image enhancement as well as underscoring the critical value of explicitly modeling graph-based priors in addressing the complex, spatially varying distortion characteristic of underwater environments.