Underwater image enhancement method, system, device and medium based on transformer network

Abstract

The present application relates to the field of image processing, and discloses a method, system, device and medium for underwater image enhancement based on a transformer network. The method includes constructing an enhancement network and training it, obtaining a shallow feature image and feeding it into the N-1th layer encoding module to obtain a deep feature image; fusing the output image of the N-1th layer encoding module with the output image of the N-1th layer decoding module through the N-1th layer feature fusion module, and outputting a residual image after inputting the N-1th layer decoding module; the feature fusion module is used to obtain a probability map of the output image of the corresponding layer encoding module, preliminarily fuse the probability map with the output image of the previous layer decoding module, and fuse the output image of the N-1th layer encoding module with the fusion result in the channel dimension; and adding the residual image to the input image element by element to obtain an enhanced underwater image. The present application can effectively solve the problems of color deviation, poor contrast, and limited enhancement effect in complex scenes in underwater images.

Description

Technical Field

[0001] The present application relates to the field of image processing, and specifically to a method, system, device and medium for underwater image enhancement based on a transformer network. Background Art

[0002] As an image processing method, underwater image enhancement technology has the potential to be widely used in the field of underwater detection and observation. Underwater images are affected by factors such as water quality and light attenuation, and often have problems such as blur and noise, which limit the recognition and analysis of underwater targets. Underwater image enhancement technology can improve the quality of underwater images by enhancing contrast and reducing noise, thereby improving the accuracy and efficiency of target recognition; at the same time, underwater image enhancement technology can be applied to marine resource exploration, underwater archaeology, underwater search and rescue and other fields, providing strong support for underwater detection work.

[0003] In recent years, the rapid development of deep learning technology has provided powerful tools for image recognition and enhancement. However, traditional deep learning methods have certain limitations in underwater image enhancement. Existing image enhancement methods have limitations in underwater image enhancement, such as large color deviation, poor contrast, and poor enhancement effect in complex scenes. Summary of the invention

[0004] The technical problem to be solved by the present application is to overcome the deficiencies of the prior art and provide an underwater image enhancement method, system, device and medium based on a transformer network, which can effectively solve the problems of color deviation, poor contrast and limited enhancement effect in complex scenes in underwater images.

[0005] To achieve the above-mentioned purpose, the first aspect of the present application provides an underwater image enhancement method based on a transformer network, constructing and training an underwater image enhancement network based on a transformer network, wherein the underwater image enhancement network based on the transformer network adopts an encoder-decoder architecture, including an N-layer encoding module, an N-layer decoding module and an N-layer feature fusion module, and constructing the enhancement network includes the following steps:

[0006] Step S1, obtaining a shallow feature image of an input image, wherein the input image is an underwater image to be enhanced;

[0007] Step S2, feeding the shallow feature image into an encoding module to obtain a deep feature image;

[0008] Step S3, fusing the output image of the encoding module of the N-1th layer and the output image of the decoding module of the N-1th layer through the feature fusion module of the N-1th layer to serve as the input image of the decoding module of the N-1th layer, and the decoding module of the N-1th layer outputs a residual image;

[0009] In the N-layer encoding module, there are N conv-transformer modules connected in sequence, and in the N-layer decoding module, there are N conv-transformer modules connected in sequence, where N is an integer greater than 1;

[0010] The N-1th layer feature fusion module is used to obtain the probability map of the output image of the coding module of the N-1th layer, perform preliminary fusion by weighting the product of the probability map and the output image of the decoding module of the previous layer, and perform channel-dimensional fusion of the output image of the coding module of the N-1th layer and the result after preliminary fusion;

[0011] The conv-transformer module is used to perform deep feature extraction on the input image features to obtain a deep feature image. The conv-transformer module extracts in the channel dimension branch and the space dimension branch respectively, and outputs after fusing the extraction results of the two dimension branches;

[0012] The N-1th layer feature fusion module performs feature fusion on the image output by the conv-transformer module in the N-1th layer encoding module and the image output by the conv-transformer module in the Nth layer decoding module;

[0013] Step S4, adding the residual image output in step S3 and the input image obtained in step S1 element by element to obtain an enhanced underwater image.

[0014] Furthermore, in step S3, the layers of the encoding module are cascaded using downsampling layers, and the encoding module is used to extract a deep feature image from the underwater image to be enhanced to obtain a deep feature image;

[0015] The layers of the decoding module are cascaded using an upsampling layer through the feature fusion module; the decoding module is used to decode the input image fused by the feature fusion module, and obtain the residual image after decoding by N layers of the decoding modules.

[0016] Furthermore, the conv-transformer module extracts the channel dimension branch features and performs the input tensor Perform layer normalization operation, which is expressed as:

[0017] ;

[0018] in, is the normalized output tensor, is the mean of all values ​​in the tensor, is the standard deviation of all values ​​in the tensor. For the input tensor , whose size is , is the height of the input image, is the width of the input image, is the channel dimension of the input image, Represents the separation of different dimensions of a tensor;

[0019] For the input tensor After layer normalization, learnable parameters are introduced and To normalize the output tensor Scaling and offset operations are performed, expressed as:

[0020] ;

[0021] in, To add the output tensor after scaling and offset operations, and are learnable parameters.

[0022] Furthermore, the process of extracting the channel dimension branch features includes obtaining a channel attention map and performing , , The components are transformed in dimension, expressed as:

[0023] ;

[0024] in, Represents the query, and the size after transforming the tensor dimension is , Represents the key value, and the size after transforming the tensor dimension is , Represents a value vector, and after transforming the tensor dimension, the size is , Represents the operation of transforming the dimension of a tensor;

[0025] The channel attention map is calculated as follows:

[0026] ;

[0027] in, is a learnable scaling parameter, express Activation function, Represents the self-attention map of the channel dimension, with a size of ;

[0028] The channel dimension branch output is expressed as:

[0029] ;

[0030] in, Represents the channel self-attention map Sum value vector The vector product of is the final output tensor of the channel dimension branch, represents the self-attention map of the channel dimension, Represents a vector of values.

[0031] Furthermore, during the fusion process of the conv-transformer module, the output image of the spatial dimension branch and the output image of the channel dimension branch are element-wise added and then fused through a point convolution operation. The fused image of the spatial dimension branch and the channel dimension branch is suppressed and output using the GDFN module.

[0032] Furthermore, the fusion process of the N-1th layer feature fusion module includes sequentially performing window segmentation, Function activation and window fusion operation to obtain the probability map of the output image, and also include pixel repetition operation on the output image of the conv-transformer module in the N-th layer decoding module, and multiplication weighted operation with the probability map to output a fused image, and also include channel dimension fusion of the output of the N-1-th layer encoding module and the fused image to output a deep fused image, and the deep fused image is used as the input image of the N-1-th layer decoding module;

[0033] Obtaining the probability map of the output image includes segmenting the output image of the conv-transformer module in the N-1th layer encoding module according to the set window size, and using The function is activated to obtain a probability map of each window, and the probability map of each window is fused to obtain the probability map of the output image.

[0034] Furthermore, the loss function is expressed as:

[0035] ;

[0036] Where n and m are the height and width of the image, For the real image Line The value of the column pixel, is the first Line The value of the column pixel, The loss value of the calculated loss function.

[0037] To achieve the above-mentioned purpose, the second aspect of the present application provides an underwater image enhancement system based on a transformer network, the enhancement system comprising:

[0038] An acquisition module, used to acquire a shallow feature image of an input image, wherein the input image is an underwater image to be enhanced;

[0039] N-layer encoding module, used for performing deep feature extraction on the shallow feature image;

[0040] The feature fusion module is used to obtain the probability map of the output image of the encoding module of the corresponding layer, perform preliminary fusion by weighting the product of the probability map and the output image of the decoding module of the previous layer, and perform channel dimension fusion on the output image of the encoding module and the result after preliminary fusion;

[0041] N conv-transformer modules in the N-layer encoding module are connected in sequence, and N conv-transformer modules in the N-layer decoding module are connected in sequence, where N is an integer greater than 1;

[0042] The conv-transformer module is used to perform deep feature extraction on the input image features to obtain a deep feature image. The conv-transformer module extracts in the channel dimension branch and the space dimension branch respectively, and outputs after fusing the extraction results of the two dimension branches;

[0043] N-layer decoding module, used to decode and output the fused feature map and the image output by the previous layer decoding module;

[0044] The output module is used to add the residual image output by the feature fusion module to the input image in the acquisition module element by element, and output an enhanced underwater image.

[0045] To achieve the above-mentioned objectives, the third aspect of the present application provides an underwater image enhancement device based on a transformer network, including a processor and a memory, wherein a computer program is stored in the memory, and when the computer program is executed by the processor, the underwater image enhancement method based on the transformer network as described above is implemented.

[0046] To achieve the above-mentioned objectives, the fourth aspect of the present application provides a computer-readable storage medium storing a computer program. When the computer program is executed by a processor, it implements the underwater image enhancement method based on the transformer network as described above.

[0047] After adopting the above technical solution, the present application has the following beneficial effects compared with the prior art:

[0048] In the present application, the conv-transformer module of each layer is used to extract deep features, which can effectively obtain the information contained in the input image. The conv-transformer module divides the self-attention into two branches, the channel dimension and the spatial dimension, and calculates them separately and then merges them. The channel dimension enhances the network's perception of the color channel, and the spatial dimension enhances the network's attention to local features. The combination of the two not only retains the ability of the transformer network to capture long-distance dependencies but also integrates the convolution's fine processing of local details; the feature fusion module obtains the local detail texture information of the image and integrates it into the upsampling process to reduce the loss of detail information during the upsampling process, and then fuses the output of the encoding module with the upsampled features in the channel dimension, which can retain the overall semantic information of the image. The feature fusion module is used to replace upsampling and jump connections, which can effectively reduce information loss during network transmission.

[0049] The specific implementation methods of the present application are further described in detail below in conjunction with the accompanying drawings. BRIEF DESCRIPTION OF THE DRAWINGS

[0050] The accompanying drawings are part of this application and are used to provide a further understanding of this application. The illustrative embodiments of this application and their descriptions are used to explain this application, but do not constitute an improper limitation on this application. Obviously, the drawings described below are only some embodiments. For ordinary technicians in this field, other drawings can be obtained based on these drawings without creative work.

[0051] In the drawings of the specification:

[0052] Figure 1 is a schematic diagram of the overall network in this specific implementation mode;

[0053] Figure 2 is a schematic diagram of a conv-transformer module in this specific implementation mode;

[0054] Figure 3 is a schematic diagram of a feature fusion module in this specific implementation mode;

[0055] Figure 4 It is an underwater image with blur problem and the enhancement effect of different methods;

[0056] Figure 5 It is an underwater image with color cast problem and the enhancement effect of different methods;

[0057] Figure 6 It is an underwater image with complex scenes and the enhancement effects of different methods. DETAILED DESCRIPTION

[0058] In order to make the purpose, technical solutions and advantages of the embodiments of the present application clearer, the technical solutions in the embodiments will be clearly and completely described below in conjunction with the drawings in the embodiments of the present application. The following embodiments are used to illustrate the present application but are not used to limit the scope of the present application.

[0059] At present, underwater image enhancement technology is mainly divided into three categories: underwater enhancement algorithm based on physical model, underwater enhancement algorithm based on non-physical model and underwater enhancement algorithm based on deep learning.

[0060] Underwater enhancement algorithms based on physical models regard underwater image enhancement as the inverse problem of image degradation. These methods establish an underwater image degradation model, estimate parameters based on prior knowledge, and invert this degradation problem to obtain enhanced images. However, the underwater environment is complex, a single physical model is difficult to describe the real underwater situation, and changes in water quality make it difficult to obtain its parameters.

[0061] Underwater enhancement algorithms based on non-physical models mainly rely on statistical and data-driven methods to improve image quality by adjusting the pixel values ​​of images. However, this method does not take into account the imaging mechanism of underwater images and often easily causes excessive contrast enhancement and color distortion.

[0062] The underwater enhancement algorithm based on deep learning enhances underwater images by using technologies such as deep convolutional neural networks (CNNs) and generative adversarial networks (GANs). However, the receptive field of deep convolutional neural networks is limited to local areas, and their understanding of global context is weak, and it is difficult to capture long-distance dependent features. Generative adversarial networks are somewhat insufficient in collecting global information, and the enhancement effect is limited.

[0063] Based on this, see Figure 1 The present application provides an underwater image enhancement method based on a transformer network, constructs and trains an underwater image enhancement network based on a transformer network, and the underwater image enhancement network based on the transformer network adopts an encoder-decoder architecture, including an N-layer encoding module, an N-layer decoding module, and an N-layer feature fusion module. Constructing the enhancement network includes the following steps:

[0064] Step S1, obtaining a shallow feature image of an input image, where the input image is an underwater image to be enhanced;

[0065] Step S2, feeding the shallow feature image into the encoding module to obtain the deep feature image;

[0066] Step S3, fusing the output image of the N-1th layer encoding module and the output image of the N-1th layer decoding module through the N-1th layer feature fusion module to serve as the input image of the N-1th layer decoding module, and the N-1th layer decoding module outputs a residual image;

[0067] There are N conv-transformer modules connected in sequence in the N-layer encoding module, and there are N conv-transformer modules connected in sequence in the N-layer decoding module, where N is an integer greater than 1;

[0068] The N-1th layer feature fusion module is used to obtain the probability map of the output image of the N-1th layer encoding module, perform preliminary fusion by weighting the product of the probability map and the output image of the previous layer decoding module, and perform channel-dimensional fusion of the output image of the N-1th layer encoding module and the result of preliminary fusion;

[0069] The conv-transformer module is used to extract deep features from the input image features to obtain deep feature images. The conv-transformer module extracts in the channel dimension branch and the spatial dimension branch respectively, and outputs the fusion of the extraction results of the two dimension branches;

[0070] The feature fusion module performs feature fusion on the image output by the conv-transformer module in the N-1th layer encoding module and the image output by the conv-transformer module in the Nth layer decoding module;

[0071] Step S4, adding the residual image output in step S3 to the input image obtained in step S1 element by element to obtain an enhanced underwater image.

[0072] It should be noted that the execution subject of the enhancement method in this embodiment is an underwater image enhancement device based on a transformer network, which may be an electronic device, a component in an electronic device, an integrated circuit, or a chip. The electronic device may be a mobile electronic device or a non-mobile electronic device. Exemplarily, the mobile electronic device may be a mobile phone, a tablet computer, a laptop computer, a PDA, a vehicle-mounted electronic device, a wearable device, etc., and the non-mobile electronic device may be a server and a personal computer, etc., which are not specifically limited in this application. The following takes the execution subject as an example of a server to describe the underwater image enhancement method based on a transformer network in this embodiment.

[0073] In addition, the terms "first" and "second" are used for descriptive purposes only and should not be understood as indicating or implying relative importance or implicitly indicating the number of technical features indicated. Thus, a feature defined as "first" or "second" may explicitly or implicitly include one or more of the features. In the description of the embodiments of the present application, the meaning of "multiple" is two or more, unless otherwise clearly and specifically defined.

[0074] In this embodiment, an underwater image to be enhanced is acquired and preprocessed.

[0075] Specifically, we obtain the Underwater Image Enhancement Benchmark Dataset (UIEB) and the Enhanced Underwater Visual Perception Dataset (EVP). The UIEB dataset contains 890 image pairs, of which 800 image pairs are used as training sets and the remaining 90 images are used as test sets; the EUVP dataset contains image pairs of three scenes: underwater darkness, underwater image network, and underwater scenes. We use 2185 image pairs in underwater scenes, of which 2000 image pairs are used as training sets and the remaining 185 image pairs are used as test sets. All images are resized to the same size before training.

[0076] This embodiment is described based on a three-layer encoding module-decoding module architecture.

[0077] In step S3, the layers of the encoding module are cascaded using downsampling layers, and the encoding module is used to extract deep feature images of the underwater image to be enhanced to obtain deep feature images; the layers of the decoding module are cascaded using upsampling layers through a feature fusion module, and the decoding module is used to decode the input image fused by the feature fusion module, and obtain a residual image after decoding by N layers of decoding modules.

[0078] Specifically, the input is of size The image is first shallowly extracted through a convolution with a 3×3 convolution kernel, and the extracted features are output through a 3-layer encoding module-decoding module architecture.

[0079] The encoding modules of each layer are connected by downsampling. The downsampling adopts the pixel-unshuffling method, which obtains the original feature map of the input image, decomposes the spatial dimension in the feature map, and rearranges the channels to compress the information in the space into the channel to obtain a lower resolution but information-intensive feature map.

[0080] The decoded output of each layer of the decoding module is combined with the output of the encoding module of the corresponding layer, and the upsampling and jump connection are achieved through the feature fusion module. The feature fusion module obtains the local probability map output by the encoding module to collect the detailed texture information in the features, expands the resolution of the decoding module output and compresses the channel, restores the information in the channel to the spatial dimension, and multiplies the local probability map output by the encoding module with the output of the decoding module after the resolution is expanded to make the two fully integrated, and then output after fusion with the encoding module output in the channel dimension.

[0081] The output of the feature fusion module is channel compressed by a 3×3 convolution operation to obtain a size of The residual is added to the original input image to get the enhanced image.

[0082] It should be noted that each layer of encoding modules and each layer of decoding modules are composed of conv-transformer modules, and the conv-transformer modules are used to extract deep features of images input into the encoding modules and decoding modules of each layer to obtain deep feature images.

[0083] See also Figure 2 In this embodiment, in step S3, the conv-transformer module performs feature extraction in the channel dimension branch and the spatial dimension branch respectively, and outputs the extracted results of the two dimension branches after fusion.

[0084] Please refer to Figure 2 ,The feature extraction process of the channel dimension branch is described in detail below.

[0085] Specifically, first input tensor Perform layer normalization to alleviate the gradient vanishing and exploding problems during training. The layer normalization operation is expressed as follows:

[0086] ;

[0087] in, is the normalized output tensor, is the mean of all values ​​in the tensor, is the standard deviation of all values ​​in the tensor. For the input tensor , whose size is , is the height of the input image, is the width of the input image, is the channel dimension of the input image, Represents the separation of different dimensions of a tensor.

[0088] Specifically, for the input tensor After layer normalization, learnable parameters are introduced and To normalize the output tensor Scaling and offset operations are performed to maintain the expressiveness of the model, as shown below:

[0089] ;

[0090] in, To add the output tensor after scaling and offset operations, and are learnable parameters.

[0091] Specifically, the normalized values ​​are processed by point convolution and depth convolution to emphasize the context in space and channel, respectively, to obtain , , and Four output tensors; represented as:

[0092] ;

[0093] in, denotes convolution, and its superscript , , and Corresponding to different output tensors, subscript and Represents point convolution and deep convolution. , , As the channel dimension branch, channel self-attention calculation is performed. Convolution is performed as a branch in the spatial dimension.

[0094] In an achievable implementation, the channel dimension branch feature extraction process includes obtaining a channel attention map, , , The components are transformed in dimension, expressed as:

[0095] ;

[0096] in, Represents the query, and the size after transforming the tensor dimension is , Represents the key value, and the size after transforming the tensor dimension is , Represents a value vector, and after transforming the tensor dimension, the size is , Represents the operation of transforming the dimension of a tensor;

[0097] The channel attention map is calculated as follows:

[0098] ;

[0099] in, is a learnable scaling parameter, express Activation function, Represents the self-attention map of the channel dimension, with a size of .

[0100] In another achievable implementation, the channel dimension branch output is expressed as:

[0101] ;

[0102] in, Represents the channel self-attention map Sum value vector The vector product of is the final output tensor of the channel dimension branch, represents the self-attention map of the channel dimension, Represents a vector of values.

[0103] See also Figure 2 ,The feature extraction fusion process of the channel dimension branch and the spatial dimension branch is described in detail below.

[0104] In this embodiment, during the fusion process of the conv-transformer module, the output image of the spatial dimension branch and the output image of the channel dimension branch are element-wise added and then fused through a point convolution operation. The fused image of the spatial dimension branch and the channel dimension branch is suppressed and output using the GDFN module.

[0105] Please continue to see Figure 2 , during the spatial dimension branch extraction process, the output tensor After deep convolution operation, the image is output and the local features of the spatial dimension are extracted. During the fusion process of the conv-transformer module, the output image of the spatial dimension branch and the output image of the channel dimension branch are added element by element, and then fused through point convolution. For the image after element addition operation of the two dimensional branches and point convolution fusion, the GDFN module is used to suppress features with less information, and the image is output after focusing on more detailed image attributes.

[0106] It should be noted that the full English name of the GDFN module is Gated-Dconv Feed-Forward Network, which is translated into Chinese as Gated Deep Convolutional Feedforward Network, which can effectively improve the performance of the transformer module.

[0107] See also Figure 1 and Figure 3 In this embodiment, in step S3, the feature fusion module performs feature fusion on the image output by the conv-transformer module in the N-1th layer encoding module and the image output by the conv-transformer module in the Nth layer decoding module.

[0108] Specifically, the output of the first layer encoding module is characterized , the output feature size is , the output of the second layer encoding module is the feature , the output feature size is , the output of the third layer encoding module is characterized , the output feature size is .

[0109] Specifically, the output of the first layer decoding module is characterized , the output feature size is , the output of the second layer decoding module is characterized , the output feature size is , the output of the third layer decoding module is characterized , the output feature size is .

[0110] Specifically, the third layer encoding module and decoding module pass through a conv-transformer module, and the output is the feature The third layer feature fusion module combines the features and Features As input, feature fusion is performed and the output is used as the input of the third layer decoding module; the second layer feature fusion module takes the feature and Features As input, feature fusion is performed and the output is used as the input of the second layer decoding module; the first layer feature fusion module takes the feature and Features It is used as input for feature fusion and the output is used as the input of the first layer decoding module, which outputs the residual image.

[0111] Please refer to Figure 3In this embodiment, the feature fusion process of the N-1 layer feature fusion module includes sequentially performing window segmentation, Function activation and window fusion operations are performed to obtain the probability map of the output image, and pixel repetition operation is performed on the output image of the conv-transformer module in the N-th layer decoding module, and the fused image is output by multiplying the output image with the probability map by weighted operation, and the channel dimension fusion of the output of the N-1 layer encoding module and the fused image is performed to output a deep fused image, and the deep fused image is used as the input image of the N-1-th layer decoding module.

[0112] Obtain the probability map of the output image, including segmenting the output image of the conv-transformer module in the encoding module according to the set window size, and using the image of each window The function is activated to obtain the probability map of each window, and the probability maps of all windows are fused to obtain the probability map of the output image.

[0113] Specifically, in order to obtain the local probability map output by the encoding module, it is necessary to perform window cutting on the output image of the conv-transformer module in the encoding module, and divide the original image into multiple small blocks according to the set window size. The tensor of is re-partitioned, and the output size is ,in and is the position code of a single window, and the size of a single window is , which is used in the window fusion process. For a single window, use The activation function calculates the probability of each value in the tensor, obtains the probability map of a single window, that is, the local probability map of the entire input, and fuses all the local probability maps into windows based on the position encoding, and re-fuse them into a size of Probability plot of .

[0114] Specifically, for the features input by the decoding module, firstly, the pixel is repeated to convert the size The input feature size becomes , and the probability map output by the encoding module is multiplied and weighted to obtain the details and texture features of the image in the encoding module, and the two are fully fused through point convolution. Finally, the fusion result is fused with the output of the encoding module in the channel dimension, and point convolution is used to reduce the channel dimension size.

[0115] In this embodiment, the loss function is expressed as:

[0116] ;

[0117] Where n and m are the height and width of the image, For the real image Line The value of the column pixel, is the first Line The value of the column pixel, The loss value of the calculated loss function.

[0118] Specifically, this embodiment uses MSE as the loss function. MSE loss measures the quality of the model by calculating the square of the distance between the predicted value and the actual value. The closer the predicted value and the true value are, the smaller the mean square error between the two. The predicted value refers to the enhanced underwater image, and the true value refers to the true value of the original underwater image. The purpose of using MSE as the loss function is that the function curve of MSE is smooth, continuous, and can be differentiated everywhere, which is convenient for using the gradient descent algorithm. Moreover, as the error decreases, the gradient also decreases, and it can converge to the minimum value faster.

[0119] Based on the same inventive concept, the present application also provides an underwater image enhancement system based on a transformer network, the enhancement system comprising:

[0120] An acquisition module, used for acquiring a shallow feature image of an input image, where the input image is an underwater image to be enhanced;

[0121] N-layer encoding module, used to extract deep features from shallow feature images;

[0122] Conv-transformer module: The Conv-transformer module is used to extract the deep feature images of the input encoding modules of each layer. The Conv-transformer module extracts in the channel dimension branch and the spatial dimension branch respectively, and fuses the extraction results of the two dimension branches and outputs them;

[0123] The feature fusion module is used to obtain the probability map of the output image of the encoding module of the corresponding layer, and fuse the probability map with the output image of the decoding module of the previous layer in the channel dimension;

[0124] N conv-transformer modules are connected in sequence in the N-layer encoding module, and N conv-transformer modules are connected in sequence in the N-layer decoding module, where N is an integer greater than 1;

[0125] The conv-transformer module is used to extract deep features from the input image features to obtain deep feature images. The conv-transformer module extracts in the channel dimension branch and the spatial dimension branch respectively, and outputs the fusion of the extraction results of the two dimension branches;

[0126] The feature fusion module performs feature fusion on the image output by the conv-transformer module in the N-1th layer encoding module and the image output by the conv-transformer module in the Nth layer decoding module;

[0127] N-layer decoding module, used to decode and output the fused feature map and the image output by the previous layer decoding module;

[0128] The output module is used to add the residual image output from the feature fusion module to the input image in the acquisition module element by element, and output an enhanced underwater image.

[0129] Specifically, the enhanced network in this embodiment is built using the Pytorch library, and the training process uses the AdamW optimizer, combined with the cosine annealing learning strategy, and the initial learning rate is , gradually reduced to , Set to (0.9, 0.999). The number of transformer blocks used in the conv-transformer module in the 1-3 layer decoder-encoder is [4, 6, 6, 8], and the number of attention heads used is [1, 2, 4, 8]. During network training, Batchsize is set to 1 and epoch is set to 120, where Batchsize refers to the number of samples contained in a batch of data, and epoch refers to the process of completely training the model once on the entire training dataset.

[0130] To verify the enhancement effect of this embodiment, this embodiment is compared with two traditional underwater image enhancement methods and five underwater image enhancement methods based on deep learning. The two traditional underwater image enhancement methods are: Generalized Contrast (GC) and Dark Channel Prior (DCP). The underwater image methods based on deep learning are: WaterNet, U-shape, FUNie-GAN, SCNet and Ucolor.

[0131] Figure 4This is an experimental effect diagram of the embodiment and other algorithms after enhancing the blurred underwater image. In the underwater environment, especially in deeper waters, the light attenuation is serious, resulting in reduced image contrast. Image blur will make the edges of objects unclear and details lost. Figure 4 It can be seen that the GC algorithm enhances the contrast of the image, but over-enhancement will introduce halos and reduce image quality. The DCP algorithm effectively restores the clarity of the image, but it is overall dark and is not ideal in solving the color cast problem of the image. The WaterNet algorithm cannot effectively enhance the clarity of the image when facing blurred images. The U-shape algorithm cannot effectively improve the clarity of the image when facing complex texture images, and the detail information is seriously lost. The FUNie-GAN algorithm effectively enhances the clarity of the image, but when facing images with higher brightness, the brightness will be over-enhanced, resulting in loss of image details. The SCNet algorithm enhances the contrast of the image and improves the clarity of the image, but there is an over-enhancement phenomenon, resulting in color distortion. The Ucolor algorithm enriches the color of the image, but there is an over-enhancement situation, and the deblurring effect is not good. The present invention has achieved remarkable results in increasing image clarity and removing image artifacts.

[0132] Figure 5 The following is an experimental result diagram of the enhancement of underwater images with color shift by this embodiment and other algorithms. The color shift in underwater environment is caused by underwater lighting conditions and water characteristics. Color shift will increase the complexity of image processing and affect the naturalness and authenticity of the image. Figure 5 It can be seen that the GC algorithm enhances the image contrast, but excessive enhancement will make the overall image too bright and reduce the image quality. The DCP algorithm enhances the clarity of the image, but the improvement of the color cast problem of the image is not ideal, and introduces the problem of the overall image being too dark. The WaterNet algorithm cannot effectively improve the color cast problem and cannot effectively enhance the clarity of the image when facing blurred images. The U-shape algorithm cannot effectively process image details, especially when facing scenes with more complex texture features, and the detail information is seriously lost. The FUNie-GAN algorithm improves the clarity of the image, but its brightness adjustment is too strong, and there will be excessive noise in the bright areas. The SCNet algorithm cannot effectively improve the color cast problem of the image in details. The Ucolor algorithm has the problem of over-enhancing the image color, which often introduces new color deviations. The present invention is better in dealing with color cast problems, is more in line with reality, and the image is more natural.

[0133] Figure 6The following is an experimental result of enhancing underwater images of complex scenes using this embodiment and other algorithms. Complex scenes refer to underwater environments with multiple visual elements, different lighting conditions, complex structures and details. They are more difficult to process due to problems such as uneven lighting, color cast, dynamic water flow, and background interference. Figure 6 It can be seen that the GC algorithm enhances the contrast too high and has poor restoration of image details. The DCP algorithm has a darker overall image and cannot effectively enhance the dark parts of the image. The WaterNet algorithm cannot effectively enhance areas with complex textures. The U-shape algorithm cannot effectively perform comprehensive enhancement when facing complex scenes. The FUNie-GAN algorithm cannot effectively enhance the bright parts of the image. The SCNet algorithm cannot effectively improve the image's detail restoration problem. The Ucolor algorithm has the problem of excessive local color enhancement. This embodiment has a good underwater image enhancement effect for complex scenes, and has good handling of problems such as inconsistent brightness and color cast in the image.

[0134] Table 1: Objective index evaluation of different underwater image enhancement algorithms on UIEB and EUVP datasets

[0135]

[0136] It is worth noting that Table 1 is an objective indicator evaluation of different underwater image enhancement algorithms on the UIEB and EUVP datasets, where the Structural Similarity Index Measure (SSIM) simulates the human eye's perception of images and is used to measure the similarity of image structures. The Peak Signal-to-Noise Ratio (PSNR) measures the difference between the original image and the enhanced image. The higher these two indicators are, the more similar the images are. As can be seen from Table 1, this embodiment as a whole shows the highest value in the objective evaluation indicators, indicating that the underwater images enhanced by this embodiment have higher image quality.

[0137] Based on the same inventive concept, the present application also provides an underwater image enhancement device based on a transformer network, including a processor and a memory, wherein a computer program is stored in the memory. When the computer program is executed by the processor, the underwater image enhancement method based on the transformer network as described above is implemented.

[0138] Based on the same inventive concept, the present application also provides a computer-readable storage medium storing a computer program. When the computer program is executed by a processor, the underwater image enhancement method based on the transformer network as described above is implemented.

[0139] The program product of the present application for implementing the above method may adopt a portable compact disk read-only memory and include program code, and may be run on a terminal device, such as a personal computer. However, the program product of the present application is not limited thereto. In the present application, a readable storage medium may be any tangible medium containing or storing a program, which may be used by or in combination with an instruction execution system, apparatus, or device.

[0140] It should be noted that a computer-readable storage medium may include a data signal propagated in a baseband or as part of a carrier wave, which carries a readable program code. Such propagated data signals may take a variety of forms, including but not limited to electromagnetic signals, optical signals, or any suitable combination of the above. The readable storage medium may also be any readable medium other than a readable storage medium, which may send, propagate, or transmit a program for use by or in conjunction with an instruction execution system, device, or device. The program code contained on the readable storage medium may be transmitted using any appropriate medium, including but not limited to wireless, wired, optical cable, RF, etc., or any suitable combination of the above.

[0141] The above are only preferred embodiments of the present application, and are not intended to limit the present application in any form. Although the present application has been disclosed as above with preferred embodiments, it is not intended to limit the present application. Any technician familiar with the present application can make some changes or modifications to equivalent embodiments of equivalent changes using the above-mentioned technical contents without departing from the scope of the technical solution of the present application. The implementation schemes in the above embodiments can also be further combined or replaced. However, any simple modifications, equivalent changes and modifications made to the above embodiments based on the technical essence of the present application without departing from the content of the technical solution of the present application are still within the scope of the solution of the present application.

Claims

1. A method for underwater image enhancement based on transformer network, characterized in that: Constructing and training an underwater image enhancement network based on a transformer network, wherein the underwater image enhancement network based on a transformer network adopts an encoder-decoder architecture, including an N-layer encoding module, an N-layer decoding module, and an N-layer feature fusion module. Constructing the enhancement network includes the following steps: Step S1, obtaining a shallow feature image of an input image, wherein the input image is an underwater image to be enhanced; Step S2, feeding the shallow feature image into an encoding module to obtain a deep feature image; Step S3, fusing the output image of the encoding module of the N-1th layer and the output image of the decoding module of the N-1th layer through the feature fusion module of the N-1th layer to serve as the input image of the decoding module of the N-1th layer, and the decoding module of the N-1th layer outputs a residual image; The encoding module of the N layers is connected in sequence by N conv-transformer modules, and the decoding module of the N layers is connected in sequence by N conv-transformer modules, where N is an integer greater than 1; The N-1th layer feature fusion module is used to obtain the probability map of the output image of the coding module of the N-1th layer, perform preliminary fusion by weighting the product of the probability map and the output image of the decoding module of the previous layer, and perform channel-dimensional fusion of the output image of the coding module of the N-1th layer and the result after preliminary fusion; The conv-transformer module is used to perform deep feature extraction on the input image features to obtain a deep feature image; the conv-transformer module extracts in the channel dimension branch and the space dimension branch respectively, and outputs after fusing the extraction results of the two dimension branches; The N-1th layer feature fusion module performs feature fusion on the image output by the conv-transformer module in the N-1th layer encoding module and the image output by the conv-transformer module in the Nth layer decoding module; Step S4, adding the residual image output in step S3 and the input image obtained in step S1 element by element to obtain an enhanced underwater image.

2. The underwater image enhancement method based on transformer network according to claim 1, characterized in that: In step S3, the layers of the encoding module are cascaded using downsampling layers, and the encoding module is used to extract a deep feature image from the underwater image to be enhanced to obtain a deep feature image; The layers of the decoding module are cascaded using an upsampling layer through the feature fusion module; the decoding module is used to decode the input image fused by the feature fusion module, and obtain the residual image after decoding by N layers of the decoding modules.

3. The underwater image enhancement method based on transformer network according to claim 1, characterized in that: The conv-transformer module extracts the channel dimension branch features and performs the input tensor Perform layer normalization operation, which is expressed as: ； in, is the normalized output tensor, is the mean of all values ​​in the tensor, is the standard deviation of all values ​​in the tensor. For the input tensor , whose size is , is the height of the input image, is the width of the input image, is the channel dimension of the input image, Represents the separation of different dimensions of a tensor; For the input tensor After layer normalization, learnable parameters are introduced and To normalize the output tensor Scaling and offset operations are performed, expressed as: ； in, To add the output tensor after scaling and offset operations, and are learnable parameters.

4. The underwater image enhancement method based on transformer network according to claim 3, characterized in that: The process of extracting the channel dimension branch features includes obtaining the channel attention map and performing , , The components are transformed in dimension, expressed as: ； in, Represents the query, and the size after transforming the tensor dimension is , Represents the key value, and the size after transforming the tensor dimension is , Represents a value vector, and after transforming the tensor dimension, the size is , Represents the operation of transforming the dimension of a tensor; The channel attention map is calculated as follows: ； in, is a learnable scaling parameter, express Activation function, Represents the self-attention map of the channel dimension, with a size of ; The channel dimension branch output is expressed as: ； in, Represents the channel self-attention map Sum value vector The vector product of is the final output tensor of the channel dimension branch, represents the self-attention map of the channel dimension, Represents a vector of values.

5. The underwater image enhancement method based on transformer network according to claim 4, characterized in that: During the fusion process of the conv-transformer module, the output image of the spatial dimension branch and the output image of the channel dimension branch are element-wise added and then fused through a point convolution operation. The fused image of the spatial dimension branch and the channel dimension branch is suppressed and output using the GDFN module.

6. The underwater image enhancement method based on transformer network according to claim 1, characterized in that: The fusion process of the N-1 layer feature fusion module includes sequentially performing window segmentation, Function activation and window fusion operation to obtain the probability map of the output image, and also include pixel repetition operation on the output image of the conv-transformer module in the N-th layer decoding module, and multiplication weighted operation with the probability map to output a fused image, and also include channel dimension fusion of the output of the N-1-th layer encoding module and the fused image to output a deep fused image, and the deep fused image is used as the input image of the N-1-th layer decoding module; Obtaining the probability map of the output image includes segmenting the output image of the conv-transformer module in the N-1th layer encoding module according to the set window size, and using The function is activated to obtain a probability map of each window, and the probability map of each window is fused to obtain the probability map of the output image.

7. The underwater image enhancement method based on transformer network according to claim 1, characterized in that: The loss function is expressed as: ； Where n and m are the height and width of the image, For the real image Line The value of the column pixel, is the first Line The value of the column pixel, The loss value of the calculated loss function.

8. An underwater image enhancement system based on a transformer network, characterized in that: The enhancement system comprises: An acquisition module, used to acquire a shallow feature image of an input image, wherein the input image is an underwater image to be enhanced; N-layer encoding module, used for performing deep feature extraction on the shallow feature image; A feature fusion module is used to obtain a probability map of the output image of the encoding module of the corresponding layer, perform a preliminary fusion by weighting the product of the probability map and the output image of the decoding module of the previous layer, and perform a channel-dimensional fusion of the output image of the encoding module and the result of the preliminary fusion; N-layer decoding module, used to decode the fused feature map and the image output by the previous layer decoding module to output a residual image; N conv-transformer modules in the N-layer encoding module are connected in sequence, and N conv-transformer modules in the N-layer decoding module are connected in sequence, where N is an integer greater than 1; The conv-transformer module is used to perform deep feature extraction on the input image features to obtain a deep feature image. The conv-transformer module extracts in the channel dimension branch and the space dimension branch respectively, and outputs after fusing the extraction results of the two dimension branches; The output module is used to add the residual image output by the decoding module to the input image in the acquisition module element by element, and output an enhanced underwater image.

9. An underwater image enhancement device based on a transformer network, characterized in that: It comprises a processor and a memory, wherein a computer program is stored in the memory, and when the computer program is executed by the processor, the underwater image enhancement method based on the transformer network as described in any one of claims 1 to 7 is implemented.

10. A computer-readable storage medium storing a computer program, characterized in that: When the computer program is executed by a processor, the underwater image enhancement method based on the transformer network as described in any one of claims 1 to 7 is implemented.

Images