A hybrid attention network with convolutional neural network and transformer for underwater image restoration

Underwater image restoration methods based on local feature extraction

In recent years, with the continuous development of deep learning technology, underwater image restoration methods based on convolutional neural networks (CNN) and generative adversarial network (GAN) (Fabbri, Islam & Sattar, 2018) have become a research hotspot. They achieve better results in underwater image restoration tasks through local feature extraction using convolutional layers and subsequent enhancement. As a classic neural network, GAN can be directly trained in the domain of transformation between underwater images and clear images. To correct and enhance the color of underwater images while maintaining the edge structure information of the originals, Li, Bai & Niu (2023) suggested an enhanced water CycleGAN technique based on the edge structure similarity loss function. However, it may produce some artifacts. WaterGAN (Li et al., 2017) was proposed based on this, which uses a weakly supervised style transfer to generate a dataset of underwater blurry images, and then uses a correction network to restore unmarked real underwater images. Supervised learning using the GAN-generated dataset can make the trained model more adaptable to underwater scenes. The EUVP underwater dataset and the FUnIE-GAN network were generated by using the same technique (Islam, Xia & Sattar, 2020), which found that a combination of supervised and unsupervised techniques yielded good performance in restoration. However, the performance of contrast enhancement is limited due to the use of labeled images in supervised learning. Chen et al. (2017) used GAN to create a large number of underwater images and trained the model to improve model accuracy. Cao, Peng & Cosman (2018) proposed underwater image restoration with neural network topologies that measure background light and scene transmission rate. The input image is decomposed into many sub-bands using the discrete wavelet transform in UIE-WD. The algorithm includes a multi-color space fusion network and a detail enhancement network, which restore underwater images with rich high-frequency information in the frequency domain. Naik, Swarnakar & Mittal (2021) proposed Shallow-UWNet—a lightweight convolutional neural network structure, which concatenates the input image with the output of each residual block through skip connections. Considering the diverse lighting conditions in underwater environments, Liu et al. (2022) proposed LANet—an underwater image enhancement network based on attention mechanisms and adaptive learning that adaptively learns important feature information by implicitly perceiving lighting features.

Underwater image restoration methods based on global feature extraction

In order to accomplish the capability of global modeling, Swin Transformer (Liu et al., 2021) was suggested as a hierarchical Transformer structure that calculates features through a moving window. Uformer-B (Wang et al., 2021b) adapted the network encoding-decoding structure to Transformer with the ability to restore underwater images effectively. Peng, Zhu & Bian (2023) proposed a U-shape Transformer for underwater image enhancement by integrating the channel multi-scale feature fusion transformation (CMSFT) module and the spatial global feature modeling transformation (SGFMT) module to improve the network’s attention for color channels and severely attenuated spatial regions. Using Transformer for deep feature extraction and color correction modules to enhance image contrast and create enhanced underwater images, Zhang et al. (2023) proposed a two-stage network based on Transformer for single underwater image enhancement. It is challenging to be able to achieve both the low computational and low latency requirements for real applications due to the nature of the Transformer.

Attention mechanism

The attention mechanism is a technique used in deep learning that allows a model to selectively focus on certain parts of its input when making predictions or generating output. In computer vision, attention mechanisms are also used in various models to improve their performance on tasks such as image restoration and object detection. Currently, there are two common attention mechanisms, channel attention (Wang et al., 2021a; Wang et al., 2023) and self-attention (Wang et al., 2022a). The channel attention mechanism selectively emphasizes or suppresses certain channels of the feature map in the neural network according to different channels of attention. For example, Zhang et al. (2018) embed the channel attention mechanism into the residual block for super-resolution network RCAN, which greatly improves the performance of the network. Zhang et al. (2019) proposed a residual non-local attention network for high-quality image restoration. The self-attention mechanism aims to apply self-attention to feature maps to selectively focus on different spatial locations in the input image, and is often used in models such as non-local neural networks (Wang et al., 2018) and Transformer (Vaswani et al., 2017), etc.

The objective function of the underwater imaging model

The majority of underwater imaging is made up of the direct component, the forward and backward scattered components, and the attenuated light that the camera receives. Underwater image model is often created using the Jaffe-MeGlamery (Trucco & Olmos-Antillon, 2006): (1)${\displaystyle I\left(x\right)=J\left(x\right)t\left(x\right)+B\left(1-t\left(x\right)\right)+\eta \left(x\right)}$ where $I\left(x\right)$ indicates the degraded image, J(x) denotes the clear image, J(x)t(x) means the direct component, B(1 − t(x))denotes the background scattering component, B denotes the underwater ambient light, and t(x) denotes the scene light transmittance, respectively. Another factor that commonly arises with underwater images is noise, which is denoted by a parameter $\eta \left(x\right)$, and usually additive noise.

To map the degraded image to the clean image with a deep learning approach, we typically train a neural network. Consequently, the map F:I(x) → J(x) can be defined as an objective function of a minimization problem in the following form: (2)${\displaystyle \underset{J\left(x\right)}{\text{arg min}}\frac{1}{2}\left|\right|J\left(x\right)-F\left(I\left(x\right)\right)|{|}_2^2}$

The above problem is equivalent to the following problem: (3)${\displaystyle minL\left(F_{Net}\left(y;\theta \right),x\right)}$ where F_Net is the network model that we suggested, L is the loss function, y is the network’s input of a degraded image, x is the clean image, and θ is the learning parameter.

Hybrid attention network combining CNN and Transformer

We propose a hybrid attention network (MA-CTNet) incorporating CNN and Transformer for underwater image restoration. Figure 1 presents an overall of the network architecture. The global feature extraction component, the local feature extraction component, and the feature fusion strengthening component make up the entirety of the network structure. Two residual groups (RGS) with a convolutional layer, residual channel attention block (RCAB), and channel attention (CA) mechanism are used to extract local features and enhance shallow features of input images. As local feature augmentation cannot rectify the overall color bias of damaged underwater images, the Transformer Blocks structure is integrated into the global feature extraction phase to replicate the global dependency of input image block sequences utilizing Self-Attention (S-A). Then, it extracts the details of the global feature. The improved local and global features are then combined together. The Supervised Attention (SA) mechanism is used to achieve feature transmission prior to the reinforcement phase. The SA module uses the real image as the supervision constraint, transfers locally learned features, fuses them with globally learned features, and inputs them into the reinforcement phase to further mine the image details, and then bolsters and reestablishes the deep features.

Local feature extraction

Figure 1 illustrates the components of the local feature extraction, which include a convolution layer (Conv), two residual groups (RGS), and skip connections. Two residual channel attention blocks (RCABs) with CA mechanisms are present in each residual group. CA focuses on significant channel characteristics and utilizes excellent recognition ability to extract and enhance deep features and improve deep network features during restoring texture and details (Zhang et al., 2018). Utilizing the interdependence between channels through skip connections, the useful local feature information can be recovered even in the deeper network by incorporating CA in the local feature extraction portion. As illustrated in Fig. 2, CA is made up of an average pooling layer, two convolutional layers, ReLU and Sigmoid activation functions, and residual connection. To create the compressed feature map, a pooling layer first down-samples the input feature. Dimension reduction, two convolutions, and the ReLU activation function are all used to learn the nonlinear interaction between several channels. Finally, the output of Sigmoid function is element-wise multiplied by its input to provide an output of the same magnitude when it is activated.

Global feature extraction

It can be challenging to recover the entire color of an underwater image using only local features. Deriving from the core of the Transformer mechanism of Self-Attention, we can model the interdependence of input and output image sequences without regard to their closeness in the sequence. Therefore, the global feature extraction achieved a better performance in the underwater image restoration. Thus, as shown in Fig. 1, the part for extracting global features consists of two Transformer Blocks with Self-Attention and one Patch embed (3x3 Conv layer). Through the convolution of the input images, patch embedding can provide low-level feature embedding. Then, in order to achieve global recovery, the underlying characteristics are put into Transformer Blocks to extract global features. Figure 3 depicts the composition of the transformer blocks, which are made up of Multi-DConv Head Transposed Attention (MDTA) with Self-Attention, and feed forward network (FFN) blocks. In more detail, after the underlying feature has passed through LN, the pixel-level cross-channel context is aggregated using 1 × 1 convolution, and the channel-level spatial context is then encoded using 3 × 3 depth-level convolution to produce the query (Q), key (K) and value (V) projections, denoted $Q=W_d^Q{W}_p^{Q}X$, $K=W_d^K{W}_p^{K}X$, $V=W_d^V{W}_p^{V}X$. Here, $W_p^{\left(\cdot \right)}$ is the 1 × 1 point-wise convolution, and $W_d^{\left(\cdot \right)}$ is the 3 × 3 depth-wise convolution. S-A is calculated on channels to reduce computational complexity: ${\displaystyle \hat{X}=W_{p}SA\left(Q,K,V\right)+X}$ (4)${\displaystyle SA\left(Q,K,V\right)=V\cdot Softmax\left(K\cdot Q/\alpha \right)}$

Here, α is a learnable scaling parameter used to regulate the size of the dot product and before applying the softmax function, X and $\hat{X}$ are the input and output feature maps. Additionally, W_p is the weight of the 1 × 1 point-by-point convolution. In order to acquire the required number of channels, a 1 × 1 convolution layer is employed after the normalized Softmax coefficients are converted back to feature V. The channels are separated into ‘heads’, and several attention graphs are learned concurrently for each head.

FFN uses two 1 × 1 convolutions: one to extend the feature channel and the remaining one to reduce the channel back to the original input dimension; and nonlinear GeLU activation function in the hidden layer to effectively recover from adjacent pixel positions. FFN is expressed as follows: (5)${\displaystyle \hat{X}=W_p^0\left[\varphi \left(W_d^1{W}_p^1\left(LN\left(X\right)\right)\right)⨀W_d^2{W}_p^2\left(LN\left(X\right)\right)\right]+X}$ where the operator ⊙depicts element-wise multiplication, ϕ is nonlinear GeLU activation function, LN(⋅) isthe normalization layer, $W_p^{\left(\cdot \right)}$ is the 1 × 1 point-wise convolution, and $W_d^{\left(\cdot \right)}$ is the 3 × 3 depth-wise convolution.

Feature fusion and enhancement

The supervised attention SA mechanism is integrated before the feature reinforcement part. On the one hand, focusing on using the real image as the supervision constraint in order to restore the real image; on the other hand, supervised attention SA can generate attention feature maps in order to suppress features with less information at the moment and only pass useful features to the next stage. Local and global features can be successfully fused to strengthen deep features. The enhanced component has the same structure as the local feature extraction part, and residuals are employed for deep enhancement. Figure 4 depicts the SA mechanism of supervision attention. The recovered local features are first convolution by a 3 × 3 convolution layer and superimposed on the deteriorated image to generate the intermediate output and calculate its Charbonnier loss with the ground-truth image GT. One 1 × 1 convolution layer and Sigmoid activation function are used to activate the intermediate output, which is then element-wise multiplied and overlaid with the local feature to generate the attention feature map of the output as follows:

O₁ = Conv(x_in + X) (6)${\displaystyle x_{out}=Sigmoid\left[Conv\left(O_1\right)\odot Conv\left(x_{in}\right)\right]+x_{in}}$ where the operator ⊙ represents the element-wise multiplication, Conv(⋅) represents the convolution operator.

Loss function

The hybrid attention network structure paired with CNN and Transformer has the following total loss function form: (7)${\displaystyle L=L_1\left(O,GT\right)+L_{Charbonnier}\left(O_1,GT\right)}$ Among them, the loss function contains two terms: loss L₁ and Charbonnier loss L_Charbonnier, respectively. GT represents the ground-truth image, O is the image recovered by the overall proposed network, and O₁ represents the output image obtained by the supervision and attention modules. The constraint of Charbonnier loss  (Charbonnier et al., 1994) is to extract useful features by using supervised learning, and the Charbonnier loss function is defined as follows: (8)${\displaystyle L_{Charbonnier}=\sqrt{\left|\right|O_1-GT|{|}^2+{\varepsilon }^2}.}$ To prevent the training network from vanishing gradient, we have added a numerical stabilization constant ɛ and usually set it to 0.001.

Experimental setup

We implemented the performance of our method on two datasets, EUVP (Adrian et al., 2015) and UIEB (Jaffe, 2015). The Underwater Dataset EUVP consists of 5,550 paired training images, we chose 5,022 images to train and 528 images to test during the experiment. For UIEB, we randomly selected 800 images in the UIEB dataset as the training set and the remaining 90 images as the testing set.

During network training, the batch_size and patch_size parameters are set to 16 and 48, respectively. The initial learning rate is set at 10⁻⁴, and it is reduced to half every 2 × 10⁵ back-propagation iterations. The Adam optimizer was used, and network convergence occurred when the iteration epoch reached 300. Simultaneously, the best model was chosen for testing. Furthermore, Gaussian noise was reproduced on the original data set for training and testing to validate the denoising impact.

As the comparative methods, one traditional method DCP (He, Sun & Tang, 2011), four other known deep learning image restoration methods, namely CycleGAN (Zhu et al., 2017), FUnIE-GAN (Islam, Xia & Sattar, 2020), Shallow-UWnet (Naik, Swarnakar & Mittal, 2021), and RTFAN (Fu et al., 2021), are employed in this research. To ensure fairness and dependability, all methods use the same training and test sets as described in this work. We implemented the experiments on PyTorch framework with the NVIDIA GeForce RTX 3090 GPU system.

Quantitative analysis

To evaluate the performance of the methods, we use two image quality assessment metrics: peak signal-to-noise ratio (PSNR) (Huynh-Thu & Ghanbari, 2008) and structural similarity index measure (SSIM) (Kotevski & Mitrevski, 2010). PSNR characterized how close a processed version of an image is to the ground truth, by calculating the mean square error (MSE) between the original image and the processed image: (9)${\displaystyle PSNR\left(x,y\right)=10log10\left[\frac{255^2}{MSE\left(x,y\right)}\right]}$ where, x and y are the reference image and the image to be evaluated respectively. MSE() represents the mean square error operation, and the larger the PSNR value, the better the image quality. In the meanwhile, SSIM takes into account the brightness, contrast, and structure of two images, therefore, closer to human perception as Formula (10). (10)${\displaystyle SSIM\left(x,y\right)=\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)}}$ where, μ_x and μ_y are the mean, ${\sigma }_x^2$ and ${\sigma }_y^2$ are the variance, σ_xydenotes the covariance, c₁ and c₂ are the constant that’s used to maintain stability.

In addition, UCIQE (Yang & Sowmya, 2015) is used to quantitatively evaluate underwater images’ non-homogeneous color shift, blurring, and low contrast. It is a linear combination of color intensity, saturation, and contrast as follows: (11)${\displaystyle UCIQE=c_1{\sigma }_c+c_{2}co{n}_l+c_3{\mu }_s}$ where σ is the standard deviation of the image, and it can represent the average of saturation.

All the methods are tested on EUVP test set and test sets with different Gaussian noises with zero mean and noise variance σ² with values 5, 10, and 15, respectively. Table 1 presents the quantitative results of the proposed method and various underwater image restoration algorithms.

It can be shown that the proposed method has higher performance in PSNR and SSIM scores, which can be increased by 0.13 dB in the best scenario. Furthermore, we show the exponential curve of the data during training and use it to quantify and analyze convergence. Figure 5 (left) depicts the trend of the loss curve. Figure 5 (right) depicts the trend of the PSNR curve. From Fig. 5 (left), it can be seen that the loss function of our method gradually tends to a downward trend, and when the epochs is 300, it gradually becomes flat. Figure 5 (right) shows the trend of increasing PSNR of the validation set as the number of iterations increases.

To fully validate the effectiveness of our method for different degraded images, we quantitatively compare with a traditional method and the four competing deep learning-based methods. we computed the PSNR, SSIM, and UCIQE results on the UIEB dataset in Table 2.

From Table 2, it can be seen that it is the best in PSNR, SSIM, and UCIQE metrics, although our method is slightly inferior in running time. Therefore, the proposed method achieves the best quantitative evaluation results on two datasets.

Qualitative analysis

When the noise variance is 5, five different test images are selected on the EUVP test set for comparison, and the restored visual results are shown in Fig. 6.

As shown in Fig. 6, unlike existing convolutional neural network methods that locally enhance image characteristics such as unevenly distributed color areas, the proposed method focuses on overall image consistency restoration, particularly for some vividly colored images. The proposed method suits to resolve the issue of image color bias and restore image features. The proposed method suits to resolve the issue of image color bias and restore image features. In addition, the proposed method eases image color bias (line 1 and line 5) and recovers more background details (lines 2–4) than other underwater image restoration methods.

To verify that our method has a good recovery effect on different types of data, we show the visualization results of the UIEB test set as shown in Fig. 7. In particular, we selected and processed degraded images containing green to ensure restoration. Due to space constraints, we choose the traditional enhancement algorithm DCP and one deep learning methods for underwater image restoration for comparison.

It can be seen from Fig. 7 that when the degraded underwater image appears green, the restoration results of our method have obvious visual advantages compared with other underwater image restoration methods. For example, DCP cannot handle this kind of image with a severe color cast. RTFAN will have an overall yellowish effect, and the restoration effect is not outstanding.

Ablation experiment

To fully verify the effectiveness of the proposed method in the color cast, we conduct an ablation experiment on Transformer Blocks and Channel Attention, and the quantitative and qualitative results are shown in Table 3 and Fig. 8, respectively.

Table 3 shows the benefits of adding Transformer Blocks (TBs) to further showcase the ability in fusing local and global features. At the same time, CA is tested for improving the color cast of the image as shown in Figs. 8B and 8D.

Comparing the results with/without using Transformer Blocks and CA as in Fig. 8, Both considerably improve the restoration results of overall color and details.

In addition, we also experimented with the number of Transformer Blocks in our network, and we tested the PSNR and SSIM for the cases of Transformer Blocks 1, 2, 3, and 4 on the EUVP dataset as shown in Table 4.

Hence, according to the above quantitative results and consideration of training time, we still choose two Transformer Blocks.