A Two-Stage Framework for Distortion Information Estimation and Underwater Image Restoration

Abstract

This work introduces a two-stage framework, named the Distorted underwater image Restoration Network (DR-Net), to address the complex degradation of underwater images caused by turbulence, water flow fluctuations, and optical properties of water. The first stage employs the Distortion Estimation Network (DE-Net), which leverages a fusion of Transformer and U-Net architectures to estimate distortion information from degraded images and focuses on image distortion recovery. Subsequently, the Image Restoration Generative Adversarial Network (IR-GAN) in the second stage utilizes this estimated distortion information to deblur images and regenerate lost details. Qualitative and quantitative evaluations on both synthetic and real-world image datasets demonstrate that DR-Net outperforms traditional methods and restoration strategies from different perspectives, showcasing its broader applicability and robustness. Our approach enables the restoration of underwater images from a single frame, which facilitates the acquisition of marine resources and enhances seabed exploration capabilities.

1. Introduction

The underwater imaging quality significantly affects the retrieval of marine resources and the exploitation of seabed resources [1]. Despite substantial advancements in optical imaging technology [2] driven by rapid progress in science and technology, underwater imaging differs from aerial imaging, with relatively fewer associated perturbations. The underwater environment is more complex, as water exerts a significant influence on light propagation [3]. Moreover, the presence of numerous suspended particles in water often results in underwater images suffering from quality degradation issues such as blurriness, low contrast, and color distortion. Additionally, the fluidity of water leads to frequent underwater turbulence [4,5] and water fluctuations [6], which further complicate the image degradation process by introducing unpredictable and intricate patterns, thereby increasing the difficulty of image restoration. In summary, three primary factors contribute to the prevalence of underwater image distortion: the optical properties of water, underwater turbulence (thermal perturbations in the water body under extreme conditions), and water fluctuations.

In this paper, we investigate image restoration techniques aimed at improving the clarity of underwater images affected by irregular distortion and complex blurring caused by water fluctuations and underwater turbulence. Traditional approaches involve using long-time image sequences or videos captured in static scenes to calculate mean and median frames from multiple frames [7,8], or adapting the “lucky imaging” method from astronomical imaging [9,10] to identify minimally distorted image patches for fusion, which can potentially yield undistorted images. Image alignment methods [8,11] can then utilize these obtained potentially less distorted images to iteratively align and restore the original image. The work in [12] proposed a model-based tracking approach that utilizes a fluctuation model and attempts to fit the fluctuations in image sequences to obtain an undistorted image; however, it struggles with more intricate distortions [13] and blurring degradation induced by the water environment. Li et al. [14] introduced a deep learning-based restoration method that can enhance distorted images from a single frame, demonstrating better applicability.

Additionally, the vision transformer (ViT) [15] has shown exceptional performance in image processing and restoration. Restormer [16] proposes a multi-Dconv head transposed attention (MDTA) module and a gated-Deconv feedforward network (GDFN), which effectively balance local context modeling and global dependency capture to handle high-resolution restoration tasks such as deraining and deblurring. Histoformer [17] employs histogram self-attention, sorting spatial features into intensity-based bins to capture long-range dependencies among weather-degraded pixels, and excels in unified adverse weather removal and restoration. TransMRSR [18] combines convolutional and Transformer blocks for super-resolution and image restoration, using a two-stage strategy with generative prior and self-distilled truncation. While these methods perform well in their respective domains and general fields, they are not designed to address underwater-specific challenges such as turbulence-induced distortions and optical scattering.

Building upon the aforementioned methodologies, this paper introduces a two-stage deep learning-based framework named DR-Net for restoring underwater distorted images. As illustrated in Figure 1, the framework assigns distinct restoration objectives to its two stages. The first stage, DE-Net (Distortion Estimation Network), primarily focuses on distortion correction: it leverages the global feature extraction capability of the ViT model to estimate a distortion map, which is principally represented as a displacement vector field derived from the input degraded image and subsequently generates a distortion-corrected image based on this computed distortion information. The second stage, IR-GAN (Image Restoration Generative Adversarial Network), addresses image blurring—a common artifact in underwater imagery. Notably, pixel blurring is inevitably introduced during the distortion correction process of the first stage; thus, this stage enhances image clarity through the adversarial generative network while recovering partially lost image details.

Our network was trained and evaluated using the dataset established by Li et al. [14], with additional validation performed on other real-world and synthetic datasets [12,19,20,21]. Comparative experiments against classical methods [8,11,12,14,19,22] and restoration methods [16,17,18] from various perspectives demonstrate that our approach consistently outperforms alternatives in image restoration. A key advantage is its enhanced versatility: unlike traditional methods that rely on multi-frame sequences or extended video inputs from the same scene, our method only requires a single frame for effective underwater image restoration. Furthermore, quantitative evaluations using multiple metrics confirm that our approach achieves state-of-the-art performance across most scenarios.

To summarize, the main contributions of our work are as follows:

• We propose DR-Net, a novel two-stage deep learning framework specifically designed to address underwater image degradation caused by turbulence and water fluctuations, integrating the strengths of Transformer and GAN architectures.
• The staged design of DR-Net enables targeted handling of distinct degradation factors: DE-Net focuses on distortion estimation and correction, while IR-GAN specializes in deblurring and detail recovery, effectively mitigating the cumulative artifacts introduced during sequential processing.
• Our method achieves superior performance on both synthetic and real-world datasets, outperforming classical and state-of-the-art approaches in both qualitative and quantitative evaluations. And DR-Net enhances practical applicability in scenarios by enabling high-quality restoration from a single frame.

2. Related Work

Distorted underwater image restoration. Restoration of distorted underwater images presents a formidable challenge in image processing. Various perspectives have been explored to address the restoration of degraded underwater images:

(1) Early methods, including using the mean or median frame of a static scene image sequence, or calculating water surface refraction concurrently [7] with recording water surface images, have proven effective for recovering simple distortions. (2) In the image combination perspective, the lucky imaging [9,10] method is utilized to select the least distorted patches of an image sequence and fuse them, with selection methods such as k-means clustering [23,24], bicoherence techniques [25], streaming algorithms [10], and bispectral analysis [26] employed.

(3) With regard to the image alignment approach, where multiple frames of a static scene are aligned to a reference image for recovery, techniques like averaging frames [27] or sequential frame-by-frame alignment [28] are commonly used. Researchers like Oreifej et al. [8] proposed a two-stage restoration method incorporating denoising post-alignment, while authors like Halder et al. [11] implement mean-filtered denoising in the second stage, and Yang et al. [29] combine non-rigid alignment and lucky block fusion. (4) The wave model perspective addresses water surface fluctuation-induced distortions, aiming to estimate the fluctuation structure and image distortion by analyzing consecutive image frames to reconstruct the scene image, with researchers like Tian et al. [12] pioneering a fluctuation equation-based modeling approach.

Even recently, (5) developments in pixel point tracking and compressive sensing areas focus on estimating the global pixel displacement vector, tracking multiple feature points in the image sequence, and employing compressed sensing and sparse representation methods [19,30] to dewarp the images. Additionally, sparse or dense optical flow methods [31,32] can be employed for restoration. (6) Furthermore, from a deep learning perspective, Li et al. [14] utilize adversarial generative networks to correct distortions in a single-frame underwater image. Our method integrates global feature extraction to address distortions, blurring, and details, producing high-quality restored images from a single frame.

End-to-end image translation. In image processing tasks, an end-to-end image translation methodology emerges as a straightforward and practical framework, commonly featuring an encoder–decoder architecture tailored for images. The encoder efficiently captures essential image features, while the decoder deciphers and reconstructs the desired image. One of the pioneering architectures is the U-Net [33], initially devised for medical image segmentation and promptly adapted to diverse image processing applications. Over the past decade, the versatile U-Net framework has garnered widespread adoption and demonstrated remarkable outcomes.

The integration of the end-to-end structure with Generative Adversarial Networks [34] for image-to-image translation and synthesis has proven highly successful across various visual tasks, including image denoising, deblurring [35,36,37], style migration [38,39], restoration [40,41], and super-resolution [42,43]. This approach involves training two networks in tandem through adversarial training, where the discriminator and generator engage in a game-like dynamic to enhance the realism of generated images. Isola et al. [22] introduced a comprehensive framework for image translation, using L1 loss and adversarial loss to train end-to-end network for source–target image pairs. In our study, we extend this concept by applying GANs to underwater imaging in our second phase, harnessing their potent generative capabilities to enhance image details effectively.

ViT-based image processing. In deep learning, various robust backbone networks exist, among which Transformer has been widely adopted in computer vision due to its powerful global feature modeling capability [44]. Specifically, Dosovitskiy et al. [15] pioneeringly introduced Transformers to computer vision tasks, demonstrating its significant value in visual domains and enabling its expansion to diverse visual tasks. ViT enhances operational efficiency by segmenting images into sequential patches, embedding them into a Transformer network, and effectively extracting more discriminative feature information to improve performance.

Cong et al. [45] proposed an encoding–decoding framework based on ViT, which employs Transformer modules for feature extraction and is trained with a multi-loss combination. This optimization effectively enhances color fidelity and contrast in underwater images. Similarly, Peng et al. [46] designed a U-shaped Transformer network integrating two modules: the channel-based multi-scale feature fusion Transformer (CMSFFT) and the spatial-based global feature modeling Transformer (SGFMT), strengthening the processing capacity for attentional attenuation in color channels and spatial regions. Xue et al. [47] developed a lightweight cross-gated Transformer algorithm for image restoration, aiming to achieve superior results with fewer parameters and lower computational costs. Additionally, Zheng et al. [48] built a multi-scale modulation module atop a ViT to dynamically select representative features in the feature space of underwater images. More recently, Gao et al. [49] combined a Transformer with a CNN for local feature enhancement, effectively improving image details, while Xu et al. [50] proposed a semi-supervised framework using an Agg Transformer to enhance underwater images.

Transformer-based image restoration. In recent years, there has been a proliferation of Transformer-based image restoration techniques that have excelled in various restoration tasks. For instance, SwinIR [51] utilizes the Swin Transformer as a robust baseline model for image restoration. TransMRSR [18], on the other hand, targets single-image super-resolution reconstruction by introducing a Transformer-based self-distillation generation prior network, which has demonstrated notable performance in relevant domains. Additionally, Restormer [16] and Histoformer [17] have introduced innovative approaches to the realm of image restoration. Restormer enhances the foundational elements of Transformer models, such as multi-head attention and feedforward networks, to facilitate the capture of long-range pixel interactions while maintaining adaptability to large images. Histoformer, on the other hand, centers around a novel histogram self-attention mechanism, which intelligently focuses on spatial features across a dynamic range, allowing for adaptive emphasis on deteriorated features exhibiting similar patterns. Zhu et al. [52] proposed FDNet, a Fourier transform-guided dual-channel underwater image enhancement diffusion network, outperforming existing learning-based methods on multiple underwater datasets. Transformer-based methodologies have been proposed for diverse tasks, including super-resolution, deblurring, and image enhancement.

3. Method: DR-Net

Consider a scenario where a camera is positioned above the water surface, capturing an underwater scene through the air–water medium. The optical properties of water cause light absorption and scattering [53]. Additionally, light refraction occurs at the air–water interface, exacerbated by surface fluctuations, resulting in complex distortions in the final image. Further, fluid dynamics of water and underwater turbulence introduce intricate blurring effects. Let $\mathcal{I}$ denote the captured distorted underwater image, and $\mathcal{J}$ represent the ground-truth undistorted image of the underwater scene. Our goal is to estimate a de-distorted image $\widehat{\mathcal{J}}$ from $\mathcal{I}$ such that $\widehat{\mathcal{J}}$ closely approximates $\mathcal{J}$.

Mathematically, the de-distortion process can be modeled as a transformation $\mathcal{T}$ that maps the distorted image to the de-distorted one: $\widehat{\mathcal{J}}=\mathcal{T}\left(\mathcal{I}\right)$. To optimize this transformation, we aim to find the optimal ${\mathcal{T}}^{\ast }$ that minimizes the loss between the transformed result and the ground-truth:

$${\mathcal{T}}^{\ast }=arg\underset{\mathcal{T}}{min}\mathcal{L}\left(\mathcal{T}\left(\mathcal{I}\right),\mathcal{J}\right)$$

(1)

where $\mathcal{L}(·,·)$ denotes a suitable loss function that quantifies the discrepancy between the estimated de-distorted image $\mathcal{T}\left(\mathcal{I}\right)$ and the ground-truth image $\mathcal{J}$.

To address this underwater image restoration task, we propose DR-Net, a two-stage framework that decomposes the task into a distortion correction stage (DE-Net) and a detail recovery and deblurring stage (IR-GAN). Figure 2 illustrates the following:

• The degraded captured image is first fed into DE-Net, which integrates a U-Net architecture with a ViT for global feature extraction at high dimensions, supplemented by a dual-channel attention block (DA-Block). Specifically, DE-Net initially employs U-Net to capture multi-scale features, enabling adaptation to diverse distortion factors in the image. For high-dimensional features, the DA-Block is utilized to extract effective features, which are then processed by ViT for global feature modeling. This process ultimately yields valid global distortion information (or a distortion map), which is used to generate a distortion-corrected image.
• Subsequently, the distortion-corrected image, combined with the estimated distortion information as guidance, is input to IR-GAN. Leveraging the powerful generative modeling capability of GANs, IR-GAN further deblurs the image and regenerates lost details, completing the overall restoration process.

3.1. IR-GAN: GAN-Based Restoration

Underwater image degradation arises from multiple factors, including water’s optical properties, particulate matter, surface fluctuations, and turbulence. A straightforward restoration approach involves directly mapping distorted images to clear ones via a single model. Initially, we explored an end-to-end network leveraging nonlinear fitting capabilities to simultaneously address distortion and blurriness. We modeled the intricate distortion in underwater images using a system function H, where

$$\mathcal{I}=H\left(\mathcal{J}\right(x,t\left)\right)$$

(2)

Here, x denotes pixel spatial position and t is the imaging timestamp.

For image-to-image translation, we adopted a GAN framework (building on pix2pix success) given its strong generative modeling capabilities. Our IR-GAN comprises a generator $\mathcal{G}$ (outputting restored images) and a discriminator $\mathcal{D}$ (distinguishing restored images from ground-truth undistorted ones). The goal is to optimize $\mathcal{G}$ such that $\mathcal{D}$ cannot differentiate its outputs from real undistorted images. As shown in Figure 2, IR-GAN takes distorted underwater images (incorporating distortion information) as input and outputs restored results directly.

To stabilize training and mitigate mode collapse, we employed Least Squares GAN (LSGAN) [55], which uses least squares loss to minimize the ${\chi }^2$ divergence between generated and target data. The objective functions are

$$\begin{array}{cc}\hfill \underset{\mathcal{D}}{min}{V}_{\mathrm{LSGAN}}\left(\mathcal{D}\right)& ={\displaystyle \frac{1}{2}}{\mathbb{E}}_{\mathcal{I}\sim p_{\mathrm{daa}}\left(\mathcal{I}\right)}\left[{\left(\mathcal{D}\left(\mathcal{I}\right)-1\right)}^2\right]\hfill \\ \hfill & +{\displaystyle \frac{1}{2}}{\mathbb{E}}_{\mathcal{J}\sim p_{\mathcal{J}}\left(\mathcal{J}\right)}\left[{\left(\mathcal{D}\left(\mathcal{G}\right(\mathcal{J}\left)\right)\right)}^2\right]\hfill \end{array}$$

(3)

$$\underset{\mathcal{G}}{min}{V}_{\mathrm{LSGAN}}\left(\mathcal{G}\right)={\displaystyle \frac{1}{2}}{\mathbb{E}}_{J\sim p_J\left(J\right)}\left[{(\mathcal{D}\left(\mathcal{G}\left(J\right)\right)-1)}^2\right]$$

(4)

The generator architecture (Figure 3) includes downsampling layers, residual blocks, and upsampling layers. After initial downsampling (three convolutional modules) for feature extraction, eight residual blocks refine features, followed by three upsampling modules for reconstruction. To preserve details efficiently, we adopted the lightweight dynamic upsampler DySample [56] for upsampling. The discriminator follows a design similar to [57], i.e., all convolutional layers (except the final one) are paired with Batch Normalization and LeakyReLU, while the final layer uses a Sigmoid activation to output discriminative probabilities.

3.2. DE-Net: Distortion Information Estimation

The preceding section described IR-GAN’s end-to-end reconstruction, which directly maps distorted images to restored ones via a generator $\mathcal{G}$:

$$\mathcal{I}\stackrel{\mathcal{G}}{\to }\widehat{\mathcal{J}}$$

(5)

However, evaluations of IR-GAN (see Section 4.3) reveal that this direct GAN-based approach, simultaneously handling distortion correction and deblurring, fails to achieve optimal performance. While deblurring results are relatively satisfactory, distortion correction across varying scales remains unstable.

This limitation motivates us to decouple the process, first focusing on accurate distortion estimation to lay a robust foundation for subsequent restoration. We model the underwater distortion phenomenon as

$$\mathcal{I}(x,t)=\mathcal{J}(x+w(x,t\left)\right)$$

(6)

where $w(x,t)$ denotes the distortion vector at spatial position x and time t. For clarity, this formulation isolates distortion (ignoring blurring) to specifically address spatial perturbations in the image.

To precisely capture such distortion characteristics, DE-Net integrates Transformer and U-Net architectures, leveraging their complementary strengths: U-Net extracts local hierarchical features via convolutional layers, while a Transformer encoder enhances global feature modeling. This hybrid design ensures seamless feature mapping between shallow convolutional layers and U-Net’s sampling processes, maximizing the utility of multi-scale features. Further, we incorporate a Dual-Attention Block (DA-Block) [54] into U-Net’s skip connections (Figure 4) to strengthen feature representation and utilization. The DA-Block combines a Position Attention Module (PAM) and Channel Attention Module (CAM), which respectively refine spatial- and channel-wise feature interactions. This dual attention mechanism enhances feature parsing capabilities, facilitating efficient transmission of critical information—enabling DE-Net to accurately estimate fine-grained underwater distortion details. Notably, placing the DA-Block before Transformer input also improves feature embedding quality.

In the processing pipeline, DE-Net takes camera-captured distorted images as input, estimates distortion information $w(x,t)$, and applies a Spatial Transformer Network (STN) to perform distortion correction, yielding a de-distorted intermediate image. Through iterative supervised training against ground-truth images, the model is refined to output precise distortion information, laying the groundwork for subsequent blurring removal.

3.3. Network and Training Object

While DE-Net effectively corrects spatial distortions, relying solely on it introduces issues like smooth blurring and color inconsistencies. Additionally, the spatial transformation process cannot eliminate pre-existing blur in the input image. To address these limitations, we introduce this two-stage cascaded framework: DE-Net first estimates and rectifies distortions, and IR-GAN subsequently refines details and removes blurring, leveraging GAN’s strength in recovering lost information during degradation. Experimental results confirm that this integration yields superior restoration performance. The complete architecture of DR-Net is illustrated in Figure 2.

Training Object. Three key loss functions are employed to train the network: the content loss, the adversarial loss, and the perceptual loss:

• The Content Loss: We use L1 loss to enforce pixel-wise similarity between the restored image and ground-truth:

  $${\mathcal{L}}_{con}=\sum _x\left|\mathcal{J}\left(x\right)-G\left(\mathcal{I}\right(x),x)\right|$$

  (7)

  where $\mathcal{J}$ denotes the ground-truth image, and $\mathcal{G}\left(\mathcal{I}\right(x),x)$ is the generator output. However, L1 loss alone may result in insufficient high-frequency details, motivating the inclusion of adversarial loss to enhance detail restoration and even recover information lost due to significant distortions.
• The Adversarial Loss: Adopting LSGAN’s loss function stabilizes training and mitigates mode collapse:

  $${\mathcal{L}}_{adv}={\displaystyle \frac{1}{2}}{E}_{z\sim p_z\left(z\right)}\left[{\left(D\left(G\left(z\right)\right)-1\right)}^2\right]$$

  (8)
• The Perceptual Loss: To reduce artifacts introduced by GANs, we incorporate perceptual loss, which measures feature-level similarity using a pre-trained network. It is defined as

  $${\mathcal{L}}_{per}=\sum _x\left|\psi \left(\mathcal{I}\left(x\right)\right)-\psi \left(G(\widehat{\mathcal{J}}\left(x\right),x)\right)\right|$$

  (9)
  Here, $\psi$ represents the output of an intermediate feature layer from a pre-trained network.

The perceptual loss is computed by separately feeding the generated image and the ground-truth image into the pre-trained network, extracting features from corresponding layers, and calculating the L1 or MSE distance between these feature maps. By measuring the content discrepancy between high-level features of the generated and ground-truth images, this loss enhances the generalization of the restoration results. In this study, we utilize the outputs of the conv4_3 layer from a pre-trained VGG16 model for this purpose.

Thus, the total loss function is a weighted combination of the three components:

$$\mathcal{L}={\mathcal{L}}_{con}+{\lambda }_{adv}{\mathcal{L}}_{adv}+{\lambda }_{per}{\mathcal{L}}_{per}$$

(10)

where ${\lambda }_{\mathrm{adv}}$ and ${\lambda }_{\mathrm{per}}$ are scaling factors. We initially set empirical weights and adjust them during training based on restoration quality, ultimately adopting ${\lambda }_{\mathrm{per}}=0.6$ and ${\lambda }_{\mathrm{adv}}=0.2$. For standalone DE-Net training, ${\lambda }_{\mathrm{adv}}=0$ (since adversarial loss is specific to IR-GAN).

4. Experiments

In this section, we present the evaluation of our model DR-Net using both synthetic datasets and real datasets. We conduct comparisons with the classical approach and perform ablation studies on our proposed model.

4.1. Network Training

Data preparation and metrics. For training our proposed underwater image restoration model DR-Net, we use the LiSet obtained from Li et al. [14]. This dataset is a comprehensive collection of underwater images based on the ImageNet dataset [58], where the researchers utilized a Cannon 5D Mark IV, (Beijing, China), camera to capture images by passing ImageNet images through a water tank at a distance exceeding 1.5 m. In total, the dataset comprises 324,452 images from 1000 ImageNet classes, with 5 images retained per class for a validation set totaling 5000 images.

Given the varied sensitivity of individuals to image attributes such as brightness, blur, and distortion, it is imperative to employ appropriate image evaluation metrics to subjectively assess the results of image restoration. Traditional metrics like Peak Signal-to-Noise Ratio (PSNR) [59] and Structural Similarity Index (SSIM) [60,61] are commonly used for assessing image restoration quality and similarity. This study also includes a comparison of these fully referenced evaluation metrics.

Traditional metrics may lack the ability to capture intricate perceptual differences due to their simplistic pixel-based evaluation functions. Recent advancements in deep learning have introduced more effective evaluation metrics such as Learned Perceptual Image Patch Similarity (LPIPS) [62] and Deep Image Structure and Texture Similarity (DISTS) [63], which have demonstrated superior performance in image quality assessment by incorporating perceptual features and complex image information. LPIPS, particularly adept at assessing local details that align with human visual perception, provides consistent evaluations across diverse distortion types. In this work, we utilize the pre-trained LPIPS model as described in [64]. DISTS employs deep neural networks to extract both structure and texture features, enabling a comprehensive analysis of complex image characteristics. Additionally, this paper utilizes reference-free assessment metrics, Blind/Referenceless Image Spatial Quality Evaluator (BRISQUE) [65], which considers contrast, brightness, and texture aspects.

Implementation details. Our model DR-Net is formulated and trained using the PyTorch1.10 framework as a two-stage architecture with three phases involving separate and cascading training. Initially, the Distortion Information Estimation Network (DE-Net) is trained using the Adam optimizer with a learning rate of 0.0002 and a batch size of 16 for both training and testing. A pre-trained vision transformer is utilized during training to accelerate convergence. The subsequent phase involves training the image restoration model, IR-GAN, based on Generative Adversarial Networks (GANs). The distortion information estimated by DE-Net, along with the distorted image corrected through bilinear interpolation or Spatial Transformer Networks (STN), serves as inputs for the IR-GAN model during the second phase of training. Supervision is enforced through the loss function delineated in Section 3.3, leveraging the Adam optimizer with a learning rate of 0.0002. The final phase encompasses fine-tuning the cascaded training of the two networks to enhance the cascade performance, leading to improved restoration outcomes through interactive refinement within the two-stage network structure. Detailed algorithmic training specifications are presented in Algorithm 1.

Algorithm 1 Training process of our network

Input: Training data $(\mathcal{I},\mathcal{J})$ Output:  $G^{\left(K_2\right)}{\left(M^{\left(K_1\right)}\right)}_f$ Phase 1: Training DE-Net M for  $k\leftarrow 1$  to  $K_1$  do        Train network M with input $\mathcal{I}$ and target $\mathcal{J}$ to obtain parameters $M^{\left(K_1\right)}$        Train Loss: Equations (7) and (9) end for Phase 2: Training IR-GAN G for  $k\leftarrow 1$  to  $K_2$  do        Compute M outputs $\widehat{W}$ and ${\mathcal{I}}_M$$=M\left(\mathcal{I}\right)$        Train network G with inputs $(\widehat{W},{\mathcal{I}}_M)$ and target $\mathcal{J}$ to get parameters $G^{\left(K_2\right)}$        Train Loss: Equations (7)–(9) end for Phase 3: Joint Fine-Tuning of M and G, DR-Net for  $k\leftarrow 1$  to  $K_f$  do        Fine-tune M and G together with input $\mathcal{I}$ and target $\mathcal{J}$ to obtain fine-tuned parameters $G^{\left(K_2\right)}{\left(M^{\left(K_1\right)}\right)}_f$        Train Loss: Equations (7)–(9) end for

4.2. Comparison with the State-of-The-Arts

To provide a more comprehensive evaluation of our model’s restoration performance, we conducted a comparative analysis with established underwater image restoration methods [8,11,12,14,16,17,18,19,22]. Specifically, Tian et al. [12], Oreifej et al. [8], and Halder et al. [11] all employed conventional image restoration approaches. Tian et al. implemented a model-based tracking restoration method using the fluctuation equation, while Oreifej et al. and Halder et al. primarily utilized image alignment for underwater image restoration. In our comparative study, we replace the non-rigid alignment method employed in Halder et al.’s approach with the optical flow method [66] for experimental evaluation. James et al. incorporated a compressed perception technique into the underwater image restoration process to enhance the quality of restoration outcomes.

Moreover, Isola et al. [22], Li et al. [14], TransMRSR [18], Restormer [16], and Histoformer [17] all employed deep learning approaches for image restoration tasks. Isola et al.’s method excels in pixel-to-pixel tasks, demonstrating strong performance in style transfer through the generative model. Li et al. concentrated on underwater image restoration, utilizing a dataset obtained from ImageNet capturing underwater distortions to conduct restoration-related research. TransMRSR is a Transformer-based self-distillation prior network that utilizes pre-trained Generative Adversarial Networks (GANs) to establish a robust prior for enhancing image super-resolution reconstruction. Both Restormer and Histoformer enhance the fundamental components of Transformer models to proficiently address image restoration tasks and yield positive outcomes across multiple image restoration scenarios.

Testing datasets. To conduct a comprehensive evaluation of our model, we perform simultaneous tests and evaluations on Li et al.’s test dataset along with various synthetic and real datasets: (1) LiSet [14], an extensive underwater image dataset curated by Li et al. derived from the ImageNet dataset, comprising 5000 images from 1000 classes for testing; (2) TianSet [12], an authentic dataset consisting of underwater image sequences compiled by Tian and Narasimhan from ground-truth images; (3) JamesSet [19], a diverse dataset akin to Tian’s dataset, featuring 14 image sequence categories encompassing blue lines, bricks, cartoon images, checkerboards, and more; (4) ThapaSet [20], a collection of synthetic images with 20 categories synthesized by Thapa et al.; and (5) NianSet [21], a synthetic dataset comprising images of both animals and plants assembled by Nian et al. Our evaluation of image restoration outcomes encompasses both qualitative and quantitative analyses.

Comparison results. As the approaches of Tian et al., Oreifej et al., Halder et al., and James et al. [8,11,12,19] necessitate an image sequence input, image restoration evaluations could not be conducted on the LiSet. Initially, we compare the image restoration outcomes of Isola’s method, Li’s method, Restormer, TransMRSR, Histoformer, and our model based on the LiSet test dataset. Figure 5 presents a comparison of the image restoration results on the LiSet, with the first column displaying the captured perturbed degraded image as the model input. Subsequently, the seventh and eighth columns depict the results of the image restoration via our method and the actual image, and the second to sixth columns showcase the restoration outcomes utilizing other methods, respectively.

As illustrated in Figure 5, our method, DR-Net, demonstrates enhanced recovery performance compared to other methods across various levels of distortion. This is attributed to our model’s comprehensive consideration of both image distortion and detail restoration aspects, as evidenced by superior management of image details, patterns, and deblurring. While Isola and Li’s method, based on adversarial generative networks, shows favorable outcomes in image details and restored contrast, it lacks effectiveness in addressing different scales of distortion restoration. On the other hand, the recent image restoration methods, such as Restormer, TransMRSR, and Histoformer, do not separately account for distortion degradation, thereby impacting their overall image restoration performance to some extent.

Notably, our model demonstrates enhanced robustness against partial occlusions within the image, facilitated by the guidance of the perturbation map estimated in the initial stage and the utilization of the adversarial generative model.

Table 1. Quantitative comparisons on LiSet for underwater distorted image restoration. Note: The upward arrows (↑) indicate that larger values of the corresponding evaluation metrics are better, while the downward arrows (↓) indicate that smaller values are better, where the same applies hereinafter.

Methods	PSNR ↑	SSIM ↑	FSIM ↑	GMSD ↓	LPIPS ↓	DISTS ↓	BRISQUE ↓
Isola et al.  [22]	19.1077	0.6504	0.7487	0.2016	0.4163	0.2280	14.7516
Li et al.  [14]	20.2547	0.7058	0.7859	0.1790	0.2826	0.1512	20.4872
Restormer [16]	21.9327	0.7659	0.7802	0.1709	0.3260	0.2087	46.0945
TransMRSR [18]	19.5300	0.6812	0.7391	0.1924	0.4135	0.2709	57.0114
Histoformer [17]	21.8466	0.7527	0.7689	0.1739	0.3524	0.2333	50.4283
Ours	22.5133	0.7828	0.8177	0.1554	0.2243	0.1321	28.6267

outlines the quantitative evaluation outcomes on the LiSet, indicating the favorable performance of our methodologies. Moreover, Figure 6 illustrates the model’s outstanding stability and recovery efficacy.

In addition, Figure 7 displays the test results for the image datasets by Tian, James, Thapha, and Nian. The performance of the methods proposed by Tian et al. [12], Oreifej et al. [8], Halde et al. [11], and James et al. [19] are compared using 11 frames of image input (denoted as Tian-11, Oreifej-11, Halde-11, and James-11).

The analysis in Figure 7 reveals that the methods Tian-11, Oreifej-11, Halder-11, and James-11 exhibit distortion to some extent owing to the limited number of frames in the input image sequence. Our approach, regarding visual quality, demonstrates superior performance in distortion correction and image sharpness, yielding the most visually appealing results among the showcased methods.

Table 2. Quantitative comparison of synthetic and real datasets for underwater distorted image restoration. Orange and blue indicate the best and second best performance, respectively. The top 3 results are marked as gray. Note: The upward arrows (↑) indicate that larger values of the corresponding evaluation metrics are better, while the downward arrows (↓) indicate that smaller values are better, where the same applies hereinafter.

Datasets	Metrics	Tian-61	Tian-11	Oreifej-61	Oreifej-11	Halder-61	Halder-11	Isola et al.	Li et al.	James-61	James-11	Restormer	TransMRSR	Histoformer	Ours
	PSNR↑	20.3856	18.1961	22.2898	17.9748	17.9410	16.4482	16.6478	16.8975	19.7365	18.6611	19.3051	18.2819	19.4668	20.4482
	SSIM↑	0.6285	0.4694	0.7802	0.5154	0.6822	0.4438	0.4257	0.5128	0.6426	0.5384	0.5402	0.4272	0.5341	0.6799
TianSet	LPIPS↓	0.3615	0.3647	0.1618	0.2359	0.2571	0.3406	0.5162	0.2258	0.1994	0.2322	0.3327	0.5520	0.3792	0.1611
	DISTS↓	0.2651	0.2193	0.1679	0.1662	0.2583	0.2637	0.3485	0.1563	0.1648	0.1663	0.2543	0.3765	0.3115	0.1179
	BRISQUE↓	41.7878	43.9543	56.4836	53.6872	46.4372	45.4426	48.5798	71.6325	50.0134	50.9611	44.1252	62.7634	53.6525	57.2292
	PSNR↑	21.4906	18.4396	19.994	18.6198	14.9126	14.6094	15.2982	17.1299	22.7130	18.7089	20.2811	18.8414	18.9443	20.1782
	SSIM↑	0.6222	0.4280	0.5854	0.4881	0.7093	0.3935	0.3528	0.5158	0.7574	0.4759	0.5494	0.3831	0.4292	0.5688
JamesSet	LPIPS↓	0.2879	0.2774	0.1532	0.1787	0.2085	0.3081	0.5094	0.1958	0.1297	0.1907	0.3679	0.3855	0.4290	0.1439
	DISTS↓	0.3020	0.2351	0.1715	0.1688	0.2955	0.3255	0.4071	0.1843	0.1835	0.1893	0.3500	0.3861	0.3834	0.1368
	BRISQUE↓	27.8425	26.2433	13.4870	13.0378	38.2839	25.0573	22.6950	19.5195	35.4909	32.2424	52.1540	67.9440	53.3034	15.9079
	PSNR↑	26.7043	25.8777	19.6219	20.4880	23.2897	23.6934	17.1358	20.5535	27.9386	26.2761	29.5838	24.5663	29.7377	32.8341
	SSIM↑	0.9565	0.9401	0.7853	0.8099	0.9185	0.9177	0.6389	0.8429	0.9653	0.9467	0.9728	0.9115	0.9740	0.9815
ThapaSet	LPIPS↓	0.0410	0.0355	0.1502	0.0970	0.0980	0.0737	0.4069	0.1067	0.0479	0.0427	0.0192	0.0689	0.0193	0.0193
	DISTS↓	0.0338	0.0251	0.1340	0.0761	0.1098	0.0908	0.2705	0.0902	0.0614	0.0423	0.0124	0.1025	0.0143	0.0146
	BRISQUE↓	44.1472	39.5778	46.2971	40.5773	40.7131	38.0676	27.5812	40.4162	39.1398	39.2585	41.6130	41.8464	42.5070	39.8312
	PSNR↑	23.6908	22.7268	22.6756	21.0862	18.0662	17.7436	18.8707	20.1866	25.5396	24.3508	21.2622	21.2562	21.4424	21.5932
	SSIM↑	0.6747	0.6390	0.8070	0.7478	0.8648	0.8466	0.8088	0.8538	0.7742	0.6947	0.8768	0.8694	0.8637	0.6043
NianSet	LPIPS↓	0.3753	0.3834	0.2887	0.3278	0.4119	0.4334	0.4891	0.3016	0.3116	0.3420	0.4025	0.4534	0.3754	0.2767
	DISTS↓	0.2329	0.2171	0.1703	0.2013	0.2414	0.2595	0.2660	0.1613	0.2144	0.2078	0.2225	0.2738	0.2300	0.1832
	BRISQUE↓	55.5353	42.0099	53.9865	54.5593	48.1926	47.1852	0.7258	39.8342	53.5041	49.2360	53.4011	67.9567	58.6706	45.0817

presents a comparative evaluation of our method’s quantitative test outcomes against those of other methods across various datasets.

4.3. Ablation Studies

In order to evaluate the efficacy of individual components within our two-stage restoration model DR-Net (as shown in Figure 8 and

Table 3. Ablation study of our method, DR-Net, to evaluate the efficacy of individual components.

Method	PSNR ↑	SSIM ↑	LPIPS ↓	DISTS ↓	BRISQUE ↓
IR-GAN	19.2210	0.6231	0.3602	0.1914	14.8061
CNNWarp	21.2186	0.7325	0.2908	0.1820	34.1997
CNNWarp+IR-GAN	21.3630	0.7406	0.2889	0.1683	25.8712
DE-Net	21.8322	0.7523	0.2687	0.1665	33.6816
DE-Net+IR-GAN	22.3290	0.7760	0.2310	0.1365	27.8354
DR-Net (Ours)	22.5133	0.7828	0.2243	0.1321	28.6267

), for image restoration, we performed ablation experiments to assess the quality of restored images under varying conditions, considering both qualitative and quantitative aspects. The dataset by Li et al. was utilized for both training and testing in these experiments.

Effect of the distorted information estimation. A comparative analysis of IR-GAN, CNNWarp, and DE-Net reveals that the network estimating distortion information for restoration exhibits advantages in distortion correction, yielding superior outcomes from both quantitative and qualitative standpoints. Moreover, our DE-Net model incorporating a Transformer structure demonstrates enhanced performance. The utilization of global features within the Transformer encoder enables a more effective extraction of distortion details in the degradation process of underwater images, supporting improved results in the subsequent stage. Figure 9 illustrates the distortion information estimated in the initial phase, showcasing the enhanced feature representation and more precise estimation of underwater image perturbations achieved through the self-attention mechanism of the Transformer and the integration of a two-channel attention mechanism in U-Net.

Effect of the GAN-based image restoration model. Analysis of IR-GAN restoration outputs reveals that the dynamic “game process” inherent in GAN facilitates continuous enhancement of the generator’s ability to produce realistic image data. This advancement enables the GAN-based image restoration model to effectively restore image intricacies, enhance color and contrast, and recover potential lost details. A comparative assessment between the restoration outcomes of the model utilizing IR-GAN and the model solely estimating perturbation information demonstrates the enhanced performance of incorporating a GAN-based image restoration model in the second stage. Notably, this integration proves more adept at mitigating image blurring and achieving superior quality outcomes.

Effect of the training strategy. In addition to assessing the efficacy of our model training strategy, we systematically oversee image restoration outcomes across stages. Progressing from the initial deployment of DE-Net, through the subsequent DE-Net+IR-GAN phase, and culminating in the cascade fine-tuning with DR-Net, the quantitative metrics of image restoration consistently exhibit enhancement. Comparative evaluations of image outputs further confirm a continual refinement in both image details and quality. Every component within our two-stage underwater image restoration framework, combined with the model training strategy, proficiently contributes to advancing the quality of image restoration procedures.

5. Conclusions

Our proposed DR-Net is a two-stage network designed for the restoration of underwater images afflicted by distortions caused by water surface fluctuations and turbulence-induced blurring. The initial stage involves leveraging DE-Net, a fusion network combining U-Net and Transformer architectures, to predict perturbation information in underwater images. Subsequently, utilizing the perturbation information in conjunction with the initial restoration results, IR-GAN, an image restoration model employing adversarial generative networks, addresses image blurring and detail restoration. Experimental findings demonstrate the superior effectiveness of our method in restoring degraded underwater images, showcasing promising prospects due to its efficacy as a single-frame image restoration technique.

Practical considerations to guide refinements include the following: the Transformer-GAN integration, while effective for complex patterns, entails higher computational demands (common in accuracy-prioritizing models), posing challenges for resource-constrained real-time deployment; like other data-driven methods, performance is influenced by training data diversity, with minor overfitting risks in rare distortion cases (mitigable via richer datasets); and in extreme conditions beyond the training scope, restoration quality may slightly degrade (though robust in typical scenarios). To address these aspects, our future work will focus on optimizing efficiency via lightweight architectures, enhancing generalization with advanced data augmentation and integrating underwater physical priors to strengthen adaptability—efforts aimed at resolving these limitations and further boosting DR-Net’s real-world utility.