Underwater image enhancement method based on frequency domain analysis and visual Mama

Abstract

The invention relates to an underwater image enhancement method based on frequency domain analysis and visual Mama, comprising the following steps: constructing a multistage wavelet transform network, the network being an encoder-decoder structure, the encoder comprising a space-frequency feature fusion module and discrete wavelet transform; the decoder comprises a space-frequency feature fusion module and inverse discrete wavelet transform; an underwater image is decomposed into a low-frequency part and a high-frequency part through discrete wavelet transform in an encoder, spatial feature extraction and frequency feature extraction are conducted on the low-frequency part and the high-frequency part through corresponding space-frequency feature fusion modules based on Mama and frequency domain convolution respectively, and splicing is conducted on the low-frequency part and the high-frequency part along a channel in the decoder according to the sequence of low frequency and high frequency. Reconstruction is carried out through inverse discrete wavelet transform; and inputting the underwater image data into the multistage wavelet transform network to obtain enhanced underwater image data. Compared with the prior art, the method has the advantages of improving the enhancement quality, reducing the complexity, quickly enhancing and the like.

Description

Technical Field

[0001] The present invention relates to the technical field of underwater image enhancement, and in particular to an underwater image enhancement method based on frequency domain analysis and visual Mamba. Background Art

[0002] With the continuous development of human beings in the ocean, underwater detection equipment such as underwater automatic cruisers are widely used in underwater research work such as underwater biological detection, underwater environmental monitoring and underwater resource exploration. Among them, underwater imaging is an important task. Underwater imaging contains three elements, namely direct light, forward scattering and backscattering. Direct light is the light directly reflected from underwater objects to the camera, forward scattering is the light reflected from underwater objects but scattered at a small angle, and backscattering is the light reflected from underwater tiny objects. Forward scattering and backscattering will affect the quality of underwater imaging. Forward scattering will cause the structure of underwater objects to be blurred, and backscattering will cover the details, textures and boundaries of underwater objects. In addition, light of different wavelengths will attenuate with depth in water. Among them, the red light with the longest wavelength attenuates the fastest, while blue light and green light attenuate slowly, which will cause the underwater image to present a blue-green hue. Therefore, people need an underwater image enhancement method to restore the image details and texture and correct the color.

[0003] At present, there are three main methods for underwater image enhancement, namely traditional methods, methods based on physical models, and methods based on deep learning. However, the existing methods have the following problems:

[0004] 1) Traditional methods mainly adjust image pixels directly, such as histogram equalization, image fusion, and gamma correction. These methods ignore the principles of underwater imaging and are prone to under / over-enhancement.

[0005] 2) The physical model-based method mainly estimates the parameters in the physical imaging model (atmospheric scattering model, etc.) through physical priors, and then infers the model to obtain clear underwater images, such as dark channel priors, minimum information loss, blur priors, etc. However, it is challenging to accurately estimate the parameters in the physical imaging model. At the same time, the physical prior is sensitive to the scene and is prone to under / over-enhancement and artifacts. At the same time, the time complexity of this method is high, resulting in a long processing time.

[0006] 3) Deep learning-based methods mainly use supervised learning, and the network structure is mainly based on encoder-decoder. By learning the nonlinear relationship between the degraded image and the reference image in the paired data set, end-to-end underwater image enhancement is achieved. As a representation learning method, underwater image enhancement based on deep learning, like other image processing tasks, effective global modeling is important for the quality of image enhancement. Existing deep learning-based methods often use convolution or self-attention to extract features, which have two problems: 1) Convolution is good at extracting local features, but lacks the ability to extract global features, resulting in poor generalization; 2) Although self-attention can extract global features, it has a high computational complexity and requires a lot of computing resources, which is not suitable for underwater environments with limited computing resources.

[0007] Existing methods are mainly studied in the spatial domain, and relatively few are studied in the frequency domain.

[0008] How to achieve underwater image enhancement quickly, with less computing resources and high quality has become a technical problem that needs to be solved. Summary of the invention

[0009] The purpose of the present invention is to overcome the defects of the above-mentioned prior art and provide an underwater image enhancement method based on frequency domain analysis and visual Mamba.

[0010] The purpose of the present invention can be achieved by the following technical solutions:

[0011] According to one aspect of the present invention, there is provided an underwater image enhancement method based on frequency domain analysis and visual Mamba, the method comprising:

[0012] Step S1, acquiring underwater image data;

[0013] Step S2, constructing a multi-level wavelet transform network based on frequency domain analysis and visual Mamba, the network is an encoder-decoder structure, wherein the encoder includes a space-frequency feature fusion module and a discrete wavelet transform; the decoder includes a space-frequency feature fusion module and an inverse discrete wavelet transform;

[0014] In the encoder, discrete wavelet transform decomposes underwater images into low-frequency and high-frequency parts. The low-frequency and high-frequency parts are extracted by the corresponding space-frequency feature fusion modules based on Mamba and frequency domain convolution respectively. The space-frequency feature fusion module is used to process the data in the decoder, and the data are spliced ​​along the channel in the order of low frequency first and high frequency later, and then reconstructed by inverse discrete wavelet transform.

[0015] Step S3, inputting the underwater image data into the multi-level wavelet transform network constructed in step S2 to obtain enhanced underwater image data.

[0016] Preferably, the space-frequency feature fusion module includes a parallel space feature extraction part and a frequency feature extraction part;

[0017] The data processed by discrete wavelet transform firstly passes through the linear layer to expand the channel and is divided into two equal parts along the channel, and then passes through the spatial feature extraction part and the frequency feature extraction part to extract the spatial and frequency features respectively;

[0018] The spatial feature extraction part includes layer normalization and Mamba-based spatial feature extraction submodules;

[0019] The frequency feature extraction part includes fast Fourier transform, two frequency feature extraction submodules based on convolution and inverse fast Fourier transform, wherein the two frequency feature extraction submodules based on convolution respectively extract features of the amplitude and phase of the spectrum diagram.

[0020] More preferably, the Mamba-based spatial feature extraction part includes a visual Mamba submodule and a Mamba-based channel attention submodule;

[0021] The visual Mamba submodule selectively scans the expanded input data along four different directions to obtain global dependencies;

[0022] The Mamba-based channel attention submodule includes adaptive average pooling and Sigmoid activation function. Mamba is used to selectively scan the results of adaptive average pooling in two directions, and then the Sigmoid activation function is used to generate a channel attention map.

[0023] More preferably, the space-frequency feature fusion module further includes a feed-forward neural network for splicing the space feature extraction part and the frequency feature extraction part and performing data screening;

[0024] The feedforward neural network includes layer normalization and gated attention submodules.

[0025] More preferably, the gated attention submodule includes a depthwise separable convolution, a SiLU activation function, and a linear layer for channel adjustment; the input data is subjected to a depthwise separable convolution to expand the channel and is divided into two equal parts along the channel, one of which is subjected to a SiLU activation function to obtain an attention map, and the other is inner-producted with the attention map to output the result.

[0026] Preferably, the multi-level wavelet transform network is an n-layer encoder-decoder structure, where n is greater than or equal to 3;

[0027] The first layer of the encoder includes a convolution module and a discrete wavelet transform; the second to nth layers of the encoder include multiple space-frequency feature fusion modules, convolution modules and discrete wavelet transform;

[0028] The first to nth layers of the decoder include multiple space-frequency feature fusion modules, convolution modules and inverse discrete wavelet transform; the last layer of the decoder is a refinement layer, which includes two convolution modules and multiple space-frequency feature fusion modules.

[0029] More preferably, skip connections are established between the 1st to nth layers of the encoder and the 1st to nth layers of the decoder respectively.

[0030] More preferably, the process of the decoder reconstructing the enhanced underwater image includes:

[0031] The nth layer of the decoder uses the corresponding space-frequency feature fusion module to process the low-frequency and high-frequency parts respectively, splices the processed results along the channel in the order of low frequency first and high frequency later, and reconstructs the spliced ​​results using inverse discrete wavelet transform;

[0032] The convolution result in the nth layer of the encoder and the output of the nth layer of the decoder are concatenated along the channel as the input of the n-1th layer of the decoder;

[0033] The n-1 layer of the decoder first uses convolution to reduce the number of channels of the input data, and then uses the space-frequency feature fusion module to process the result after the channel reduction, and splices the processing result with the high-frequency processing result of the n-1 layer of the encoder along the channel. The spliced ​​result is input to the inverse discrete wavelet transform for reconstruction, and the reconstruction result is jump-connected with the convolution output result of the n-1 layer of the encoder and used as the input of the n-2 layer of the decoder;

[0034] The processing flow of the decoder from layer n-2 to layer 1 is consistent with that of the decoder from layer n-1;

[0035] The refinement layer of the decoder first uses convolution to reduce the number of channels of the input data, then uses the spatial-frequency feature fusion module to refine the convolution result, and finally enhances the underwater image by the convolution output.

[0036] Preferably, the number of space-frequency feature fusion modules in each layer in the encoder and the decoder is different.

[0037] Preferably, the loss function L of the multi-level wavelet transform network is total The calculation is as follows:

[0038] L total =λ1L L1 +λ2L edge +θ3L fft +θ4Ladv

[0039] Among them, L L1 L1 loss is used to calculate the absolute value of the difference between the predicted value and the true value; L edge is the boundary loss, which is used to calculate the difference between the boundaries of two images; L fft is the frequency loss, which is used to calculate the difference between the phase and amplitude of the two images; L adv It is the adversarial loss, which is used to measure the difference between the data generated by the generator and the real data, and the ability of the discriminator to distinguish between the real data and the generated data; λ1, λ2, λ3, and θ4 are all hyperparameters.

[0040] Compared with the prior art, the present invention has the following beneficial effects:

[0041] 1) The present invention constructs a multi-level wavelet transform underwater image enhancement network based on frequency domain analysis and visual Mamba, uses wavelet transform to separate high-frequency and low-frequency parts, and fuses features from both spatial domain and frequency domain to enhance underwater images, involving a small number of parameters, wherein Mamba is used to process spatial domain data, and convolution is used to process frequency domain data, the frequency domain data includes the phase and amplitude of the spectrum graph, and the degraded area in the underwater image is better enhanced according to the phase and amplitude in the spectrum graph; Mamba extracts the global dependency of the image with linear complexity, so that the model can effectively improve the visual quality of the underwater image with less computing resources; thereby greatly improving the image visual quality, computing complexity, and computing resources.

[0042] 2) The present invention reflects the degradation of underwater images in both low-frequency and high-frequency parts, processes the low-frequency and high-frequency parts separately, extracts spatial features to restore detail information, and also extracts frequency domain features to restore structural information, and fuses the two features to enhance the quality of underwater images.

[0043] 3) The present invention introduces the visual Mamba submodule and the channel attention submodule of Mamba, selectively scans the expanded input data along different directions to obtain global dependencies, and selectively scans the global information sequence of adaptive average pooling from two directions, and then uses Sigmoid to output the channel attention map, and performs inner product with the input to emphasize important channels and suppress unimportant channels, thereby solving the channel redundancy problem of Mamba and improving the underwater image enhancement capability. BRIEF DESCRIPTION OF THE DRAWINGS

[0044] Figure 1 It is a flow chart of the underwater image enhancement method based on frequency domain analysis and visual Mamba in the present invention;

[0045] Figure 2It is a schematic diagram of the structure of the multi-level wavelet transform network in the present invention;

[0046] Figure 3 It is a structural schematic diagram of the space-frequency feature fusion module in the present invention;

[0047] Figure 4 It is a schematic diagram of the spatial feature extraction submodule based on Mamba in the present invention;

[0048] Figure 5 It is a schematic diagram of the structure of the visual Mamba submodule in the present invention;

[0049] Figure 6 It is a schematic diagram of the structure of the Mamba channel attention submodule in the present invention;

[0050] Figure 7 Schematic diagram of the structure of the gated attention submodule in the present invention;

[0051] Figure 8 It is a frequency feature extraction module based on convolution in the present invention;

[0052] In the accompanying figure, AvgPool: adaptive average pooling, Conv: convolution, ConvBlock: frequency feature extraction based on convolution, Channel: channel attention submodule based on Mamba, DWConv: depthwise separable convolution, FFT: fast Fourier transform, iFFT: inverse fast Fourier transform, Gated Attention: gated attention submodule, Leaky ReLU: LeakyReLU activation function, Linear: linear layer, LayerNorm: layer normalization, Sigmoid: Sigmoid activation function, SiLU: SiLU activation function, Spatial Mamba: visual Mamba submodule, VSSM: spatial feature extraction submodule based on Mamba. DETAILED DESCRIPTION

[0053] The following will be combined with the drawings in the embodiments of the present invention to clearly and completely describe the technical solutions in the embodiments of the present invention. Obviously, the described embodiments are part of the embodiments of the present invention, not all of the embodiments. Based on the embodiments of the present invention, all other embodiments obtained by ordinary technicians in this field without creative work should fall within the scope of protection of the present invention.

[0054] In view of the problems of poor generalization, high computational complexity and slow speed of existing methods, first of all, we believe that the generalization of deep learning-based methods is stronger than that of traditional methods and physical model-based methods, and the time complexity of deep learning-based methods is less than or equal to that of physical model-based methods and traditional methods. Therefore, with the main purpose of improving the generalization ability of deep learning-based methods and reducing their time complexity, we use Mamba and convolution from both the spatial domain and the frequency domain to achieve effective and fast underwater image enhancement with fewer parameters and lower time complexity, thereby improving the visual quality of underwater images.

[0055] This embodiment relates to an underwater image enhancement method based on frequency domain analysis and visual Mamba, such as Figure 1 , including the following steps:

[0056] S1. Obtain and construct underwater image dataset

[0057] S2, establish a multi-level wavelet transform network based on frequency domain analysis and visual Mamba;

[0058] S3. Design the loss function of the multi-level wavelet transform network;

[0059] S4, train the model using underwater image dataset;

[0060] S5. Test the model using underwater image dataset;

[0061] S6. Use the model to enhance real underwater images.

[0062] The multi-level wavelet transform network based on frequency domain analysis and visual Mamba is a three-layer encoder-decoder structure. The first layer of the encoder includes a 3×3 convolution and discrete wavelet transform; the second layer of the encoder includes multiple space-frequency feature fusion modules, a 3×3 convolution and discrete wavelet transform; the third layer of the encoder includes multiple space-frequency feature fusion modules, a 3×3 convolution and discrete wavelet transform.

[0063] The first layer of the decoder includes multiple space-frequency feature fusion modules, one 3×3 convolution and inverse discrete wavelet transform; the second layer of the decoder includes multiple space-frequency feature fusion modules, one 3×3 convolution and inverse discrete wavelet transform; the third layer of the decoder includes multiple space-frequency feature fusion modules and one inverse discrete wavelet transform. The last layer of the decoder is the refinement layer, which includes two 3×3 convolutions and multiple space-frequency feature fusion modules. The number of space-frequency feature fusion modules in each layer is different.

[0064] A skip connection is established between the first layer of the encoder and the first layer of the decoder; a skip connection is established between the second layer of the encoder and the second layer of the decoder; and a skip connection is established between the third layer of the encoder and the third layer of the decoder.

[0065] The discrete wavelet transform decomposes the input underwater image into low-frequency and high-frequency data and downsamples the input data. The high-frequency and low-frequency data are processed separately by the corresponding space-frequency feature fusion modules. The inverse discrete wavelet transform reconstructs and upsamples the input low-frequency and high-frequency data.

[0066] The spatial-frequency feature fusion module is divided into three parts, including spatial feature extraction, frequency feature extraction and feedforward neural network. Spatial feature extraction includes layer normalization and Mamba-based spatial feature extraction submodules; frequency feature extraction includes fast Fourier transform, two convolution-based frequency feature extraction submodules and inverse fast Fourier transform; feedforward neural network includes layer normalization and gated attention submodules. The data processed by discrete wavelet transform first expands the channel through a linear layer and is divided into two equal parts along the channel, and then undergoes spatial feature extraction and frequency feature extraction respectively. Finally, the feedforward neural network performs data screening on the splicing results of spatial feature extraction and frequency feature extraction.

[0067] The Fast Fourier Transform (FFT) transforms the spatial domain data into a spectrogram, and the Inverse Fast Fourier Transform (IFFT) transforms the spectrogram into the spatial domain.

[0068] The Mamba-based spatial feature extraction module contains two linear layers for adjusting channels, one 3×3 depthwise separable convolution, one SiLU activation function, one visual Mamba submodule, one Mamba-based channel attention submodule, and one layer normalization. The input underwater image data first passes through the linear layer to expand the channel, then passes through the depthwise separable convolution, SiLU activation function, visual Mamba submodule, Mamba-based channel attention submodule, and layer normalization. Finally, the linear layer restores the number of channels to the original size.

[0069] The visual Mamba submodule selectively scans the expanded input data along four different directions to obtain global dependencies.

[0070] The Mamba-based channel attention submodule includes an adaptive average pooling and a Sigmoid activation function. Mamba is used to selectively scan the results of the adaptive average pooling in two directions, and then the Sigmoid activation function is used to generate a channel attention map.

[0071] The gated attention submodule includes a 3×3 depthwise separable convolution, a SiLU activation function, and a channel-adjusted linear layer. The input data is expanded through the depthwise separable convolution and divided into two equal parts along the channel. One part is activated through the SiLU activation function to obtain the attention map, and the other part is inner-producted with the attention map to obtain the result.

[0072] The convolution-based frequency feature extraction module includes two 1×1 convolutions and two LeakyReLU activation functions. The input data passes through 1×1 convolution, LeakyReLU activation function, 1×1 convolution, and LeakyReLU activation function in sequence.

[0073] Step S1: Obtain and construct an underwater image dataset. Obtain an underwater image dataset from the Internet, including a paired underwater image dataset and an unpaired underwater image dataset. Currently commonly used underwater image datasets include EUVP, UIEB, LSUI, OceanDark, RUIE, etc. The underwater image dataset is obtained by randomly selecting a training set and a test set. The training set is a paired underwater image dataset, and the test set is a paired underwater image dataset and an unpaired underwater image dataset.

[0074] Step S2: Establish a multi-level wavelet transform network based on frequency domain analysis and visual Mamba.

[0075] Figure 2The overall structure of the multi-level wavelet transform network based on frequency domain analysis and visual Mamba is presented. In the first layer of the encoder, the input image is subjected to 3×3 convolution to extract shallow features and double the channels. Then, the feature map is decomposed into 1 equal low-frequency subband and 3 equal high-frequency subbands using discrete Haar wavelet transform. This allows the data to be downsampled and the computational resources required at the initial resolution to be reduced with little information loss. In the second layer of the encoder, different numbers of spatial-frequency feature fusion modules are set according to the number of input data channels to process the input high-frequency and low-frequency data respectively. Since the degradation of underwater images is reflected in both low-frequency and high-frequency, we believe that if the data is processed directly without high- and low-frequency separation, the restored structure and detail information will be insufficient. Therefore, we process the high-frequency and low-frequency separately. Afterwards, the channels of the low-frequency processing results are doubled using 3×3 convolution to enhance the representation ability of the model, and the low-frequency data is further decomposed using discrete Haar wavelet transform. The processing flow of the third layer of the encoder is consistent with that of the second layer of the encoder. The third layer of the decoder uses different numbers of spatial-frequency feature fusion modules to process low-frequency and high-frequency data, splices the processed results along the channel in the order of low frequency first and high frequency later, and reconstructs the spliced ​​results using the inverse discrete Haar wavelet transform. The result after the inverse discrete Haar wavelet transform belongs to low-frequency data, and the jump connection can assist the decoder in reconstruction. Therefore, a jump connection is introduced in the third layer of the decoder, and the result after convolution in the third layer of the encoder and the output of the third layer of the decoder are spliced ​​along the channel as the input of the second layer of the decoder. The second layer of the decoder first uses 3×3 convolution to reduce the number of channels of the input data by 1 times, and then uses multiple spatial-frequency feature fusion modules to process the result after the channel reduction, splices the processing result with the high-frequency processing result of the second layer of the encoder along the channel, and inputs the spliced ​​result into the inverse discrete Haar wavelet transform for reconstruction. The reconstructed result is jump-connected with the convolution output result of the second layer of the encoder and used as the input of the first layer of the decoder. The processing flow of the first layer of the decoder is consistent with that of the second layer of the decoder. The last layer of the decoder (after the first layer) is the refinement layer, which first uses a 3×3 convolution to reduce the number of channels of the input data by a factor of 1, then uses multiple spatial-frequency feature fusion modules to refine the convolution result, and finally outputs the result through a 3×3 convolution.

[0076] Figure 3The structure of the spatial-frequency feature fusion module is shown. The input data is passed through a linear layer to double the number of channels, which is convenient for subsequent channel division. The advantage of spatial domain processing is that it can restore detail information, while the advantage of frequency domain processing is that it can restore structural information. We believe that combining the two can enhance the representation learning ability of the model and improve the enhancement effect. Therefore, we input the equally divided results into the Mamba-based spatial feature extraction module and the convolution-based frequency feature extraction module respectively. Then, the output results of the two modules are spliced ​​along the channel and input into the gated attention submodule. Among them, the fast Fourier transform (FFT) is used to transform the data from the spatial domain to the frequency domain, and then the absolute value and arctan() of the spectrum graph are calculated to obtain the phase and amplitude. Studies have shown that the amplitude can reflect the degradation of the image to a certain extent. Therefore, we input the phase and amplitude into the convolution-based feature extraction module respectively. Considering that the numerical range of the phase is [-Π~Π], the convolution-based feature extraction module corresponding to the phase removes the last LeakyReLU activation function. The feature map is then converted into a spectrum map, and the spectrum map is then converted into the spatial domain by the inverse fast Fourier transform. Finally, the spatial domain results are input to the layer normalization and gated attention submodules. In order to speed up training and prevent gradient disappearance, we introduce two residual connections in the module. The formula flow of the spatial-frequency feature fusion module is as follows:

[0077] x_1,x_2=Chunk(Linear(Input)) (1)

[0078] x_2^Freq=FFT(x_2) (2)

[0079] x_2^Phase=abs(x_2^Freq) (3)

[0080] x_2^Amplitude=arctan(x_2^Freq) (4)

[0081] x_2^Real=ConvBlock(x_2^Amplitude)×cos(x_2^Phase) (5)

[0082] x_2^Imag=ConvBlock(x_2^Amplitude)×sin(x_2^Phase) (6)

[0084] 〖Output〗_mid=Linear(Cat[LN(VSSM(x_1)),iFFT(x_2^real+x_2^imagi)])+Input (7)

[0085] Output=LN(GatedAttention(〖Output〗_mid))+〖Output〗_mid (8)

[0086] Where Input represents the decomposed data of the input underwater image, Linear represents the linear layer, Chunk represents the channel division, x_1 and x_2 represent the results after division, FFT represents the fast Fourier transform, x_2^Freq represents the spectrum of x2, abs represents the absolute value, x_2^Phase represents the phase, x_2^Amplitude represents the amplitude, ConvBlock represents the frequency feature extraction module based on convolution, x_2^Real represents the real part of the complex number, x_2^Imag represents the imaginary part of the complex number, VSSM represents the spatial feature extraction module based on Mamba, LN represents layer normalization, Cat represents the splicing operation, iFFT represents the inverse fast Fourier transform, Output mid Represents the intermediate result, GatedAttention represents the gated attention submodule, and Output represents the output.

[0087] Figure 4 The structure of the spatial feature extraction module based on Mamba is shown. The input data first passes through a linear layer to expand the number of channels, and then is input into a depthwise separable convolution to extract local features and a SiLU activation function to increase nonlinearity. Mamba is a selective state-space model that can effectively extract the global dependencies of data with linear complexity. With the help of Mamba's powerful global modeling capabilities, we introduce the visual Mamba submodule. Figure 5 As shown in the figure, the input data is selectively scanned in the four directions (upper left to lower right, lower right to upper left, upper right to lower left, and lower left to upper right) as shown in the figure, and the scan results in the four directions are added to get the output result. Considering the channel redundancy problem of Mamba, we designed channel attention based on Mamba. Figure 6 As shown in the figure, we first use adaptive average pooling to extract the global information of the data, then use Mamba to selectively scan the global information sequence in two directions (from front to back and from back to front), and finally use Sigmoid to output the channel attention map. The attention map is inner-producted with the input to emphasize important channels and suppress unimportant channels. After the visual Mamba submodule and the Mamba-based channel attention processing, the results are output through layer normalization and linear layers. The formula flow of the Mamba-based spatial feature extraction module is as follows:

[0088] Output=Linear(LN(ChannelaMamba(SpatialMamba(SiLU(DWConv(Linear(Input)))))))) (9)

[0091] Among them, DWConv stands for depthwise separable convolution, SiLU stands for SiLU activation function, SpatialMamba stands for visual Mamba submodule, and ChannelMamba stands for Mamba-based channel attention module.

[0092] Figure 7 The structure of the gated attention submodule is shown. Gated attention extracts local features through depthwise separable convolution, and then divides it into two equal parts. One of them is activated by SiLU function to generate an attention map, and the attention map is inner-producted with the other part, thus achieving control of information flow. The formulaic process of gated attention is as follows:

[0093] x1,x2=Chunk(DWConv(Input)) (10)

[0094] Output = SiLU (x1) × x2 (11)

[0095] Figure 8 The structure of the frequency feature extraction module based on convolution is shown. Since the information at each position of the spectrum graph represents global information, we use 1×1 convolution to process the spectrum graph, and then use the LeakyReLU activation function to increase nonlinearity. The formula process of the frequency feature extraction module based on convolution is as follows:

[0096] Output=LeakyReLU(Conv(LeakyReLU(Conv(Input)))) (12)

[0097] Where Conv stands for convolution and LeakyReLU stands for the LeakyReLU activation function.

[0098] Step 3: Design a loss function. Training a deep neural network requires a loss function to constrain the network.

[0099] The loss function L of the multi-level wavelet transform network based on frequency domain analysis and visual Mamba total as follows:

[0100] L total =λ1L L1 +λ2L edge +λ3L fft +λ4L adv (1)

[0101] Where LL1 is the L1 loss, L edge is the boundary loss, L fft is the frequency loss, L adv To combat the loss, λ1, λ2, λ3, and λ4 are hyperparameters.

[0102] L1 loss measures the absolute value of the difference between the predicted value and the true value. Using L1 loss can reduce the impact of noise and outliers on the reconstruction results.

[0103] Boundary loss calculates the difference between the boundaries of two images. First, the image is extracted based on the Sobel operator, and then the L1 loss of the boundary map of the reconstructed image and the reference image is calculated. Boundary loss allows the model to focus on the texture of the image, so that the reconstruction result retains more texture information.

[0104] Frequency loss calculates the difference between the phase and amplitude of the two images. The image is converted to the frequency domain using the Fast Fourier Transform, the phase and amplitude of the spectrogram are calculated separately, and finally the L1 loss between the phase and amplitude of the reconstructed image and the reference image is calculated. Frequency loss can be used to optimize the frequency representation of the image, thereby more effectively repairing the degraded area.

[0105] Adversarial loss is used to measure the difference between the data generated by the generator and the real data, as well as the ability of the discriminator to distinguish between real data and generated data. PatchGAN acts as a discriminator. Through adversarial loss, the data generated by the generator will be more realistic.

[0106] Step 4: Train the model using the underwater image dataset

[0107] We take the original images in the paired underwater image dataset as input and the reference images as the objects of supervised learning.

[0108] The number of iterations for training the model is set to 300,000.

[0109] The optimizer for training the model is the AdamW optimizer, the learning rate is initially set to 0.0003, and the betas are set to (0.9, 0.999).

[0110] The learning rate scheduler for the training model is the cosine annealing learning rate scheduler with an adjustment interval of [92000, 208000], which decreases the learning rate from 0.0003 to 0.000001 within the interval.

[0111] Step 5: Test the network using the underwater image dataset. Use the test set in the underwater image dataset to test the enhancement effect of the model, and use multiple evaluation indicators to evaluate the enhanced images.

[0112] The evaluation indicators used in the present invention include reference evaluation indicators and non-reference evaluation indicators. The reference evaluation indicators include peak signal-to-noise ratio (PSNR), structural similarity (SSIM) and mean square error (MSE), and the non-reference evaluation indicators include underwater image quality measurement (UIQM), underwater color image quality evaluation index (UCIQE) and normalized image quality evaluation index (NIQE).

[0113] Peak signal-to-noise ratio (PSNR) is a reference evaluation indicator used to measure the quality and denoising effect of an image. It evaluates the quality of image enhancement by comparing the difference between the enhanced image and the reference image. The higher the PSNR value, the closer the enhanced image is to the reference image and the smaller the distortion.

[0114] Structural Similarity (SSIM) is a reference evaluation metric that aims to measure the structural similarity between two images. Structural similarity takes into account the brightness, contrast, and structural changes of the image, so it can better reflect the human visual system's perception of image quality. The SSIM index ranges from 0 to 1, where 1 means that the two images are completely similar and 0 means that they are completely different. The higher the SSIM value, the better the image quality.

[0115] The mean square error (MSE) is a reference evaluation index used in image enhancement to quantify the difference between the enhanced image and the reference image. The image quality is evaluated by calculating the mean square of the difference in pixel values ​​between the two images. The lower the MSE value, the closer the enhanced image is to the reference image, and the better the image quality.

[0116] The underwater image quality measurement index (UIQM) is a reference-free evaluation index. It uses the color measurement index (UICM), the clarity measurement index (UISM), and the contrast measurement index (UIConM) as the evaluation basis for the degradation mechanism and imaging characteristics of underwater images, and expresses UIQM as a linear combination of the three. The higher the value of UIQM, the better the quality of the enhanced image.

[0117] The Underwater Color Image Quality Evaluation Index (UCIQE) is a non-reference evaluation index. UCIQE linearly combines color concentration, saturation and contrast to quantitatively evaluate the color cast, blur and low contrast of underwater images. The higher the UCIQE value, the better the quality of the enhanced image.

[0118] The Normalized Image Quality Evaluation (NIQE) is a non-reference evaluation metric. NIQE evaluates poor visual quality in an image, such as noise, blur, and distortion, without relying on specific features of the image content. NIQE values ​​are usually between 0 and 1, with lower values ​​indicating better image quality and higher values ​​indicating worse image quality.

[0119] Step 6: Use the model to enhance real underwater images.

[0120] An underwater robot or diver uses a camera to capture underwater images, inputs the underwater images into a trained model, and outputs enhanced underwater images.

[0121] The above is only a specific embodiment of the present invention, but the protection scope of the present invention is not limited thereto. Any technician familiar with the technical field can easily think of various equivalent modifications or replacements within the technical scope disclosed by the present invention, and these modifications or replacements should be included in the protection scope of the present invention. Therefore, the protection scope of the present invention shall be based on the protection scope of the claims.

Claims

1. An underwater image enhancement method based on frequency domain analysis and visual Mamba, characterized in that: The method comprises: Step S1, acquiring underwater image data; Step S2, constructing a multi-level wavelet transform network based on frequency domain analysis and visual Mamba, the network is an encoder-decoder structure, wherein the encoder includes a space-frequency feature fusion module and a discrete wavelet transform; the decoder includes a space-frequency feature fusion module and an inverse discrete wavelet transform; In the encoder, discrete wavelet transform decomposes underwater images into low-frequency and high-frequency parts. The low-frequency and high-frequency parts are extracted by the corresponding space-frequency feature fusion modules based on Mamba and frequency domain convolution respectively. The space-frequency feature fusion module is used to process the data in the decoder, and the data are spliced ​​along the channel in the order of low frequency first and high frequency later, and then reconstructed by inverse discrete wavelet transform. Step S3, inputting the underwater image data into the multi-level wavelet transform network constructed in step S2 to obtain enhanced underwater image data.

2. The underwater image enhancement method based on frequency domain analysis and visual Mamba according to claim 1, characterized in that: The space-frequency feature fusion module includes a parallel space feature extraction part and a frequency feature extraction part; The data processed by discrete wavelet transform firstly passes through the linear layer to expand the channel and is divided into two equal parts along the channel, and then passes through the spatial feature extraction part and the frequency feature extraction part to extract the spatial and frequency features respectively; The spatial feature extraction part includes layer normalization and Mamba-based spatial feature extraction submodules; The frequency feature extraction part includes fast Fourier transform, two frequency feature extraction submodules based on convolution and inverse fast Fourier transform, wherein the two frequency feature extraction submodules based on convolution respectively extract features of the amplitude and phase of the spectrum diagram.

3. The underwater image enhancement method based on frequency domain analysis and visual Mamba according to claim 2 is characterized in that: The Mamba-based spatial feature extraction part includes a visual Mamba submodule and a Mamba-based channel attention submodule; The visual Mamba submodule selectively scans the expanded input data along four different directions to obtain global dependencies; The Mamba-based channel attention submodule includes adaptive average pooling and Sigmoid activation function. Mamba is used to selectively scan the results of adaptive average pooling in two directions, and then the Sigmoid activation function is used to generate a channel attention map.

4. The underwater image enhancement method based on frequency domain analysis and visual Mamba according to claim 2 is characterized in that: The space-frequency feature fusion module also includes a feed-forward neural network for splicing the space feature extraction part and the frequency feature extraction part and performing data screening; The feedforward neural network includes layer normalization and gated attention submodules.

5. The underwater image enhancement method based on frequency domain analysis and visual Mamba according to claim 4 is characterized in that: The gated attention submodule includes a depthwise separable convolution, a SiLU activation function, and a linear layer for channel adjustment; the input data is subjected to a depthwise separable convolution to expand the channel and is divided into two equal parts along the channel, one of which is subjected to a SiLU activation function to obtain an attention map, and the other is inner-producted with the attention map to output the result.

6. The underwater image enhancement method based on frequency domain analysis and visual Mamba according to claim 1, characterized in that: The multi-level wavelet transform network is an n-layer encoder-decoder structure, where n is greater than or equal to 3; The first layer of the encoder includes a convolution module and a discrete wavelet transform; the second to nth layers of the encoder include multiple space-frequency feature fusion modules, convolution modules and discrete wavelet transform; The first to nth layers of the decoder include multiple space-frequency feature fusion modules, convolution modules and inverse discrete wavelet transform; the last layer of the decoder is a refinement layer, which includes two convolution modules and multiple space-frequency feature fusion modules.

7. The underwater image enhancement method based on frequency domain analysis and visual Mamba according to claim 6, characterized in that: The 1st to nth layers of the encoder respectively establish jump connections with the 1st to nth layers of the decoder.

8. The underwater image enhancement method based on frequency domain analysis and visual Mamba according to claim 6, characterized in that: The decoder reconstructs the enhanced underwater image through the following steps: The nth layer of the decoder uses the corresponding space-frequency feature fusion module to process the low-frequency and high-frequency parts respectively, splices the processed results along the channel in the order of low frequency first and high frequency later, and reconstructs the spliced ​​results using inverse discrete wavelet transform; The convolution result in the nth layer of the encoder and the output of the nth layer of the decoder are concatenated along the channel as the input of the n-1th layer of the decoder; The n-1th layer of the decoder first uses convolution to reduce the number of channels of the input data, and then uses the space-frequency feature fusion module to process the result after the channel reduction, and splices the processing result with the high-frequency processing result of the n-1th layer of the encoder along the channel, and inputs the splicing result into the inverse discrete wavelet transform for reconstruction. The reconstructed result is combined with the encoder n The convolution output of layer -1 is skip-connected and used as the input of the decoder layer n-2; The processing flow of the decoder from layer n-2 to layer 1 is consistent with that of the decoder from layer n-1; The refinement layer of the decoder first uses convolution to reduce the number of channels of the input data, then uses the spatial-frequency feature fusion module to refine the convolution result, and finally enhances the underwater image by the convolution output.

9. The underwater image enhancement method based on frequency domain analysis and visual Mamba according to claim 1, characterized in that: The number of spatial-frequency feature fusion modules per layer is different in the encoder and decoder.

10. The underwater image enhancement method based on frequency domain analysis and visual Mamba according to claim 1, characterized in that: The loss function L of the multi-level wavelet transform network is total The calculation is as follows: L total =λ1L L1 +λ2L edge +λ3L fft +λ4L adv Among them, L L1 L1 loss is used to calculate the absolute value of the difference between the predicted value and the true value; L edge is the boundary loss, which is used to calculate the difference between the boundaries of two images; L fft is the frequency loss, which is used to calculate the difference between the phase and amplitude of the two images; L adv It is the adversarial loss, which is used to measure the difference between the data generated by the generator and the real data, and the ability of the discriminator to distinguish between the real data and the generated data; λ1, λ2, λ3, and λ4 are all hyperparameters.

Images