An effective transformer based on dual attention fusion for underwater image enhancement

Underwater image formation model

There are similarities between foggy weather images and blurry underwater images. As a result, the imaging in underwater environments is traditionally described by the atmospheric scattering model (Frenkel, 1924). However, it is only applicable in some scenes which shallow water scenes with low backscattering as the atmospheric scattering model is not suitable for all the diverse and complex underwater scenes in the real world. This results in a significant deviation from the actual underwater imaging model. Therefore, most existing approaches usually use an improved model (Chiang & Chen, 2011) to describe the underwater imaging process. Compared with the atmospheric scattering model, the selective attenuation of different wavelengths of light in water is considered by it, thus improving its generalization. The improved model of underwater imaging can be expressed as

(1) $${\mathrm{U}}_{\lambda }(x)={\mathrm{I}}_{\lambda }(x)\cdot T_{\lambda }(x)+B_{\lambda }\cdot (1-T_{\lambda }(x))$$where $x$ represents a specific point in the underwater scene and $\lambda$ denotes the wavelength of the light at that point, including red, green and blue colors. The underwater image captured by the camera is denoted as ${\mathrm{U}}_{\lambda }(x)$, while ${\mathrm{I}}_{\lambda }(x)$ represents the clear scene image. $B_{\lambda }$ denotes the global background light component. $T_{\lambda }(x)$ denotes the remaining energy ratio of the light radiation from the scene points reaching the camera, resulting in light scattering and color shift, which is expressed as

(2) $$T_{\lambda }(x)={\displaystyle \frac{E_{\lambda }^{\mathrm{r}\mathrm{e}\mathrm{s}\mathrm{i}\mathrm{d}\mathrm{u}\mathrm{a}\mathrm{l}}(x)}{E_{\lambda }^{\mathrm{i}\mathrm{n}\mathrm{i}\mathrm{t}\mathrm{i}\mathrm{a}\mathrm{l}}(x)}=\mathrm{N}\mathrm{r}\mathrm{e}\mathrm{r}(\lambda {)}^{d(x)}}$$where $E_{\lambda }^{\mathrm{r}\mathrm{e}\mathrm{s}\mathrm{i}\mathrm{d}\mathrm{u}\mathrm{a}\mathrm{l}}$ and $E_{\lambda }^{\mathrm{i}\mathrm{n}\mathrm{i}\mathrm{t}\mathrm{i}\mathrm{a}\mathrm{l}}$ represent the residual energy and initial energy, respectively, of the light with wavelength $\lambda$ that at the camera and the scene point. while $d(x)$ is the distance traveled by light from scene point $x$. $\mathrm{N}\mathrm{r}\mathrm{e}\mathrm{r}(\lambda )$ is the normalized residual energy ratio. It describes the ratio of the residual energy of light propagating per unit distance in the medium, which is related to the wavelength of the light. For water media, the red light with longer wavelengths attenuates more rapidly than the shorter wavelengths of blue and green, resulting in underwater images that usually appear blue or green in style.

Underwater image enhancement

Physical model-based methods first assume an imaging model. Then, physical parameters within the imaging model are estimated using prior methods. Finally, the imaging model is reversed to get a clear image. The core problem is to estimate the two parameters of transmittance and global background light by using priors. Existing priors include underwater dark channel priors (Drews et al., 2013), red channel prior (Galdran et al., 2015), minimum information prior (Li et al., 2016), underwater light attenuation prior (Song et al., 2018), etc. The underwater dark channel prior method (Drews et al., 2013) improves the DCP method (He, Sun & Tang, 2010) to better adapt to the underwater scene, but it only considers the blue and green channels, while the information in the red channel is ignored. The method of image blurriness and light absorption, suggested by Peng & Cosman (2017), can estimate depth maps for underwater images. The underwater light attenuation prior (Song et al., 2018) can effectively estimate scene depth, which assumes that it is closely related to the difference between the red channel and the maximum value of the channels blue and green. These physical model-based methods can produce good results when the imaging model and priors correspond with the underwater scene. But it does not generalize well as the difficulty of adapting to varied underwater scenes. Moreover, the resulting image has poor visual effects because it ignores the vision perception characteristics of the human eye.

Non-physical model-based methods enhance the image primarily through adjustments to its brightness and contrast. They include conventional methods used earlier, such as histogram equalization (Ghani & Isa, 2015; Hitam et al., 2013) to improve contrast and white balance (Iqbal et al., 2010, 2007) to correct color. In addition, there are fusion based methods (Ancuti et al., 2012) and Retinex based methods (Fu et al., 2014; Zhang et al., 2017). Ancuti et al. (2012) proposes an image fusion method to enhance underwater images by fusing two images with corrected color and improved contrast using different weight maps. Zhang et al. (2017) proposed a multi-scale Retinex method utilizing the characteristics of the three channels of the CIELab color space. The non-physical model-based method has a significant effect on improving the brightness and contrast of underwater images. However, since underwater images have spatially inconsistent attenuation, it is not reasonable to globally correct the color of the images. Because the physical imaging process is not considered, the effectiveness of its enhancement is limited. Moreover, it often introduces artificial enhancement effects.

Since outstanding results have been seen for computer vision tasks through the application of convolutional neural networks (CNNs), several works have employed CNNs to enhance underwater images (Anwar, Li & Porikli, 2018; Fabbri, Islam & Sattar, 2018; Islam, Xia & Sattar, 2020; Li et al., 2019a). Anwar, Li & Porikli (2018) proposed the UWCNN network, which consists of densely connected convolutional layers. In addition, they based on the underwater imaging model and different water-type parameters to synthesize an underwater image dataset. Li et al. (2019a) proposed a Water-Net to enhance underwater images, which incorporates a gated fusion architecture. Three images are used as inputs to the network, which are obtained by applying gamma correction, histogram equalization and white balance to the underwater images respectively. It combines the input images according to the learned confidence map to construct the enhanced image. In addition, some methods utilize the concept of game theory to propose a generative adversarial network. Islam, Xia & Sattar (2020) propose a GAN approach to enhance underwater images. Its generator is based on U-Net, and the discriminator uses the PatchGAN (Isola et al., 2017) structure, which has fewer computational parameters and can better capture high-frequency features. Fabbri, Islam & Sattar (2018) proposed UGAN, which uses synthetic data for network training. The training images are generated by CycleGAN (Zhu et al., 2017) converting underwater images from clear to blurry.

Although the above methods achieve better results compared to physical and non-physical model-based methods, they also have some shortcomings. The enhancement results generated by the GAN-based model are unstable, although it has excellent visual perception. The convolution has parameter sharing, which makes it inflexible to adapt to input features. Furthermore, it cannot model long-range dependencies because of limited receptive fields. The Transformer (Vaswani et al., 2017) is a Seq2Seq model devised for processing sequences within the domain of natural language processing (NLP), which was later utilized for computer vision. Its self-attention mechanism can effectively capture global dependencies between data. Moreover, unlike convolution, the self-attentive mechanism allows the image content to interact with the attention weights. With these advantages, it is extensive applications in various visual tasks including image denoising (Liang et al., 2021), image segmentation (Strudel et al., 2021), image recognition (Dosovitskiy et al., 2020) and object detection (Carion et al., 2020). Dosovitskiy et al. (2020) reshape an image into a sequence of image patches and feed it into a pure Transformer model. Its excellent performance on image classification tasks shows that the dependence on computer vision on CNN is unnecessary. Liu et al. (2021) proposed a Swin Transformer network, which uses a shifted window strategy to limit the attention calculation range. It not only has higher computational efficiency but also has the locality of the convolution operation. Peng, Zhu & Bian (2023) proposed a U-shape network that introduces Transformer into a GAN for underwater image enhancement. The color channels with severely attenuated are strengthened by the inter-channel self-attention of the network. It also computes spatial global attention between image patches to focus on the regions with severe degradation. Compared to previous underwater image enhancement methods, our network focuses not only on more severely attenuated spatial areas and color channels but also on local modeling that is beneficial for image recovery to reconstruct refined details.

Proposed method

Our goal is to develop an efficient Transformer model for underwater image enhancement with good robustness. First, the suggested overall architecture is presented. After that, the important components of the network are introduced, including the channel attention block, local attention block and adaptive feature fusion block. Finally, the loss function for training is introduced.

Network architecture

The overall architecture of the proposed network is presented in Fig. 1. Components of the overall pipeline including encoders, bottlenecks, and decoders. The decoder receives information through skip-connection from the encoder.

In the encoder part, input the given raw underwater image $\mathrm{I}\in {\mathbb{R}}^{H\times W\times 3}$, first perform $3\times 3$ convolution to extract low-level features $F_l\in {\mathbb{R}}^{H\times W\times C}$, where $H\times W$ is the spatial dimension, and $C$ is the number of channels. Then, the patch embedding layer maps it to the embedding features $F_e\in {\mathbb{R}}^{(H/2)\times (W/2)\times C}$. Following the design of U-net, the embedded features pass through $K$ encoder stages. Each stage has a sequence of proposed FABs and a downsample layer. The encoder extracts image features at different scales and finally maps the input to deeper feature domains $F_d\in {\mathbb{R}}^{(H/16)\times (W/16)\times 8C}$ to strengthen semantic representation. The downsample layer reduces spatial dimensions and increases channels.

The bottleneck is located after the encoder, and it has a sequence of Fusion Attention blocks for increasing the aggregation of global information. This stage gets the depth feature output of the encoder $F_d\in {\mathbb{R}}^{(H/16)\times (W/16)\times 8C}$, capturing global pixel dependencies and avoiding the result of excessive or insufficient enhancement.

The decoder also contains $K$ stages, and each stage has a upsample layer and a sequence of proposed fusion attention blocks. Each stage receives information from both the previous stage and the skip-connection information of the corresponding encoder. It receives low-resolution depth features and gradually expands the resolution to reconstruct clear underwater images $F_e\in {\mathbb{R}}^{(H/2)\times (W/2)\times C}$. After that, the image reconstruction block reconstructs refined underwater images $F_r\in {\mathbb{R}}^{H\times W\times C}$ and recovers spatial dimensions by upsampling. Finally, the network outputs an enhanced image ${\mathrm{I}}^{\prime }\in {\mathbb{R}}^{H\times W\times 3}$ by using a $3\times 3$ convolution layer to restore the scale of channels.

Fusion attention block

The inconsistently attenuation of underwater images on the space lead to the different degraded regions of the image should have different weights. Obviously, the convolution with the same treatment on the space does not accommodate this difference characteristic well.

The original Swin Transformer block is included in our proposed fusion attention block, where the self-attention mechanism allows the attention weights to interact with the image content, enabling flexible focus on any location in the input. STB limits the growth of computational complexity by using cyclic shift operations and masked MSA. The self-attention is computed within the non-overlapping local windows divided from the image. Two consecutive STBs are used as a whole unit. The first STB calculates self-attention within the window, and the second STB computes the self-attention within the shifted window. This strategy of shifting window enables the establishment of cross-window dependencies.

The calculation of the Swin Transformer block is expressed as

$$\begin{array}{rl}& {\hat{S}}_l=\mathrm{W}-\mathrm{M}\mathrm{S}\mathrm{A}(\mathrm{L}\mathrm{N}(S_{l-1}))+S_{l-1}\\ & S_l=\mathrm{M}\mathrm{L}\mathrm{P}(\mathrm{L}\mathrm{N}({\hat{S}}_l))+{\hat{S}}_l\end{array}$$where ${\hat{S}}_l$ and $S$ are the output features of the W-MSA module and $\mathrm{M}\mathrm{L}\mathrm{P}$. $\mathrm{L}\mathrm{N}$ represents Layer Norm, while W-MSA and $\mathrm{M}\mathrm{L}\mathrm{P}$ refer to window-based multi-head self-attention and multilayer perceptron, respectively.

Our network is a hierarchical architecture, computing window self-attention on small size feature maps in deep networks, enabling modeling of long-range dependencies. The relative position embedding $B\in {\mathbb{R}}^{(2M-1)\times (2M-1)}$ is added to the attention calculation, where $M$ represents the window size. The dot product attention calculation is expressed as

$$\text{Attention}(Q,K,V)=\text{ SoftMax}\left(\frac{{\text{QK}}^{\text{T}}}{\sqrt{\text{C}}}\text{+B}\right)\cdot \text{V}$$where $Q,K,V\in {\mathbb{R}}^{C\times HW}$ represent the query, key, and value respectively. $C$ refers to the channel dimension of the input feature, while $\mathrm{S}\mathrm{o}\mathrm{f}\mathrm{t}\mathrm{M}\mathrm{a}\mathrm{x}$ is multiple logistic regression.

Underwater images usually suffer from color shift since the water medium is biased to absorb certain wavelengths of light. Pixels in certain color channels of the images may have extremely low values, resulting in varying weights for different channel information. If we treat features of different channels equally, it cannot flexibly adapt to degraded underwater images.

We introduce a channel attention block (CAB) (Zhang et al., 2018) as shown in Fig. 2, which allows the model to capture weak degraded features. For the input feature map $x\in {\mathbb{R}}^{H\times W\times C}$, apply a global average pooling layer to generate a series of local descriptors $z\in {\mathbb{R}}^{H\times W\times C}$. Then, the first convolution compresses the scale of the channel, and after activation by the ReLU function, the second convolution restores the channels scale. The Sigmoid function activates multiple channel features to generate channel statistics $s\in {\mathbb{R}}^{1\times 1\times C}$. The input is multiplied with this channel statistics to produce the feature map with weights in the channel dimension. The channel attention block is expressed as

$$\begin{array}{rl}& z=\mathrm{G}\mathrm{A}\mathrm{P}(x)\\ & r=\mathrm{R}\mathrm{e}\mathrm{L}\mathrm{U}(\mathrm{C}\mathrm{o}\mathrm{n}{\mathrm{v}}_D(z))\\ & s=\mathrm{S}\mathrm{i}\mathrm{g}\mathrm{m}\mathrm{o}\mathrm{i}\mathrm{d}(\mathrm{C}\mathrm{o}\mathrm{n}{\mathrm{v}}_U(r))\\ & {\hat{x}}_c=s_c\cdot x_c\end{array}$$where $x$ represents the input feature map, $\mathrm{G}\mathrm{A}\mathrm{P}$ refers to the global average pooling layer, and $\mathrm{C}\mathrm{o}\mathrm{n}{\mathrm{v}}_D$ and $\mathrm{C}\mathrm{o}\mathrm{n}{\mathrm{v}}_U$ denote two convolution layers. $\mathrm{R}\mathrm{e}\mathrm{L}\mathrm{U}$ and $\mathrm{S}\mathrm{i}\mathrm{g}\mathrm{m}\mathrm{o}\mathrm{i}\mathrm{d}$ both represent activation functions. In the $c$-th channel, $s_c$ denotes the scaling factor, while $x_c$ represents the feature map.

Although Transformer can capture long-range dependencies well, it has insufficient ability to capture local features (Li et al., 2021; Wu et al., 2021) due to the lack of inductive bias of convolution. However, local context information is essential for image detail recovery. The local domain information of degraded pixels contributes to reconstructing clear images.

Therefore, we propose a local attention block (LAB), which is placed in another branch to capture useful local context information. The roles of the two branches are distinguished and exploit their respective advantages to enhance underwater images effectively. In Fig. 3, the local attention block comprises four residual blocks, and the calculation of every residual block as

$$X_l=\mathrm{C}\mathrm{o}\mathrm{n}\mathrm{v}(\mathrm{R}\mathrm{e}\mathrm{L}\mathrm{U}(\mathrm{C}\mathrm{o}\mathrm{n}\mathrm{v}(X_{l-1})))+X_{l-1}$$where $\mathrm{C}\mathrm{o}\mathrm{n}\mathrm{v}$ denotes the convolution layer of $3\times 3$. The input is activated by the $\mathrm{R}\mathrm{e}\mathrm{L}\mathrm{U}$ layer after passing the first convolution, then enters the second convolution, and finally obtains a residual output. The input of the $l$-th residual block is $X_{l-1}$, and the output is $X_l$. We set $l$ to 4.

We use Swin Transformer block (STB) to model remote dependencies and perform global attention while letting local attention modules capture local contextual information. To enable the adaptive collaboration between the global attention branch and local attention branch, we introduce an adaptive feature fusion block (AFFB) (Li et al., 2019b) that is learnable and allows dynamic adjusting weights of two branches.

The adaptive feature fusion block is illustrated in Fig. 4. First, incorporate the features of two branches by element addition, and then the global pooling layer is employed to obtain the global feature vector $v$. Following that, the dimensions are scaled down by the first linear layer, and then the two linear layers restore the scale and generate two weight vectors $W_1$ and $W_2$ to guide the selection of the two branches. Finally, the weights are multiplied with the input features separately, and the output results fuse the information on both branches. It allows the outputs of the FAB can be chosen flexibly between the two branches according to their features. The calculation of the AFFB is expressed as

$$\begin{array}{rl}& v=\mathrm{G}\mathrm{A}\mathrm{P}(X_{\mathrm{S}\mathrm{T}\mathrm{L}}+X_{\mathrm{L}\mathrm{A}})\\ & \hat{v}=\mathrm{F}{\mathrm{C}}_{\mathrm{D}}(v)\\ & W_1,W_2=\mathrm{S}\mathrm{o}\mathrm{f}\mathrm{t}\mathrm{M}\mathrm{a}\mathrm{x}([\mathrm{F}{\mathrm{C}}_{\mathrm{U}1}(\hat{v}),\mathrm{F}{\mathrm{C}}_{\mathrm{U}2}(\hat{v})])\\ & F=X_{\mathrm{S}\mathrm{T}\mathrm{L}}\cdot W_1+X_{\mathrm{L}\mathrm{A}}\cdot W_2\end{array}$$where $X_{\mathrm{S}\mathrm{T}\mathrm{L}}$ and $X_{\mathrm{L}\mathrm{A}}$ denotes the output of STB and local attention block. $\mathrm{F}{\mathrm{C}}_{\mathrm{D}}$ denotes the Linear layer with reduced dimensionality. $\mathrm{F}{\mathrm{C}}_{\mathrm{U}1}$ and $\mathrm{F}{\mathrm{C}}_{\mathrm{U}2}$ denotes two independent Linear layers. $\mathrm{G}\mathrm{A}\mathrm{P}$ stands for global average pooling layer, and $\mathrm{S}\mathrm{o}\mathrm{f}\mathrm{t}\mathrm{M}\mathrm{a}\mathrm{x}$ is a multiple logistic regression. $W_1$ and $W_2$ represent the weights on the two branches, respectively.

In summary, the process of FAB is defined as

$$\begin{array}{rl}& \hat{X}=\mathrm{C}\mathrm{A}(X)\\ & Y_1,Y_2=\mathrm{S}\mathrm{T}\mathrm{B}(\hat{X}),\mathrm{L}\mathrm{A}(\hat{X})\\ & Z=\mathrm{A}\mathrm{F}\mathrm{F}\mathrm{B}(Y_1,Y_2)\end{array}$$where $X$ is the input of FAB, $\mathrm{C}\mathrm{A}$ denotes channel attention block, $\mathrm{S}\mathrm{T}\mathrm{B}$ and $\mathrm{L}\mathrm{A}$ refer to Swin Transformer block and local attention block, and $\mathrm{A}\mathrm{F}\mathrm{F}\mathrm{B}$ represents adaptive feature fusion block. $Y_1$ and $Y_2$ represent global attention and local attention, respectively, while $Z$ represents the fused features outputted by the FAB. Our proposed FAB serves as the basic building block of the network, which has channel attention to strengthen the channel information with color degradation. To focus on degraded image regions and recover image details, it performs global attention and local attention in two parallel branches and adaptively fuses the attention results.

Loss function

To train our network, we employed a loss function that comprises two parts, $L_{\mathrm{c}\mathrm{h}\mathrm{a}}$ and $L_{\mathrm{p}\mathrm{e}\mathrm{r}}$. We use the robust Charbonnier loss (Lai et al., 2018) to learn the global similarity of image contents. Compared to L1 loss, it allows the function differentiable at zero by introducing a constant, which avoids the drawbacks of L1 loss. The Charbonnier loss calculation is shown below

$$L_{\mathrm{c}\mathrm{h}\mathrm{a}}=\sqrt{{\Vert I_u-X_u\Vert }^2+{ϵ}^2}$$where $I_u$ represents the enhanced underwater image, $X_u$ is the ground-truth image, and $\epsilon =10^{-3}$ refers to a constant value.

However, due to the neglect of the visual characteristics of the human eye by the Charbonnier loss, the generated results fail to maintain visual consistency, resulting in suboptimal visual effects. To enhance the subjective visual effects of images on human eyes, we introduce a perceptual loss in training. The perceptual structure of an image is represented by the pre-trained network VGG-19 (Simonyan & Zisserman, 2014). The perceptual loss $L_{\mathrm{p}\mathrm{e}\mathrm{r}}$ represents the distance between the two features of the enhanced image $I_u$ and the ground truth image $X_u$ output by the pre-training model, which is defined as

$$L_{\mathrm{p}\mathrm{e}\mathrm{r}}=\sum _{j=1}^N{\displaystyle \frac{1}{C_j{H}_j{W}_j}{\Vert {\varphi }_j(I_u)-{\varphi }_j(X_u)\Vert }_2^2}$$where $N$ denotes the batch size. ${\varphi }_j(\cdot )$ indicates the output features of the $j$th layer VGG-19 network and $C_j{H}_j{W}_j$ denotes the dimensions feature map, namely the channel, height and width.

In summary, Charbonnier loss exhibits good robustness to outliers and possesses smooth penalties, making it suitable for image restoration tasks. While perceptual loss is based on extracted features, it can capture the semantic information of an image, making it suitable for learning the content of the image. The overall loss function integrating $L_{\mathrm{c}\mathrm{h}\mathrm{a}}$ and $L_{\mathrm{p}\mathrm{e}\mathrm{r}}$ is expressed as

$$L_{\mathrm{s}\mathrm{u}\mathrm{m}}={\lambda }_1{L}_{\mathrm{c}\mathrm{h}\mathrm{a}}+{\lambda }_2{L}_{\mathrm{p}\mathrm{e}\mathrm{r}}$$where ${\lambda }_1$ and ${\lambda }_2$ denote the importance coefficients, they are set to 1 and 0.1. They represent the balance between overall performance and visual effects.

Underwater image synthesis

In contrast to some advanced vision tasks with large datasets available, underwater image enhancement tasks rely on synthetic datasets. The datasets used in the current related research can be classified into three categories. Firstly, using the unsupervised method of CycleGAN to generate images with underwater style. Images generated in this way need to select good results, and the dataset lacks the characteristics of various water types. Secondly, use publicly available datasets. However, these datasets are usually small, which can constrain the model's generalization ability if used for training. The last category is using indoor scene images with depth maps to synthesize simulated underwater images in accordance with the underwater image formation model.

We follow Zhou & Yan (2020) to synthesize the dataset for the training network. It describes three typical water types and provides the attenuation and background light coefficients. To simulate complex underwater scenes, the parameters provided have random values that increase the diversity of the dataset. The indoor dataset provided by Sun rgb-d (Song, Lichtenberg & Xiao, 2015) is adopted to synthesize underwater images in accordance with the underwater image degradation process defined in Eqs. (1) and (2).

Due to imperfections in some of the depth maps, we exclude flawed depth maps and select 7,680 indoor images with corresponding depth maps to synthesize a dataset. The dataset contains three types of underwater scenes consisting of 23,040 simulated underwater images and corresponding ground truth images.

Table 1 shows the normalized residual energy ratio $\mathrm{N}\mathrm{r}\mathrm{e}\mathrm{r}$ and background light $B_{\lambda }$ for three types of water. We set the scene depth $d(x)$ from the nearest 0.5 m to the farthest 15 m (Anwar, Li & Porikli, 2018). To increase the computational efficiency of network training, we resized the synthesized underwater images to $256\times 256$. Overall, the synthesized underwater image dataset is based on underwater image formation model and conforms to the principles of underwater imaging.

Implement details

The proposed model is trained using the synthetic dataset and set the image size of the input network to $256\times 256$. The network training utilizes 22,610 images from the synthetic dataset, while the validation set contains 430 images. Our network is implemented on Ubuntu20 with Intel(R) Xeon E7-4809 v3 CPU, NVIDIA RTX 3090 GPU, and PyTorch framework. Training with Adam (Kingma & Ba, 2014) optimizer, which ${\beta }_1$ and ${\beta }_2$ set to 0.9 and 0.999. We configured the batch size at 2 and set the learning rate at 0.0005 and a decay rate of 0.5 every 50 epochs, training 200 epochs.

Synthetic underwater images evaluation

First, our proposed method is evaluated on synthetically underwater images. To demonstrate its superior performance, we compared proposed network with seven competition methods. It includes three traditional methods (UDCP (Drews et al., 2013), Fusion (Ancuti et al., 2012), RGHS (Huang et al., 2018)), three CNN-based methods (UWCNN (Anwar, Li & Porikli, 2018), FUnIE (Islam, Xia & Sattar, 2020), WaterNet (Li et al., 2019a)) and one Transformer-based method (U-shape (Peng, Zhu & Bian, 2023)). The results of the comparison method are obtained from executing the code released by the corresponding author.

Our synthetic datasets were subjected to comprehensive experiments. In the dataset, test images are denoted as Test-S and training images are denoted as Train-U. The PSNR and SSIM quality measurement are employed for full-reference evaluation on synthetic datasets with reference images. The full-reference evaluation metric indicates its proximity to the reference image. PSNR evaluates the image quality from the perspective of the loss of signal. A higher PSNR value indicates a higher quality reconstructed image with more similar image content. The SSIM measurement takes into account luminance, contrast and structure, which is comparable to the human visual system. It has a value in the range of 0 to 1, with higher values meaning a more similar image structure.

The quantitative evaluations as shown in Table 2, while qualitative comparisons results on Test-S are demonstrated in Fig 5. Our model has the best results on PSNR and SSIM metrics, as indicated in Table 2. Its 20.46 PSNR and 0.82 SSIM on Test-S are higher than the suboptimal method by 0.1 PSNR and 0.01 SSIM. Moreover, the proposed method has a moderate number of parameters, and low computational cost during training.

The enhanced outcomes of synthetic underwater images of the test set, which includes images of degraded underwater scenes in dark green, are displayed in Fig. 5. Compared to competing methods, the images enhanced by our method exhibit the highest similarity to the ground truth. The traditional methods have limited ability to correct the color, which is only for images with color shifts slightly. Fusion overcorrects image colors resulting in unnatural tones, while RGHS cannot correct degraded image colors. The deep learning approach enhances images by extracting image features, which obtain better results than traditional methods. Since UWCNN-I uses one type of water dataset for training, it leads to a model with insufficient generalization ability and produces poor image enhancement results. Water-Net and FUnIE remove most color shift effects. However, the resultant image of Water-Net has poor details and low color saturation, while FUnIE recovers the colors incompletely to produce color artifacts and low image contrast. U-shape removes the color shift effect but causes over-enhancement of colors and a lack of ability to reconstruct image detail. The proposed method has sharper enhancement results with better color and detail reconstruction.

Real underwater images evaluation

Furthermore, our model was evaluated on real-world underwater images to validate its generalization. Likewise, our model is compared with the seven competing approaches to observe enhancement effects on real-world underwater images. These images utilized for the evaluation are taken from the available public datasets.

We perform evaluations on real underwater images Test-C60 (Li et al., 2019a) and Test-M to verify actual effectiveness of our model. Test-C60 Contains 60 challenge images, which are difficult to enhance by existing methods. Test-M contains 515 underwater images including underwater marine scenes and marine animals. Since some of the images did not have corresponding reference images, we have employed non-reference quality evaluation measure such as UIQM (Panetta, Gao & Agaian, 2015), UCIQE (Yang & Sowmya, 2015), NIQE (Mittal, Soundararajan & Bovik, 2012b) and BRISQUE (Mittal, Moorthy & Bovik, 2012a). Among them, higher UIQM and UCIQE scores indicate that the result has better visual perception. UIQM evaluates image quality from aspects of color, sharpness and contrast. Also, UCIQE measurement from aspects of image hue, saturation and contrast. A lower NIQE score indicates higher visual quality, which measures the difference in image features over a multivariate distribution constructed from features of natural images. BRISQUE evaluates the quality of an image based on the natural scene statistics present in the image, including the gradient magnitude, contrast, and brightness.

The quantitative evaluations as shown in Table 3, while qualitative comparisons are demonstrated in Figs. 6 and 7. According to the statistical results, the FUnIE achieved the highest UIQM and NIQE scores on Test-C60. However, the results of FUnIE display drawbacks with low color saturation, color artifacts and insufficient detail recovery. Due to the inability of certain evaluation metrics to differentiate perceptual differences, even if better evaluation values are achieved, the subjective visual effects of the results may still exhibit color distortions. This could be attributed to the sensitivity of the evaluation metrics to the types of distortions present. For instance, UIQM and UCIQE are insensitive to color artifacts and have biases for certain features (Berman et al., 2020). Moreover, in the no-reference evaluation of underwater images, non-data-driven methods generally outperform data-driven methods. From this perspective, the proposed method performs relatively well among deep learning-based methods, with the exception of the FUnIE results. In summary, the quantitative evaluation cannot determine the final performance, and a visual comparison is necessary.

Figure 6 shows the visual effects of the comparison approaches on Test-C60. It is evident that the traditional approaches produce unsatisfactory enhancement results. UDCP introduces unnatural colors and the image has low brightness. While the results of Fusion and RGHS have color artifacts. The data-driven methods are better than the traditional methods in general, but there are also some defects. The results of UWCNN-I have the problem of color casts. FUnIE and Water-Net result in images with low color contrast. The method based on Transformer obtains enhanced images with better color restoration. However, U-shape produces relatively blurry results, while the proposed method has better image detail reconstruction ability and image color contrast.

Figure 7 shows the visual effects of the different methods on Test-M. As can be seen, the traditional methods cannot effectively deal with the color shift problem. For background areas with more backscattering, Fusion and RGHS cannot correct the color. And UDCP over-enhances and introduces unnatural colors. FUnIE corrects the color to some extent, but color artifacts remain. The dehazing effects of UWCNN-I are not satisfactory, while the enhanced results of Water-Net lack color saturation. The Transformer-based method has better visual results since it solves the color shift problem more effectively than other methods. Compared to the U-shape, the results of our method have better color saturation and sharper details.

In our network, the fusion attention block can strengthen the attention to degraded spatial regions and color channels, thus obtaining results without color bias and artifacts. It also extracts image features by utilizing global and local attention to reconstruct the images with high-quality details.

Ablation study

In order to prove the effectiveness of every block in our model, we conducted an ablation study. The factors we consider include networks without channel attention block (CAB), networks without local attention block (LAB) and adaptive feature fusion block (AFFB) and networks without perceptual loss $L_{\mathrm{p}\mathrm{e}\mathrm{r}}$.

All experiments were trained by Train-U. Table 4 reports the average statistics on Test-S. Where full model suggested in this article achieves the best results. Therefore, this reflects the effectiveness of the fusion attention block of our model. Table 5 reports the influence of different training losses. As shown in Table 5, with the addition of the loss function, the performance of the proposed network has improved, indicating the effectiveness of introducing perceptual loss.

As shown in Fig. 8, the result of only CAB has more noise due to the lack of local detail reconstruction ability of LAB. While AFFB + LAB avoids more noise generation but produces results with hue shifting. Overall, our full model produces the highest-quality images, which are closest to the ground truth image. These results demonstrate that the FAB and loss function used in the proposed method significantly contribute to achieving the final visually pleasing results.