Wavelet guided Transform underwater image enhancement method based on prompt learning

Abstract

The invention relates to a wavelet guided Transform underwater image enhancement method based on prompt learning. The method comprises the following steps: S1, obtaining a degraded underwater image; s2, preprocessing the degraded underwater image serving as the training set to obtain training sample data; s3, inputting the sample into a wavelet guided Transform network based on prompt learning for training; and S4, after training is completed, performing image enhancement on the test sample to obtain a result. The performance of the method and the system is superior to that of other underwater image enhancement methods, the details of the image can be recovered more finely, and the method has advantages in the aspect of avoiding artifacts compared with other methods.

Description

Technical Field

[0001] This invention belongs to the field of deep learning technology and relates to a wavelet-guided Transformer underwater image enhancement method based on cue learning. Background Technology

[0002] Underwater imagery is a crucial medium for acquiring information about the marine environment, playing an indispensable role in marine resource exploration, underwater visible light communication, and navigation of autonomous underwater vehicles (AUVs). However, due to the complex underwater optical environment, light undergoes severe non-uniform attenuation during transmission, resulting in images generally exhibiting low contrast. Simultaneously, the scattering effects of suspended particles and dissolved organic matter in the water introduce complex noise interference. The coupling effect of these degradation factors significantly reduces the quality of the original image, often failing to meet the needs of practical applications. Therefore, researching effective underwater image enhancement techniques is of significant practical value for improving image visibility, ensuring the reliable operation of underwater vision systems, and advancing marine scientific research.

[0003] Traditional underwater image enhancement methods focus on accurately estimating the transport characteristics of the medium, reversing the physical imaging model to recover a clear image by estimating the transport map and other key underwater imaging parameters. However, due to the complexity and diversity of real underwater environments, the applicability of these methods is often limited, and their effectiveness is usually confined to relatively simple scenes. In recent years, Convolutional Neural Networks (CNNs) have achieved significant success in underwater image enhancement due to their powerful feature extraction capabilities. However, the limited receptive field of CNNs restricts their ability to capture global information, thus limiting the performance of these methods. Recently, Transformer-based algorithms, based on a self-attention mechanism, have achieved global image modeling, effectively capturing long-range dependencies and significantly improving the performance of underwater image enhancement. Although existing methods have achieved some success, they mainly focus on spatial domain processing and fail to fully explore the feature information contained in the frequency domain. The frequency domain can effectively characterize the hierarchical structure of an image: high-frequency components correspond to detail and texture information, while low-frequency components carry overall contour and background information; while the spatial domain focuses more on capturing the relationships between pixels. The synergistic analysis of these two domains is crucial for accurately understanding complex underwater scenes. In recent years, research has shown that introducing degradation information as a cue signal can effectively guide models to perform enhancement tasks. This method provides valuable supervised guidance for underwater image enhancement by establishing semantic associations between degradation features and enhancement targets. However, current cue-based learning methods are still mainly limited to spatial domain processing strategies, lacking a systematic exploration of high-frequency and low-frequency components of images. This single-domain processing paradigm inevitably limits the model's adaptability to complex underwater degradation conditions and also restricts its potential for further performance improvement. Summary of the Invention

[0004] In view of this, the purpose of this invention is to provide a wavelet-guided Transformer underwater image enhancement method, WFPFormer, based on cue learning. This method selectively extracts and processes frequency components through a frequency-guided Transformer module (FGTM), thereby achieving more refined and information-rich feature representations. It employs a wavelet-guided attention refinement module (WRAM) to integrate wavelet-based frequency domain information into channel and spatial modeling, mapping low-frequency components to the spatial domain to enhance the model's adaptability to complex underwater scenes. Furthermore, it utilizes a wavelet-guided cueing module (WGPM) to fuse contextual and frequency features to generate cue information, establishing a dynamic correlation between degraded features and the reconstruction target. While maintaining theoretical rigor, this method and system achieve continuous performance optimization through the organic integration of deep feature learning and physical knowledge, giving it advantages over other underwater image enhancement methods. This method not only enhances images but also possesses stronger detail structure recovery capabilities and exhibits significant advantages in artifact suppression compared to other methods.

[0005] To achieve the above objectives, the present invention provides the following technical solution: A wavelet-guided Transformer underwater image enhancement method based on cue learning is disclosed. The method includes the following steps: S1: acquiring degraded underwater images; S2: preprocessing the degraded underwater images used as the training set to obtain training sample data; S3: inputting the samples into a wavelet-guided Transformer network based on cue learning for training; S4: after training, performing image enhancement on the test samples to obtain the results.

[0006] Furthermore, in step S1, the underwater robot is used to acquire degraded underwater images, specifically including: S11, deploying the underwater robot to underwater environments with diverse water scenes, water types, lighting conditions, and categories to collect data.

[0007] Furthermore, in step S2, the degraded images used as the training set are preprocessed. First, the degraded images are subjected to data augmentation operations, including random rotation and flipping. Then, the images are processed using the best existing underwater image augmentation methods. Finally, the data-augmented underwater images, depth maps, gradient maps, and their corresponding images processed by the underwater image augmentation methods constitute the training sample pairs.

[0008] Furthermore, in step S3, the input image obtains multi-scale features through three coding layers, each containing a 3×3 convolution and four FGTMs. For each coding layer, the input first undergoes a 3×3 convolution, as shown in the following formula:

[0009] Where X is the input feature map. Indicates the filter size is Convolution operation, This indicates that the filter size is... The result of the convolution operation.

[0010] Furthermore, after the 3×3 convolution, it is immediately followed by an FGTM. The FGTM is divided into three branches: the first branch performs flat pooling, the second branch performs wavelet convolution, and the third branch performs max pooling, as shown in the following formula:

[0011] Where X is the input feature map. For the first The output of each branch, Represents average pooling operation. Represents wavelet convolution operation, This represents the max pooling operation.

[0012] In the second branch, the input undergoes wavelet convolution to obtain one low-frequency component and three high-frequency components, as shown in the following formula:

[0013] in, As input features, , , , These represent the low-frequency components and the high-frequency components in the vertical, horizontal, and diagonal directions, respectively. This represents the wavelet convolution operation.

[0014] For low-frequency components, they are first input into the low-frequency enhancement module for spatial and frequency domain processing. In spatial processing, they are first divided into two equal parts along the channel dimension, and then further divided into query value, key value, and attribute value through three 1×1 convolutions, as shown in the following formula:

[0015] in, For the input feature map, For splitting operations along the channel, This is the first channel of the input features after channel partitioning. part, , , Representing the Partial query values, key values, and attribute values.

[0016] Then, exchange and Multi-head attention calculation is performed using the following formula:

[0017] in, , , Indicates the first Part 1 The query value, key value, and attribute value of the head. Indicates the first Part 1 Self-attention based on size Represents the softmax equation, This represents each head dimension.

[0018] Then, the calculated attention is concatenated along the channel dimension, as shown in the following formula:

[0019] in, Representing the Partial self-attention, This represents a splicing operation along the channel dimension.

[0020] Then, the two self-attention results are concatenated along the channel dimension, and the spatial processing output is obtained through a 1×1 convolution, as shown in the following formula:

[0021] in, This represents the output of spatial processing.

[0022] Then, in the frequency domain processing, the amplitude and phase are first obtained through Fourier transform, as shown in the following formula:

[0023] in, For the input feature map, This indicates the Fourier transform operation. , Amplitude and phase, respectively.

[0024] Then, the amplitude and phase are refined by inputting 1×1 convolution, LeaklyReLU activation function, and 1×1 convolution respectively, as shown in the following formula:

[0025] in, Representing LeaklyReLU, , These represent the refined amplitude and phase characteristics, respectively.

[0026] Then, these features are combined and subjected to an inverse Fourier transform to obtain the output of the frequency domain processing, as shown in the following formula:

[0027] in, Represents the output of frequency domain processing. This represents the inverse Fourier transform operation.

[0028] Subsequently, after processing in the spatial and frequency domains respectively, the features from the two domains are concatenated along the channels, and after a 1×1 convolution, added to the input features to obtain the output of the low-frequency enhancement module, as shown in the following formula:

[0029] in, For the input feature map, The output of spatial processing, The output of frequency domain processing, This represents the output of the low-frequency enhancement module.

[0030] For high-frequency components, they are first concatenated along the channel dimension, as shown in the following formula:

[0031] in, This represents the result of splicing together three high-frequency components.

[0032] Subsequently, the high-frequency components are input into the high-frequency enhancement module, and after being separable convolved along three 3×3 depths with hole ratios of 1, 2, and 3 respectively, they are summed to obtain the output of the high-frequency enhancement module, as shown in the following formula:

[0033] in, For the input feature map, For the output of the high-frequency enhancement module, Indicates the filter size is void ratio Depth-separable convolution operations.

[0034] Then, the output of the first branch is added to the output of the low-frequency enhancement module, and the output of the third branch is added to the output of the high-frequency enhancement module to obtain the enhanced high and low frequency characteristics, as shown in the following formula:

[0035] in, For the input feature map, This represents a strengthened low-frequency characteristic. High-frequency features that represent enhancement.

[0036] Then, the enhanced high- and low-frequency features are input into the cross-frequency interaction attention module. For low-frequency features, they are first divided into query values, key values, and attribute values ​​through three 1×1 convolutions, as shown in the following formula:

[0037] in, This represents a strengthened low-frequency characteristic. , , These represent query values, key values, and attribute values ​​that emphasize low-frequency characteristics.

[0038] Then, multi-head attention is calculated using the following formula:

[0039] in, , , Indicating low-frequency characteristics The query value, key value, and attribute value of the head. Indicates the first Part 1 Attention score based on height.

[0040] Then, for high-frequency features, a 1×1 convolution is first applied to transform them into attribute values, as shown in the following formula:

[0041] in, High-frequency features representing enhancement The attribute value represents the enhancement of high-frequency features.

[0042] Then, the low-frequency attention score is multiplied by the high- and low-frequency attribute values ​​respectively, concatenated along the channel dimension, and added to the original high- and low-frequency features to obtain the high- and low-frequency interaction features, as shown in the following formula:

[0043] in, Representing high-frequency characteristics The attribute value of size. , These represent the low-frequency and high-frequency characteristics after frequency interaction, respectively.

[0044] Then, the low-frequency and high-frequency features after frequency interaction are deconvolved using wavelet to obtain the output of the cross-frequency interaction attention module, as shown in the following formula:

[0045] in, This represents the result of the cross-frequency interaction attention module. This represents the wavelet deconvolution operation.

[0046] Then, the output of the cross-frequency interaction attention module is fed into a gated convolutional feedforward neural network for non-linear interaction. The input features are first passed through two 1×1 convolutions, one result is multiplied by the other using the GELU activation function, and then added to the original input via another 1×1 convolution to obtain the output of the FGTM, as shown in the following formula:

[0047] in, As input features, Represents the output of FGTM, This represents the GELU activation function.

[0048] Furthermore, after three levels of encoding, the features are input into the WRAM, along with wavelet channel pyramid attention and wavelet spatial interaction attention. In the wavelet channel pyramid attention, the features are first subjected to bilinear interpolation, and then high- and low-frequency features are separated through wavelet convolution, as shown in the following formula:

[0049] in, As input features, , , , These represent the low-frequency components and the high-frequency components in the vertical, horizontal, and diagonal directions, respectively. This represents the bilinear interpolation operation.

[0050] Then, the low-frequency components are processed through three scales of adaptive average pooling and adaptive max pooling to obtain pooling pyramid features. The adaptive average pooling and adaptive max pooling at each corresponding scale are added together and concatenated along the channel dimension. After passing through a multilayer perceptron network, the features are multiplied by the original input using the sigmoid equation, as shown in the following formula:

[0051] in, As input features, Low-frequency components, and These represent compressing the features to a size of [size value]. Adaptive average pooling and adaptive max pooling operations, Represents a multilayer perceptron network. Represents the sigmoid equation. The output represents the wavelet channel pyramid attention.

[0052] Then, in the wavelet spatial interactive attention, the features are also subjected to bilinear interpolation, and then high- and low-frequency features are separated by wavelet convolution. After extracting the low-frequency features, the results of average pooling and max pooling in the channel dimension are added together and concatenated along the channel. After passing through a 7×7 convolution and the sigmoid equation, it is multiplied with the original input, as shown in the following formula:

[0053] in, As input features, Low-frequency components, and These represent the operations performed along the channel dimension after average pooling and max pooling, respectively. The output represents the wavelet space interactive attention.

[0054] Then, the wavelet channel pyramid attention and the wavelet spatial interaction attention are added together to obtain the output of the wavelet refinement attention module, as shown in the following formula:

[0055] in, This represents the output of WRAM.

[0056] Further, after WRAM, decoding begins step by step. The decoding process is identical to the encoding process, except for adding a WGPB before PGTM. In PGTM, the features are first separated into a low-frequency component and three high-frequency components in different directions through wavelet convolution. The three high-frequency components in different directions are then concatenated along the channel and subjected to a 1×1 convolution, as shown in the following formula:

[0057] in, As input features, , , , These represent the low-frequency components and the high-frequency components in the vertical, horizontal, and diagonal directions, respectively. This represents the result of fusing three high-frequency components.

[0058] Then, the high- and low-frequency components interact with two sets of learnable cues, respectively. Taking the high-frequency component as an example, it first undergoes global average pooling and a linear layer, then the attention weights are obtained through a softmax equation. These weights are then multiplied by the high-frequency cues and summed, as shown in the following formula:

[0059] in, The representative is the input feature. This represents the global average pooling operation. For linear layers, Representing the A tip, This indicates a frequency-guided prompt.

[0060] Then, the frequency-guided cue is divided into key values ​​and attribute values ​​through two 1×1 convolutions. The input features are then converted into query values ​​through a 1×1 convolution, as shown in the following formula:

[0061] in, The representative is the input feature. This represents a strengthened low-frequency characteristic. , , Represents query value, key value, and attribute value.

[0062] Then, multi-head attention is calculated using the following formula:

[0063] in, , , Indicates the first The query value, key value, and attribute value of the head. Indicates the first Self-attention based on size.

[0064] Then, the calculated attention is concatenated along the channel dimension and added to the input, as shown in the following formula:

[0065] in, The representative is the input feature. Represents frequency interaction characteristics.

[0066] Then, the high-frequency interaction features are convolved with the low-frequency interaction features through wavelet deconvolution, concatenated with the original features along the channel dimension, and then convolved again with 1×1, as shown in the following formula:

[0067] in, The representative is the input feature. and These represent low-frequency interaction features and high-frequency interaction features, respectively. This represents the output of WGPM.

[0068] Furthermore, just as three scales of encoded features are obtained during encoding, three scales of decoded features are also obtained during decoding. During decoding, encoded and decoded features at the same scale are concatenated along the channels and then subjected to a 1×1 convolution to obtain skip connection features, as shown in the following formula:

[0069] in, Representing the Skip connectivity features at various scales, Representing the Encoding features at each scale, Representing the Decoding features at each scale.

[0070] Furthermore, the decoded features at the last scale are added to the input image to obtain the enhanced result of WFPFormer, as shown in the following formula:

[0071] in, As input features, This is the decoding feature for the last scale. This represents the output of WFPFormer.

[0072] Furthermore, the loss function for the overall network is set as follows:

[0073] in, Indicates enhanced results, This represents a clear and realistic underwater image. A loss function is calculated, and the model parameters of the underwater image enhancement framework are optimized based on the loss function and the backpropagation process. After training, a trained underwater image enhancement framework is obtained. The input samples are then enhanced using the trained underwater image enhancement framework, and the enhanced underwater image is output.

[0074] The beneficial effects of this invention are as follows: This invention proposes a Frequency-Guided Transformer Module (FGTM) that introduces a frequency learning mechanism into the Transformer architecture, enabling the learning of more refined image features. The proposed Wavelet Refinement Attention Module (WRAM) maps clean low-frequency components to the feature space, enhancing the model's adaptability to complex underwater environments. The proposed Wavelet Guided Cueing Module (WGPM) effectively combines frequency learning and cueing learning, achieving a dynamic association between degradation information and enhancement targets. This invention proposes an end-to-end cueing-based wavelet-guided Transformer network, WFPFormer, for underwater image enhancement. This network combines the advantages of frequency learning and cueing learning, enabling the model to learn more robust features. Experimental results on publicly available datasets demonstrate that the proposed WFPFormer outperforms state-of-the-art underwater image enhancement methods.

[0075] Other advantages, objectives, and features of the invention will be set forth in part in the description which follows, and in part will be apparent to those skilled in the art from the following examination, or may be learned from practice of the invention. The objectives and other advantages of the invention can be realized and obtained through the following description. Attached Figure Description

[0076] To make the objectives, technical solutions, and advantages of the present invention clearer, the preferred embodiments of the present invention will be described in detail below with reference to the accompanying drawings, wherein: Figure 1 This is a flowchart of the method of the present invention; Figure 2 This is a schematic diagram of underwater image acquisition. Figure 3 A diagram of a cue-based wavelet-guided Transformer network (WFPFormer) for underwater image enhancement; Figure 4 This is a structural diagram of the Frequency Guided Transformer module (FGTM) of the present invention; Figure 5 This is a structural diagram of the wavelet refinement attention module (WRAM) of the present invention; Figure 6 This is a structural diagram of the wavelet guidance prompting module (WGPM) of the present invention; Figure 7 This is a scene from the UIEB underwater image enhancement dataset, where (a) is a degraded underwater image, (b) is the underwater image enhancement result of the method described in this invention, and (c) is the ground truth value; Figure 8 Visualizations of different methods on an underwater image augmentation dataset, including (a) WaterNet, (b) SUWNet, (c) PUIE, (d) DICAM, (e) PUGAN, (f) WPFNet, (g) WFPFormer, and (h) ground truth maps. Detailed Implementation

[0077] The technical solution of the present invention will now be described in detail with reference to the accompanying drawings.

[0078] Figure 1 This is a flowchart of the method of the present invention. The present invention provides an underwater image enhancement method based on cue-learning wavelet-guided Transformer. As shown in the figure, in the image acquisition stage, an underwater robot is used to acquire underwater images. The deep learning network used for underwater image enhancement is as follows: Figure 3As shown, he is able to combine the advantages of frequency learning and cue learning to capture details and contextual information in the data to obtain high-quality underwater images.

[0079] The network mainly consists of a Frequency Guided Transformer module (FGTM), a Wavelet Refinement Attention module (WRAM), and a Wavelet Guided Cueing module (WGPM). The input image is enhanced after passing through a multi-stage encoder-decoder structure.

[0080] Each encoding layer consists of a 3×3 convolution and four FGTMs, and each decoding layer consists of a 3×3 convolution, four FGTMs, and one WGPM. Each encoder-decoder layer includes a 3×3 convolution and four FGTMs to capture multi-scale features and global contextual relationships. The FGTM designed in this invention achieves efficient global modeling by incorporating frequency learning into the Transformer. The WRAM designed in this invention improves the model's adaptability to complex environments through wavelet channel pyramid attention and wavelet spatial interaction attention. The WGPM designed in this invention promotes the dynamic relationship between degraded and reconstructed targets by combining frequency learning with cue learning. Specifically, the technical solution of this invention includes the following: 1. Data Acquisition: Underwater image acquisition is achieved using underwater robots. These robots are deployed to diverse underwater environments with varied water scenes, water types, lighting conditions, and types of data collection. (See diagram below.) Figure 2 .

[0081] 2. Data Preprocessing: The degraded images used as the training set are preprocessed. First, data augmentation operations, including random rotation and flipping, are performed on the degraded images. Then, the depth map and gradient map corresponding to each image are obtained through dark channel priors. Finally, the images are processed using the best existing underwater image augmentation methods. The data-augmented underwater images, depth maps, and gradient maps, along with their corresponding images processed by the underwater image augmentation methods, constitute the training sample pairs.

[0082] 3. The input image is processed through three coding layers to obtain multi-scale features. Each coding layer contains one 3×3 convolution and four FGTMs. For each coding layer, the input first undergoes a 3×3 convolution, as shown in the following formula:

[0083] Where X is the input feature map. Indicates the filter size is Convolution operation, This indicates that the filter size is... The result of the convolution operation.

[0084] 4. After the 3×3 convolution, it is immediately passed through an FGTM. The FGTM has three branches: the first branch performs flat pooling, the second branch performs wavelet convolution, and the third branch performs max pooling. The formulas are as follows:

[0085] Where X is the input feature map. For the first The output of each branch, Represents average pooling operation. Represents wavelet convolution operation, This represents the max pooling operation.

[0086] In the second branch, the input undergoes wavelet convolution to obtain one low-frequency component and three high-frequency components, as shown in the following formula:

[0087] in, As input features, , , , These represent the low-frequency components and the high-frequency components in the vertical, horizontal, and diagonal directions, respectively. This represents the wavelet convolution operation.

[0088] For low-frequency components, they are first input into the low-frequency enhancement module for spatial and frequency domain processing. In spatial processing, they are first divided into two equal parts along the channel dimension, and then further divided into query value, key value, and attribute value through three 1×1 convolutions, as shown in the following formula:

[0089] in, For the input feature map, For splitting operations along the channel, This is the first channel of the input features after channel partitioning. part, , , Representing the Partial query values, key values, and attribute values.

[0090] Then, exchange and Multi-head attention calculation is performed using the following formula:

[0091] in, , , Indicates the first Part 1 The query value, key value, and attribute value of the head. Indicates the first Part 1 Self-attention based on size Represents the softmax equation, This represents each head dimension.

[0092] Then, the calculated attention is concatenated along the channel dimension, as shown in the following formula:

[0093] in, Representing the Partial self-attention, This represents a splicing operation along the channel dimension.

[0094] Then, the two self-attention results are concatenated along the channel dimension, and the spatial processing output is obtained through a 1×1 convolution, as shown in the following formula:

[0095] in, This represents the output of spatial processing.

[0096] Then, in the frequency domain processing, the amplitude and phase are first obtained through Fourier transform, as shown in the following formula:

[0097] in, For the input feature map, This indicates the Fourier transform operation. , Amplitude and phase, respectively.

[0098] Then, the amplitude and phase are refined by inputting 1×1 convolution, LeaklyReLU activation function, and 1×1 convolution respectively, as shown in the following formula:

[0099] in, Representing LeaklyReLU, , These represent the refined amplitude and phase characteristics, respectively.

[0100] Then, these features are combined and subjected to an inverse Fourier transform to obtain the output of the frequency domain processing, as shown in the following formula:

[0101] in, Represents the output of frequency domain processing. This represents the inverse Fourier transform operation.

[0102] Subsequently, after processing in the spatial and frequency domains respectively, the features from the two domains are concatenated along the channels, and after a 1×1 convolution, added to the input features to obtain the output of the low-frequency enhancement module, as shown in the following formula:

[0103] in, For the input feature map, The output of spatial processing, The output of frequency domain processing, This represents the output of the low-frequency enhancement module.

[0104] For high-frequency components, they are first concatenated along the channel dimension, as shown in the following formula:

[0105] in, This represents the result of splicing together three high-frequency components.

[0106] Subsequently, the high-frequency components are input into the high-frequency enhancement module, and after being separable convolved along three 3×3 depths with hole ratios of 1, 2, and 3 respectively, they are summed to obtain the output of the high-frequency enhancement module, as shown in the following formula:

[0107] in, For the input feature map, For the output of the high-frequency enhancement module, Indicates the filter size is void ratio Depth-separable convolution operations.

[0108] Then, the output of the first branch is added to the output of the low-frequency enhancement module, and the output of the third branch is added to the output of the high-frequency enhancement module to obtain the enhanced high and low frequency characteristics, as shown in the following formula:

[0109] in, For the input feature map, This represents a strengthened low-frequency characteristic. High-frequency features that represent enhancement.

[0110] Then, the enhanced high- and low-frequency features are input into the cross-frequency interaction attention module. For low-frequency features, they are first divided into query values, key values, and attribute values ​​through three 1×1 convolutions, as shown in the following formula:

[0111] in, This represents a strengthened low-frequency characteristic. , , These represent query values, key values, and attribute values ​​that emphasize low-frequency characteristics.

[0112] Then, multi-head attention is calculated using the following formula:

[0113] in, , , Indicating low-frequency characteristics The query value, key value, and attribute value of the head. Indicates the first Part 1 Attention score based on height.

[0114] Then, for high-frequency features, a 1×1 convolution is first applied to transform them into attribute values, as shown in the following formula:

[0115] in, High-frequency features representing enhancement The attribute value represents the enhancement of high-frequency features.

[0116] Then, the low-frequency attention score is multiplied by the high- and low-frequency attribute values ​​respectively, concatenated along the channel dimension, and added to the original high- and low-frequency features to obtain the high- and low-frequency interaction features, as shown in the following formula:

[0117] in, Representing high-frequency characteristics The attribute value of size. , These represent the low-frequency and high-frequency characteristics after frequency interaction, respectively.

[0118] Then, the low-frequency and high-frequency features after frequency interaction are deconvolved using wavelet to obtain the output of the cross-frequency interaction attention module, as shown in the following formula:

[0119] in, This represents the result of the cross-frequency interaction attention module. This represents the wavelet deconvolution operation.

[0120] Then, the output of the cross-frequency interaction attention module is fed into a gated convolutional feedforward neural network for non-linear interaction. The input features are first passed through two 1×1 convolutions, one result is multiplied by the other using the GELU activation function, and then added to the original input via another 1×1 convolution to obtain the output of the FGTM, as shown in the following formula:

[0121] in, As input features, Represents the output of FGTM, This represents the GELU activation function.

[0122] 5. After three levels of encoding, the features are input into the WRAM, including wavelet channel pyramid attention and wavelet spatial interaction attention. In wavelet channel pyramid attention, the features are first subjected to bilinear interpolation, and then high- and low-frequency features are separated by wavelet convolution, as shown in the following formula:

[0123] in, As input features, , , , These represent the low-frequency components and the high-frequency components in the vertical, horizontal, and diagonal directions, respectively. This represents the bilinear interpolation operation.

[0124] Then, the low-frequency components are processed through three scales of adaptive average pooling and adaptive max pooling to obtain pooling pyramid features. The adaptive average pooling and adaptive max pooling at each corresponding scale are added together and concatenated along the channel dimension. After passing through a multilayer perceptron network, the features are multiplied by the original input using the sigmoid equation, as shown in the following formula:

[0125] in, As input features, Low-frequency components, and These represent compressing the features to a size of [size value]. Adaptive average pooling and adaptive max pooling operations, Represents a multilayer perceptron network. Represents the sigmoid equation. The output represents the wavelet channel pyramid attention.

[0126] Then, in the wavelet spatial interactive attention, the features are also subjected to bilinear interpolation, and then high- and low-frequency features are separated by wavelet convolution. After extracting the low-frequency features, the results of average pooling and max pooling in the channel dimension are added together and concatenated along the channel. After passing through a 7×7 convolution and the sigmoid equation, it is multiplied with the original input, as shown in the following formula:

[0127] in, As input features, Low-frequency components, and These represent the operations performed along the channel dimension after average pooling and max pooling, respectively. The output represents the wavelet space interactive attention.

[0128] Then, the wavelet channel pyramid attention and the wavelet spatial interaction attention are added together to obtain the output of the wavelet refinement attention module, as shown in the following formula:

[0129] in, This represents the output of WRAM.

[0130] 6. After WRAM, decoding begins step by step. The decoding process is identical to the encoding process, except for adding a WGPB before PGTM. In PGTM, the features are first separated into a low-frequency component and three high-frequency components in different directions through wavelet convolution. The three high-frequency components in different directions are then concatenated along the channel and subjected to a 1×1 convolution, as shown in the following formula:

[0131] in, As input features, , , , These represent the low-frequency components and the high-frequency components in the vertical, horizontal, and diagonal directions, respectively. This represents the result of fusing three high-frequency components.

[0132] Then, the high- and low-frequency components interact with two sets of learnable cues, respectively. Taking the high-frequency component as an example, it first undergoes global average pooling and a linear layer, then the attention weights are obtained through a softmax equation. These weights are then multiplied by the high-frequency cues and summed, as shown in the following formula:

[0133] in, The representative is the input feature. This represents the global average pooling operation. For linear layers, Representing the A tip, This indicates a frequency-guided prompt.

[0134] Then, the frequency-guided cue is divided into key values ​​and attribute values ​​through two 1×1 convolutions. The input features are then converted into query values ​​through a 1×1 convolution, as shown in the following formula:

[0135] in, The representative is the input feature. This represents a strengthened low-frequency characteristic. , , Represents query value, key value, and attribute value.

[0136] Then, multi-head attention is calculated using the following formula:

[0137] in, , , Indicates the first The query value, key value, and attribute value of the head. Indicates the first Self-attention based on size.

[0138] Then, the calculated attention is concatenated along the channel dimension and added to the input, as shown in the following formula:

[0139] in, The representative is the input feature. Represents frequency interaction characteristics.

[0140] Then, the high-frequency interaction features are convolved with the low-frequency interaction features through wavelet deconvolution, concatenated with the original features along the channel dimension, and then convolved again with 1×1, as shown in the following formula:

[0141] in, The representative is the input feature. and These represent low-frequency interaction features and high-frequency interaction features, respectively. This represents the output of WGPM.

[0142] 7. During encoding, three scales of encoded features are obtained; similarly, three scales of decoded features are obtained during decoding. During decoding, encoded and decoded features at the same scale are concatenated along the channels and then subjected to a 1×1 convolution to obtain skip connection features, as shown in the following formula:

[0143] in, Representing the Skip connectivity features at various scales, Representing the Encoding features at each scale, Representing the Decoding features at each scale.

[0144] 8. The decoded features of the last scale are added to the input image to obtain the enhanced result of WFPFormer, as shown in the following formula:

[0145] in, As input features, This is the decoding feature for the last scale. This represents the output of WFPFormer.

[0146] 9. The loss function for the overall network is set as follows: ,in, Indicates enhanced results, This represents a clear and realistic underwater image. A loss function is calculated, and the model parameters of the underwater image enhancement framework are optimized based on the loss function and the backpropagation process. After training, a trained underwater image enhancement framework is obtained. The input samples are then enhanced using the trained underwater image enhancement framework, and the enhanced underwater image is output.

[0147] like Figure 8 The experimental results of the WFPFormer underwater image enhancement network described in this invention on an open-source underwater image enhancement dataset show that underwater details are well restored. The reconstruction effect of this invention can be further illustrated through comparative experiments. The method of this invention is compared with other existing methods such as WaterNet, SUWNet, PUIE, DICAM, PUGAN, and WPFNet on the underwater image enhancement dataset. Peak signal-to-noise ratio (PSNR) and structural similarity (SSIM) are calculated respectively. A higher PSNR indicates relatively less distortion and higher hyperspectral image quality; a higher SSIM indicates that the image structure, brightness, and contrast are more similar to the real image, and the image quality is higher. Table 1 shows a comparison of the enhancement results of different methods on the UIEB dataset: Table 1. Comparison of WFPFormer and various methods on the UIEB dataset.

[0148] The method of this invention achieves optimal accuracy on this dataset. Figure 8 Visualized detection results of different methods in a single scene are presented, demonstrating that the method described in this invention outperforms other underwater image enhancement methods. The proposed method can finely restore image details and has advantages over other methods in avoiding artifacts.

[0149] Finally, it should be noted that the above embodiments are only used to illustrate the technical solutions of the present invention and are not intended to limit it. Although the present invention has been described in detail with reference to preferred embodiments, those skilled in the art should understand that modifications can be made to the technical solutions of the present invention without departing from the spirit and scope of the present invention, and all such modifications should be covered within the scope of the claims of the present invention.

Claims

1. A wavelet-guided Transformer underwater image enhancement method based on cue learning, characterized in that: The method includes the following steps: S1: Acquire degraded underwater images; S2: Preprocess the degraded underwater images used as the training set to obtain training sample data; S3: Input the sample into the wavelet-guided Transformer network based on cue learning for training; S4: After training, perform image enhancement on the test samples to obtain the results; In step S1, the underwater robot is used to acquire degraded underwater images, specifically including: deploying the underwater robot to underwater environments with diverse water scenes, water types, lighting conditions, and categories to collect data; In step S2, the degraded images used as the training set are preprocessed. First, the degraded images are subjected to data augmentation operations, including random rotation and flipping. Then, the images are processed using the best existing underwater image augmentation methods. Finally, the data-augmented underwater images and their corresponding images processed by the underwater image augmentation methods constitute the training sample pairs. In step S3, the input image obtains multi-scale features through three coding layers. Each coding layer contains a 3×3 convolution and four FGTMs. For each coding layer, the input first passes through a 3×3 convolution, as shown in the following formula: Where X is the input feature map. The filter size is Convolution operation, This indicates that the filter size is... The result of the convolution operation; Furthermore, after 3×3 convolution, it passes through the frequency-guided Transformer module FGTM. FGTM is divided into three branches: the first branch is average pooling, used to simulate low-frequency features; the second branch is wavelet convolution, used to separate high- and low-frequency features, and then further enhance the high- and low-frequency features; the third branch is max pooling, used to simulate high-frequency features. The low-frequency features simulated by the first branch and the high-frequency features simulated by the third branch are used to further supplement the high- and low-frequency features obtained by the second branch.

2. The wavelet-guided Transformer underwater image enhancement method based on cue learning according to claim 1, in FGTM, is divided into three branches: the first branch is a flat pooling operation, the second branch is a wavelet convolution operation, and the third branch is a max pooling operation, as shown in the following formula: ,in, X is the input feature map. For the first The output of each branch, Represents average pooling operation. Represents wavelet convolution operation. This represents the max pooling operation; In the second branch, the input undergoes wavelet convolution to obtain one low-frequency component and three high-frequency components, as shown in the following formula: ,in, For input features, , , , These represent the low-frequency components and the high-frequency components in the vertical, horizontal, and diagonal directions, respectively. Represents wavelet convolution operation; For low-frequency components, they are first input into the low-frequency enhancement module for spatial and frequency domain processing. In spatial processing, they are first divided into two equal parts along the channel dimension, and then divided into query value, key value, and attribute value through three 1×1 convolutions, as shown in the following formula: ,in, For the input feature map, For splitting operations along the channel, This is the first channel of the input features after channel partitioning. part, , , Representing the Partial query values, key values, and attribute values; Then, exchange and Multi-head attention calculation is performed using the following formula: ,in, , , Indicates the first Part 1 The query value, key value, and attribute value of the head. Indicates the first Part 1 The result of self-attention in terms of size Represents the softmax equation, Represents each head dimension; Then, the calculated attention results are concatenated along the channel dimension, as shown in the following formula: ,in, Representing the Partial self-attention, This represents a splicing operation along the channel dimension; Then, the two self-attention results are concatenated along the channel dimension, and the spatial processing output is obtained through a 1×1 convolution, as shown in the following formula: ,in, The output represents spatial processing; Then, in the frequency domain processing, the amplitude and phase are first obtained through Fourier transform, as shown in the following formula: ,in, For the input feature map, This indicates the Fourier transform operation. , Amplitude and phase respectively; Then, the amplitude and phase are refined by inputting 1×1 convolution, LeaklyReLU activation function, and 1×1 convolution respectively, as shown in the following formula: ,in, Representing the Leakly ReLU equation, , These represent the refined amplitude and phase characteristics, respectively. Then, these features are combined and subjected to an inverse Fourier transform to obtain the output of the frequency domain processing, as shown in the following formula: ,in, Represents the output of frequency domain processing. This represents the inverse Fourier transform operation. Subsequently, after processing in the spatial and frequency domains respectively, the features from the two domains are concatenated along the channels, and after a 1×1 convolution, added to the input features to obtain the output of the low-frequency enhancement module, as shown in the following formula: ,in, For the input feature map, The output of spatial processing, The output of frequency domain processing, This represents the output of the low-frequency enhancement module; For high-frequency components, they are first concatenated along the channel dimension, as shown in the following formula: ,in, This represents the result of splicing together three high-frequency components. Subsequently, the high-frequency components are input into the high-frequency enhancement module, and after being separable convolved along three 3×3 depths with hole ratios of 1, 2, and 3 respectively, they are summed to obtain the output of the high-frequency enhancement module, as shown in the following formula: ,in, For the input feature map, For the output of the high-frequency enhancement module, Indicates the filter size is void ratio Depth-separable convolution operations; Then, the output of the first branch is added to the output of the low-frequency enhancement module, and the output of the third branch is added to the output of the high-frequency enhancement module to obtain the enhanced high and low frequency characteristics, as shown in the following formula: ,in, For the input feature map, This represents a strengthened low-frequency characteristic. High-frequency features representing enhancement; Then, the enhanced high- and low-frequency features are input into the cross-frequency interaction attention module. For low-frequency features, they are first divided into query values, key values, and attribute values ​​through three 1×1 convolutions, as shown in the following formula: ,in, This represents a strengthened low-frequency characteristic. , , Query values, key values, and attribute values ​​that emphasize low-frequency characteristics; Then, multi-head attention is calculated using the following formula: ,in, , , Indicating low-frequency characteristics The query value, key value, and attribute value of the head. Indicates the first Part 1 Attention score based on height; Then, for high-frequency features, a 1×1 convolution is first applied to transform them into attribute values, as shown in the following formula: ,in, High-frequency features representing enhancement Attribute values ​​that represent enhanced high-frequency features; Then, the low-frequency attention score is multiplied by the high- and low-frequency attribute values ​​respectively, concatenated along the channel dimension, and added to the original high- and low-frequency features to obtain the high- and low-frequency interaction features, as shown in the following formula: ,in, Representing high-frequency characteristics The attribute value of size. , These represent the low-frequency and high-frequency characteristics after frequency interaction, respectively. Then, the low-frequency and high-frequency features after frequency interaction are deconvolved using wavelet to obtain the output of the cross-frequency interaction attention module, as shown in the following formula: ,in, This represents the result of the cross-frequency interaction attention module. Represents wavelet deconvolution operation; Then, the output of the cross-frequency interaction attention module is input into a gated convolutional feedforward neural network for non-linear interaction; the input features are first passed through two 1×1 convolutions, one result is multiplied by the other through the GELU activation function, and then added to the original input through another 1×1 convolution to obtain the output of FGTM, as shown in the following formula: ,in, For input features, Represents the output of FGTM, This represents the GELU activation function.

3. The wavelet-guided Transformer underwater image enhancement method based on cue learning according to claim 1, after three-level encoding, the features are input to the wavelet refinement attention module WRAM, including wavelet channel pyramid attention and wavelet spatial interaction attention; in the wavelet channel pyramid attention, the features are first subjected to bilinear interpolation, and then high- and low-frequency features are separated by wavelet convolution, as shown in the following formula: ,in, For input features, , , , These represent the low-frequency components and the high-frequency components in the vertical, horizontal, and diagonal directions, respectively. This represents the bilinear interpolation operation; Then, the low-frequency components are processed through three scales of adaptive average pooling and adaptive max pooling to obtain pooling pyramid features. The adaptive average pooling and adaptive max pooling at each corresponding scale are added together and concatenated along the channel dimension. After passing through a multilayer perceptron network, the features are multiplied by the original input using the sigmoid equation, as shown in the following formula: ,in, For input features, Low-frequency components, and These represent compressing the features to a size of [size value]. Adaptive average pooling and adaptive max pooling operations, Represents a multilayer perceptron network. Represents the sigmoid equation. The output representing the wavelet channel pyramid attention; Then, in the wavelet spatial interactive attention, the features are also subjected to bilinear interpolation, and then high- and low-frequency features are separated by wavelet convolution. After extracting the low-frequency features, the results of average pooling and max pooling in the channel dimension are added together and concatenated along the channel. After passing through a 7×7 convolution and the sigmoid equation, it is multiplied with the original input, as shown in the following formula: ,in, For input features, Low-frequency components, and These represent the operations performed along the channel dimension after average pooling and max pooling, respectively. The output represents the wavelet spatial interactive attention; Then, the wavelet channel pyramid attention and the wavelet spatial interactive attention are added together to obtain the output of the wavelet refinement attention module, as shown in the following formula: ,in, This represents the output of WRAM.

4. The wavelet-guided Transformer underwater image enhancement method based on cue learning as described in claim 1, after WRAM, progressive decoding begins. The decoding process is identical to the encoding process except for the addition of a wavelet-guided cue module WGPM before PGTM. In WGPM, the features are first separated into a low-frequency component and three high-frequency components in different directions through wavelet convolution. The three high-frequency components in different directions are then concatenated along the channel and subjected to a 1×1 convolution, as shown in the following formula: ,in, For input features, , , , These represent the low-frequency components and the high-frequency components in the vertical, horizontal, and diagonal directions, respectively. This represents the result of fusing three high-frequency components; Then, the high- and low-frequency components interact with two sets of learnable cues respectively. Taking the high-frequency component as an example, it first goes through global average pooling and a linear layer, then uses the softmax equation to obtain the attention weights, which are then multiplied by the high-frequency cues and added together, as shown in the following formula: ,in, Represented by input features, This represents the global average pooling operation. For linear layers, Representing the A tip, This indicates a frequency-guided prompt; Then, the frequency-guided cue is divided into key values ​​and attribute values ​​through two 1×1 convolutions. The input features are then converted into query values ​​through a 1×1 convolution, as shown in the following formula: ,in, Represented by input features, This represents a strengthened low-frequency characteristic. , , Query values, key values, and attribute values; Then, multi-head attention is calculated using the following formula: ,in, , , Indicates the first The query value, key value, and attribute value of the head. Indicates the first Self-attention in terms of size; Then, the calculated attention is concatenated along the channel dimension and added to the input, as shown in the following formula: ,in, Represented by input features, This represents the self-attention after being spliced ​​along the channel dimension. Represents frequency interaction characteristics; Then, the high-frequency interaction features are convolved with the low-frequency interaction features through wavelet deconvolution, concatenated with the original features along the channel dimension, and then convolved again with 1×1, as shown in the following formula: ,in, Represented by input features, and These represent low-frequency interaction features and high-frequency interaction features, respectively. Represents the output of WGPM; During the encoding process, three scales of encoded features are obtained, and the decoding process also yields three scales of decoded features. During decoding, the encoded features and decoded features at the same scale are concatenated along the channels and then subjected to a 1×1 convolution to obtain skip connection features, as shown in the following formula: ,in, Representing the Skip connectivity features at various scales, Representing the Encoding features at each scale, Representing the Decoding features at each scale; The final scale's decoded features are added to the input image to obtain the enhanced WFPFormer result, as shown in the following formula: ,in, For input features, This is the decoding feature for the last scale. This is the final output of the model; The loss function for the overall network is set as follows: ,in, Indicates enhanced results, The system represents a clear underwater image; calculates a loss function, optimizes the model parameters of the underwater image enhancement framework based on the loss function and the backpropagation process, and obtains a trained underwater image enhancement framework after training; enhances the input samples using the trained underwater image enhancement framework and outputs an underwater image enhancement effect image.

Images