Underwater image enhancement method based on Mama multi-feature enhancement fusion

Abstract

The invention provides an underwater image enhancement method combining high-frequency enhancement, space-frequency mixed feature extraction and a Mama enhancement fusion technology. The method aims at solving the problems of low contrast ratio, color distortion, detail loss and the like caused by absorption and scattering effects in the light propagation process of the underwater image, the image quality is improved, and the visual perception effect is improved. The method specifically comprises the following four main steps: firstly, acquiring and preprocessing initial data sets of different underwater scenes; secondly, extracting texture detail information of the image through a texture enhancement module; thirdly, further enhancing texture details by using a space-frequency mixed feature extraction module; and finally, in combination with the original image, the texture information and the spatial frequency domain enhancement result, realizing feature fusion of the image through a Mama enhancement fusion module, thereby effectively recovering details of the underwater image, correcting color deviation and improving the definition and visual effect of the image. According to the method, the enhancement performance of the underwater image in a complex environment can be improved, and powerful support is provided for applications such as underwater detection, target detection, object recognition and image segmentation.

Description

Technical Field

[0001] The invention belongs to the field of underwater image enhancement, and relates to an underwater image enhancement method combining high frequency enhancement, space-frequency mixed feature extraction and Mamba enhancement fusion. Background Art

[0002] The ocean covers about 70% of the Earth's surface and has rich biological and mineral resources. It plays a vital role in the Earth's ecosystem and human society. With the increasing demand for underwater research, resource exploration and engineering operations, underwater image enhancement (UIE) technology plays a key role in autonomous underwater vehicles (AUVs), remotely operated vehicles (ROVs) and other fields. These systems are highly dependent on high-quality visual data for navigation, target detection and environmental analysis. However, underwater imaging faces many challenges, especially the particularity of underwater light propagation characteristics, which significantly affects the quality of the image.

[0003] The main degradation problems of underwater imaging come from the absorption and scattering effects of light, which lead to low contrast, color distortion, loss of detail and blurring of images. The absorption characteristics of water cause long-wave light (such as red light) to attenuate significantly, while short-wave light (such as blue-green light) is easily scattered, causing the image to appear blue or green. In addition, factors such as dissolved organic matter, suspended particles and microorganisms in the water cause additional scattering and absorption during the propagation of light, resulting in underwater images that are not only distorted in color, but also often accompanied by high noise, further reducing the clarity and effectiveness of the image.

[0004] In order to solve these problems, researchers have proposed a variety of underwater image enhancement techniques, including contrast enhancement, frequency domain enhancement, color constancy restoration and image fusion. However, these traditional methods usually focus on optimizing single image features, and lack robustness in applications with different water qualities and underwater scenes, making it difficult to effectively cope with complex and changeable underwater environments. With the rapid development of deep learning technology, more and more studies have applied deep learning methods to underwater image enhancement, effectively improving image quality through automated feature extraction and nonlinear mapping, especially in dehazing, denoising and color correction. However, due to the high complexity of the underwater environment, the problem of underwater image degradation has not been completely solved. Therefore, improving the quality of underwater images, especially for problems such as color distortion, detail recovery and image blur, is still a major challenge facing the current field of image enhancement. Improving the quality of underwater images is of vital importance for applications such as underwater detection, target detection, object recognition and image segmentation.

[0005] In response to the above challenges, this paper proposes an underwater image enhancement method that combines high-frequency enhancement, spatial-frequency hybrid feature extraction, and Mamba enhancement fusion technology. This method aims to improve the image enhancement performance of the model in complex underwater environments, effectively address the problem of underwater image degradation, and improve the visual quality and detail performance of the image. Summary of the invention

[0006] In view of this, the technical problem to be solved by the present invention is to provide an underwater image enhancement method that combines high-frequency enhancement, spatial-frequency hybrid feature extraction, and Mamba enhancement fusion. The method can enhance the detail information of degraded underwater images, correct color deviation, and improve visual perception effects.

[0007] The technical solution adopted by the present invention to solve the technical problem is as follows:

[0008] An underwater image enhancement method of high frequency enhancement, spatial-frequency hybrid feature extraction and Mamba enhancement fusion, including:

[0009] Step 1) obtaining an initial data set of different underwater scenes and preprocessing and dividing the images in the data set;

[0010] Step 2) inputting the preprocessed image into a high-frequency enhancement module to extract high-frequency detail information of the original underwater image;

[0011] Step 3) inputting the extracted high-frequency details into a space-frequency hybrid feature extraction module to further enhance the high-frequency detail information;

[0012] Step 4) inputting the original image information, high-frequency detail information and spatial-frequency domain enhancement information into the Mamba enhancement fusion module for feature fusion;

[0013] Further, the step 1) specifically includes the following steps:

[0014] Step 11) Underwater dataset division: The public underwater image dataset is divided into training set, validation set and test set in a ratio of 8:1:1. At the same time, the provided no-reference images are used as challenging images to verify the generalization of the model.

[0015] Further, the step 2) specifically includes the following steps:

[0016] Step 21) extracting high-frequency details: first, the initial underwater image is decomposed into high-frequency and low-frequency information by wavelet transform, and then the high-frequency information in the horizontal, vertical and diagonal directions is spliced ​​with the low-frequency information respectively, and finally the stacking operation in the channel dimension is used to extract the high-frequency details of the image.

[0017] Step 22) High-frequency detail enhancement: Repeat step 21) three times to enhance the high-frequency detail information of the original image.

[0018] Further, the step 3) specifically includes the following steps:

[0019] Step 31) Multi-scale local features: First, the feature map containing high-frequency details is segmented in the channel dimension, and 1x1, 3x3, 5x5, and 7x7 convolutions are used to extract multi-scale local features, and then a stacking operation in the channel dimension is performed.

[0020] Step 32) frequency domain operation, Fourier transform the feature map containing local features in step 31) into the frequency domain, use 1x1 convolution to enhance the expression of information in the frequency domain, then perform inverse Fourier transform to the spatial domain, and finally perform residual connection to supplement the original information of the local features.

[0021] Further, the step 4) specifically includes the following steps:

[0022] Step 41) Global features: The original image is expanded in channel dimension using 1x1 convolution and input into the Mamba enhanced fusion module as global features.

[0023] Step 42) local features, concatenate the high-frequency details in step 22) and the frequency domain enhanced local features in step 32) in the channel dimension, and adjust the stacked local features to the resolution of the global features by upsampling.

[0024] Step 43) Global-local feature fusion: concatenate the global features and local features in the channel dimension and perform feature fusion through 3x3 convolution. BRIEF DESCRIPTION OF THE DRAWINGS

[0025] In order to make the purpose, technical solutions and advantages of the present invention more clear, the present invention will be further described in detail below with reference to the accompanying drawings, in which:

[0026] Figure 1 This is the flow chart of the method;

[0027] Figure 2 This is the overall model framework diagram of this method;

[0028] Figure 3 This is the high frequency extraction structure diagram proposed by this method;

[0029] Figure 4 This is the spatial-frequency hybrid feature extraction structure diagram proposed in this method;

[0030] Figure 5 Diagram of the Mamba structure proposed for this method;

[0031] Figure 6This is the multi-scale fusion structure diagram proposed by this method; DETAILED DESCRIPTION

[0032] The preferred embodiments of the present invention will be described in detail below with reference to the accompanying drawings.

[0033] The present invention provides an underwater image enhancement method based on Mamba multi-feature enhancement fusion, the method flow chart is as follows: Figure 1 As shown in the figure, the overall model framework diagram proposed by this method is as follows Figure 2 As shown, the method comprises the following steps:

[0034] Step 1) divide the underwater image enhancement dataset;

[0035] Step 2) constructing a high-frequency feature extraction module for extracting the original image;

[0036] Step 3) constructing a hybrid feature extraction module in the space-frequency domain;

[0037] Step 4) construct the Mamba multi-scale feature fusion module;

[0038] Step 5) training the underwater image enhancement network;

[0039] Step 6) testing the underwater image enhancement network;

[0040] Further, step 1) specifically includes the following steps:

[0041] Step 11) Obtain an open source underwater image enhancement dataset, which contains single object underwater scene images, complex underwater scene images, and underwater images with different color shift conditions.

[0042] Step 12) Divide the obtained data set into a training set, a test set, and a validation set in a ratio of 8:1:1.

[0043] Further, step 2) specifically includes the following steps:

[0044] Step 21) Figure 3 As shown, firstly, discrete wavelet transform (DWT) is performed on the original image to obtain low-frequency component yL and high-frequency component yH, and horizontal detail coefficient y_HL, vertical detail coefficient y_LH and diagonal detail coefficient y_HH are extracted from the high-frequency component respectively, and the low-frequency component yL is concatenated with the high-frequency components in three directions along the channel dimension, and the concatenated tensor is reduced in the channel dimension. The two-dimensional wavelet transform of the image can be expressed as formula (1), and the Haar wavelet is simplified to formula (2).

[0045] DWT(x)=(LL,HL,LH,HH) (1)

[0046] Among them, LL represents the low-frequency part, HL, LH and HH are the high-frequency parts, which contain different detailed information of the image.

[0047]

[0048] The Haar wavelet decomposes the signal by a simple piecewise linear function.

[0049] Step 22) The original image is subjected to a third discrete wavelet transform (DWT) to obtain local features.

[0050] Further, step 3) specifically includes the following steps:

[0051] Step 31) Figure 4 As shown in the figure, the spatial-frequency domain hybrid feature extraction module is the core module of detail enhancement, which includes multi-scale local feature extraction and frequency domain operations. First, the tensor is input to the multi-scale local feature extraction module, which divides the tensor into 4 parts in the channel dimension, and performs depth-separable convolution using 1x1, 3x3, 5x5, and 7x7 convolution kernels respectively. Finally, the tensors after convolution are spliced ​​together.

[0052] Step 32) The spliced ​​tensor is transformed by two-dimensional Fourier transform to obtain the real part and the imaginary part, and then the real part and the imaginary part are merged, and then more useful features are further extracted through convolution layer, batch normalization and activation function to enhance the expression ability of the network. The two-dimensional Fourier transform is expressed as formula (3)

[0053]

[0054] x(m,n) is the image pixel value in the time domain.

[0055] X(u,v) are the Fourier coefficients in the frequency domain.

[0056] M and N are the height and width of the input image respectively.

[0057] is the rotation factor and i is the imaginary unit.

[0058] X(u,v)=X real (u,v)+·X imag (u,v) (4)

[0059] X real (u,v) is the real part of the complex number.

[0060] X imag (u,v) is the imaginary part of the complex number.

[0061] Further, step 4) specifically includes the following steps:

[0062] Step 41) Figure 5 As shown in the figure, the input tensor is first adjusted in dimension to obtain branch one and branch two. The temporal features are extracted by convolution of branch one and nonlinearity is added by activation. Then, the Selective SSM (state space model) is input, as shown in formula (5) (6). Branch two is activated as a residual connection and multiplied with the output of branch one, and finally the dimensionality is reduced and output.

[0063] Step 42) first adjusts the dimension of the input tensor, uses the module of step 41) to extract forward features and backward features from the input tensor respectively, adds the forward features to the backward features, and finally adjusts the tensor dimension output.

[0064] h′(t)=Ah(t)+Bx(t) (5)

[0065] y(t)=Ch(t)+Dx(t) (6)

[0066] Step 43) Figure 6 As shown, the outputs of step 2 and step 3 are spliced ​​together and upsampled to the size of the original image, the output is spliced ​​together with the output of step 4, feature fusion is achieved through convolution, and then the dimensionality is reduced.

[0067] Further, step 5) specifically includes the following steps:

[0068] Step 51) Configure the corresponding environment, set training parameters, and carry out training tasks;

[0069] Step 52) The training environment configured is: Ubuntu 22.04, CUDA 11.8, Python 3.10.14 and Pytorch 2.3.0. The hardware equipment includes CPU: Intel(R) Xeon(R) CPU E5-2680 v4@2.40GHz, memory is 32GB; GPU: NVIDIARTX 3090, memory is 24GB.

[0070] Step 53) The training parameters are configured as follows: the input image size is 256×256, the batch size is set to 4, the learning rate is 0.0001, and the number of training rounds is 500.

[0071] Further, step 6) specifically includes the following steps:

[0072] Step 61) Load the trained weight file into the underwater image enhancement network, select pictures from the test data set for enhancement, and return the enhanced underwater image.

[0073] Finally, it should be noted that the above embodiments are only used to illustrate the technical solutions of the present invention rather than to limit it. Although the present invention has been described in detail through the above embodiments, those skilled in the art should understand that various changes can be made in form and details without departing from the scope defined by the claims of the present invention.

Claims

1. An underwater image enhancement method based on Mamba multi-feature enhancement fusion. Characterized by: The following steps are involved: Step 1: Collect and preprocess relevant underwater image enhancement datasets: Download open source underwater image enhancement datasets from public websites and divide them into training set, test set and validation set in a ratio of 8:1:

1. The scenes include underwater complex environment map, underwater biological cluster map and underwater biological individual map. Pair the initial dataset, that is, one underwater degraded image corresponds to one underwater reference image. Step 2, build a high-frequency detail enhancement module: The core of this module is the wavelet transform decomposition operation, which divides the image into high-frequency and low-frequency parts. The high-frequency component contains the details and edge information of the image, while the low-frequency component represents the basic contour and lighting information. On this basis, the module enhances the image detail performance by splicing high-frequency information and low-frequency information in the horizontal, vertical and diagonal directions. Then, the channel dimension stacking operation is used to fuse the spliced ​​information to further enrich the image detail features, thereby improving the image clarity and detail perception ability. This process effectively enhances the high-frequency details in the original image and significantly improves the detail recovery and visual effects of underwater images. Step 3, construct a spatial-frequency domain information enhancement module: The core of this module is the combination of multi-scale feature extraction and frequency domain information enhancement. First, by segmenting the feature map containing high-frequency details in the channel dimension, 1x1, 3x3, 5x5 and 7x7 convolution kernels are used to extract local features of different scales. Small convolution kernels are used to capture tiny local details, while large convolution kernels help capture local features in a larger range. Then, through the channel dimension stacking operation, the feature maps of the two scales are fused to form a richer multi-scale local feature representation. Then, the feature map enters the frequency domain through Fourier transformation, and the frequency domain features are further enhanced by 1x1 convolution to optimize the expression of its frequency components. Subsequently, the inverse Fourier transform is performed to convert the frequency domain features back to the spatial domain, and the original local features are fused with the enhanced frequency domain features through residual connection to enhance the image details. By effectively combining spatial and frequency domain information, this module not only improves the expression of local features of the image, but also enhances the global perception ability and optimizes the overall detail performance of the image. Step 4, construct a global-local feature fusion module: The core of this module is the effective fusion of global features and local features. First, the original image is expanded in channel dimension through 1x1 convolution to extract global features, and then input into the Mamba enhanced fusion module for further processing. Next, the high-frequency detail features extracted in step 2 and the frequency domain enhanced local features in step 3 are spliced ​​in the channel dimension, and then the spliced ​​local features are adjusted to the same resolution as the global features through upsampling operation to ensure that the global and local features are fused at the same scale. Finally, the global features are fused with the local features using 3x3 convolution to further improve the accuracy and overall expressiveness of image detail recovery. Through this fusion of global and local features, the module can effectively enhance the image detail recovery capability, improve the image performance quality, and thus optimize the visual effect of underwater images. Step 5, training the underwater image enhancement network model: This step uses paired original underwater images and reference underwater images as network input data, and adopts a weighted fusion strategy to reasonably match the structural loss, gradient loss and pixel loss to supervise the training of the underwater image enhancement network. In this process, by combining different weights of various loss functions, the performance of the model is optimized, and it is expected that the model weights with high-quality enhancement effects can be obtained through training, thereby improving the network's ability to enhance underwater images. Step 6, test the underwater image enhancement network model: After the model training is completed, the obtained weight file is loaded into the underwater image enhancement network, and the image enhancement test is performed using a paired test set. During the test, the model will return the enhanced underwater image and calculate relevant performance indicators, such as structural similarity (SSIM) and peak signal-to-noise ratio (PSNR). In order to further verify the generalization ability of the model, this step also includes the processing of non-reference images, and the image quality evaluation indicators (such as color, saturation and contrast) are used to comprehensively evaluate the enhancement effect and performance of the model in complex underwater environments.

2. An underwater image enhancement method based on Mamba multi-feature enhancement fusion according to patent claim 1. It is characterized in that: The high-frequency detail enhancement module constructed in step 2 processes underwater images in the following steps: First, the underwater image is decomposed into high-frequency and low-frequency information through wavelet transform. High-frequency information mainly contains the details and edges of the image, while low-frequency information reflects the basic outline and illumination changes of the image. Next, the high-frequency information and low-frequency information in the horizontal, vertical and diagonal directions are spliced ​​to enhance the detail performance of the image. Subsequently, the spliced ​​high-frequency and low-frequency information are merged through the stacking operation of the channel dimension to further enhance the expressiveness of the image details. Finally, the module effectively enhances the high-frequency details in the original image in this way, significantly improves the clarity and detail perception of the image, and thus improves the visual effect of underwater image enhancement.

3. According to the underwater image enhancement method based on Mamba multi-feature enhancement fusion described in patent claim 1, it is characterized in that: The spatial-frequency information enhancement module designed in step 3 integrates the techniques of multi-scale feature extraction and frequency domain enhancement. First, the feature map containing high-frequency details is segmented into channels, and 1x1, 3x3, 5x5 and 7x7 convolution kernels are used to extract local features of different scales. The 1x1 and 3x3 convolution kernels focus on capturing subtle local information, while the 5x5 and 7x7 convolution kernels extract features in a larger range. Then, the feature maps of these two scales are fused through the stacking operation of the channel dimension to form a richer multi-scale local feature representation. Subsequently, the feature map is converted to the frequency domain, processed in the frequency domain by Fourier transform, and then the information in the frequency domain is enhanced by 1x1 convolution to optimize the performance of its frequency components. After the inverse Fourier transform, the frequency domain features are reconverted to the spatial domain, and the enhanced frequency domain information is combined with the original local features through residual connection to further enhance the image detail performance. Through this interaction between the spatial domain and the frequency domain, the module not only enhances the expressiveness of image details, but also improves the overall perception of the image.

4. The underwater image enhancement method based on Mamba multi-feature enhancement fusion according to patent claim 1 is characterized in that: The global-local feature fusion module constructed in step 4 is designed to effectively integrate global features and local features. First, the original image is expanded in its channel dimension through a 1x1 convolution operation to extract global features, and is input into the Mamba enhanced fusion module for further processing. Next, the high-frequency detail features extracted in step 2 and the local features enhanced in the frequency domain in step 3 are spliced ​​in the channel dimension, and the spliced ​​local features are adjusted to the same resolution as the global features through an upsampling operation, thereby ensuring that the global and local features are fused at a unified scale. Finally, the global and local features are further fused through a 3x3 convolution to optimize the detail recovery and overall performance of the image. Through this organic combination of global and local features, the module significantly improves the detail expression and recovery effect of the image, ensuring that the image is enhanced in multiple dimensions.

Images