Underwater image enhancement method based on transformer and generative adversarial network

Abstract

The present application relates to an underwater image enhancement method based on a Transformer and a generative adversarial network, which solves the defect that it is difficult to perform enhancement processing on underwater images compared with the prior art.The present application comprises the following steps: acquisition and preprocessing of underwater images; construction of an underwater image enhancement model; training of the underwater image enhancement model; acquisition of the underwater image to be enhanced; enhancement processing of the underwater image.The present application uses a window-based double local enhancement module to make up for the deficiency of the Transformer in extracting local features, and further eliminates blur; a generative adversarial network containing a double-branch discriminator is used to generate more realistic colors, so that the enhanced image is more consistent with the human visual system.

Description

Technical Field

[0001] This invention relates to the field of underwater image processing technology, specifically to an underwater image enhancement method based on Transformer and generative adversarial networks. Background Technology

[0002] As a crucial carrier and presentation form of underwater information, underwater images play an irreplaceable role in the exploration of water bodies. Unlike ordinary images, underwater images are susceptible to various degradation problems due to the complex and diverse underwater environment. According to the imaging principles of underwater images, degradation is mainly caused by the absorption of light by water during propagation, as well as forward and backscattering caused by encountering suspended particles in the water.

[0003] Current underwater image processing mainly falls into three categories: physical model-based restoration methods that consider the underwater image degradation process; enhancement methods that directly enhance visual effects without considering the underwater image degradation process; and data-driven deep learning-based methods. Among these, enhancement methods do not consider the image degradation process and reallocate the pixel values ​​of a given image to enhance contrast and correct colors.

[0004] Furthermore, from an artificial intelligence perspective, data-driven deep learning methods are now mainly divided into three categories: convolutional neural networks (CNNs), generative adversarial networks (GANs), and Transformers. Recently, GANs and Transformers have been successfully applied to image-to-image translation tasks with excellent results.

[0005] Image restoration methods typically rely on prior knowledge to reconstruct degraded images; however, inaccurate prior knowledge often leads to significant estimation errors. The lack of reliable prior knowledge regarding underwater images has become a major obstacle to research in this area.

[0006] Enhancement methods directly utilize image processing techniques, subjectively adjusting pixel values ​​in the image to eliminate noise, improve edge blur, enhance target object features, and reduce the influence of irrelevant environmental features on the target. However, because they do not consider underwater optical imaging models, they introduce additional noise, causing oversaturation in different areas of the image. Furthermore, existing deep learning-based methods sometimes fail to thoroughly correct color casts in images with severe color distortion, and the enhanced images may exhibit blurring, poor contrast, and over-enhancement.

[0007] Therefore, how to use deep learning technology for underwater image enhancement has become an urgent technical problem to be solved. Summary of the Invention

[0008] The purpose of this invention is to address the shortcomings of existing technologies in enhancing underwater images by providing an underwater image enhancement method based on Transformer and Generative Adversarial Networks.

[0009] To achieve the above objectives, the technical solution of the present invention is as follows:

[0010] An underwater image enhancement method based on Transformer and Generative Adversarial Networks includes the following steps:

[0011] Underwater image acquisition and preprocessing: Acquire underwater images and perform preprocessing;

[0012] Construction of an underwater image enhancement model: An underwater image enhancement model is constructed based on Transformer and Generative Adversarial Network;

[0013] Training of the underwater image enhancement model: Input the preprocessed underwater image into the underwater image enhancement model for training;

[0014] Acquisition of underwater images to be enhanced: Acquire underwater images to be enhanced and preprocess them, adjusting the size of all images to 256×256 using bicubic interpolation;

[0015] Underwater image enhancement: The preprocessed underwater image to be enhanced is input into the trained underwater image enhancement model to obtain the enhanced underwater image.

[0016] The construction of the underwater image enhancement model includes the following steps:

[0017] The underwater image enhancement model consists of three parts: a dual local enhancement module, a generator network structure, and a dual-branch discriminator structure.

[0018] Constructing a window-based dual local enhancement module: Building a window-based dual local enhancement module suitable for underwater image enhancement;

[0019] Constructing the generator network structure: Based on the Encoder-Decoder framework, Inception, Bottlenect, and Fusion modules are added to construct the generator network structure;

[0020] Construct a two-branch discriminator structure: Construct a two-branch discriminator containing feature branches and color branches.

[0021] The construction of the window-based dual local enhancement module includes the following steps:

[0022] The window-based dual local enhancement module consists of two parts: the first part is a self-attention module for feature extraction and a self-attention feature map enhancement module; the second part is a local enhancement module, whose input is the original image and whose output is the feature map fused by the self-attention feature map enhancement module and the local enhancement module.

[0023] The self-attention module for feature extraction is configured as follows:

[0024] Using an 8×8 window as a unit, the feature map is stretched to a 1D shape. After layer normalization, it is input into three different learnable parameter matrices to obtain query Q, key K, and value V vectors. Then, query Q and key K are multiplied, relative position encoding B is added, and zero-mean normalization is performed to obtain the attention matrix Attn. Finally, Attn is activated by Softmax and multiplied with value V to output the attention feature map.

[0025] The self-attention feature map enhancement module is configured as follows:

[0026] The self-attention feature map is normalized by layer, then linearly projected using a 1×1 convolutional kernel to reshape it into a 2D shape. After passing through a 3×3 convolutional layer, it is stretched into a 1D shape. Finally, it is linearly projected to the same dimension as the input feature using a 1×1 convolutional kernel to output the self-attention enhanced feature map.

[0027] The local enhancement module is configured as follows:

[0028] The input to the local enhancement module is the original image, which is processed by two Conv+BatchNorm+GELU modules, where the size of the convolution kernel is 3×3, and the output is a local enhancement feature map.

[0029] The self-attention enhancement feature map and the local enhancement feature map are fused to obtain a self-attention feature map with dual local enhancement.

[0030] The construction of the generator network structure includes the following steps:

[0031] The generator network structure consists of four parts: the first part is the Inception part, which extracts the original features; the second part is the Encoder-Decoder framework; the third part is the Bottlenect part, which extracts global information; and the fourth part is the Fusion part, which integrates global information into each scale. Its input is the original image, and its output is the image enhanced by the generator network.

[0032] The input image first passes through the Inception part, where the original features are extracted and then output to the Encoder part in the second part. The Encoder part extracts and encodes the features and then outputs them to the Bottlenect part. Each layer of the Encoder part has skip connections to the Decoder part. The Bottlenect part extracts global information and outputs it to the Fusion part and the Decoder part. The Fusion part incorporates the global information into each scale of the Decoder part. The Decoder part decodes and reconstructs the features and then outputs the enhanced image.

[0033] The Inception configuration is as follows:

[0034] A window-based dual local enhancement module is used to extract features from the original image and output the extracted original features.

[0035] Configure the Encoder-Decoder framework as follows:

[0036] The Encoder-Decoder framework consists of an Encoder part and a Decoder part. The Encoder part has 5 encoding layers, which perform multi-scale feature extraction on the raw features initially extracted by the Inception part. The first layer contains only one convolutional layer, and each of the other layers contains a Conv+BatchNorm+ReLU module. The parameters of all convolutional layers are size=4×4, stride=2, padding=1, which serves to downsample while extracting features.

[0037] Finally, the extracted 512×8×8 shape is output to the Bottlenect part, and the features extracted from each layer are passed to the corresponding layer of the Decoder part through skip connections;

[0038] The Decoder part has 5 decoding layers, which receive global information extracted from the Bottlenect part. The structural design is the same as that of the Encoder part, except that only the first layer adds the Tanh activation function, and all the convolutional layers are replaced with transposed convolutional layers, which can reconstruct features while upsampling.

[0039] The output image after five decoding layers is an enhanced image with a shape of 3×256×256.

[0040] The Bottlenect configuration is as follows:

[0041] Two DleWin modules were used. When the features extracted by the Inception part were downsampled to an 8×8 size by the Encoder, which is the same as the window size of the DleWin module, the Transformer module extracted the global information.

[0042] The Fusion settings are as follows:

[0043] Global information extracted from Bottlenect, such as overall lighting and layout, is integrated into each scale.

[0044] The global information first passes through a 1×1 convolutional layer to adjust the channels of the global information to correspond with the Decoder. Then, by copying and reshaping the global information, the shape of the output fused information is finally the same as the feature map of the corresponding layer of the Decoder.

[0045] The construction of the dual-branch discriminator structure includes the following steps:

[0046] The dual-branch discriminator structure consists of two parts: the first part is a feature branch discriminator that preserves the image content features, and the second part is a color branch discriminator. Its inputs are the enhanced image and the ground truth image, and its outputs are the adversarial map and underwater index map corresponding to the enhanced image and the ground truth image, respectively.

[0047] The feature branch discriminator is configured as follows:

[0048] PatchGAN is used. After preserving the image content through a convolutional layer, three layers of Conv+BatchNorm+Leaky-ReLU modules and one convolutional layer are stacked to distinguish between real and fake images. A 30×30 adversarial graph is generated for evaluation. Finally, the average of the adversarial graphs is used as the generator loss to promote the generator to generate realistic images.

[0049] The color branch discriminator is set as follows: PatchGAN is used, and 5 layers of CBL modules and one layer of convolution are directly stacked to distinguish whether the image belongs to the underwater scene. A 15×15 underwater index map is generated to evaluate the intensity of underwater attributes, which helps the generator to generate colors that match the image in the air.

[0050] The adversarial graph generated by the feature branch discriminator and the underwater index graph generated by the color branch discriminator respectively guide the generator training.

[0051] The training of the underwater image enhancement model includes the following steps:

[0052] The objective function of the dual-branch discriminator is set using Wasserstein GAN with gradient penalty, which is used to eliminate the gradient vanishing problem of standard GAN and guide stable training of feature branches. Underwater exponential loss Uloss is introduced to guide color branch training.

[0053] The objective function of the generator is set as feature branch discriminator loss to preserve image content and underwater exponential loss to generate more realistic colors. Finally, L1 loss is used to eliminate blur.

[0054] In the initial training phase, the generator is not given underwater index loss, but the color branch discriminator is fully trained. That is, the generator is given underwater index loss after epoch=30.

[0055] In the training of the underwater image enhancement model, the dual-branch discriminator is trained first, and then the generator is trained. During the training process, each iteration first trains the dual-branch discriminator and then trains the generator.

[0056] Training the bi-branch discriminator:

[0057] The original image is input into the generator to generate an enhanced image. Then, the enhanced image and the ground truth are concatenated with the original image by channel to obtain two input images, which are then input into a dual-branch discriminator. After passing through the feature branch and the color branch, the corresponding adversarial map and underwater index map are output.

[0058] The gradient penalty term of the feature branch discriminator is calculated for the two adversarial graphs output by the feature branch. The loss is calculated based on the adversarial graphs and the gradient penalty term. The feature branch discriminator is then trained and its parameters are updated by backpropagation.

[0059] The underwater index loss is directly calculated based on the underwater index map, and the color branch discriminator is trained and updated by backpropagation of error.

[0060] Training generator:

[0061] After the original image is input into the generator, the original features are extracted by the Inception part, which contains one window-based dual local enhancement module. Then, the image features are extracted layer by layer by the Encoder part, which contains five coding layers. The output is sent to the Bottlenect part, which contains two window-based dual local enhancement modules, to provide global information. The Fusion part integrates the global information into each scale of the Decoder part. Finally, the Decoder part, which contains five decoding layers, receives the global information and the fused information, and reconstructs the features layer by layer to output the enhanced image.

[0062] The enhanced image is input into the discriminator to obtain the adversarial map and the underwater index map. The total loss of the generator is calculated by combining the L1 loss, and the generator is trained and updated by backpropagation of error.

[0063] The training was performed using 6000 images, with 32 images per batch. All images were used for training in each epoch, for a total of 188 batches. The training ended after 200 epochs.

[0064] Beneficial effects

[0065] The underwater image enhancement method based on Transformer and Generative Adversarial Network of the present invention, compared with the prior art, utilizes a window-based dual local enhancement module to compensate for the shortcomings of Transformer in extracting local features and further eliminates blur; it uses a generative adversarial network with a dual-branch discriminator to generate more realistic colors, making the enhanced image more consistent with the human visual system.

[0066] This invention proposes a window-based dual local enhancement (DleWin) module, which is more suitable for underwater image enhancement tasks, and a Transformer module, which is also better suited for underwater image enhancement tasks. The DleWin module implements a self-attention mechanism, which can effectively capture long-range dependencies; on the other hand, local features are crucial for underwater image enhancement tasks, so we introduce CNNs in the DleWin module in both serial and parallel ways for local enhancement.

[0067] This invention also proposes a GAN with a dual-branch discriminator, which includes a feature branch and a color branch. The feature branch is used to preserve image features and enhance contrast, while the color branch performs color correction to generate more realistic colors.

[0068] Meanwhile, this invention proposes a method that combines convolutional neural networks and Transformers in parts. Since Transformers are good at capturing long-range dependencies and extracting original information, while CNNs are good at extracting local features, our generator is based on the CNN-based UNet network, and the DleWinTransformer module is used in the part that extracts original information and global information. Attached Figure Description

[0069] Figure 1 This is a sequence diagram of the method of the present invention;

[0070] Figure 2 This is a network structure diagram of the present invention;

[0071] Figure 3 This is a diagram of the self-attention module of the present invention;

[0072] Figure 4 This is a network structure diagram of the Encoder and Decoder parts in the generator of this invention;

[0073] Figure 5 This is a comparison chart of the present invention with other methods on paired datasets;

[0074] Figure 6 This is a comparison chart of the present invention with other methods on an unpaired dataset;

[0075] Figure 7 This is a demonstration image showing the results of feature point matching using the SIFT algorithm on images enhanced by the present invention and several representative methods, as well as their flipped images.

[0076] Figure 8 This is a demonstration image showing the edge detection results of the image enhanced by the present invention and several representative methods on the Canny edge detection algorithm;

[0077] Figure 9 This is a demonstration image showing the results of underwater target detection using the Yolov5 model on images enhanced by the present invention and several representative methods.

[0078] Figure 10 This is a diagram illustrating the application test data of the present invention. Detailed Implementation

[0079] To provide a better understanding of the structural features and effects achieved by the present invention, a detailed description is provided below, accompanied by preferred embodiments and accompanying drawings:

[0080] like Figure 1 and Figure 2 As shown, the underwater image enhancement method based on Transformer and Generative Adversarial Networks of this invention includes the following steps:

[0081] The first step is the acquisition and preprocessing of underwater images: acquiring underwater images and performing preprocessing.

[0082] The second step involves constructing an underwater image enhancement model: This model is built based on Transformer and Generative Adversarial Networks (GANs). The underwater image enhancement model incorporates a novel, meticulously designed generator network structure and a two-branch discriminator. Then, a window-based Dual Local Enhancement (DleWin) module, better suited for underwater image enhancement, is proposed, using WGAN-GP loss, Uloss, and L1 loss to guide our network training.

[0083] The construction of an underwater image enhancement model includes the following steps:

[0084] (1) The underwater image enhancement model is defined as consisting of three parts: a dual local enhancement module, a generator network structure, and a dual-branch discriminator structure.

[0085] We propose a window-based Dual Local Enhancement (DleWin) module, a Transformer module better suited for underwater image enhancement tasks. The DleWin module implements a self-attention mechanism, effectively capturing long-range dependencies. Furthermore, local features are crucial for underwater image enhancement; therefore, we incorporate CNNs into the DleWin module using both serial and parallel approaches for local enhancement. The generator is built upon this DleWin module. Secondly, we propose a method combining convolutional neural networks and Transformer components. Since Transformers excel at capturing long-range dependencies and extracting original information, while CNNs excel at extracting local features, the generator is based on a CNN-based UNet network, employing the DleWin Transformer module for extracting both original and global information. Finally, we propose a GAN with a dual-branch discriminator, including a feature branch for preserving image features and enhancing contrast, and a color branch for color correction to generate more realistic colors. Finally, we implement the discriminator using stacked convolutional layers, and employ Wasserstein GAN (WGAN-GP) loss with gradient penalty to guide the training of the feature branch and underwater exponential loss (Uloss) to guide the training of the color branch. Based on these three designs, our proposed method achieves state-of-the-art performance in underwater image enhancement tasks.

[0086] (2) Construct a window-based dual local enhancement module: Construct a window-based dual local enhancement module suitable for underwater image enhancement;

[0087] A1) The window-based dual local enhancement module consists of two parts: the first part is a self-attention module for extracting features and a self-attention feature map enhancement module; the second part is a local enhancement module, whose input is the original image and whose output is the feature map fused by the self-attention feature map enhancement module and the local enhancement module.

[0088] A2) The self-attention module for feature extraction is set as follows:

[0089] like Figure 3As shown, the feature map is stretched to a 1D shape using an 8×8 window as the unit. After layer normalization, it is input into three different learnable parameter matrices to obtain query Q, key K, and value V vectors. Then, query Q and key K are multiplied, relative position encoding B is added, and zero-mean normalization is performed to obtain the attention matrix Attn. Finally, Attn is activated by Softmax and multiplied with value V to output the attention feature map.

[0090] A3) The self-attention feature map enhancement module is set as follows:

[0091] The self-attention feature map is normalized by layer, then linearly projected using a 1×1 convolutional kernel to reshape it into a 2D shape. After passing through a 3×3 convolutional layer, it is stretched into a 1D shape. Finally, it is linearly projected to the same dimension as the input feature using a 1×1 convolutional kernel to output the self-attention enhanced feature map.

[0092] A4) Configure the local enhancement module as follows:

[0093] The input to the local enhancement module is the original image, which is processed by two Conv+BatchNorm+GELU modules, where the size of the convolution kernel is 3×3, and the output is a local enhancement feature map.

[0094] A5) The self-attention enhancement feature map and the local enhancement feature map are fused to obtain a self-attention feature map with dual local enhancement.

[0095] (3) Constructing the generator network structure: Based on the Encoder-Decoder framework, Inception, Bottlenect and Fusion modules were added to construct the generator network structure.

[0096] We effectively combine Transformer and Convolutional Neural Networks (CNNs). We believe that CNNs are superior to Transformers in multi-scale feature extraction. Therefore, we use CNNs in the Encoder and Decoder parts of multi-scale feature extraction and reconstruction, which effectively reduces edge blurring and preserves more details. We use Transformers in the Inception and Bottlenecks parts because they excel at extracting both raw and global information from the image. Integrating the fully extracted global information from the Transformer into each feature scale is particularly effective in addressing color cast issues in underwater image degradation.

[0097] B1) The generator network structure is set to include four parts: the first part is the Inception part which extracts the original features; the second part is the Encoder-Decoder framework; the third part is the Bottlenect part which extracts global information; and the fourth part is the Fusion part which fuses global information into each scale. Its input is the original image, and its output is the image enhanced by the generator network.

[0098] The input image first passes through the Inception part, where the original features are extracted and then output to the Encoder part in the second part. The Encoder part extracts and encodes the features and then outputs them to the Bottlenect part. Each layer of the Encoder part has skip connections to the Decoder part. The Bottlenect part extracts global information and outputs it to the Fusion part and the Decoder part. The Fusion part incorporates the global information into each scale of the Decoder part. The Decoder part decodes and reconstructs the features and then outputs the enhanced image.

[0099] B2) The Inception section is configured as follows:

[0100] A window-based dual local enhancement module is used to extract features from the original image and output the extracted original features.

[0101] B3) Configure the Encoder-Decoder framework as follows:

[0102] like Figure 4 As shown, the Encoder-Decoder framework consists of an Encoder part and a Decoder part. The Encoder part has 5 encoding layers, which perform multi-scale feature extraction on the original features initially extracted by the Inception part. The first layer contains only one convolutional layer, and each of the other layers contains a Conv+BatchNorm+ReLU module. The parameters of all convolutional layers are size=4×4, stride=2, and padding=1, which can downsample while extracting features.

[0103] Finally, the extracted 512×8×8 shape is output to the Bottlenect part, and the features extracted from each layer are passed to the corresponding layer of the Decoder part through skip connections;

[0104] The Decoder part has 5 decoding layers, which receive global information extracted from the Bottlenect part. The structural design is the same as that of the Encoder part, except that only the first layer adds the Tanh activation function, and all the convolutional layers are replaced with transposed convolutional layers, which can reconstruct features while upsampling.

[0105] The output image after five decoding layers is an enhanced image with a shape of 3×256×256.

[0106] B4) The Bottlenect section is set as follows:

[0107] Two DleWin modules were used. When the features extracted by the Inception part were downsampled to an 8×8 size by the Encoder, which is the same as the window size of the DleWin module, the Transformer module extracted the global information.

[0108] B5) The Fusion section is set as follows:

[0109] Global information extracted from Bottlenect, such as overall lighting and layout, is integrated into each scale.

[0110] The global information first passes through a 1×1 convolutional layer to adjust the channels of the global information to correspond with the Decoder. Then, by copying and reshaping the global information, the shape of the output fused information is finally the same as the feature map of the corresponding layer of the Decoder.

[0111] (4) Construct a dual-branch discriminator structure: Construct a dual-branch discriminator containing feature branches and color branches.

[0112] The feature branch preserves image features and enhances contrast, while the color branch corrects colors, generating more realistic colors and mitigating the domain shift problem between the synthesized underwater image and the real underwater image. Using a dual-branch discriminator allows the generator to focus on different aspects of the image.

[0113] C1) The dual-branch discriminator structure consists of two parts: the first part is a feature branch discriminator that preserves the image content features, and the second part is a color branch discriminator. Its inputs are the enhanced image and the ground truth image, and its outputs are the adversarial map and underwater index map corresponding to the enhanced image and the ground truth image, respectively.

[0114] C2) The feature branch discriminator is set as follows:

[0115] PatchGAN is used. After preserving the image content through a convolutional layer, three layers of Conv+BatchNorm+Leaky-ReLU modules and one convolutional layer are stacked to distinguish between real and fake images. A 30×30 adversarial graph is generated for evaluation. Finally, the average of the adversarial graphs is used as the generator loss to promote the generator to generate realistic images.

[0116] C3) The color branch discriminator is set as follows: PatchGAN is used, and 5 layers of CBL modules and one layer of convolution are directly stacked to distinguish whether the image belongs to the underwater scene. A 15×15 underwater index map is generated to evaluate the intensity of underwater attributes and promote the generator to generate colors that match the image in the air.

[0117] The adversarial graph generated by the C4 feature branch discriminator and the underwater index graph generated by the color branch discriminator respectively guide the generator training.

[0118] The third step is training the underwater image enhancement model: input the preprocessed underwater image into the underwater image enhancement model for training.

[0119] (1) Set the objective function of the dual-branch discriminator to use Wasserstein GAN with gradient penalty, which is used to eliminate the gradient vanishing problem of standard GAN and guide the stable training of feature branches. Introduce underwater exponential loss Uloss to guide the training of color branches.

[0120] (2) The objective function of the generator is set as feature branch discriminator loss to preserve image content and underwater index loss to generate more realistic colors. Finally, L1 loss is used to eliminate blur.

[0121] In the initial training phase, the generator is not given underwater index loss, but the color branch discriminator is fully trained, i.e., underwater index loss is added to the generator after epoch=30.

[0122] (3) In the training of the underwater image enhancement model, the dual-branch discriminator is trained first, and then the generator is trained. In the training process, the dual-branch discriminator is trained first and then the generator is trained in each iteration.

[0123] (4) Training the two-branch discriminator:

[0124] The original image is input into the generator to generate an enhanced image. Then, the enhanced image and the ground truth are concatenated with the original image by channel to obtain two input images, which are then input into a dual-branch discriminator. After passing through the feature branch and the color branch, the corresponding adversarial map and underwater index map are output.

[0125] The gradient penalty term of the feature branch discriminator is calculated for the two adversarial graphs output by the feature branch. The loss is calculated based on the adversarial graphs and the gradient penalty term. The feature branch discriminator is then trained and its parameters are updated by backpropagation.

[0126] The underwater index loss is directly calculated based on the underwater index map, and the color branch discriminator is trained and updated by error backpropagation.

[0127] (5) Training the generator:

[0128] After the original image is input into the generator, the original features are extracted by the Inception part, which contains one window-based dual local enhancement module. Then, the image features are extracted layer by layer by the Encoder part, which contains five coding layers. The output is sent to the Bottlenect part, which contains two window-based dual local enhancement modules, to provide global information. The Fusion part integrates the global information into each scale of the Decoder part. Finally, the Decoder part, which contains five decoding layers, receives the global information and the fused information, and reconstructs the features layer by layer to output the enhanced image.

[0129] The enhanced image is input into the discriminator to obtain the adversarial graph and the underwater index graph. The total loss of the generator is calculated by combining the L1 loss, and the generator is trained and updated by backpropagation of error.

[0130] (6) Use 6000 images for training, with 32 images per batch. Use all images for training once per epoch, for a total of 188 batches and 200 epochs.

[0131] The fourth step is to acquire underwater images to be enhanced: acquire underwater images to be enhanced and preprocess them, adjusting the size of all images to 256×256 using bicubic interpolation.

[0132] The fifth step is underwater image enhancement: the preprocessed underwater image to be enhanced is input into the trained underwater image enhancement model to obtain the enhanced underwater image.

[0133] like Figure 5 As shown, it presents a comparison graph of the present invention with other methods on paired datasets. From Figure 5 As can be seen, using ground truth as a reference, some methods used for comparison offer limited quality improvement, while others, although showing significant quality improvement, may result in over-enhancement or incorrect color correction. The image color balance enhanced by our proposed model exhibits higher contrast and better visual effects. In paired image tests, this invention demonstrates excellent performance, almost identical to the ground truth. Table 1 shows that our invention achieves the best MSE, PSNR, and SSIM on the validation set Val test group and the underwater_dark and underwater_imagenet test groups, respectively. Furthermore, it also achieves the best average performance across all paired image test groups. This is attributed to the powerful learning capability of our model's dual local enhancement module.

[0134] like Figure 6 As shown, it presents a comparison graph of the present invention with other methods on unpaired datasets. From Figure 6As can be seen, the present invention achieves the best results in terms of both color correction and blur removal, resulting in enhanced images with richer colors, higher contrast, and clearer details.

[0135] The numerical comparison results are shown in Table 2. As can be seen from the table, our method performs better. The average results of the four test groups show that our proposed method outperforms all other methods in UIQM, NIQE, BRISQUE, and FRIQUEE, with UIQM at least 0.3175 higher, NIQE at least 5.9267 lower, BRISQUE at least 8.1168 lower, and FRIQUEE at least 2.4478 higher. Our proposed method's performance in UCIQE, Entropy, and underwater index needs improvement, as it only achieved third place in all of them.

[0136] Table 1. Test results of the present invention and other methods on paired image test sets using reference metrics.

[0137]

[0138]

[0139] Table 2 shows the test results of this invention and other methods on an unpaired image test set using non-reference metrics.

[0140]

[0141]

[0142] Table 3 shows the runtime of our invention compared to other methods. As can be seen from the table, our proposed method requires less time and processes images faster. Only GC, FUnieGAN, and UWCNN are faster than our method, which is because these three methods primarily prioritize time performance. Our proposed method achieves the best enhancement effect while also reaching a satisfactory level of efficiency. Furthermore, its runtime meets the requirements of real-time tasks.

[0143] Table 3 shows runtime on Intel(R) Core i5-9th CPU and GeForce RTX 3090 GPU.

[0144]

[0145] In this section, we used several downstream application test cases, including the SIFT keypoint matching algorithm, the Canny edge detection algorithm, and underwater target detection, to test the proposed model from the perspectives of basic features and target detection. We only selected a few representative methods for comparison. Figure 7 As shown, the SIFT algorithm only matches a small number of keypoints in the original image, while other enhancement methods have improved the number of matched keypoints. In the image enhanced by our proposed model, salient features are extracted and a large number of exact matches appear. Figure 8 As shown, the Canny edge detection algorithm detects relatively few edges in the original image, while more edge information is detected in the image enhanced by other methods. In the image enhanced by the method presented in this paper, Canny detects even more edge information. This demonstrates that underwater image enhancement processing is very helpful for feature extraction tasks. Figure 9 As shown, we use the Yolov5 model for underwater target detection, trained on a dataset containing 300 labeled images. It can be seen that compared to the original image and other methods, the proposed enhancement algorithm effectively improves target detection accuracy.

[0146] Application testing improvement effect, such as Figure 10 As shown, this paper demonstrates improvements in all methods, including SIFT feature point matching, Canny edge detection, and underwater target detection using the Yolov5 model, with better results than other methods.

[0147] The foregoing has shown and described the basic principles, main features, and advantages of the present invention. Those skilled in the art should understand that the present invention is not limited to the above embodiments. The embodiments and descriptions in the specification are merely principles of the invention. Various changes and modifications can be made to the invention without departing from its spirit and scope, and all such changes and modifications fall within the scope of the claimed invention. The scope of protection claimed by the appended claims and their equivalents is defined.

Claims

1. An underwater image enhancement method based on Transformer and Generative Adversarial Networks, characterized in that, Includes the following steps: 11) Acquisition and preprocessing of underwater images: Acquire underwater images and perform preprocessing; 12) Construction of underwater image enhancement model: An underwater image enhancement model is constructed based on Transformer and generative adversarial network; The construction of the underwater image enhancement model includes the following steps: 121) The underwater image enhancement model is defined as consisting of three parts: a dual local enhancement module, a generator network structure, and a dual-branch discriminator structure; 122) Construct a window-based dual local enhancement module: Construct a window-based dual local enhancement module suitable for underwater image enhancement; 123) Constructing the generator network structure: Based on the Encoder-Decoder framework, the Inception, Bottlenect, and Fusion modules are added to construct the generator network structure; 124) Construct a two-branch discriminator structure: Construct a two-branch discriminator containing a feature branch and a color branch; 13) Training the underwater image enhancement model: Input the preprocessed underwater image into the underwater image enhancement model for training; 14) Acquisition of underwater images to be enhanced: Acquire underwater images to be enhanced and preprocess them. Adjust the size of all images to 256×256 using bicubic interpolation. 15) Underwater image enhancement: Input the preprocessed underwater image to be enhanced into the trained underwater image enhancement model to obtain the enhanced underwater image.

2. The underwater image enhancement method based on the Transformer and the generative adversarial network according to claim 1, characterized in that, The construction of the window-based dual local enhancement module includes the following steps: 21) The window-based dual local enhancement module consists of two parts: the first part is a self-attention module for extracting features and a self-attention feature map enhancement module; the second part is a local enhancement module, whose input is the original image and whose output is the feature map fused by the self-attention feature map enhancement module and the local enhancement module. 22) The self-attention module for feature extraction is set as follows: Using an 8×8 window as a unit, the feature map is stretched to a 1D shape. After layer normalization, it is input into three different learnable parameter matrices to obtain query Q, key K, and value V vectors. Then, query Q and key K are multiplied, relative position encoding B is added, and zero-mean normalization is performed to obtain the attention matrix Attn. Finally, Attn is activated by Softmax and multiplied with value V to output the attention feature map. 23) The self-attention feature map enhancement module is set as follows: The self-attention feature map is normalized by layer, then linearly projected using a 1×1 convolutional kernel to reshape it into a 2D shape. After passing through a 3×3 convolutional layer, it is stretched into a 1D shape. Finally, it is linearly projected to the same dimension as the input feature using a 1×1 convolutional kernel to output the self-attention enhanced feature map. 24) The local enhancement module is set as follows: The input to the local enhancement module is the original image, which is processed by two Conv+BatchNorm+GELU modules, where the size of the convolution kernel is 3×3, and the output is a local enhancement feature map. 25) The self-attention enhancement feature map and the local enhancement feature map are fused to obtain a self-attention feature map with dual local enhancement.

3. The underwater image enhancement method based on the Transformer and the generative adversarial network according to claim 1, characterized in that, The construction of the generator network structure includes the following steps: 31) The generator network structure is set to include four parts: the first part is the Inception part which extracts the original features; the second part is the Encoder-Decoder framework; the third part is the Bottlenect part which extracts global information; and the fourth part is the Fusion part which fuses global information to each scale. Its input is the original image and its output is the image enhanced by the generator network. The input image first passes through the Inception part, where the original features are extracted and then output to the Encoder part in the second part. The Encoder part extracts and encodes the features and then outputs them to the Bottlenect part. Each layer of the Encoder part has skip connections to the Decoder part. The Bottlenect part extracts global information and outputs it to the Fusion part and the Decoder part. The Fusion part incorporates the global information into each scale of the Decoder part. The Decoder part decodes and reconstructs the features and then outputs the enhanced image. 32) The Inception section is set up as follows: A window-based dual local enhancement module is used to extract features from the original image and output the extracted original features. 33) Configure the Encoder-Decoder framework as follows: The Encoder-Decoder framework consists of an Encoder part and a Decoder part. The Encoder part has 5 encoding layers, which perform multi-scale feature extraction on the raw features initially extracted by the Inception part. The first layer contains only one convolutional layer, and each of the other layers contains a Conv+BatchNorm+ReLU module. The parameters of all convolutional layers are size=4×4, stride=2, and padding=1, which can downsample while extracting features. Finally, the extracted 512×8×8 shape is output to the Bottlenect part, and the features extracted from each layer are passed to the corresponding layer of the Decoder part through skip connections; The Decoder part has 5 decoding layers, which receive global information extracted from the Bottlenect part. The structural design is the same as that of the Encoder part, except that the first layer adds the Tanh activation function and all convolutional layers are replaced with transposed convolutional layers, which can reconstruct features while upsampling. The output image after five decoding layers is an enhanced image with a shape of 3×256×256. 34) The Bottlenect part is set as follows: Two DleWin modules were used. When the features extracted by the Inception part were downsampled to an 8×8 size by the Encoder, which is the same as the window size of the DleWin module, the Transformer module extracted the global information. 35) Configure the Fusion section as follows: Global information extracted from Bottlenect, such as overall lighting and layout, is integrated into each scale. The global information first passes through a 1×1 convolutional layer to adjust the channels of the global information to correspond with the Decoder. Then, by copying and reshaping the global information, the shape of the output fused information is finally the same as the feature map of the corresponding layer of the Decoder.

4. The underwater image enhancement method based on the Transformer and the generative adversarial network according to claim 1, characterized in that, The construction of the dual-branch discriminator structure includes the following steps: 41) The dual-branch discriminator structure consists of two parts: the first part is a feature branch discriminator that preserves the image content features, and the second part is a color branch discriminator. Its inputs are the enhanced image and the ground truth image, and its outputs are the adversarial map and underwater index map corresponding to the enhanced image and the ground truth image, respectively. 42) The feature branch discriminator is set as follows: PatchGAN is used. After preserving the image content through a convolutional layer, three layers of Conv+BatchNorm+Leaky-ReLU modules and one convolutional layer are stacked to distinguish between real and fake images. A 30×30 adversarial graph is generated for evaluation. Finally, the average of the adversarial graphs is used as the generator loss to promote the generator to generate realistic images. 43) The color branch discriminator is set as follows: PatchGAN is used, and 5 layers of CBL modules and one layer of convolution are directly stacked to distinguish whether the image belongs to the underwater scene. A 15×15 underwater index map is generated to evaluate the intensity of underwater attributes and promote the generator to generate colors that match the image in the air. 44) The adversarial graph generated by the feature branch discriminator and the underwater index graph generated by the color branch discriminator respectively guide the generator training.

5. The underwater image enhancement method based on the Transformer and the generative adversarial network according to claim 1, characterized in that, The training of the underwater image enhancement model includes the following steps: 51) Set the objective function of the dual-branch discriminator to use Wasserstein GAN with gradient penalty, which is used to eliminate the gradient vanishing problem of standard GAN and guide the stable training of feature branches. Introduce underwater exponential loss Uloss to guide the training of color branches. 52) Set the objective function of the generator to feature branch discriminator loss to preserve image content and underwater index loss to generate more realistic colors, and finally use L1 loss to eliminate blur; In the initial training phase, the generator is not given underwater index loss, but the color branch discriminator is fully trained, i.e., underwater index loss is given to the generator after epoch=30. 53) In the training of the underwater image enhancement model, the dual-branch discriminator is trained first, and then the generator is trained. During the training process, the dual-branch discriminator is trained first and then the generator is trained in each iteration. 54) Training the bi-branch discriminator: The original image is input into the generator to generate an enhanced image. Then, the enhanced image and the ground truth image are concatenated with the original image by channel to obtain two input images, which are then input into a dual-branch discriminator. After passing through the feature branch and the color branch, the corresponding adversarial map and underwater index map are output respectively. The gradient penalty term of the feature branch discriminator is calculated for the two adversarial graphs output by the feature branch. The loss is calculated based on the adversarial graphs and the gradient penalty term. The feature branch discriminator is then trained and its parameters are updated by backpropagation. The underwater index loss is directly calculated based on the underwater index map, and the color branch discriminator is trained and updated by backpropagation of error. 55) Training generator: After the original image is input into the generator, the original features are extracted by the Inception part, which contains one window-based dual local enhancement module. Then, the image features are extracted layer by layer by the Encoder part, which contains five coding layers. The output is sent to the Bottlenect part, which contains two window-based dual local enhancement modules, to provide global information. The Fusion part integrates the global information into each scale of the Decoder part. Finally, the Decoder part, which contains five decoding layers, receives the global information and the fused information and reconstructs the features layer by layer to output the enhanced image. The enhanced image is input into the discriminator to obtain the adversarial map and the underwater index map. The total loss of the generator is calculated by combining the L1 loss, and the generator is trained and updated by backpropagation of error. 56) Use 6000 images for training, with 32 images per batch. Use all images for training once per epoch, for a total of 188 batches, and end training after 200 epochs.

Images