An underwater image enhancement method based on transformer and multi-scale spatial domain feature fusion

Abstract

The present application relates to the field of industrial vision enhancement, more specifically, it relates to an underwater image enhancement method based on Transformer and multi-scale spatial domain feature fusion, the present application introduces the mechanism of Transformer, uses self-attention and linear attention mechanism to capture the complex spatial relationship in the image and improve the image quality, at the same time, the multi-scale feature convolution module is designed to extract information of different scales and levels, which can reduce the complexity of the network and improve the calculation efficiency, in addition, an air domain feature fusion module is designed to calculate the attention in the local neighborhood of pixels, thereby improving the calculation efficiency and feature fusion, the method has the advantages of good generalization ability, high robustness and high enhancement efficiency, and can be widely applied in underwater pipeline detection and maintenance, underwater environment monitoring, underwater biological research and other fields, and provides strong support for scientific research and engineering practice in related fields.

Description

Technical Field

[0001] This invention relates to the field of industrial visual enhancement, and more specifically, to an underwater image enhancement method based on Transformer and multi-scale spatial domain feature fusion. Background Technology

[0002] Underwater image enhancement technology, as an image processing technique, has the potential for widespread application in underwater exploration and observation. Underwater images are often affected by factors such as water quality and light attenuation, resulting in blurriness and noise, which limits the identification and analysis of underwater targets. Underwater image enhancement technology can improve underwater image quality by enhancing contrast and reducing noise, thereby improving the accuracy and efficiency of target identification. Furthermore, underwater image enhancement technology can be applied to fields such as marine resource exploration, underwater archaeology, and underwater search and rescue, providing strong support for underwater exploration 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 suffer from poor generalization ability, low robustness, and low enhancement efficiency in underwater image enhancement. Therefore, developing an image enhancement algorithm with good generalization ability, high robustness, and high enhancement efficiency is of great significance for underwater exploration work. Summary of the Invention

[0004] To address the shortcomings of existing technologies, the present invention aims to provide an underwater image enhancement method based on Transformer and multi-scale spatial domain feature fusion, which has the advantages of good generalization ability, good robustness, and high enhancement efficiency.

[0005] The above-mentioned technical objective of this invention is achieved through the following technical solution: an underwater image enhancement method based on Transformer and multi-scale spatial domain feature fusion, comprising the following steps:

[0006] S1. An underwater illumination source emits light towards an underwater target, and an underwater imaging device receives the light from the underwater target to form an underwater image, thereby constructing an underwater image dataset;

[0007] S2. Build a deep learning network model. Use an underwater RGB image of size N×H×W into the deep learning network model. Then, use a multi-scale shallow convolution module to extract information of different scales and levels in the feature encoding part to enhance the model's understanding ability and its ability to process contextual information in the underwater environment.

[0008] Where N is the number of channels, H is the image height, and W is the image width, the deep learning network model includes a feature extraction network, a feature encoding network, a feature decoding network, a feature compression network, a feature reconstruction network, a multi-scale feature convolution module, a spatial domain feature fusion module and a transformer module, a multi-scale deep convolution fusion module and an attention convolution module;

[0009] S3. Set the training parameters, including setting the learning rate to lr, the loss function to L, the batch size to h, the training set and test set to be divided into m and n, the optimizer to D, and the total training cycle to A. Then train the deep learning network model.

[0010] S4. Enhance the image using the trained deep learning network model. Input the image to be tested and output an enhanced underwater image. Remove noise from the image, correct color deviation caused by underwater reflection, and finally obtain the feature-enhanced output image.

[0011] In one embodiment, the underwater image dataset is represented as: ;

[0012] Among them, dataset The total number of elements in the middle is The image size is , For image channels, Image height, The image width is specified, and the image file format is PNG.

[0013] In one embodiment, S2 further includes the following steps:

[0014] Let x be the given input feature map of the multi-scale feature convolution module, and let its three feature maps be respectively... , , It is generated by three convolutional blocks with different kernel sizes of 1×1, 3×3 and 5×5, and is represented as follows:

[0015]

[0016] in, This represents a convolution operation with kernel size Z = 1, 3, 5;

[0017] The input features then pass through the channel attention module, whose output feature map Represented as:

[0018]

[0019] in, Represents the channel attention function;

[0020] After channel attention, the feature maps are combined with the fully connected layer to output the feature maps separately. , and ,in,

[0021]

[0022] Indicates a fully connected layer;

[0023] Finally, the above feature maps are summed using addition to obtain the output feature map. , represented as , × indicates element-wise multiplication.

[0024] In one embodiment, the operation of the spatial feature fusion module in S2 further includes:

[0025] Input features , will the Features are processed through the multi-scale deep convolutional fusion module and the attention convolutional module;

[0026] In the spatial domain feature fusion module, the spatial domain feature map is converted to the frequency domain by fast Fourier transform, and the feature map in the frequency domain is converted back to the spatial domain by inverse fast Fourier transform, so as to realize feature map reconstruction.

[0027] In one embodiment, the feature information processing process of the spatial feature fusion module can be represented as follows:

[0028]

[0029]

[0030] in, , and These respectively represent the multi-scale deep convolutional fusion module, the attention convolutional module, and the transformer module. and The size component and phase component of the feature are specifically represented respectively. , and These are represented as 1×1 convolution, Fourier domain convolution, and inverse fast Fourier domain transform, respectively. The hyperparameter α is a constant used to control the fusion ratio of spatial frequency domain information, thereby enhancing the network's sensing field and quality.

[0031] In one embodiment, S2 further includes the following steps:

[0032] Let the input feature map be Multiple attention heads are connected and examined through a feedforward network. Assume the feature map size of the input feature map is... The transformer module enhances the input feature map by capturing important information in the image through an attention mechanism and uses residual connections to preserve information from the original image. Subsequently, the feature map is processed by convolutional layers to obtain an enhanced output feature map. The feature map extraction process of the transformer module can be represented as follows:

[0033]

[0034]

[0035] in , and These represent Batch Normalization, ReLU activation function, and 3×3 convolution, respectively. This represents multi-head attention with dimension x, where x=2, 3. The two formulas above represent the transformer encoder and transformer decoder, respectively.

[0036] The underwater image enhancement method described above, based on Transformer and multi-scale spatial domain feature fusion, has the advantages of good generalization ability, high robustness, and high enhancement efficiency. It can be widely used in underwater pipeline inspection and maintenance, underwater environmental monitoring, underwater biological research and other fields, providing strong support for scientific research and engineering practice in related fields. Attached Figure Description

[0037] Figure 1 This is a schematic diagram of the underwater imaging device in this embodiment;

[0038] Figure 2 This is a schematic diagram of the overall architecture of the enhanced network in this embodiment;

[0039] Figure 3 For this embodiment Figure 2 Schematic diagram of the network architecture of the output module;

[0040] Figure 4 For this embodiment Figure 2 Schematic diagram of the network architecture of the basic module;

[0041] Figure 5 For this embodiment Figure 2 Schematic diagram of the network architecture of the convolutional module;

[0042] Figure 6 For this embodiment Figure 7 Schematic diagram of the network architecture of the multi-scale feature convolution module;

[0043] Figure 7 For this embodiment Figure 2 Schematic diagram of the network architecture of multi-scale deep convolutional modules and multi-scale shallow convolutional modules;

[0044] Figure 8 For this embodiment Figure 2 Schematic diagram of the network architecture of the spatial feature fusion module;

[0045] Figure 9 For this embodiment Figure 2 A schematic diagram of the network architecture of the Transformer decoder and Transformer encoder modules.

[0046] The following are the labeling elements in the figure:

[0047] 101. Computer; 102. Underwater imaging equipment; 103. Underwater lighting source; 104. Underwater target. Detailed Implementation

[0048] The present invention will now be described in detail with reference to the accompanying drawings and embodiments.

[0049] In the description of this invention, it should be understood that the terms "center," "longitudinal," "lateral," "length," "width," "thickness," "upper," "lower," "front," "rear," "left," "right," "vertical," "horizontal," "top," "bottom," "inner," "outer," "clockwise," "counterclockwise," "axial," "radial," and "circumferential" indicate the orientation or positional relationship based on the orientation or positional relationship shown in the accompanying drawings. They are used only for the convenience of describing this invention and are not intended to indicate or imply that the device or element referred to must have a specific orientation, or be constructed and operated in a specific orientation. Therefore, they should not be construed as limitations on this invention.

[0050] Furthermore, the terms "first" and "second" are used for descriptive purposes only and should not be construed as indicating or implying relative importance or implicitly specifying the number of indicated technical features. Thus, a feature defined as "first" or "second" may explicitly or implicitly include at least one of that feature. In the description of this invention, "a plurality of" means at least two, unless otherwise explicitly specified.

[0051] In this invention, unless otherwise explicitly specified and limited, the terms "installation," "connection," "linking," and "fixing," etc., should be interpreted broadly. For example, they can refer to a fixed connection, a detachable connection, or an integral part; they can refer to a mechanical connection, an electrical connection, or a connection that allows communication between them; they can refer to a direct connection or an indirect connection through an intermediate medium; they can refer to the internal communication of two components or the interaction between two components, unless otherwise explicitly limited. Those skilled in the art can understand the specific meaning of the above terms in this invention according to the specific circumstances.

[0052] In this invention, unless otherwise explicitly specified and limited, "above" or "below" the second feature can mean that the first and second features are in direct contact, or that they are in indirect contact through an intermediate medium. Furthermore, "above," "over," and "on top" of the second feature can mean that the first feature is directly above or diagonally above the second feature, or simply that the first feature is at a higher horizontal level than the second feature. "Below," "below," and "under" the second feature can mean that the first feature is directly below or diagonally below the second feature, or simply that the first feature is at a lower horizontal level than the second feature.

[0053] An underwater image enhancement method based on Transformer and multi-scale spatial domain feature fusion includes the following steps:

[0054] S1. The underwater illumination source 103 emits light towards the underwater target 104, and the underwater imaging device 102 receives the light from the underwater target 104 to form an underwater image, thereby constructing an underwater image dataset. , where the dataset The total number of elements in the middle is =890, image size is 3×256×256. The image file format is PNG.

[0055] In the above steps, the computer 101 is connected to the underwater imaging device 102. The computer 101 acquires the underwater target and underwater image through deep learning image algorithms and performs image enhancement.

[0056] S2. Construct a deep learning network model, which includes a feature extraction network, a feature encoding network, a feature decoding network, a feature compression network, a feature reconstruction network, a multi-scale feature convolution module, a spatial feature fusion module, a transformer module, a multi-scale deep convolution fusion module, and an attention convolution module.

[0057] In one example, an underwater RGB image with a size of 64×128×128 is used as input into the deep learning network model. Initially, as... Figure 4As shown, the image feature map undergoes basic module processing in the feature extraction section, the purpose of which is to learn the basic features of the image and perform simple information extraction. Subsequently, the image passes through the feature encoding section. In this section, as... Figure 7 As shown, a flexible multi-scale shallow convolutional module is proposed to extract information at different scales and levels, enhancing the model's ability to understand and process contextual information in underwater environments. This process produces a shallow embedding feature map with 64 channels. These shallowly embedded features are fed into the proposed transformer encoder to capture global information and improve the model's generalization ability. Next, as... Figure 7 As shown, the shallow features extracted from the transformer encoder are fed into a multi-scale deep convolutional module, which uses contextual information to fuse features to generate deep embedding feature maps. Subsequently, these features are transformed into deeper layers. Addressing the issue that most current degraded image restoration methods heavily rely on spatial domain processing, while traditional convolutional methods often ignore the rich global information in the Fourier domain, such as... Figure 8 As shown, we propose a spatial feature fusion module, whose output depth feature map is... ;

[0058] In practical applications, whitening can minimize the correlation of image features and ensure that each segment has the same variance. Therefore, whitening is introduced in the feature compression part to reduce redundancy between image features, decrease the correlation between components, and balance their variance. To mitigate the impact of degraded diversity and strengthen the connection between feature encoding and feature decoding, four convolutions are merged on the skip connection line. These convolutions help enhance global dependencies, passing through the bottom two feature extraction modules (spatial domain feature fusion module, such as...) in the feature encoding and feature decoding parts. Figure 8 (as shown) and multi-scale deep convolutional modules (such as...) Figure 7 After the output feature map is recorded as shown, the output feature map is recorded as follows: The output feature map size remains unchanged. The whitening operation eliminates redundant information in the input data, improving the efficiency of core information extraction.

[0059] like Figure 2 As shown, similar to the feature decoding part, the transformer decoder is designed on the skip connections of the feature decoding part. By flexibly balancing global and local information in different frequency domains, the transformer decoder improves the preservation of both overall and local quality. This enhances the detail of underwater images and improves the adaptability of the proposed network to complex underwater environments.

[0060] After four reconstruction modules (multi-scale deep convolutional modules, such as...), Figure 7 As shown), convolutional modules (such as...) Figure 5 (as shown) and two basic modules (such as) Figure 4 As shown), the shallow feature map is then recorded as follows: Finally, the feature map is convolved with the output of the feature reconstruction part (e.g., ...). Figure 9 (As shown) After that, by embedding the feature map Projecting onto the RGB channels restores a clear underwater image.

[0061] Specifically, step S2 also includes the following steps:

[0062] like Figure 5 As shown in the example, in one instance, the given input feature map of the multi-scale feature convolution module is set to x, and three feature maps are used. , , Generated by three convolutional blocks with different kernel sizes of 1×1, 3×3, and 5×5, it is represented as follows:

[0063]

[0064] in, This represents convolution operations with kernel sizes Z=1, 3, and 5. This structure helps the network better adapt to information at different scales and levels by automatically extracting and fusing key features from feature maps at three different scales, thus improving model performance. Furthermore, the input features also pass through a channel attention module, and the output feature map... Represented as:

[0065]

[0066] in Let represent the channel attention function. By introducing channel attention, multi-scale convolutional modules can reduce their sensitivity to noise or irrelevant features while enhancing the attention between different channels in the input feature map x. This helps the network capture key features in the image more effectively. After channel attention, we combine the feature map with the fully connected layer, and the output feature maps are represented as follows: , and :

[0067]

[0068] in This represents a fully connected layer. Finally, the above feature maps are summed using an addition operation to obtain the output feature map Y, represented as:

[0069]

[0070] Where × represents element-wise multiplication.

[0071] Furthermore, the spatial feature fusion module operation process in step S2 also includes:

[0072] First, given the input features are Initially, the input features are processed through a multi-scale deep convolutional fusion module and an attention convolutional module (such as...). Figure 6 (As shown) is processed.

[0073] Multi-scale deep convolutional fusion modules can fuse features from different levels to improve the model's expressive power and better capture information from input features. Attention-based convolutional modules can finely allocate and process input features.

[0074] Therefore, by summing the features output by the multi-scale deep convolutional fusion module and the attention convolutional module, image details can be extracted more effectively. Furthermore, the transformer encoder can capture information biases from the input features, improving the network's generalization ability and reducing overfitting to the training data.

[0075] In the spatial domain feature fusion module, the Fast Fourier Transform (FFT) is used to transform the spatial domain feature map into the frequency domain, facilitating the utilization of its frequency components and spectral features. The Inverse Fast Fourier Transform (IFFT) is then used to transform the feature map from the frequency domain back into the spatial domain to achieve feature map reconstruction. The feature information processing procedure of the spatial domain feature fusion module can be represented as follows:

[0076]

[0077]

[0078] in , and These represent the multi-scale deep convolutional fusion module, the attention convolutional module, and the transformer module, respectively. Additionally, and The size component and phase component of the feature are specifically represented respectively. , and These are represented as 1×1 convolution, Fourier domain convolution, and inverse fast Fourier transform, respectively. The hyperparameter α in SDFFM is a constant that controls the fusion ratio of spatial frequency domain information, enhancing the network's sensing field and quality.

[0079] Furthermore, S2 also includes the following steps:

[0080] Assume the input feature map is The transformer allows the model to jointly infer attention from different representation subspaces, connecting multiple attention heads and further examining them through a feedforward network. Here, the MHSA mechanism uses three heads. Assuming the feature map size of the input feature map is given, the transformer enhances the input feature map by capturing important information in the image through the attention mechanism. Residual connections are used to preserve information from the original image. Subsequently, the feature map is processed by convolutional layers to obtain an enhanced output feature map. The feature map extraction process of the transformer is represented as follows:

[0081]

[0082]

[0083] in , and These represent Batch Normalization, ReLU, and 3×3 convolution, respectively. This represents multi-head attention, with dimension x, where x=2, 3. The two formulas above represent the transformer encoder and transformer decoder, respectively.

[0084] S3. Set the training learning rate to 0.0001 and the loss function to MSELOSS. Configure the training parameters as follows: learning rate is 0.0001, batch size is 4, training set and test set are divided into 700 and 190 images respectively, the optimizer is Adam, and the total training cycle is 300.

[0085] S4. Use the trained network for enhancement. Input the test image, and output the enhanced underwater image. The image size is input into the network. The input image is enhanced by the network to remove noise, correct color deviations caused by underwater reflections, and finally obtain a feature-enhanced output image.

[0086] Considering the above solutions and existing technologies, current image enhancement algorithms generally struggle to simultaneously improve image denoising speed and accurately remove color deviations across multiple scales. This invention, however, introduces a transformer mechanism, utilizing self-attention and linear attention to capture complex spatial relationships in images and improve image quality. Simultaneously, a multi-scale feature convolution module is designed to extract information at different scales and levels, reducing network complexity and improving computational efficiency. Furthermore, a spatial feature fusion module is designed to calculate attention within the local neighborhood of pixels, thereby enhancing computational efficiency and feature fusion.

[0087] The above methods have the advantages of good generalization ability, high robustness, and high enhancement efficiency. They can be widely used in underwater pipeline inspection and maintenance, underwater environmental monitoring, underwater biological research and other fields, providing strong support for scientific research and engineering practice in related fields.

[0088] The embodiments described above are merely illustrative of several implementations of the present invention, and while the descriptions are specific and detailed, they should not be construed as limiting the scope of the present invention. It should be noted that those skilled in the art can make various modifications and improvements without departing from the concept of the present invention, and these modifications and improvements all fall within the scope of protection of the present invention. Therefore, the scope of protection of this patent should be determined by the appended claims.

Claims

1. An underwater image enhancement method based on Transformer and multi-scale spatial domain feature fusion, characterized in that, Includes the following steps: S1. An underwater illumination source emits light towards an underwater target, and an underwater imaging device receives the light from the underwater target to form an underwater image, thereby constructing an underwater image dataset; S2. Construct a deep learning network model by using an underwater RGB image of size N×H×W into the deep learning network model. The deep learning network model includes: a feature extraction network, a feature encoding network, a feature decoding network, a feature compression network, a feature reconstruction network, a multi-scale shallow convolution module, a multi-scale feature convolution module, a spatial feature fusion module, a transformer module, a multi-scale deep convolution fusion module, and an attention convolution module. Where N is the number of channels, H is the image height, and W is the image width, the feature extraction network is used to extract basic features from the input underwater RGB image, the feature encoding network is provided with the multi-scale shallow convolution module and the transformer module, the multi-scale shallow convolution module extracts information of different scales and levels from the feature map output by the feature extraction network to generate a shallow embedded feature map and input it to the transformer module for global information modeling, the feature map output from the transformer module is sequentially input to the multi-scale deep convolution fusion module and the spatial domain feature fusion module for deep feature fusion and spatial domain and / or frequency domain feature enhancement, after being processed by the feature compression network and the feature decoding network, it is input to the feature reconstruction network to restore the enhanced underwater image; S3. Set the training parameters, including setting the learning rate to lr, the loss function to L, the batch size to h, the training set and test set to be divided into m and n, the optimizer to D, and the total training cycle to A. Then train the deep learning network model. S4. Enhance the image using the trained deep learning network model. Input the image to be tested, output an enhanced underwater image, remove noise from the image, correct color deviation caused by underwater reflection, and finally obtain the feature-enhanced output image. S2 also includes the following steps: Let X be the given input feature map of the multi-scale feature convolution module, and let its three feature maps be respectively... , , It is generated by three convolutional blocks with different kernel sizes of 1×1, 3×3 and 5×5, and is represented as follows: in, This represents a convolution operation with kernel size Z = 1, 3, 5; The input features then pass through the channel attention module, whose output feature map Represented as: in, Represents the channel attention function; After channel attention, the feature maps are combined with the fully connected layer to output the feature maps separately. , and ,in, Indicates a fully connected layer; Finally, the above feature maps are summed using addition to obtain the output feature map. , represented as × indicates element-wise multiplication, and the output feature map Y serves as the input feature map for subsequent feature encoding and feature fusion modules in the deep learning network model.

2. The underwater image enhancement method based on Transformer and multi-scale spatial domain feature fusion according to claim 1, characterized in that: The underwater image dataset is represented as follows: ； Among them, dataset The total number of elements in the middle is The image size is , For image channels, Image height, The image width is specified, and the image file format is PNG.

3. The underwater image enhancement method based on Transformer and multi-scale spatial domain feature fusion according to claim 1, characterized in that, The operation of the spatial feature fusion module in S2 also includes: Input features , will the Features are processed through the multi-scale deep convolutional fusion module and the attention convolutional module; In the spatial domain feature fusion module, the spatial domain feature map is converted to the frequency domain by fast Fourier transform, and the feature map in the frequency domain is converted back to the spatial domain by inverse fast Fourier transform, so as to realize feature map reconstruction.

4. The underwater image enhancement method based on Transformer and multi-scale spatial domain feature fusion according to claim 3, characterized in that, The feature information processing process of the spatial feature fusion module can be represented as follows: First to The size component of the frequency domain feature is obtained by performing a fast Fourier transform. and phase components The above two frequency domain feature components are used to calculate the output spatial domain feature map. ; in, , and These respectively represent the multi-scale deep convolutional fusion module, the attention convolutional module, and the transformer module. and The size component and phase component of the feature are specifically represented respectively. , and These are represented as 1×1 convolution, Fourier domain convolution, and inverse fast Fourier transform, respectively. The spatial domain feature map output by the spatial domain feature fusion module is shown. The hyperparameter α is a constant used to control the fusion ratio of spatial frequency domain information, thereby enhancing the network's sensing field and quality.

5. The underwater image enhancement method based on Transformer and multi-scale spatial domain feature fusion according to claim 1, characterized in that, S2 further includes the following step: Let the input feature map be... Multiple attention heads are connected and examined through a feedforward network; the feature map size of the input feature map is... The transformer module enhances the input feature map by capturing important information in the image through an attention mechanism and uses residual connections to preserve information from the original image. Subsequently, the feature map is processed by convolutional layers to obtain an enhanced output feature map. The feature map extraction process of the transformer module can be represented as follows: in , and These represent Batch Normalization, ReLU activation function, and 3×3 convolution, respectively. The above formulas represent multi-head attention with dimension x, where x = 2 and 3. The two formulas above represent the transformer encoder and transformer decoder, respectively, and their output feature maps are used as input feature maps for subsequent feature fusion and feature reconstruction in the deep learning network model.

Images