MFP-PAINet: Enhancing Underwater Images Through Multi-Dimensional Feature Fusion and Probabilistic Uncertainty Modeling

Abstract

In supervised training for underwater image enhancement (UIE), many deep learning methods rely on approximate reference images; however, they tend to neglect the inherent uncertainty of these references. To address this problem, this study proposes MFP-PAINet, an underwater image enhancement network, which can integrate multi-dimensional feature fusion with probabilistic adaptive uncertainty modeling. Specifically, histogram equalization, white balance correction, and gamma correction were first implemented to generate preliminary enhanced inputs so as to improve contrast and color balance. Subsequently, a multi-dimensional feature mapping module was used to extract high-dimensional representations and produce multiple feature maps, which were dynamically weighted and adaptively fused by a confidence generation module to initially handle the uncertainty at the feature level. Finally, a probabilistic module was introduced to further model the uncertainty and refine image details. Extensive experiments conducted on eight publicly available datasets demonstrate that the proposed MFP-PAINet outperforms both traditional and deep learning-based UIE methods. For example, on the LSUI dataset, MFP-PAINet achieved SSIM, PSNR, and UIQM scores of 0.854, 23.044, and 3.056, respectively. Furthermore, MFP-PAINet exhibited certain promising performance in image deraining and image dehazing tasks, proving that it can effectively retrieve image details and improve image clarity.

1. Introduction

Underwater Image Enhancement (UIE) has become a focal research direction in such fields as image processing and underwater vision over recent years [1,2]. High-quality underwater images are of significant application value in many tasks, and they play a crucial role in such disciplines as marine biology, archeology, and marine ecology [3,4,5], while becoming essential in practical operations like underwater engineering, search and rescue, and resource exploration. Nevertheless, given the complexity of underwater environments, such as variable lighting conditions and the scattering and absorption characteristics of water, underwater images are often compromised by certain factors, such as reduced contrast and color distortion, thus severely degrading their quality and usability.

The root cause of underwater images’ degradation lies in the wavelength absorption and scattering that light would inescapably encounter as it propagates through water. Typically, such images would exhibit a distinct blue or green hue [6], because as light travels through water, red wavelengths would be absorbed first, enabling blue and green light to penetrate farther. In a word, underwater images degrade due to wavelength absorption and scattering, including both forward and backward scattering [7,8]: the former causes images to blur, while the latter results in reduced brightness, thus further degrading image quality [9,10,11]. Moreover, the attenuation of wavelengths [10,12], scattering, and other optical properties of water bodies can cause irregular nonlinear distortions [13,14,15], which lead to low image contrast and are often accompanied by blurring and color degradation. Although traditional UIE methods can improve image quality to some extent, they usually rely on specific assumptions and prior knowledge, which limits their adaptability to different water environments. As a result, these methods often suffer from color deviation or over-enhancement in complex underwater scenes, reflecting their limited generalization capability.

Over recent years, significant advancements have been made with deep learning in the UIE field, and the methods adopted typically rely on paired distorted and reference images so as to learn the mapping between them [6,16]. However, due to the complexity and variability of underwater environments, it becomes a challenging task to take perfect, undistorted underwater images (i.e., the reference images). To address this issue, the currently mainstream solutions often involve generating approximate reference images, which are used as undistorted, clear underwater images. For example, the solution in reference [17] enhances distorted images by using various advanced algorithms, and then the best-quality image is selected as the reference image. However, this approach is influenced by multiple factors, such as human visual perception, imaging equipment, and algorithmic parameters, therefore leading to uncertainty in the reference image. This means that there may be multiple potential solutions to the same distorted image. If directly training a model with such distorted images and their corresponding reference images, poor model generalization performance may be the result due to the uncertainty of such reference images [18]. Moreover, similar challenges related to uncertainty and unstable supervision have also been dealt with in other signal enhancement domains. For instance, Siddique et al. [19,20] found that improving robustness under noise and uncertain signal patterns requires enhanced representations and reliable feature learning, a principle that is also essential for underwater image enhancement for ambiguous or approximate reference images.

To address the aforementioned issues, this study proposes a novel UIE network, called MFP-PAINet, which combines multi-dimensional feature fusion with probabilistic adaptive uncertainty distribution to generate high-quality underwater images. The main contributions of this study include

(1)

“Fusion-then-Modeling” Hierarchical UIE Framework: A hierarchical framework was constructed to integrate multi-dimensional feature fusion with probabilistic uncertainty modeling. By fusing multi-dimensional features prior to performing probabilistic modeling, the framework could obtain more stable and robust feature representations.

(2)

Confidence-Guided Multi-Dimensional Feature Fusion Strategy: A collaborative design for the multi-dimensional feature mapping module and the confidence generation module is proposed. By virtue of dynamic weighting and adaptive fusion of multi-dimensional features, this design strategy can preliminarily handle uncertainty at the feature level.

(3)

Experimental Validation and Cross-Task Adaptability: On all the eight public UIE datasets, the proposed method achieved superior performance over the existing UIE approaches. In addition, the method exhibits a certain degree of cross-task adaptability in image deraining and image dehazing tasks, demonstrating promising generalization capability and practical application potential.

2. Related Work

2.1. Traditional UIE Methods

Traditionally, UIE methods were based on physical approaches [21], and prior knowledge was utilized to estimate the parameters of underwater imaging models for improving image quality. Such priors include underwater dark channel prior [22], minimum information prior [23], haziness prior [24], etc. For example, Ancuti et al. [25] proposed a method to deliver color correction and contrast enhancement through weighted mapping and image fusion techniques, ultimately improving the visual impact of images. By using local adaptive contrast enhancement techniques while minimizing color loss, Zhang et al. [26] sharpened image clarity and contrast. He et al. [27] introduced the Dark Channel Prior (DCP) algorithm to estimate the transmission of light by calculating the “dark channel” in images, effectively enhancing image quality. Based on DCP, Peng et al. [28] proposed the Generalized Dark Channel Prior (GDCP) method, which assumes dark channel information in underwater images to restore image clarity. On the basis of dehazing algorithms, Chiang et al. [29] developed a UIE method to improve image clarity by removing haze from images. Galdran et al. [30] introduced an improved dark channel prior method, which used red channel information to estimate the depth of underwater images, thus improving image quality. Li et al. [31] created a physics-based UIE method to estimate an intermediary transmission map and recover image visibility and color, and this method was found to perform well in various water bodies and adverse weather conditions, thus demonstrating broad applicability.

Generally, traditional methods tend to implement physical and statistical techniques to enhance underwater image quality. Although they improve the visual appeal of images to some extent, they often rely on specific assumptions and prior knowledge. Unfortunately, they may encounter limitations when dealing with complex and unstable underwater environments.

2.2. Deep Learning-Based UIE Methods

Deep learning-based methods for UIE are designed to leverage automatic feature extraction and utilize extensive paired or unpaired training data to learn enhancement mappings. As a UIE network based on medium transmission guidance embedded in multiple color spaces, UColor [32] can significantly enhance the visual quality and quantitative metrics of underwater images by integrating feature representations across different color spaces and utilizing medium transmission information from physical models. Li et al. [17] introduced a UIE network called Water-Net, which, after being trained on a large-scale underwater image dataset, could significantly enhance visual quality and detail retention by combining multi-scale features with adversarial learning strategies. Shallow-UWNet [33] is a lightweight neural network structure to reduce the number of parameters and enhance a model’s compression rate and speed while maintaining performance comparable to the existing state-of-the-art models on multiple synthetic and real datasets. Water-CycleGAN [34] is a weakly supervised color transfer method for underwater image color correction, and it is designed with a multi-loss function, including adversarial loss, cycle consistency loss, and Structural Similarity Index Measure (SSIM) loss; this method can correct image colors to render them as taken in air, while preserving the content and structure of the input images. By using Generative Adversarial Networks (GAN) for UIE, Fabbri et al. [6] proposed a method called UGAN, which improved the visual quality of underwater scenes through generative adversarial networks and enhanced the performance of autonomous underwater robot visual-driven behaviors, delivering significant effects in multiple tests. Jiang et al. [35] introduced a target-oriented perceptual adversarial fusion network called TOPAL by combining multi-scale dense enhancement modules with deep esthetic rendering modules; through a dual-channel attention mechanism and a global-local adversarial mechanism, it achieved high-quality enhancement for underwater images.

In a word, through advanced neural network architectures and optimization tricks, the deep learning approaches above performed excellently in complex underwater environments, significantly improving image clarity, color accuracy, and contrast. Nevertheless, although these methods overcame the limitations of traditional approaches, they all failed to give proper consideration to the uncertainties in the image enhancement process. In contrast, by combining multi-dimensional feature fusion strategies with probabilistic adaptive distribution, the MFP-PAINet proposed in this study can learn and handle uncertainties in a better manner in the image enhancement process, thereby generating high-quality underwater images.

3. The Proposed Model

3.1. The Principle and Overall Framework of the Model

This study proposes a UIE network, called MFP-PAINet, based on multi-dimensional feature fusion and probabilistic adaptive uncertainty distribution. As shown in Figure 1 for the overall architecture of the network, it includes a multi-dimensional feature mapping module, a confidence generation module, and a probabilistic module, which is composed of a Conditional Variational Autoencoder (CVAE) [36,37] and Probabilistic Adaptive Instance Normalization (PAdaIN) [18].

After being inputted, the distorted underwater images would undergo certain preprocessing operations, including histogram equalization (HE), white balance correction (WB) and gamma correction (GC), so as to preliminarily enhance the contrast and color balance of the images. Specifically, HE is implemented at the luminance channel in the HSV color space to enhance the overall image contrast; WB adopts the Perfect Reflector method, purposed to normalize the maximum values of the RGB channels and eliminate color deviation; and GC employs a power-law transformation with $\gamma =2.0$, so as to enhance the dark-region details and optimize the overall visual quality. The preprocessed images are then fed into the multi-dimensional feature mapping module, which generates various multi-dimensional feature mappings of the images while capturing and integrating more advanced features to improve naturalness and detail quality. The confidence generation module would, based on the output of the multi-dimensional feature mapping module, assign corresponding confidence maps to these multi-dimensional feature mappings and dynamically fuse them. This module can preliminarily resolve the uncertainties by synthesizing multiple potential solutions so as to improve the final image quality. The merged results are then inputted into the probabilistic module, which uses the PAdaIN module to encode the uncertainties in a deeper manner. In addition, this module can dynamically adjust the mean and standard deviation of the image features and introduce random variables to express uncertainties, thus further ensuring the consistency and stability of the enhancement results.

Proposed by Fuet et al. [18], PUIE-Net first introduced uncertainty modeling into UIE and designed the PAdaIN module for probabilistic feature encoding. Subsequently, Saleh et al. [38] extended this concept by constructing an unsupervised enhancement model, called UDnet, to further broaden the application of uncertainty modeling. Similarly to PUIE-Net and UDnet, the method used in this study also adopts PAdaIN for feature-level uncertainty encoding. However, there is a key distinction: this study implements a multi-dimensional feature fusion and confidence-based dynamic weighting mechanism prior to uncertainty modeling. This mechanism adaptively integrates results from different enhancement paths (e.g., HE, WB, and GC), thereby explicitly mitigating early-stage uncertainties at the feature level. The differences in the workflow among PUIE-Net, UDnet, and the proposed method are illustrated in Figure 2. This “fusion-then-modeling” strategy enables the model proposed in this study to obtain more stable and reliable feature representations before probabilistic distribution learning, eventually resulting in enhanced consistency and structural fidelity in the final enhanced images.

3.2. Multi-Dimensional Feature Mapping Module

Images of different types are inputted into the multi-dimensional feature mapping module to generate corresponding feature mappings. This module integrates a multi-branch structure, Atrous Spatial Pyramid Pooling (ASPP) [39], and SE channel attention mechanisms, and its structure is illustrated in Figure 3. The multi-branch module captures features from local to global scales by using different convolution kernel sizes and then fuses these features along the channel dimension. Subsequently, the SE [40] module enhances feature expression capabilities, and residual blocks further optimize feature extraction effects. After multiple iterations, the feature maps are inputted into the ASPP module, which uses convolution layers with varying dilation rates to extract multi-scale contextual information and ultimately produce the feature mappings of the inputted images. By virtue of multi-dimensional feature fusion and residual learning, the network enables multi-dimensional feature mappings to capture and integrate more advanced features, thus significantly enhancing the naturalness and detail representation of the images. The ASPP module is introduced to exploit convolutions with different dilation rates for extracting multi-scale contextual information, thereby enhancing the network’s global perception of complex underwater scenes. The SE attention mechanism is employed to adaptively adjust channel-wise feature weights so as to highlight the critical information while suppressing redundant features, purposed to improve the effectiveness of feature representation.

3.3. Confidence Generation Module

The Confidence Generation Module, a key component in the UIE framework, is responsible for generating confidence maps for different feature mappings of the multi-dimensional feature mapping module and dynamically weighting and merging these feature mappings to preliminarily address uncertainties in the enhancement process. As shown in Figure 4, the structure of this module is based on the U-Net architecture [41] and enhanced with residual blocks and SE (Squeeze-and-Excitation) modules, purposed to improve the capability and precision of feature extraction. The multi-dimensional feature mappings outputted from the multi-dimensional feature mapping module are sequentially fed into the encoder of the Confidence Generation Module, where the encoder, working with residual blocks, progressively downsamples to extract high-dimensional feature representations. In the decoding phase, skip connections are used to fuse the features extracted in the encoding phase, with them gradually restored during decoding, so as to ensure the comprehensive integration of the multi-dimensional features. The SE module in the decoder is used to adaptively adjust weights between channels so as to enhance key information specific to each feature mapping, thereby highlighting high-confidence features. The formula for dynamically weighting and fusing the multi-dimensional feature mappings with the confidence maps is as follows:

$$I_{en}=R_{WB}\odot C_{WB}+R_{HE}\odot C_{HE}+R_{GC}\odot C_{GC}$$

(1)

where $R_{HE}$, $R_{WB}$, and $R_{GC}$ are the three feature mappings outputted from the multi-dimensional feature mapping module; $C_{HE}$, $C_{WB}$, and $C_{GC}$ are the confidence maps for these feature mappings; $I_{en}$ is the feature map after dynamically weighted fusion; and $\odot$ represents element-wise multiplication.

3.4. Probabilistic Module

To further address the uncertainties, the probabilistic module introduces a latent variable $z$ to represent uncertainties in the enhancement process. By integrating a CVAE [36,37] and PAdaIN [18], this module generates diverse enhancement results. The objective is to learn the mapping from the degraded image $x$ to the clean image $y$, while accounting for the uncertainty represented by $z$. This process can be formulated as:

$$p\left(y|x\right)\approx {\displaystyle \frac{1}{S}}\sum _{s=1}^{S}p\left(y\left|z^{\left(S\right)}\right.,x\right), z^{\left(S\right)}~p\left(z|x\right)$$

(2)

where $p\left(z|x\right)$ denotes the uncertainty distribution, and $S$ represents the number of samples. In this way, the model cannot only generate multiple enhancement results, but can capture the diversity across different enhancement strategies. To identify the optimal enhancement result, the maximum likelihood estimation is implemented, and the sample with the highest probability among the generated samples is selected as the final output:

$$p\left(y|x\right)\approx p\left(y\left|z_{max}\right.,x\right), z_{max}~p\left(z|x\right)$$

(3)

where $z_{max}$ represents the latent variable sample with the highest probability.

Assume that $x$ is a distorted image and $y$ is a reference image. In the CVAE process, the feature extractor maps the input image into feature representations to construct a probability distribution, from which random samples are drawn and fed into the PAdaIN module. The structure of the feature extractor is illustrated in Figure 5. PAdaIN can be expressed as:

$$PAdaIN\left(x\right)=b\left({\displaystyle \frac{x-\mu \left(x\right)}{\sigma \left(x\right)}}\right)+a$$

(4)

where $a$ and $b$ indicate the random variables sampled from the mean and standard deviation distributions; and $\mu \left(x\right)$ and $\sigma \left(x\right)$ indicate the mean and standard deviation. During training, the feature extractor computes the posterior distribution; during testing, it computes the prior distribution. In the training phase, $a$ and $b$ can be represented as:

$$a~{\mathcal{N}}_m\left(\mu \left(y,x\right) ,{\sigma }^2\left(y,x\right)\right)$$

(5)

$$b~{\mathcal{N}}_s\left(m\left(y,x\right) ,v^2\left(y,x\right)\right)$$

(6)

In the testing phase, $a$ and $b$ can be represented as:

$$a~{\mathcal{N}}_m\left(\mu \left(x\right) ,{\sigma }^2\left(x\right)\right)$$

(7)

$$b~{\mathcal{N}}_s\left(m\left(x\right) ,v^2\left(x\right)\right)$$

(8)

where ${\mathcal{N}}_m$ and ${\mathcal{N}}_s$ denote N-dimensional Gaussian distributions for the mean and standard deviation, respectively.

3.5. Loss Function

This study combines the Mean Squared Error (MSE) loss [42], VGG16 perceptual loss [43], KL divergence loss [44] and edge detection loss to synergistically enhance the model’s image enhancement effect.

Mean Squared Error loss is used to measure the pixel difference between the enhanced image and the reference image, with the following formula:

$$L_{MSE}={\displaystyle \frac{1}{N}}\sum _{i=1}^N{\left(y_i-{\widehat{y}}_i\right)}^2$$

(9)

where $y_i$ represents the pixel values of the reference image, ${\widehat{y}}_i$ represents the pixel values of the enhanced image, and $N$ represents the total number of pixels.

VGG16 perceptual loss is used to measure the difference between the enhanced image and the reference image in a higher feature dimension. Specifically, both the enhanced and reference images pass through the first 31 layers of the pretrained VGG16 network to extract high-level perceptual features; and then, the Mean Squared Error (MSE) between these feature representations is computed. The formula is as follows:

$$L_{VGG16}={\displaystyle \frac{1}{N}}\sum _{i=1}^N{‖\varnothing \left(y_i\right)-\varnothing \left({\widehat{y}}_i\right)‖}^2$$

(10)

where $\varnothing \left(·\right)$ represents the feature vectors extracted through the VGG16 model, and $N$ represents the total number of feature vectors.

KL divergence loss is used to measure the difference between the posterior distribution $q\left(z|x\right)$ and the prior distribution $p\left(z\right)$, so as to help the model learn more accurate representations in the latent space. The formula is as follows:

$$L_{KL}=D_{KL}\left(q\left(z|x\right)‖p\left(z\right)\right)$$

(11)

where $D_{KL}$ represents the KL divergence. Specifically, the calculations of KL divergence include the following two components:

$$L_m=D_{KL}\left({\mathcal{N}}_m\left(x\right)‖{\mathcal{N}}_m\left(y ,x\right)\right)$$

(12)

$$L_s=D_{KL}\left({\mathcal{N}}_s\left(x\right)‖{\mathcal{N}}_s\left(y ,x\right)\right)$$

(13)

where ${\mathcal{N}}_m$ and ${\mathcal{N}}_s$ represent the Gaussian distributions of means and standard deviations, respectively.

To better preserve details and edge information, the edge detection loss is introduced in this study. Edge maps of the enhanced and reference images are extracted using a Laplacian edge detector, with the edge loss calculated as follows:

$$L_{edge}=MSE\left(E_{pred} ,E_{target}\right)$$

(14)

where $L_{edge}$ represents the edge loss; $E_{pred}$ and $E_{target}$ represent the edge maps of the enhanced and reference images, respectively; and MSE represents the mean squared error.

The total loss function integrates the pixel difference, perceptual difference, distribution alignment, and edge information:

$$L=L_{MSE}+L_{VGG16}+L_m+L_s+L_{edge}$$

(15)

By minimizing the total loss, the model is optimized at the pixel, feature, and distribution levels, therefore significantly enhancing the quality and robustness of image enhancement.

4. Experimental Results and Analysis

This section would, through a series of experiments, evaluate the performance of the proposed method. Both quantitative and qualitative comparison results of the image enhancement method proposed in this study would be presented against various models on multiple public datasets. Ablation studies are also conducted to demonstrate the effectiveness of each module in this study’s approach. To illustrate the model’s generalization and effectiveness in a better manner, additional experiments were conducted on image deraining and dehazing.

4.1. Datasets

Following the experimental settings provided in [45], this study used 800 pairs of real-world underwater images from the UIEB [17] dataset (each pair consisting of a degraded image and its corresponding reference image), together with 800 pairs of synthetic un-derwater images generated by CycleGan [46]. Therefore, a total of 1600 image pairs were utilized for the training. In the testing phase, eight datasets covering diverse scenes and task types were employed, including: (1) EUVP [47], LSUI [48], UFO-120 [49], and UIEB [17]: they are four mainstream UIE benchmark datasets with available reference images; (2) RUIE [16]: a no-reference UIE dataset collected from real-world underwater scenes; (3) AUDD [50] and DUO [51]: they are underwater object detection datasets; and (4) UOT32 [52]: an underwater object tracking dataset.

Table 1. Summary of the datasets used in this study (including the numbers of image pairs).

Phase	Dataset	Reference Type	Images
Training	Synthetic	Paired	800
	UIEB	Paired	800
Testing—Reference	EUVP	Paired	500
	LSUI	Paired	500
	UFO-120	Paired	120
	UIEB	Paired	90
Testing—No-reference	AUDD	Unpaired	100
	DUO	Unpaired	100
	RUIE	Unpaired	100
	UOT32	Unpaired	100

summarizes the detailed statistics of these datasets, and Figure 6 displays certain samples from each dataset.

Figure 7 compares the RGB channel distributions between the synthetic and real-world underwater datasets. As can be observed, the synthetic data exhibit smoother and more idealized color distributions, whereas the real-world underwater images show case greater variability and imbalance among color channels due to complex illumination and scattering effects. These findings indicates a clear domain gap between the synthetic and real-world data in terms of color distribution and contrast. Therefore, it is necessary to introduce of probabilistic modeling and adaptive fusion, in order to enhance cross-domain generalization.

4.2. Evaluation Metrics

To comprehensively evaluate the performance of the proposed model, both full-reference and no-reference evaluation metrics were employed. Full-reference metrics include Peak Signal-to-Noise Ratio (PSNR) and Structural Similarity Index (SSIM) [53]. A higher PSNR indicates that the output image is perceptually closer to the reference image, while a higher SSIM suggests greater structural similarity between them. The no-reference metric used is the Underwater Image Quality Measure (UIQM) [54], which combines Colorfulness Measure (UICM), Sharpness Measure (UISM), and Contrast Measure (UIConM). A higher UIQM indicates better visual quality.

4.3. Implementation Details

This study’s model was trained on images with a resolution of 256 × 256 pixels. First, the shortest side of the images was scaled to 256 pixels, and then regions of 256 × 256 were randomly cropped to ensure diversity and consistency in the input images so as to enhance the model’s generalization capability. The training process was conducted over 400 epochs, with a batch size of 4, on an NVIDIA GeForce RTX 3090 GPU with 24 GB of RAM (Nvidia Corporation, Santa Clara, CA, USA). The model construction and training were based on the Pytorch framework (version 1.11.0) [55], using the ADAM optimizer with a learning rate of $1\times {10}^{-4}$.

4.4. Comparison Methods

To comprehensively evaluate the performance of the proposed UIE method, four traditional unsupervised methods and seven deep learning methods were selected for comparison, covering the main research advancements in the field of UIE. For the traditional methods, UDCP [56], Fusion [25], Retinex [57], and Haze-line [14] were selected. UDCP used the dark channel prior to removing blue-green tones from underwater images; Fusion enhanced image quality through weighted mapping of multiple enhancement results; Retinex improved the contrast based on the separation of illumination and reflection; and Haze-line restored image colors by simplifying the attenuation ratio. For deep learning methods, this study selected UDnet [38], FUNIE-GAN [47], UGAN [6], CycleGan [46], U-Shape [48], PUGAN [58], and Mamba-UIE [59]. Specifically, UDnet is an unsupervised UIE method based on uncertainty modeling; FUNIE-GAN, UGAN, CycleGan, and PUGAN are UIE methods built on generative adversarial networks (GANs); U-Shape is a UIE method designed with a Transformer-based architecture; and Mamba-UIE is a recently developed UIE framework based on the Mamba architecture.

For a fair comparison, all deep learning models were retrained and evaluated on the same datasets. Each baseline deep learning model was strictly reproduced and trained following the parameter settings and training strategies given in their original publication, purposed to ensure that all models would be evaluated under their optimal performance conditions.

4.5. Quantitative Comparison

The quantitative comparison results on the eight datasets are summarized in

Table 2. Quantitative comparison on EUVP and LSUI datasets (red indicates the best performance, blue the second-best, and green the third-best values; $\uparrow$ denotes that higher values correspond to better performance; the same notation applies to remaining tables).

Method	EUVP			LSUI
	$\mathbf{SSIM} \uparrow$	$\mathbf{PSNR} \uparrow$	$\mathbf{UIQM} \uparrow$	$\mathbf{SSIM} \uparrow$	$\mathbf{PSNR} \uparrow$	$\mathbf{UIQM} \uparrow$
UDCP	0.723	19.099	2.306	0.682	13.398	2.326
Fusion	0.726	17.087	3.043	0.773	17.632	2.941
Retinex	0.538	11.272	3.009	0.727	13.981	3.015
Haze-line	0.696	16.863	2.289	0.697	14.554	2.386
UDnet	0.722	19.278	3.170	0.829	20.712	3.119
FUNIE-GAN	0.701	18.356	3.045	0.779	19.146	3.083
UGAN	0.721	18.659	3.026	0.842	22.625	3.091
CycleGan	0.722	19.730	2.866	0.787	20.484	2.963
U-Shape	0.723	18.712	2.914	0.846	21.961	2.998
PUGAN	0.666	16.523	2.929	0.766	19.638	3.003
Mamba-UIE	0.729	18.326	2.911	0.877	22.967	2.949
Ours	0.735	19.501	3.051	0.854	23.044	3.056

,

Table 3. Quantitative comparison on UFO-120 and UIEB datasets.

Method	UFO-120			UIEB
	SSIM  $\uparrow$	PSNR  $\uparrow$	UIQM  $\uparrow$	$\mathbf{SSIM} \uparrow$	$\mathbf{PSNR} \uparrow$	$\mathbf{UIQM} \uparrow$
UDCP	0.723	18.727	2.324	0.663	14.658	2.287
Fusion	0.733	17.351	2.976	0.755	17.262	2.834
Retinex	0.632	12.714	2.956	0.539	10.461	2.805
Haze-line	0.691	15.310	2.204	0.693	16.170	2.304
UDnet	0.728	18.844	3.113	0.773	18.962	3.000
FUNIE-GAN	0.712	18.727	2.992	0.768	18.573	2.968
UGAN	0.726	18.788	2.966	0.812	21.288	3.005
CycleGan	0.726	19.388	2.836	0.751	19.779	2.765
U-Shape	0.718	18.813	2.841	0.811	21.000	2.947
PUGAN	0.666	16.613	2.79	0.731	18.381	2.785
Mamba-UIE	0.731	18.120	2.845	0.850	21.781	2.843
Ours	0.741	19.359	2.991	0.810	20.569	2.980

and

Table 4. Quantitative comparison on AUDD, DUO, RUIE, and UOT32 datasets.

Method	AUDD	DUO	RUIE	UOT32
	UIQM $\uparrow$	$\mathbf{UIQM} \uparrow$	$\mathbf{UIQM} \uparrow$	$\mathbf{UIQM} \uparrow$
UDCP	2.674	3.072	2.617	3.055
Fusion	3.161	3.484	3.303	3.355
Retinex	3.227	3.417	3.196	3.084
Haze-line	2.671	2.582	2.628	2.907
UDnet	3.283	3.453	3.338	3.483
FUNIE-GAN	3.164	3.413	3.155	3.207
UGAN	3.151	3.473	3.213	3.289
CycleGan	2.943	3.154	3.073	3.171
U-Shape	3.209	3.207	3.150	3.069
PUGAN	3.102	3.444	3.194	2.677
Mamba-UIE	3.085	3.384	3.085	3.092
Ours	3.234	3.467	3.342	3.502

, where the red color indicates the best performance, the blue the second-best, and the green the third-best values for each metric. Given that no reference images were available for the datasets reported in Table 4, only the no-reference metric UIQM is used for the quantitative evaluation. It can be found that overall, the proposed MFP-PAINet achieved the best comprehensive performance across all the eight datasets, indicating its superior robustness and strong generalization ability under varying water types and illumination conditions.

On the EUVP and UFO-120 datasets, MFP-PAINet attained the highest SSIM values (0.735 and 0.741, respectively), and its PSNR and UIQM scores could be consistently ranked second. For example, on EUVP, its UIQM score of 3.051 was only slightly lower than that of UDnet (3.170). On LSUI dataset, MFP-PAINet achieved the highest PSNR (23.044), and its SSIM (0.854) was marginally lower than that of Mamba-UIE (0.877). Although its UIQM score on UIEB ranked third (2.980), its SSIM and PSNR values remained at a comparatively high level, demonstrating stable performance on this challenging dataset, which is representative of diverse underwater scenes and degradation patterns. On the no-reference datasets (Table 4), MFP-PAINet achieved the highest UIQM scores on RUIE and UOT32 (3.342 and 3.502, respectively), and remained close to the best performance on AUDD and DUO, indicating that even under significant distribution shifts and across heterogeneous underwater scenarios, the proposed model could maintain high perceptual quality, thus reflecting its strong generalization capability and robustness.

Furthermore, the computational complexity and parameter counts of different methods are reported in

Table 5. Comparison of MACs (Multiply–Accumulate Operations) and Parameters between Different Methods.

	UDnet	FUNIE-GAN	UGAN	CycleGan	U-Shape	PUGAN	Mamba-UIE	Ours
MACs (G)	30.150	10.695	19.820	56.864	2.985	75.400	381.569	211.204
Params (M)	1.402	7.716	57.169	11.378	22.817	101.187	6.046	17.481

. As can be observed, the parameter size of MFP-PAINet (17.481 M) lies at a moderate level among all the compared models. Although its MACs are higher than those of some lightweight models (e.g., U-Shape), they remain substantially lower than those of the Mamba-UIE model. With both enhancement performance and model size considered, MFP-PAINet achieved a reasonable trade-off between computational cost and restoration quality.

4.6. Qualitative Comparison

The qualitative comparison results on the eight datasets are shown in Figure 8, Figure 9, Figure 10, Figure 11, Figure 12 and Figure 13, and they are further supported by the corresponding SSIM and UIQM values. Figure 8, Figure 9, Figure 10 and Figure 11 present the results on the datasets with reference images, whereas Figure 12 and Figure 13 correspond to the datasets without reference images. It can be found that overall, in all the competing methods tested, the proposed MFP-PAINet could most effectively mitigate image turbidity and color distortion, while producing more visually natural colors, better balanced contrast, and clearer structural details. By contrast, traditional UIE methods often exhibited over- or under-enhancement and pronounced color bias; on the other hand, although deep learning–based methods are superior to traditional approaches, they still suffered from varying degrees of color shifts and color artifacts.

On the EUVP dataset (see Figure 8), the results of UDCP and CycleGan exhibit an overall bluish-green color cast; Retinex and Haze-line tend to over-enhance the images; UDnet and UGAN render insufficient color restoration; and the enhancement results of FUNIE-GAN, U-Shape, and PUGAN showcase noticeable color deviations. Although Mamba-UIE yielded better visual results than the aforementioned methods, it was still inferior to MFP-PAINet in terms of color fidelity and local detail preservation. On the LSUI dataset (see Figure 9), the results produced by PUGAN display obvious color bias. On the UFO-120 dataset (see Figure 10), the enhanced images generated by UGAN, CycleGan, and Mamba-UIE retained a prominent bluish-green background, which was more evident in locally zoomed regions, while U-Shape and FUNIE-GAN introduce distinct color artifacts, again clearly visible in local enlargements. It can be concluded that overall, for the datasets with reference images, MFP-PAINet is visually the closest to the ground truth (GT) in terms of global tone, local texture, and perceptual naturalness, indicating that it possesses strong color restoration and detail reconstruction capabilities in complex underwater scenes. For the no-reference datasets, the enhancement results produced by deep learning–based UIE methods generally look more natural in visual perception and exhibit richer color content, whereas traditional UIE methods, such as UDCP and Haze-line, still suffer from a pronounced bluish-green background and noticeable color distortion.

To further verify the effect of uncertainty modeling in the UIE process, multiple potential candidate enhancement results are compared with the final enhanced output (Figure 14). As illustrated in Equations (7) and (8), the model can generate several potential enhancement candidates during the testing phase (six of which are shown in this study, denoted as Samples 1–6), while the final enhancement result (Sample 7) corresponds to the latent variable sample with the highest probability, as determined by Equation (3). To visually demonstrate the differences among these enhancement results, the PSNR heatmaps were computed and visualized for all the samples. As shown in Figure 14, distinct variations exist across different candidate samples, particularly in local regions. The final enhancement result, obtained based on Equation (3), achieved the highest PSNR value (24.60), indicating superior performance in terms of structural preservation and detail restoration. These findings validate that the proposed probabilistic modeling mechanism can effectively evaluate and sort out the optimal enhancement outcome, thereby improving the entire model’s stability and robustness under complex underwater conditions.

4.7. Ablation Studies

To validate the effectiveness of each module in MFP-PAINet, ablation studies were conducted on the LSUI dataset, with the results presented in

Table 6. Ablation Study Results: Comparison of Metrics Across Different Baselines on the LSUI Dataset.

Multi-dimensional Feature Mapping Module	Confidence Generation Module	Probabilistic Module	$\mathbf{SSIM} \uparrow$	$\mathbf{PSNR} \uparrow$
√	√		0.841	21.963
	√	√	0.851	22.677
√		√	0.851	22.795
		√	0.848	22.754
√	√	√	0.854	23.044

for the impact of removing different modules on network performance. The LSUI dataset was chosen for ablation studies because its rich, high-quality reference images and diverse underwater scenes provide ideal conditions for assessing the model’s performance in terms of detail preservation, structural similarity, and color restoration. As can be observed, compared to other datasets, the diversity and standards of LSUI better reveal the contributions of each module in image enhancement, making the experimental results more representative and persuasive.

When the probabilistic module was removed, SSIM and PSNR dropped to 0.841 and 21.963, respectively, indicating that the probabilistic module plays a key role in handling image uncertainty, enhancing structural similarity, and improving the signal-to-noise ratio. Removing the multi-dimensional feature mapping module drove SSIM and PSNR to decrease to 0.851 and 22.677, respectively. Although with weakened performance, the ablated model still closely approached the complete model in terms of performance. When the confidence generation module was removed, there was only a slight decrease in performance (SSIM of 0.851 and PSNR of 22.795), indicating that the confidence module mainly addressed the uncertainties in image enhancement initially. Keeping only the probabilistic module resulted in an SSIM and PSNR of 0.848 and 22.754, respectively, suggesting that while the probabilistic module could handle some uncertainties, it performed best when functioning with other modules. Overall, the complete network (SSIM of 0.854 and PSNR of 23.044) delivered the highest performance, validating the critical role of the synergy in enhancing image quality for the modules.

4.8. Extended Experiments

To further assess the generalization ability of MFP-PAINet, this study additionally evaluated it on image deraining (Rain1400 [60]) and dehazing (RESIDE [61]) tasks. SSIM, PSNR, and RMSE were used as the evaluation metrics, with the quantitative results shown in

Table 7. Quantitative results for image deraining and dehazing. $\downarrow$ lowest value.

Method	Rain1400 (Deraining)			RESIDE (Dehazing)
	$\mathbf{SSIM} \uparrow$	$\mathbf{PSNR} \uparrow$	$\mathbf{RMSE} \downarrow$	$\mathbf{SSIM} \uparrow$	$\mathbf{PSNR} \uparrow$	$\mathbf{RMSE} \downarrow$
UDCP	0.653	18.419	32.594	0.704	16.964	37.842
Fusion	0.704	18.501	31.484	0.632	10.841	76.369
Retinex	0.512	8.368	99.166	0.634	11.309	71.816
Haze-line	0.697	17.529	36.190	0.636	11.308	73.775
UDnet	0.733	20.469	25.653	0.787	16.726	41.160
FUNIE-GAN	0.842	26.687	12.247	0.850	25.938	13.349
UGAN	0.821	26.282	12.732	0.819	24.471	15.774
CycleGan	0.858	28.186	10.610	0.826	21.914	22.411
U-Shape	0.820	26.881	12.024	0.844	27.843	10.640
PUGAN	0.784	24.491	15.806	0.787	20.944	24.103
Mamba-UIE	0.894	28.229	10.383	0.897	26.256	12.844
Ours	0.900	28.214	10.421	0.920	25.987	13.643

and the qualitative results in Figure 15 and Figure 16.

As can be seen, on the Rain1400 deraining task, MFP-PAINet achieved the highest SSIM (0.900), indicating that its outputs are structurally closest to the reference images. In addition, its PSNR (28.214) and RMSE (10.421) were the second best, only slightly lower than those of Mamba-UIE (28.229 and 10.383, respectively). As shown in Figure 15, the derained results produced by MFP-PAINet and Mamba-UIE are visually closest to the ground truth, effectively removing rain streaks across the entire image, while preserving fine details and restoring natural colors in a better manner. In contrast, such traditional methods as UDCP and Haze-line failed to remove most rain streaks, but introduced noticeable color casts; on the other hand, although deep learning–based methods, such as UDnet, UGAN, and PUGAN, were able to eliminate a large portion of rain streaks, they still left residual streaks and showed limited detail recovery.

On the RESIDE dehazing task, MFP-PAINet attained the highest SSIM (0.920) again, and its PSNR (25.987) and RMSE (13.643) were close to those of Mamba-UIE (26.256 and 12.844, respectively), thus demonstrating its overall superior ability over traditional methods and other GAN-based approaches. As illustrated in Figure 16, MFP-PAINet provided the best performance in terms of detail and structural reconstruction, a result consistent with its leading SSIM score. By comparison, Fusion and Retinex almost failed to remove haze effectively, while UDCP and Haze-line introduced obvious artifacts and color shifts. Other deep learning–based methods, such as FUNIE-GAN and UGAN, delivered certain dehazing effects but still suffered from color artifacts and loss of fine details, especially in the locally enlarged regions.

It is important to note that MFP-PAINet is designed primarily for UIE tasks. The goal of the extended experiments as described in this section is not to emphasize its general applicability to other image enhancement tasks but to demonstrate its adaptability across different task scenarios. These experiments validate that MFP-PAINet not only performs excellently in UIE but can deliver superior performance in cross-task applications. The extended experiments in this section validate the generalization and high performance of MFP-PAINet and further reveal its practical application potential in other domains.

5. Conclusions

This study proposes MFP-PAINet, a UIE method based on multi-dimensional feature fusion and probabilistic adaptive uncertainty distribution to address the limitations of the existing methods in handling uncertainties and visual degradation. MFP-PAINet comprises three modules: a multi-dimensional feature mapping module, a confidence generation module, and a probabilistic module. Specifically, the multi-dimensional feature mapping module is used to capture and integrate higher-level features through a multi-scale feature fusion strategy and generate multi-dimensional feature mappings; the confidence generation module preliminarily addresses uncertainties in the enhancement process by dynamically weighting and fusing various feature mappings; and the probabilistic module further manages these uncertainties through the combination of CVAE and PAdaIN. Experimental results demonstrate that the model proposed in this study significantly outperforms the existing approaches in such metrics as SSIM, PSNR, and UIQM, with SSIM (0.854) and PSNR (23.044) delivered on the LSUI dataset, a performance surpassing other methods. Ablation studies validate the contribution of each module to the model’s performance, and extended experiments in image de-raining and de-hazing tasks further demonstrate the applicability and robustness of MFP-PAINet. While delivering superior structural clarity and perceptual quality, MFP-PAINet created slight fluctuations in color metrics (e.g., UIQM), suggesting a trade-off between structural fidelity and color accuracy. Moreover, the experiments in this study were conducted on publicly available datasets, which might fail to fully map the complexity of real-world underwater environments. Therefore, future efforts are needed to be focused on improving color consistency and further test the proposed method in practical underwater scenarios.