UIR-Net: A Simple and Effective Baseline for Underwater Image Restoration and Enhancement

Abstract

Because of the unique physical and chemical properties of water, obtaining high-quality underwater images directly is not an easy thing. Hence, recovery and enhancement are indispensable steps in underwater image processing and have therefore become research hotspots. Nevertheless, existing image-processing methods generally have high complexity and are difficult to deploy on underwater platforms with limited computing resources. To tackle this issue, this paper proposes a simple and effective baseline named UIR-Net that can recover and enhance underwater images simultaneously. This network uses a channel residual prior to extract the channel of the image to be recovered as a prior, combined with a gradient strategy to reduce parameters and training time to make the operation more lightweight. This method can improve the color performance while maintaining the style and spatial texture of the contents. Through experiments on three datasets (MSRB, MSIRB and UIEBD-Snow), we confirm that UIR-Net can recover clear underwater images from original images with large particle impurities and ocean light spots. Compared to other state-of-the-art methods, UIR-Net can recover underwater images at a similar or higher quality with a significantly lower number of parameters, which is valuable in real-world applications.

1. Introduction

Underwater optical imaging is an important tool enabling humans to understand, exploit and protect the ocean. It has the advantages of intuitive detection of targets, high imaging resolution and high information content. Underwater images are applied in multiple research areas, such as marine biology, archaeology [1,2], fisheries [3], marine resource surveys, underwater target detection [4], etc. However, in high-turbidity waters such as coastal waters or lakes, scatterers in the water (including bubbles, suspended sediment, plankton, etc.) will cause reductions in and scattering of visible light. Severely scattered light enters the imaging system and causes a veiling light effect, producing a lot of noise in the image. As a result, the quality of the captured images is degraded, and the edges and contours of the scene and the targets are blurred. Thus, the features and information in degraded images cannot be effectively extracted and utilized by downstream tasks, such as underwater target detection, segmentation and recognition. This makes the enhancement and restoration of degraded underwater images necessary for marine applications.

There are a great number of solutions, including formation-model-based methods [5] and deep-learning-based methods [6], for enhancing or recovering images on land, but translating these methods for use on underwater images is not an easy task. On one hand, these two kinds of imaging models face different problems; on the other hand, there are no training samples of underwater images for these methods. For these reasons, there exists an urgent need to build relevant high-quality datasets and design enhancement and reconstruction algorithms specifically for underwater images. A number of methods have been proposed for underwater image recovery, which can be broadly classified into traditional methods and depth-based methods. Traditional methods are based on physical models and non-physical models. Non-physical image enhancement methods include white-balancing methods, histogram-equalization methods, Retinex-based methods, etc. Physical image restoration methods construct a mathematical model for the underwater image degradation process and then reverse it to obtain an undegraded image in an ideal state.

With the development and applications of deep neural networks in various scenarios, generative adversarial networks have become increasingly used for image restoration and have achieved impressive visual results [7,8,9,10,11]. Recent work on underwater image recovery using UGAN-P [12], for example, trains a CycleGAN [13] with unpaired clear and degraded underwater images and then generates a training set of degraded underwater images corresponding to the clear underwater images using CycleGAN [13]. To obtain the underwater image training dataset, Water-GAN [14] uses a generative adversarial network to generate images. Jiang [15] has designed a two-stage underwater image recovery network. The first stage is a two-line parallel encoder and decoder structure, and then a dual-channelwise attention operation is carried out. Jiao [16] used Dark channel prior to enhance the image, and refined the rough transmission mapping with the proposed adaptive non-convex non-smooth variational algorithm. The current deep-learning-based approach using generative adversarial networks has three main problems: The first is that it ignores the optical properties of underwater images, thereby easily introducing color biases and artifacts and thus not being ideal for color correction of underwater images. The second is that a CNN often constitutes a neural network with a very large number of parameters and a long training time, making it difficult to achieve a lightweight deployment. The third and most essential problem is that there is no algorithm that combines both tasks, forcing researchers to choose between image enhancement and image restoration.

In order to solve the problems mentioned above, this paper proposes a self-supervised network for underwater image restoration and image enhancement named UIR-Net. In this paper, we consider both tasks as domain-to-domain conversion problems, for which we can achieve high-quality restoration and enhancement of underwater images by mapping high-dimensional features to low-dimensional features through an encoder–decoder structure while preserving the images’ spatial and color features. Firstly, regarding the problem of lacking datasets of clear and degraded underwater images, we created the Marine Spot Impurity Removal Benchmarking (MSIRB) dataset based on the MSRB dataset [17] with scene extension. We also employed UIEBD-Snow, a UIEBD-based dataset [18] that integrates both underwater image enhancement and underwater image restoration to fuse the spots with underwater images. Secondly, to address the GAN-based underwater image restoration methods’ inability to solve the underwater image degradation well, we present a gradient-based policy update scheme and a channel residual prior to extract the channel of the image to be recovered as a prior. Thirdly, to reduce parameters and training time, we introduce a gradient strategy to make the gradient descent algorithm trainable, thus obtaining a more lightweight method. Specifically, we combine a gradient estimation strategy with proximal gradient descent (PGD) algorithm with a learnable step size. Unknown degradation issues in real-world scenarios are solved by combining flexible gradient descent and informative proximal mapping using a deep neural network. Overall, the contributions of this paper are:

1.

In this paper, we present a novel, end-to-end, lightweight underwater image restoration network named UIR-Net, which can generate context-rich and spatially accurate outputs.

2.

UIR-Net is the first network that combines underwater image enhancement and underwater image restoration, which means this can be considered a pioneering study with great significance for practical applications and deployments.

3.

This article proposes the Marine Spot Impurity Removal Benchmarking (MSIRB) dataset and UIEBD-Snow.

4.

Extensive experiments have been carried out on the underwater MSRB [17], the MSIRB dataset, and UIEBD-Snow that demonstrate the effectiveness of our UIR-Net. Our method is able to achieve better results than other state-of-the-art methods. The results of UIR-Net are shown in Figure 1.

2. Related Works

2.1. Underwater Image Restoration

Marine pollutants are one of the main causes of the degradation of underwater images. These pollutants include sinking organic particles, zooplankton remnants, floating excrement, suspended silt and other inorganic materials [19]. They vary in size, shape and transparency. When taking photos (or filming) underwater, the light is scattered by these tiny particles before being reflected back to the camera, resulting in low-contrast and blurry images. Light reflections from these underwater scatterers cause random near-white stains in various shapes, sizes, and opacities in underwater images/videos, decreasing scene perception [20]. It is worth noting that these particles fall to the seafloor in the same way that snow does in atmospheric conditions. There are many existing methods in this area. The authors of [21,22,23,24] proposed a modified median filter to remove the effects of underwater pollutants in these images. Wang et al. [25] advocated for machine vision in marine fisheries and presented an integrated two-channel model to improve underwater picture quality, along with a marine network model for impurity elimination.

DB-ResNet [26] is a “deep detail network” structure designed to remove natural raindrop patterns from captured images. A deep residual network (ResNet) is used as a parameter layer to obtain deeper picture features and to simplify the network’s structure by effectively shortening the mapping distance between input and output features. Other authors [27] proposed a progressive optimization residual network called progressive ResNet (PRN) and a progressive recurrent network (PReNet) to perform image de-raining, which can achieve good de-raining results without increasing network parameters. Maxim [28] is the latest MLP-based U-Net backbone network that implements both global and local perceptual fields and can be applied directly to high-resolution images. Restormer [29] is an efficient transformer that achieves several key improvements in the design of its improved multi-head attention and feed-forward networks MPRNet [30] is a novel collaborative design with a multi-stage structure designed to learn the recovery features of degraded inputs concurrently to decomposing the entire recovery process. Specifically, MPRNet first learns context-dependent features using an encoder–decoder architecture, and then combines it with high-resolution branches that preserve local information better. Compared to the original proximal gradient descent (PGD) algorithm, DGUNet [31] introduces a Deep-Network-based gradient estimation method into its gradient descent progress without loss of interpretability, allowing it to handle complex and real-world image degradation.

All of these methods are very competitive in natural image restoration, but land-based and underwater imaging models are not interchangeable, which means that directly translating image enhancement methods from land-based to underwater image restoration cannot achieve satisfying results; underwater imaging models are simply more complicated.

2.2. Underwater Image Enhancement

Many image enhancement and recovery methods based on marine scenes, such as Neural Networks (NNs) and Dark Channel Priors (DCP) based on Image Formation Model (IFM) class algorithms, have evolved in recent years [32]. In the machine learning literature [33], high-quality restored image data can be generated from a single picture, and the neural network’s generalization was tested using different images. A Convolutional Neural Network (CNN)-based algorithm [34] was proposed in this publication. A Convolutional Neural Network (CNN)-based framework for underwater image improvement implements a pixel interruption method for color correction and defogging. The authors of [14] utilized a Generative Adversarial Network (GAN) to produce a real IFM imaging model. The authors of [35] proposed an underwater image enhancement CNN model based on underwater scene prior, called UWCNN.The authors of [5] generated a huge training dataset encompassing real underwater photos in depth using a GAN for underwater image color correction. An IFM-based method to enhance underwater images by removing background light (BL) based on light transmission parameters is developed in the literature. The authors of [36] simplified the full underwater optical transmission model to extract the transmission map (TM) from the degraded underwater images. Optimal enhancement performance is also guaranteed based on regional background estimation. In [37], the background light is estimated by using a color variable that changes with depth. In [6], Drewa obtains the transmission map with an underwater Dark Channel Prior (UDCP) model by reducing the R channel’s effect on the dark channel. In [38] Carlevaris et al. used a Maximum Intensity Prior (MIP) model to estimate the transmission map to recover underwater images. In [39], Song offered an Underwater Light Attenuation Prior (ULAP) model in order to obtain the scene depth of underwater images. The authors of [40] proposed an Image Blurriness and Light Absorption (IBLA) algorithm, which achieves the estimation of the system parameters by selecting the background light and depth map from different scene regions. The S-model is used to fuse the parameters and estimate the system parameters to complete the image’s recovery.Recently, attention mechanisms have been successfully used in image processing tasks, such as image denoising [41], image super-resolution [42] and image segmentation [43]. The performance of the image enhancement model can be effectively improved by leveraging the focus of the attention mechanism.

However, the existing underwater image enhancement methods only achieve enhancement (color correction). Restoring underwater images degraded by visual impurities to a satisfactory level still remains a challenge.

3. Method

UIR-Net is inspired by DGUNet [31]. Achieving high-quality image restoration and color correction is an essential part of utilizing underwater optical images. The biggest advantage of self-supervised learning is that the results generated by our trained model can match the labels to a greater extent. This paper takes advantage of the self-supervised method, divides the images to be enhanced and the original ones into two domains and forms a domain-to-domain conversion to achieve image restoration and image enhancement after a lot of learning. UIR-Net conforms to the U-Net encoder and decoder principle, as shown in Figure 2. UIR-Net mainly consists of two parts: the first part is the residual prior module, which guides the subsequent underwater image recovery by subtracting the maximum and minimum features as a prior, and the second part is the strategy gradient network, which integrates PGD strategies for iterative training.

In addition, perceptual loss is added to the generator to further constrain the detail and semantic information of generated images. In the discriminator, we use a fully convolutional block-based discriminator to learn the loss of the structure. A great number of approaches combining traditional optimization algorithms with DNNs have been proposed. These methods usually require predefined degradation processes or artificial assumptions and are therefore difficult to handle complex underwater real-world applications.

The structure of UIR-Net is shown in Figure 2.

3.1. Channel Residual Prior

This part refers to Kaiming He’s dark channel defogging algorithm, which can restore the color and visibility of an image and which also uses the fog concentration to estimate the distance of an object. He’s work focuses on the statistical features of fog-free images. It was found that, in a fog-free image, every local region is likely to have shaded areas, something in black or something in pure color. Therefore, it is highly likely that there always exists at least one color channel with very low light values. He named this statistical law Dark Channel Prior or Dark Primary Color Prior. Intuitively, Dark Channel Prior assumes that there is always something very dark in every local area, and it is the essential fundamental law in their study of defogging problems. Since fog is always grayish, once the image is affected by fog, then these features that are supposed to be dark will become grayish. Moreover, based on the physical fog formation formula, they can also determine the fog’s concentration based on how gray these features are. Therefore, their proposed Dark Channel Prior can remove the effect of fog very effectively while using the concentration of features to estimate the distance (depth) of objects.

For dark tones and ocean impurity occlusion, if we think about it, we can identify a statistical rule for underwater images: for most underwater images without impurity particles, at least one of the R, G and B values of any pixel is very low (this is quite understandable, as if all the three values are all high, then that pixel obviously has a tendency to be overly white). The “darker” values (channels) are arranged in a certain way to form a dark channel map corresponding to the original image. Based on this idea and the laws of statistics, the basic model of the image can be reconstructed more clearly without any streaks in the image, but we found that the structures of some constantly changing objects in the reconstructed image are deformed. To address this problem, we propose introducing a channel residual prior (CRP) model. Furthermore, the CRP lacks spots and bubbles, which is beneficial to derive altered background information from the residuals of the maximum and minimum channel values of the picture.

Referring to Li et al. [44], we can express the color intensity of the rainband image as follows:

$$\tilde{\mathcal{G}}\left(x\right)=t{\gamma }_{rs}\left(x\right)\mathcal{A}\alpha +(T-t)\mathcal{R}\pi$$

(1)

where $\tilde{\mathcal{G}}\left(x\right)$ is the intensity of the color vector indicating the presence of color, and ${\gamma }_{rs}$ consists of the refraction coefficient, the specular reflection coefficient and the internal reflection coefficient of the spot [45]. $A={(Ar,Ag,Ab)}^T$ expresses the luminance, $\mathcal{R}=R_r+R_g+R_b·\alpha =\mathbf{A}/\mathcal{A}\mathrm{and}\pi =\mathbf{R}/\mathcal{R}$ denotes the the reflection coming from the background, $R=R_r+R_g+R_b$. $\alpha =A/\mathbf{A}\mathrm{and}\pi =R/\mathbf{R}$ are the A and R chromaticity, respectively. The exposure duration is T, and the time it takes for raindrops to flow over pixel x is t. The first term in Equation (1) is the rain streak term, and the second term is the backdrop term. We can achieve the following normalization steps when predicting $\alpha$ using any known color constancy algorithm:

$$\mathcal{G}\left(x\right)=\frac{\tilde{\mathcal{G}}\left(x\right)}{\alpha }=G_{rs}\left(x\right)\mathbf{l}+G_{bg}\left(x\right)$$

(2)

where $l=(1,1,1)$ in Equation (2), $T,G_{rs}=t{\beta }_{rs}\mathcal{A}$, and $G_{bg}=(T-t)R/\alpha$.

When we normalize the image, the chromaticity of the light deteriorates, erasing the color, and the spectrum’s sensitization impact is likewise eliminated. Therefore, according to Equation (2), given an underwater bubble image G, one can define a residual channel prior P for G:

$$P_{prior}\left(x\right)=\underset{c\in r,g,b}{max}{\mathcal{G}}^c\left(x\right)-\underset{d\in r,g,b}{min}{\mathcal{G}}^d\left(x\right).$$

(3)

3.2. Gradient Strategy

From a Bayesian maximum a posteriori perspective, the traditional model-based approach defines image reconstruction as follows:

$$\widehat{\mathbf{x}}=\underset{\mathbf{x}}{argmax}log{\mathcal{P}}_{prior}(\mathbf{x}\mid \mathbf{y})=\underset{\mathbf{x}}{argmax}log{\mathcal{P}}_{prior}(\mathbf{y}\mid \mathbf{x})+log{\mathcal{P}}_{prior}\left(\mathbf{x}\right)$$

(4)

where $log{\mathcal{P}}_{prior}(\mathbf{y}\mid \mathbf{x}$) represents the data fidelity term and $log{\mathcal{P}}_{prior}\left(x\right)$ represents the prior term. Usually, the $L_2$ loss function is used as the data fidelity term. Therefore, the above equation can be further deformed as:

$$\widehat{\mathbf{x}}=\underset{\mathbf{x}}{argmin}\frac{1}{2}{\parallel \mathbf{y}-\mathbf{B}\mathbf{x}\parallel }_2^2+\lambda J\left(\mathbf{x}\right)$$

(5)

Based on the proximal gradient descent (PGD) algorithm [31], the above equation can be developed as an iterative solution process:

$${\widehat{\mathbf{x}}}^k=\underset{\mathbf{x}}{argmin}\frac{1}{2}{∥\mathbf{x}-\left({\widehat{\mathbf{x}}}^{k-1}-\theta \nabla g\left({\widehat{\mathbf{x}}}^{k-1}\right)\right)∥}_2^2+\lambda J\left(\mathbf{x}\right)$$

(6)

The first part represents the gradient descent, and the second part represents the proximal mapping. In turn, the above equation can be decomposed into:

$$\begin{array}{c}{\mathbf{v}}^k={\widehat{\mathbf{x}}}^{k-1}-\theta {\mathbf{B}}^{\top }\left(\mathbf{B}{\widehat{\mathbf{x}}}^{k-1}-\mathbf{y}\right)\\ {\widehat{\mathbf{x}}}^k={prox}_{\lambda ,J}\left({\mathbf{v}}^k\right)\end{array}$$

(7)

where A represents the degradation matrix, and $prox$ is the known degradation matrix B. It is difficult to deal with the real scenes where the degradation is not known. In this paper, based on the PGD algorithm, we expand it with a deep network. In detail, the degradation matrix A and the proximal mapping equation are replaced by a deep network, and the step size $\theta$ of the gradient descent is set as a learnable parameter. Thus, the gradient descent process is redefined as:

$${\mathbf{v}}^k={\widehat{\mathbf{x}}}^{k-1}-{\rho }^k{\mathcal{F}}_{\mathbf{B}}^k\top \left({\mathcal{F}}_{\mathbf{B}}^k\left({\widehat{\mathbf{x}}}^{k-1}\right)-\mathbf{y}\right)$$

(8)

The proximal mapping module also incorporates inter-stage information to overcome the inter-stage information loss in the deep unfolding network:

$${\widehat{\mathbf{x}}}^k,{\mathbf{F}}^k={prox}_{{\omega }^k}\left({\mathbf{v}}^k,{\mathbf{F}}^{k-1}\right)$$

(9)

and the structure of PGD is shown in Figure 3, from the structure, the gradient is first passed through a residual layer, then by doing the difference with itself, and then summed with the prior feature, then we assigned a dynamic based on 0–0.5 parameters $\theta$ and the output feature matrix just multiplied, after the residual layer output.

3.3. Loss Function

We apply the perceptual loss [46] $L_{perctual}$ and Edge loss [30] $L_{edge}$ in the network. As the name implies, $L_{perctual}$ penalizes outcomes that are not close to the references in the sense of activation in the pretrained network by defining distance measurement. $L_{perctual}$ is defined as:

$${\mathcal{L}}_{perctual}=\mathbb{E}\left[\sum _i\frac{1}{N_i}{∥{\varphi }_i\left({\mathbf{I}}_{gt}\right)-{\varphi }_i\left({\mathbf{I}}_{out}\right)∥}_1\right]$$

(10)

where ${\varphi }_i$ is the activation map of the i-th layer of the pre-trained network. For our work, ${\varphi }_i$ corresponds to activating the VGG-19 network pre-trained on the ImageNet dataset from the layers relu1-1, relu2-1, relu3-1, relu4-1 and relu5-1. The edge loss is defined as follows:

$${\mathcal{L}}_{\mathrm{edge}}=\sqrt{{\left|\Delta \left({\mathbf{X}}_S\right)-\Delta \left(\mathbf{Y}\right)\right|}^2+{\epsilon }^2}$$

(11)

where $\epsilon$ is empirically set as ${10}^{-3}$, and $\Delta$ means Laplacian (for extracting edges). Thus, the final loss is:

$${\mathcal{L}}_{final}={\mathcal{L}}_{\mathrm{edge}}+{\lambda }_{perctual}\times {\mathcal{L}}_{perctual}$$

(12)

There the ${\lambda }_{perctual}$ equal 0.01, the ablation study of ${\lambda }_{perctual}$ is shown in Section 4.7.2.

4. Experimental

4.1. Datasets of Underwater Restoration

1.

Marine Snow Removal Benchmarking Dataset (MSRB)

In this section, we present the specifications of the ocean snow artifacts synthesized by the MSRB dataset [17], which has been mainly built for two general tasks of marine snow artifacts removal: MSR task 1 is dedicated removing small-size artifacts, while MSR task 2 is used to cope with various artifacts of different sizes.

Obviously, it would be much more difficult to handle multiple sizes of underwater snow than just focusing on small-sized ones. Corresponding to each MSR task, each sub-dataset is composed of a training set with 2300 pairs of images and a test set with 400 pairs, all with a pixel resolution of 384 × 384. Each image pair includes an original underwater collected picture and one containing synthetic marine snow artifacts, original underwater collected picture is the ground truth of the dataset. Each composite image is added with N marine snow particles, while N, which is the number of added particles, is generated by a discrete uniform distribution of U{100,600}. Based on our preliminary observations, in each synthetic image, H-type and V-type ocean snow spots are randomly generated with a probability of 0.7 for H-type and 0.3 for V-type. Most of our parameters are chosen based on our observations of real-world collected images and the corresponding artificial influence of ocean snow.

2.

Marine Spot Impurity Removal Benchmarking Dataset (MSIRB)

At present, there is no accessible underwater dataset of ocean snow images that can be used to train and test deep neural networks to eliminate ocean snow particles, which represents a major inconvenience to relevant ongoing research on marine light spot removal algorithms. In order to present the diverse morphological features of the real-world ocean snow particles and the complex composition of ocean snow scenes, we propose a new dataset called the Marine Spot Impurity Removal Benchmarking Dataset (MSIRB). The MSIRB dataset contains (1) synthetic ocean light spot images, (2) the corresponding real images and (3) an ocean light spot mask. We reference [47] and use its mask to produce underwater light spot images, each basic mask corresponds to small, medium, and large grain sizes. The dataset is shown in Figure 4.

It is important to focus on the fact that inside the MSIRB dataset’s ground truth and input binary images, there is no difference in color, only the presence or absence of light spot impurities.

3.

An underwater image enhancement benchmark dataset and beyond with Snow (UIEBD-Snow)

To explore the generalization of our work in image enhancement tasks, the same operation that was applied on the MSIRB dataset is performed in this paper on the UIEBD dataset, a dataset with its potential ground truth generated by adopting a dozen of cutting-edge image improvement methods. Various morphological features of real snow particles of the ocean are fused inside the UIEBD to make our new dataset called UIEBD-Snow, which contains 890 underwater images with corresponding reference maps.

It is worth noting that the difference between UIEBD-Snow and MSIRB is mainly in the color correction, and the dataset example can be referred to Figure 4

4.2. Training Parameter Settings

Our UIR-Net is trained end-to-end, and our training strategy is to use the same setting as MPRNet and DGUNet [30,31] in our experiments. Specifically, we use the Adam optimizer with an initialized learning rate of $2\times {10}^{-4}$. A typical warm-up strategy [43] is adopted to gradually increase the learning efficiency considering the depth of our network. Our model is trained with images which are randomly cropped into 256 × 256 image patches from our training set. Our batch size is set to 16 and $4\times {10}^5$ iterations. In image compression perception, the network is trained on 32 × 32 image patches with a learning rate of $1\times {10}^{-4}$. For 200 epochs, the batch size is set to 128. Our model training is performed on 3090 GPU and can be completed in 4 hours. Our experimental results are shown in Figure 5 and Figure 6. The source code for UIR-Net can be obtained from https://github.com/meixinkui/UIR-Net, accessed on 1 November 2022.

4.3. Evaluation Metrics

1.

Full-Reference Evaluation

We adopt the standard metrics (PSNR, SSIM and RMSE [48]) as full-reference evaluation.

PSNR is short for Peak Signal to Noise Ratio. The better the PSNR, the less the image distortion. We can obtain the PSNR as follows:

$$PSNR=10\ast {log}_{10}\left(\frac{{\left(2^n-1\right)}^2}{MSE}\right)$$

(13)

Assume that the current image I and the reference images K, H and W are the height and width of the image respectively. MSE can be calculated as follows:

$$MSE=\frac{1}{H\ast W}\sum _{i=1}^H\sum _{j=1}^W{[I(i,j)-K(i,j)]}^2$$

(14)

SSIM is short for Structural Similarity, we can obtain the SSIM of the two images (image x and image y) as follows:

$$SSIM(x,y)=\frac{\left(2{\mu }_x{\mu }_y+c_1\right)\left(2{\sigma }_{xy}+c_2\right)}{\left({\mu }_x^2+{\mu }_y^2+c_1\right)\left({\sigma }_x^2+{\sigma }_y^2+c_2\right)}$$

(15)

where ${\mu }_x$, ${\mu }_y$ are the means of x and y, and ${\sigma }_x$, ${\sigma }_y$, ${\sigma }_{xy}$ represent the variance and covariance of x and y, respectively. $C_1$ and $C_2$ are used to maintain stable constants.

RMSE is short for Root Mean Square Error. We can obtain RMSE as follows:

$$RMSE=\sqrt{\frac{1}{N}\sum _{i=1}^n{\left(Y_i-f\left(x_i\right)\right)}^2}$$

(16)

2.

Non-Reference Evaluation

UIR-Net is supervised network, there is an image pair for input and the corresponding target, but we also adopt the UCIQE [49] and UIQM [50] as non-reference evaluation, which are commonly used for underwater image quality assessment.

UCIQE is a linear combination of color intensity, saturation and contrast, used to quantitatively evaluate the non-homogeneous color shift, blurring and low contrast of underwater images. The better the UCIQE, the better the underwater image quality, we can obtain UCIQE as follows:

$$UCIQE=c_1\ast {\sigma }_c+c_2\ast {con}_l+c_3\ast {\mu }_s$$

(17)

where $c_1$,$c_2$ and $c_3$ are the weight factors of each measurement component in the linear combination, ${\sigma }_c$ is image chromaticity measurement component, ${con}_l$ is image contrast measurement component and ${\mu }_s$ is image saturation measurement component.

UIQM is short for underwater image quality measurement. We can obtain it as follows:

$$UIQM=c_1\ast UICM+c_2\ast UISM+c_3\ast UIconM$$

(18)

where $c_1$,$c_2$ and $c_3$ are the weight factors of each measurement component in the linear combination, UICM is underwater image colorfulness measure, UISM is underwater image sharpness measure and UIconM is underwater image contrast measure. The better the UIQM, the more consistent the results are with human visual perception.

4.4. Comparison with State-of-the-Art Methods on Underwater Image Restoration

We compared UIR-Net with existing methods, including natural image restoration methods: AirNet [51], DB-ResNet [26], PReNet [27], DGUNet [31], Maxim [28], Restormer [29], and MPRNet [30]. In AirNet, the network structure consists of a base layer and a detail layer. The detail layer is fed into this two-layer convolutional neural network, and then, the detail image after de-raining is obtained by deconvolution. The enhanced base layer and the enhanced detail layer are superimposed linearly to obtain the de-rained image. We think this is similar to the underwater image recovery process, and most underwater images are enhanced, so we chose this method for comparison with our proposed method.

4.5. Comparison with State-of-the-Art Methods on Underwater Image Enhancement

We compared UIR-Net with existing methods, including PUIE [52], CWR [53], IFM [54] and FUnIE-GAN [55]. PUIE [52] is a multi-color space encoder network in which the adaptively selected features of different color spaces are fused based on the attention mechanism. A decoder network based on medium transmission guidance is used to guide the model to give more attention to degraded regions. The complementary benefits between underwater imaging domain knowledge and DNNs are discussed in this work. CWR [53] introduces contrast learning into underwater image enhancement, takes the original image as a reference and restores the underwater image through TV loss. FUnIE-GAN [55] uses a structure called FUNI-Gan. The generator uses a five-layer U-net structure. The discriminator part uses PatchGAN, which can check the quality of the generated pictures in blocks. IMF [54] proposes a network consisting of two modules. One aims to estimate backscatter, and the other is used for estimating the direct transmission of light. A carefully designed reconstruction module follows the above two and uses their previous outputs to produce the enhanced underwater image.

4.6. Qualitative Evaluations

4.6.1. Quantitative Comparisons of Underwater Image Restoration on the MSRB Dataset

As we can see in Figure 7, AirNet [51] and PReNet [27] barely remove most of the small near-white spots. PReNet [27] removes the ocean snow grains while placing some dark particles near the original stained location, resulting in an overall darker picture as a result. DB-ResNet [26], maxim [28] and Restormer [29] can remove a small portion of the ocean snow grains while preserving the boundaries of other objects. However, these three methods still perform poorly when removing large ocean snow particles. Our method and MPRNet [30], DGUNet [31] can remove large particles in the ocean very well. In terms of the existing multi-stage techniques, we believe that MPRNet [30] structurally uses an encoder–decoder architecture (the encoding stage effectively propagates contextual information, but does not do much to preserve spatial details) using single scales at the same time. The spatial accuracy and semantic information are also well-preserved, and using both structures at the same time is effective for underwater image restoration.

As we can see in

Table 1. Comparison results of UIR-Net with other SOTA methods on the MSRB dataset.

Method	PSNR	SSIM	RMSE	UIQM	UICQE
Nature image restoration
AirNet [51]	24.359	0.954	6.745	3.418	0.549
DB-ResNet [26]	23.258	0.920	9.060	3.531	0.553
PReNet [27]	21.837	0.899	13.932	3.722	0.524
Maxim [28]	24.578	0.931	7.836	3.606	0.553
Restormer [29]	24.411	0.954	6.783	3.977	0.550
DGUNet [31]	33.292	0.977	3.408	3.755	0.567
MPRNet [30]	35.322	0.984	2.659	3.715	0.566
Our	34.464	0.984	2.649	3.877	0.576

, UIR-Net achieves the best results at the values of SSIM RMSE and UICQE. The values of PSNR and UIQM are second.

4.6.2. Quantitative Comparisons of Underwater Restoration on MSIRB Dataset

The results of the ocean light spot removal example are shown in Figure 8. AirNet [51], PReNet [27], maxim [28] and Restormer [29] still have defects in their ocean light spot removal tasks. We believe that the rain streak removal task is a little easier than the underwater image recovery task because there is little need to take into account lighting effects and background factors, and as clearly observed, PReNet [27] removes some of the ocean light spots by replacing them with dark patches. Although the Restormer [29] model alleviates the drawbacks of CNNs (i.e., limited perceptual field and maladaptation to the input content), it cannot be directly applied to design recovery tasks for underwater images considering its quadratically increased computational complexity with spatial resolution.

MPRNet [30] shows better performance than Restormer [29] and maxim [28] in underwater image recovery, and we believe that MRNet’s multi-stage progressive restoration is helpful in restoring underwater images. MPRNet first divides the feature map into four copies and then into two copies after restoration by U-Net. It then performs a residual network to complete the refinement. DGUNet [31] is an N-stage progressive generation network that also shows good results for underwater images with dark backgrounds. We attribute this success to its gradient strategy and the encoder and decoder components in each stage, as we can see in

Table 2. Comparison results of UIR-Net with other SOTA methods on the MSIRB dataset.

Method	PSNR	SSIM	RMSE	UIQM	UICQE
Nature image restoration
AirNet [51]	20.536	0.922	9.878	3.761	0.575
DB-ResNet [26]	20.506	0.897	10.743	3.758	0.574
PReNet [27]	20.310	0.888	14.237	4.075	0.550
Maxim [28]	23.599	0.924	8.316	3.832	0.565
Restormer [29]	20.396	0.927	9.817	3.810	0.574
DGUNet [31]	33.785	0.984	2.584	3.911	0.562
MPRNet [30]	33.561	0.985	2.592	3.725	0.554
Our	34.513	0.986	2.391	3.837	0.565

. UIR-Net achieves the best results at the values of PSNR SSIM and RMSE, in terms of UIQM and UICQE, UIR-Net approximate the best results.

4.6.3. Quantitative Comparisons of Underwater Enhancement on UIEBD-Snow Dataset

The results of the ocean light spot removal and image enhancement example are shown in Figure 9. We can see that UIR-Net achieves better results in both underwater image enhancement and underwater image restoration tasks.

As we can see from

Table 3. Comparison results of UIR-Net with other SOTA methods on the UIEBD-Snow dataset.

Method	PSNR	SSIM	RMSE	UIQM	UICQE
Underwater image enhancement
PUIE [52]	16.926	0.723	18.030	3.594	0.606
CWR [53]	15.374	0.556	23.276	4.620	0.507
IMF [54]	17.006	0.707	18.904	3.354	0.633
FUnIE-GAN [55]	15.619	0.486	24.762	5.106	0.596
Nature image restoration
MPRNet [30]	20.027	0.775	15.048	3.400	0.598
DGUNet [31]	20.711	0.795	14.068	3.363	0.597
Our	21.200	0.807	13.142	3.610	0.596

, UIR-Net achieves better performance and outperforms the other methods at PSNR SSIM and RMSE, in terms of UIQM and UICQE, UIR-Net approximate the best results. Specifically, the CWR [53] method based on contrast learning relies on the real image used and only takes the color correction problem into account. It performs poorly when confronted with fused light spot impurities in image enhancement, showing unnatural color and saturation and blurred image details. A multi-color space coding proposed by PUIE [52] proves to be effective in image enhancement alone, yet has little effect in the face of light spots, which leads use to believe that IFM [54] is relatively ineffective in image enhancement and restoration and does not improve the contrast, which is severely degraded by these colors. The objective function is also formulated to evaluate the perceptual quality based on global content, color, local texture and style. We believe that FUnIE-GAN [55] is superior to our method in image correction. The results of FUnIE-GAN are obviously different from those of other methods in terms of color, being more colorful and more consistent with the judgment of UIQM in terms of non-reference evaluation indicators, but FUnIE-GAN can only achieve image enhancement and cannot be used for image recovery. It can be seen that while most of the methods can perform contrast enhancement, some of the deficiencies cause severe visual defects such as unsatisfactory colors; our method, on the other hand, works well in all these cases, and our results look cleaner and more natural with fine-grained textures. MPRNet [30] and DGUNet [31] can achieve image enhancement to some extent, but they are not effective in image restoration. In addition, there is still a large portion of blur and missing textures in their outcomes. We believe that MPRNet and DGUNet focus on space in the task of domain-to-domain conversion and are unable to remove light spot occlusion, which is similar to image inpainting.

4.6.4. FLOPs and Params Comparisons with MPRNet and DGUNet

For a better quantitative comparison, we performed a fairer comparison with MPRNet [30] and DGUNet [31] in terms of FLOPs and Params, i.e., time complexity and space complexity, respectively. According to

Table 4. FLOPs and Params comparisons.

Method	FLOPs	Total Params	Training Time
DGUNet [31]	$8.4278\times {10}^{11}$	12,176,119	94,600 s
MPRNet [30]	$2.1946\times {10}^{12}$	3,637,249	41,400 s
UIR-Net	$2.4120\times {10}^{11}$	3,632,740	22,000 s

, we can see that UIR-Net has one-tenth the number parameters of MPRNet in terms of FLOPs, and there is not much difference in the amount of data in Params. We thus measured the training time again. We set the epoch to 200, the batch-size to 8 and the patch size to 128, and all methods adopted the warm up strategy. It can be seen that our training time is half that of MPRNet, and the inference speed is also better, so the operation method and convergence time of this paper are much better than MPRNet, which is very helpful for practical applications.

4.7. Ablation Study

4.7.1. Ablation Study of Channel Prior and Gradient Descent

As shown in Figure 10, while using only the gradient strategy, proximal gradient descent or the residual prior can obtain an improvement over the baseline, combining all of the above strategies can achieve the best performance, and the integration of these components can generate realistic results and restore details in the case of oceanic light spots and large grain impurities. The channel residual prior, based on the observation of the absorption rate of the red channel in a great many of underwater images, is suggested to be able to remain sensitive to low-illumination images.

As we can see in

Table 5. Ablation study of UIR-Net on the MSRB and MSIRB datasets.

residual prior	proximal gradient descent	PSNR	SSIM	RMSE
		33.319	0.981	2.953
√		33.319	0.981	2.953
	√	34.356	0.981	2.931
√	√	34.464	0.984	2.649
residual prior	proximal gradient descent	PSNR	SSIM	RMSE
		32.477	0.982	2.912
√		33.295	0.983	2.649
	√	34.121	0.984	2.765
√	√	34.513	0.986	2.391

, after adding the residual prior, the PSNR, SSIM and RMSE were effectively improved to 33.656, 0.983 and 2.954, respectively, enhancing the significant texture features. When only adding the gradient strategy, PSNR climbed to 34.356 and SSIM rose to 0.981, proving the superiority of our model in underwater image restoration. The process of gradient descent is essentially the same as the operation of a neural network, except that when training a neural network, the samples and the parameters to be trained are known. Gradient descent has explicit variables so that it can rely on such variables to complete the iterations when training a neural network, instead of relying purely on samples and parameters. After adding the residual prior and the gradient strategy, PSNR, SSIM and RMSE reached 34.464 and 0.984 and 2.649, respectively.

4.7.2. Ablation Study of ${\lambda }_{perctual}$

We need to analyze the paramater of the Loss Function. As we can see from Figure 11, we can obtain that when ${\lambda }_{perctual}$ equal 0.01, the information of the image is preserved more completely. Moreover we can see in

Table 6. Ablation results of paramater of ${\lambda }_{perctual}$ on the UIEBD-Snow dataset.

Paramater	PSNR	SSIM	RMSE
${\lambda }_{perctual}$ = 0.1	20.203	0.781	14.030
${\lambda }_{perctual}$ = 0.05	19.371	0.791	13.301
${\lambda }_{perctual}$ = 0.02	20.039	0.797	13.904
${\lambda }_{perctual}$ = 0.01	21.200	0.807	13.142

we can clearly deduce that when ${\lambda }_{perctual}$ equal 0.01, all of the paramaters(PSNR SSIM and RMSE) are the best, so there the ${\lambda }_{perctual}$ equal 0.01.

5. Conclusions

In this paper, we propose a simple and effective baseline named UIR-Net that integrates underwater image restoration and enhancement. For dark tones and ocean impurity occlusion, UIR-Net uses a channel residual prior to extract the channel of the image to be recovered as a prior, combined with a gradient strategy to reduce the number of parameters and training time to achieve a lightweight deployment. Our experimental results prove that UIR-Net can effectively improve images with both large ocean particle impurities and ocean light spots on these datasets (MSRB, MSIRB and UIEBD-Snow). In addition, our approach is able to achieve color correction on image enhancement, maintain the style and spatial texture of the content. In addition, the network in this paper outperforms many existing methods with less computational complexity and faster training speed, which is very useful for practical applications.