An underwater image enhancement method based on wavelet Mamba

Abstract

The application relates to a wavelet Mamba-based underwater image enhancement method, which aims at the problems of color deviation, low contrast and fuzzy details of underwater images, extracts shallow features of underwater low-quality images and inputs the multi-scale coding network, uses a wavelet Mamba unit to jointly enhance low-frequency color information and high-frequency detail features, combines layer-by-layer compression of down-sampling to obtain global semantic features, completes global modeling of color deviation correction and contrast enhancement, restores the spatial resolution through up-sampling, combines a jump connection with the wavelet Mamba unit to realize reconstruction and enhancement of texture details and color information, finally maps the features to an image domain, fuses the features with input image residuals, and generates enhanced underwater images with natural colors, clear textures and balanced contrast; the wavelet Mamba unit is used to realize joint feature modeling of the frequency domain and the spatial domain, and the color authenticity, structural clarity and detail perceptibility of the underwater images are significantly improved.

Description

Technical Field

[0001] This invention relates to the fields of computer vision and underwater image enhancement technology, specifically to an underwater image enhancement method based on wavelet Mamba, which can be applied to underwater exploration, ocean monitoring, underwater robots and other application fields. Background Technology

[0002] Underwater imaging is affected by various physical factors such as light absorption, scattering, and interference from suspended particles, resulting in common problems in images including color shift, insufficient brightness, decreased contrast, and blurred texture. These phenomena are significantly aggravated with increasing water depth and changes in lighting conditions. Especially in low-light or turbid water environments, the red light channel of the image decays rapidly, resulting in an overall blue-green or grayish-white color cast, blurred target boundaries, and severe degradation of detail. Since image quality directly affects the reliability of tasks such as underwater exploration, ocean monitoring, navigation of unmanned underwater vehicles, and underwater robotic operations, effectively enhancing underwater images has become a core research direction in the field of image restoration.

[0003] Existing technical approaches mainly fall into four categories: traditional methods based on dark channels, dehazing models, histograms, or Gamma transformations are highly sensitive to water type and lighting conditions, prone to color shifts, over-enhancement, or local artifacts, and lack context-adaptive capabilities; convolutional neural networks (CNNs) can improve local textures end-to-end, but they essentially rely on fixed-size local convolutional kernels, making it difficult to capture long-distance dependencies, resulting in poor global consistency and blurred high-frequency details; Transformers achieve global modeling through self-attention, but operate only in the spatial domain, which is insufficient for recovering periodic textures and high-frequency boundaries, and their quadratic complexity limits high-resolution or real-time applications.

[0004] Recent wavelet-depth hybrid methods only treat the frequency domain branch as an additional module, failing to form end-to-end joint optimization with the backbone, resulting in low performance ceilings. The emerging Mamba state-space model achieves long-range modeling with linear complexity, providing a new approach to image enhancement. However, current technologies still suffer from problems such as insufficient color shift and contrast, loss of texture and detail, inadequate frequency domain modeling, insufficient end-to-end optimization, and high computational complexity.

[0005] In view of this, it is very necessary to provide an underwater image enhancement method based on wavelet Mamba to solve the above-mentioned defects in the prior art. Summary of the Invention

[0006] To overcome the problems of insufficient color shift and contrast, loss of texture and detail, inadequate frequency domain modeling, insufficient end-to-end optimization, and high computational complexity in existing underwater image enhancement methods, this invention proposes an underwater image enhancement method based on wavelet Mamba. By combining the U-Net multi-scale coding and decoding architecture with the wavelet frequency domain decomposition mechanism, a Mamba state-space model is introduced at each scale to achieve efficient modeling of global dependencies, enabling simultaneous capture of low-frequency color attenuation information and high-frequency texture detail features of underwater images. Through the design of a lightweight gated feedforward network, this method achieves joint enhancement in the frequency and spatial domains, improving the network's overall performance in color restoration, contrast enhancement, and detail reconstruction, while significantly reducing model complexity and computational overhead while maintaining high enhancement quality, further improving the underwater image enhancement effect. This solves the aforementioned technical problems existing in the prior art.

[0007] To achieve the above objectives, the present invention provides the following technical solution:

[0008] In a first aspect, the technical solution of the present invention provides an underwater image enhancement method based on wavelet Mamba, specifically including the following steps:

[0009] Step S1, the step of extracting shallow features, in which:

[0010] Input a low-quality underwater image and perform shallow feature extraction; the shallow features include the image's basic color and texture information;

[0011] Step S2, the multi-scale coding and feature compression step, in which:

[0012] The extracted shallow features are input into a multi-scale coding network. During the coding process, wavelet Mamba units are used to perform wavelet frequency domain decomposition and state space modeling on the image to extract low-frequency color information and high-frequency detail features. Global semantic features are obtained by combining downsampling and layer-by-layer compression.

[0013] Step S3, the deep representation and global modeling of multi-scale fused features, in which:

[0014] Deep representation of multi-scale fusion features is performed in the bottleneck layer of the network to complete the global modeling of underwater image color shift correction and contrast enhancement.

[0015] Step S4, the step of upsampling to restore resolution and reconstructing and enhancing texture details and color information, in which:

[0016] Spatial resolution is gradually restored through upsampling, and texture details and color information are reconstructed and enhanced by combining skip connections and wavelet Mamba units.

[0017] Step S5, mapping features to the image domain and residual fusion to generate an enhanced underwater image, in which:

[0018] Features are mapped to the image domain and residually fused with the input image to generate enhanced underwater images with natural colors and clear textures.

[0019] In step S1:

[0020] Shallow feature extraction is performed on low-quality underwater images, and the shallow features include the image's basic color and texture information.

[0021] The specific implementation method of the shallow feature extraction is as follows:

[0022] Input a raw, low-quality underwater image, specifically represented as follows:

[0023]

[0024] in, Represents underwater images, Represents the real number field. The dimensions of the image are represented by H (height, number of pixels vertically), W (width, number of pixels horizontally), and 3 (number of channels, usually RGB three channels).

[0025] Furthermore, shallow features from the original low-quality underwater image are extracted using convolutions with a kernel size of 3×3, and used as input to the encoder network. This process is represented as follows:

[0026]

[0027] in, The input underwater image represents the raw data to be processed, with dimensions of [dimension number missing]. (Height × Width × 3 channels, RGB color image); This indicates a convolution operation with a kernel size of 3×3; 3×3 is the spatial size of the convolution kernel (3 pixels high and 3 pixels wide), which is a classic convolution kernel specification for extracting local features in computer vision and can effectively capture shallow details such as edges and textures of images. This represents the shallow feature map obtained after a 3×3 convolution, and is the output of the convolution operation.

[0028] Furthermore, The specific scope is as follows:

[0029]

[0030] in, This represents the shallow feature map obtained after initial processing, which serves as the input for subsequent convolution operations. This indicates that the values ​​of the feature map belong to the real number domain, that is, the feature values ​​are real numbers; These represent the height and width of the feature map, respectively, describing the spatial dimensions of the feature map; C represents the number of channels in the feature map, which reflects the number of different feature types contained in the feature map and is the basis for extracting diverse features in convolution operations.

[0031] Step S1 extracts shallow features through 3×3 convolution, which can quickly capture the basic texture and color information of the image, providing a foundation for subsequent feature extraction and enhancement. Shallow feature extraction is simple to operate and has low computational cost, which can effectively reduce the computational complexity of the model, while retaining the key information of the image, providing support for the subsequent separation of high-frequency and low-frequency features.

[0032] In step S2:

[0033] The coding network of this invention employs a multi-scale wavelet Mamba encoder based on a U-Net symmetric structure. The encoder is composed of multiple cascaded wavelet Mamba units, each of which consists of wavelet frequency domain decomposition and inverse transform, Mamba state space modeling and downsampling units, enabling the gradual extraction of deep semantic features from high resolution to low resolution.

[0034] Furthermore, the U-Net is a classic encoder-decoder architecture widely used in image processing. Its symmetrical design ensures consistency in the number of layers and modules between the encoder and decoder, facilitating feature extraction and fusion at different scales. This structure not only helps to gradually reduce the spatial resolution of the image to extract deep semantic features, but also facilitates the gradual restoration of the image's spatial resolution and the fusion of multi-scale features during the decoding stage, thereby achieving high-quality image reconstruction.

[0035] Furthermore, the encoder is constructed from multiple cascaded wavelet Mamba units. Each wavelet Mamba unit layer contains the following key operations:

[0036] The input feature map is decomposed into low-frequency components (such as color information) and high-frequency components (such as texture information) using two-dimensional discrete wavelet transform (2D-DWT). The low-frequency components are used to extract the structural and color information of the image, while the high-frequency components are used to extract the texture and detail information. After processing, the low-frequency and high-frequency components are recombined using two-dimensional inverse wavelet transform (IWT) to generate the reconstructed feature map.

[0037] In low-frequency components, a Mamba state-space model is introduced to model long-range dependencies and enhance global consistency. The Mamba state-space model, with its linear complexity and long-range modeling capabilities, effectively handles global information in images, overcoming the limitations of traditional convolutional models. In high-frequency components, 3×3 convolution operations are used for detail enhancement, improving image texture clarity and edge information.

[0038] By progressively reducing the spatial resolution of the feature map through downsampling operations (such as convolutional downsampling or pooling) while increasing the number of channels, deeper semantic features can be extracted. This process helps the network better understand the overall structure and semantic information of the image.

[0039] Furthermore, the encoder starts with a high-resolution input image and progressively reduces the spatial resolution through multiple layers of wavelet Mamba units while simultaneously extracting deep semantic features. Each wavelet Mamba unit not only handles low-frequency and high-frequency components but also enhances global consistency through the Mamba state-space model and improves texture clarity through detail enhancement operations. Through this progressive extraction method, the network can extract rich semantic features from shallow to deep layers, providing high-quality feature representations for subsequent image enhancement and reconstruction.

[0040] Through the above design, the coding network of the present invention can effectively extract deep semantic features of underwater images while maintaining global consistency and detailed information, providing a solid foundation for subsequent image enhancement and reconstruction.

[0041] Furthermore, the encoding process is as follows: Input shallow features, and extract deep features using 4 layers of wavelet Mamba units. The deep features are represented as follows:

[0042]

[0043] in, This represents the deep feature map obtained after processing by four layers of wavelet Mamba units. It is the result of shallow features after multiple layers of feature extraction and transformation. This indicates that the values ​​of the deep feature map belong to the real number domain, that is, the feature values ​​are real numbers; These represent the height and width of the deep feature map, respectively. These dimensions are the same as the input shallow feature map size. This indicates that after multiple processing steps (such as sampling), the spatial resolution of the feature map is reduced, allowing for a greater focus on extracting global semantic features. This represents the number of channels in the deep feature map; where C is the number of channels in the shallow feature map. After processing by 4 layers of wavelet Mamba units, the number of channels becomes 8 times the original, which means that the feature map contains richer and more abstract feature information, which is helpful for subsequent global modeling and image enhancement tasks.

[0044] Furthermore, in the encoder, starting with high-resolution input, the encoder gradually reduces the spatial size while expanding the channel capacity;

[0045] The specific encoding process is as follows:

[0046] The input features are decomposed in the frequency domain using wavelet Mamba units and then globally modeled using a state-space model. The input features are represented as follows:

[0047]

[0048] in, The input feature map to the wavelet Mamba unit is the object of subsequent frequency domain decomposition and global modeling operations; i indicates that the value of the feature map belongs to the real number domain; H×W represents the spatial size of the feature map; C represents the number of channels of the feature map.

[0049] Furthermore, the wavelet Mamba unit performs frequency domain decomposition on the input features, with the following specific steps:

[0050] The wavelet Mamba unit decomposes features using a two-dimensional discrete wavelet transform (2D-DWT); low-frequency components are obtained to extract the image structure and three high-frequency components, and further, texture details of the image are extracted; the discrete wavelet transform formula is as follows:

[0051]

[0052] in, Indicates the characteristics of low-frequency components. This represents the characteristics of three high-frequency components. It is a two-dimensional discrete wavelet transform. This represents the input features.

[0053] Furthermore, the wavelet Mamba unit is divided into two branches: a low-frequency branch and a high-frequency branch.

[0054] Furthermore, in the low-frequency branch, the low-frequency components are input into a 3×3 convolutional activation block, and the low-frequency components are represented as follows:

[0055]

[0056] in, This represents the feature tensor (i.e., low-frequency component) of the low-frequency branch after wavelet decomposition. "LL" is a common name for the "low-frequency-low-frequency" subband in wavelet transform, representing the result after low-frequency filtering of the original signal in both the horizontal and vertical directions. The number of elements in the low-frequency component (i.e., the size of the tensor) represents the total number of elements in the low-frequency component. Wavelet transform typically has a "downsampling" property, which reduces the resolution of the low-frequency branch; therefore, the number of elements is a key feature of the original component. 1 / 2 of W.

[0057] Furthermore, its height and width dimensions are flattened into one dimension to generate a feature tensor in sequence form; this operation integrates spatial features while preserving the channel dimension, which facilitates subsequent modeling of channel directions.

[0058] Furthermore, low-frequency information is enhanced through Mamba state space modeling and a 3×3 convolution, resulting in the enhanced low-frequency output features, denoted as:

[0059]

[0060] in, This is the final output feature tensor of the low-frequency branch after complete enhancement processing.

[0061] This process can be represented as:

[0062]

[0063] in, Representing the visual state space model, This indicates the LayerNorm operation. express Activation function This represents the final output feature tensor of the low-frequency branch after complete enhancement processing. This represents the characteristic tensor of the low-frequency branch after wavelet decomposition. This indicates a convolution operation with a kernel size of 3×3.

[0064] Furthermore, in the high-frequency branch, the three high-frequency components are directly input into a 3×3 convolutional module to obtain the enhanced high-frequency features, specifically as follows:

[0065]

[0066] in, There are three high-frequency components; High-Low components represent features in the high-frequency-low-frequency direction and high-frequency-low-frequency dimension, typically corresponding to high-frequency details such as edges and textures in a certain direction of an image. The Low-High component represents the characteristics of the low-frequency to high-frequency direction and the high-frequency to low-frequency dimension, and is also a form of high-frequency detail representation; The "High-High" component represents the characteristics of the high-frequency-high-frequency direction and the high-frequency-low-frequency dimension, containing rich high-frequency detail information.

[0067] This process can be represented as:

[0068]

[0069] in, It is a high-frequency enhancement feature. This is a 3×3 convolution operation module. This is the enhanced set of three high-frequency components obtained after processing by a 3×3 convolution module.

[0070] Furthermore, the low-frequency enhancement features and high-frequency enhancement features are input into the inverse wavelet transform (IWT) module to reconstruct the frequency domain features back into the image space, obtaining the final output features:

[0071]

[0072] in, These are the four sub-band components of the input; Inverse wavelet transform is the inverse operation of wavelet transform, used to reconstruct or synthesize the decomposed sub-band components into the original signal or image. The output result obtained after inverse wavelet transform can be a reconstructed image, signal, or feature map, which is the final representation after integrating the input subband components. The "Low-Low" component represents the features in the low-frequency-low-frequency direction. It is an approximate component of the image and contains the main structural information of the image. The "High-Low" component represents the features in the high-frequency-low-frequency direction and the high-frequency-low-frequency dimension, which usually correspond to high-frequency details such as edges and textures in a certain direction in an image; The "Low-High" component represents the characteristics of the low-frequency to high-frequency direction and the high-frequency to low-frequency dimension, and is also a form of high-frequency detail. The "High-High" component represents the features in the high-frequency-high-frequency direction and the high-frequency-low-frequency dimension, containing rich high-frequency detail information.

[0073] Furthermore, the Mamba state space modeling in the wavelet Mamba unit achieves efficient modeling of global dependency features through the Visual State Space Model (VSSB).

[0074] Furthermore, the state-space model is a mathematical structure used to describe the evolution of the implicit state within a dynamic system as the input changes, and it can model feature sequences through an implicit state recursion mechanism.

[0075] Furthermore, the state-space model maps the input signal to the latent state, and then generates the output response from the latent state. Its dynamic process can be represented by the following set of first-order linear state equations:

[0076]

[0077]

[0078] Where N represents the state size, A, B, C, and D are learnable weight matrices; x(t) is the input signal and x(t)∈R; For potential state and ; y(t) is the output response and y(t)∈R; The output signal of the system at time t is the system's response to the input. and internal state The comprehensive response, such as the filtered output in signal processing, the feedback signal in the control system, and the predicted output of the machine learning model; State vector The derivative with respect to time t describes the rate of change of the state, that is, the speed at which the system state dynamically evolves over time.

[0079] Furthermore, to enable the aforementioned continuous state-space model to be embedded in deep neural networks and trained and inferred on discrete data, it needs to be discretized. The purpose of discretization is to map continuous-time parameters to discrete parameters that can be executed at sampling time steps; let Δ be the time step parameter, used to characterize the interval of discrete computation, then the parameters A and B in the continuous form can be transformed into discrete parameters through the discretization process. , Then, the zero-order preservation method is used as the discretization strategy, and its discretization expression can be expressed as:

[0080]

[0081]

[0082]

[0083] in, To discretize the time step, , (Including the previous state) The discrete input, latent state, and output at step t are respectively: The continuous parameters A and B are the state transition matrix and input matrix after being discretized with zero order preservation, while C and D are the output matrix and direct transfer coefficients that remain unchanged after discretization, which together adapt to the training and inference of the neural network on discrete data.

[0084] Furthermore, the Mamba structure improves upon the state-space model and introduces a selective scanning mechanism, enabling the model to adaptively adjust the state transition process based on the input image content while maintaining linear computational complexity, thereby effectively enhancing the expressive power of features and global modeling capabilities.

[0085] Specifically, given the input features, they are represented as follows:

[0086]

[0087] Where H and W represent the height and width of the image, and C represents the channel; For input features.

[0088] Furthermore, depthwise separable convolution (DWConv) and the SiLU activation function are used to extract local neighborhood features.

[0089] Furthermore, a 2D selective scan mechanism (SS2D) is used to model long-range dependencies. After layer normalization of the output features, the original input features are multiplied element-wise by the SiLU activation function, and finally passed through a linear function to obtain the final state-space model output. The model output is represented as follows:

[0090]

[0091] Where H and W represent the height and width of the image, and C represents the channel; This is the output.

[0092] This process can be represented as:

[0093]

[0094] in, It is a linear layer (fully connected layer); This refers to Layer Normalization. It is a "space-channel second-order dynamic"; This is the activation function for the Sigmoid Linear Unit. This refers to Depthwise Convolution. The i-th feature input is used in step l; The output feature of the l-th path; the symbol ⊕ indicates feature fusion and residual connection.

[0095] Furthermore, the 2D selective scanning mechanism in the state space model, by simulating human eye perception, sequentially scans image features in the horizontal and vertical directions to achieve input-dependent global information aggregation.

[0096] Furthermore, after frequency domain reconstruction, the wavelet Mamba unit employs a lightweight gated feedforward network to enhance feature representation. This network structure includes depthwise separable convolution, channel splitting and gated activation, and channel attention fusion.

[0097] Furthermore, the depthwise separable convolution further includes: processing the reconstructed features through depthwise separable convolution and extracting local contextual information using depthwise convolution. This convolution method significantly reduces the amount of computation while maintaining the feature extraction capability.

[0098] Furthermore, the channel splitting and gating activation specifically includes: dividing the convolutional feature channels into two paths, one of which is processed by the gating module and then multiplied element-wise with the other feature path to achieve dynamic feature aggregation. This process enhances the expressive power of the features through the gating mechanism.

[0099] Furthermore, the channel attention fusion specifically includes: adaptively enhancing the aggregated features through a channel attention mechanism; generating attention weights through global average pooling and fully connected layers to further highlight important features and suppress unimportant features. The enhanced features are added to the input features to obtain the final feedforward output. This process not only preserves the feature information of the original input but also enhances the expressive power of the features through gating and attention mechanisms, providing richer feature representations for subsequent processing.

[0100] The specific expression is as follows:

[0101]

[0102]

[0103]

[0104]

[0105] in, This represents a depthwise convolution operation. This indicates a splitting operation. This represents the GeLU activation function. Indicates channel attention. This represents element-wise multiplication. To reconstruct the feature representation, Z represents the feedforward output, and Z represents the features after depthwise convolution. This represents the two-way features after segmentation; Features after gating activation; This is an attention fusion operation.

[0106] Step S2 can achieve the following technical effects:

[0107] By employing wavelet frequency domain decomposition, low-frequency color information and high-frequency texture details of the image are explicitly separated. Adaptive processing methods are used for features in different frequency bands—low-frequency features are globally dependently modeled using the Mamba state-space model, while high-frequency features are enhanced for detail through convolutional operations. This ensures accurate extraction of all types of features, avoiding information omissions or errors due to frequency band differences. The method supports feature extraction at various scales, covering mainstream feature types for underwater images, enabling it to flexibly handle multi-source heterogeneous image data environments and improving its versatility in practical applications. Through a standardized and precise feature extraction process, the extracted feature information is ensured to be complete and accurate, providing high-quality data input for subsequent global modeling and detail restoration using deep networks, thus guaranteeing the effectiveness of the entire image enhancement process.

[0108] In step S3:

[0109] The bottleneck layer uses a wavelet Mamba unit, specifically the fourth layer wavelet Mamba unit described in the above encoding process; the specific steps include:

[0110] The bottleneck layer receives multi-scale fused features from the last layer of the encoder. These feature maps have been processed by multiple wavelet Mamba units and contain rich low-frequency color information and high-frequency texture details. In the bottleneck layer, these multi-scale features are further fused to extract deeper semantic information.

[0111] Furthermore, the fused features are further processed through multiple convolutional layers or more complex network structures to enhance their expressive power. These operations aim to extract deeper semantic information, ensuring that the feature maps can better represent the overall structure and details of the image.

[0112] Furthermore, long-range dependency modeling of low-frequency components is performed using wavelet Mamba units to enhance global consistency, thereby achieving color shift correction in underwater images. Detail enhancement is applied to high-frequency components to improve image texture clarity and edge information, thus enhancing contrast.

[0113] Furthermore, a wavelet Mamba unit is used in the bottleneck layer to perform wavelet decomposition, enhancement, and frequency domain reconstruction on the input deep features; the features are decomposed into low-frequency components and high-frequency components; the low-frequency components are globally modeled using the Mamba state space model, and the high-frequency components are enhanced in detail through convolution operations, and the enhanced low-frequency and high-frequency components are recombined; thus completing the global modeling of underwater image color shift correction and contrast enhancement.

[0114] Furthermore, after deep representation and global modeling, the generated feature map is used as the output of the bottleneck layer, ultimately yielding the enhanced deep features, represented as:

[0115]

[0116] in, This indicates the depth features obtained after processing through 4 layers of wavelet Mamba units, which are the result of shallow features after multiple layers of feature extraction and transformation. This indicates that the numerical value of the deep feature map belongs to the real number field; These represent the height and width of the deep feature map, respectively. This represents the number of channels in the deep feature map.

[0117] Step S3 can achieve the following technical effects:

[0118] Deep representation of the fused multi-scale features at the bottleneck layer further integrates feature information at different scales, enhances the model's understanding of global semantics, and ensures the integrity and accuracy of global features. Color shift correction and contrast enhancement are completed through global modeling, which effectively solves common color distortion and insufficient contrast problems in underwater images, improves the overall visual quality of the image, and enhances the adaptability and robustness of the method under different lighting conditions and water environments. Accurate global modeling ensures the accuracy and consistency of global semantic information, providing a high-quality global semantic foundation for subsequent detail reconstruction using wavelet Mamba units, and ensuring the coherence and effectiveness of the entire image enhancement process.

[0119] In step S4:

[0120] Step S4 is the decoding process, in which: depth features are upsampled to restore spatial resolution, and texture details and color information are reconstructed and enhanced by skip connections and wavelet Mamba units.

[0121] Furthermore, the decoding process includes image restoration using three layers of wavelet Mamba units combined with an upsampling module. During decoding, image resolution is restored layer by layer, starting from depth features, and channel capacity is compressed; the specific implementation process is as follows:

[0122] The deep features are upsampled to restore spatial resolution, fused with the corresponding coding layer features, and then enhanced in the frequency domain and reconstructed using wavelet Mamba units before being output.

[0123]

[0124]

[0125]

[0126] in, Indicates an upsampling operation. This indicates a splicing operation. Represents a wavelet Mamba unit; For coding layer features, the coding stage is the first The intermediate features of the layer (the source of the skip connections) contain the local texture and detail information of that layer. They are fused with the features of the decoding layer through skip connections to supplement the detail information needed for reconstruction and avoid loss of detail during resolution restoration. The initial deep features (input) of the decoding process, and the final output from the encoding stage, are high-dimensional features that have undergone channel compression and retain global semantic information, providing basic semantic support for decoding and reconstruction. For the first The upsampled features from the layer decoding are processed by upsampling the features from the previous layer. Alternatively, spatial resolution recovery can be performed on the previously decoded features to initially restore the spatial structure of the image; For the first The first fusion feature of the layer decoding is to fuse the upsampled features with the corresponding coding layer features through the Concat (channel concatenation) operation, integrating semantic information and detail information; For the first The final fusion feature of the layer decoding is a secondary concatenation, which is to infer the supplementary fusion coding layer features by combining the context, strengthen the connection between details and semantics, and provide rich input features for subsequent wavelet Mamba units; No. The features output by the layer decoding are enhanced in the frequency domain and reconstructed in the structure to optimize the texture details and color consistency of the fused features, while compressing the channel capacity, laying the foundation for the next layer decoding or the final image output; the layer index (integer) of the decoding layer and the coding layer. To match the hierarchical correspondence between "upsampled features" and "coding layer skip connection features" during the decoding process, and to ensure that the fused features come from the same resolution or semantic level, the reconstruction accuracy is improved.

[0127] Step S4 can achieve the following technical effects:

[0128] By gradually restoring the spatial resolution of the image through upsampling operations and combining wavelet Mamba units to reconstruct and enhance texture details and color information, the image's texture clarity and color naturalness can be effectively improved, making the enhanced image closer to the real scene. By combining skip connections to fuse features from the encoding and decoding stages, the image's detail information can be effectively preserved, avoiding the loss of important details during resolution restoration. This enhances the flexibility and effectiveness of the enhancement method in detail restoration. Through precise detail reconstruction and enhancement, the integrity and accuracy of image details are ensured, providing high-quality detail information for the final image output and guaranteeing the final effect of the entire image enhancement process.

[0129] In step S5:

[0130] The enhanced features output from the decoding stage are mapped to the image domain, and the final enhanced underwater image is generated through a residual fusion mechanism; the specific implementation process is as follows:

[0131] The features output from the last layer of the decoder are denoted as:

[0132]

[0133] Where H and W represent the height and width of the image, respectively, and C represents the number of channels; This is the output feature of the last layer.

[0134] Furthermore, to map the features to the RGB image domain, a 3×3 convolution operation is performed on the output features of the last layer to obtain the residual image, specifically represented as follows:

[0135]

[0136] Where R is the residual image, For the output features of the last layer, This is a 3×3 convolution operation.

[0137] Furthermore, the residual image is added pixel by pixel to the original input image to obtain the final enhanced underwater image, specifically expressed as:

[0138] =

[0139] in, To enhance underwater images, The original input image, This is the residual image.

[0140] Furthermore, the method employs an end-to-end training approach, with a loss function comprising a weighted combination of reconstruction loss and frequency loss, used to jointly constrain the reconstruction quality of the network output in both the spatial and frequency domains, thereby improving the structural fidelity and detail representation of the enhanced image. The reconstruction loss can be expressed as:

[0141]

[0142] in, This indicates the enhanced underwater image. Represents a real image. express -norm, This indicates the losses incurred during reconstruction.

[0143] Furthermore, the total loss function Defined as:

[0144]

[0145] in, Represents the reconstruction loss, used to measure the difference between the output image and the reference image in the pixel domain. This is an adjustable non-negative weighting parameter used to balance spatial reconstruction accuracy and frequency domain feature consistency; based on empirical values, it is set to = 0.05. This represents the frequency loss, used to measure the fidelity of the enhanced image in terms of spectral structure, thereby enhancing the model's ability to recover texture details and structural information.

[0146] Step S5 can achieve the following technical effects:

[0147] By using a residual fusion mechanism to add the enhanced features to the original input image pixel by pixel, the structural information of the original image can be preserved while introducing enhanced details and color information. This results in an enhanced image that maintains a natural feel while having higher clarity and contrast. It can effectively transform the feature information learned by the model into the visual effect of the image, ensuring that the enhancement method is applicable to various underwater image scenarios and improving the universality and practicality of the method. Through precise residual fusion, the quality and stability of the final output image are ensured, providing high-quality image input for subsequent underwater image analysis, target recognition and other applications, and ensuring the effectiveness of the entire underwater image processing workflow.

[0148] Secondly, the present invention also provides an underwater image enhancement system based on wavelet Mamba, specifically including a shallow feature extraction module, a multi-scale coding and feature compression module, a deep representation and global modeling module for multi-scale fused features, an upsampling module to restore resolution and reconstruct enhanced texture details and color information, and a feature mapping module to the image domain and residual fusion module to generate enhanced underwater images.

[0149] The shallow feature extraction module further includes: inputting a low-quality underwater image and extracting shallow features; the shallow features include the image's basic color and texture information;

[0150] The multi-scale coding and feature compression module further includes: inputting the extracted shallow features into the multi-scale coding network; performing wavelet frequency domain decomposition and state space modeling on the image through wavelet Mamba units during the coding process; extracting low-frequency color information and high-frequency detail features; and combining downsampling with layer-by-layer compression to obtain global semantic features.

[0151] The deep representation and global modeling module of the multi-scale fusion features further includes: performing deep representation of the multi-scale fusion features in the bottleneck layer of the network to complete the global modeling of underwater image color shift correction and contrast enhancement.

[0152] The upsampling module for restoring resolution and reconstructing and enhancing texture details and color information further includes: gradually restoring spatial resolution through upsampling, and reconstructing and enhancing texture details and color information by combining skip connections and wavelet Mamba units;

[0153] The module for mapping features to the image domain and performing residual fusion to generate enhanced underwater images further includes: mapping features to the image domain and performing residual fusion with the input image to generate enhanced underwater images with natural colors and clear textures.

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

[0155] To address the aforementioned problems, this invention proposes an underwater image enhancement method based on wavelet Mamba. This method decomposes underwater images into low-frequency color sub-bands and high-frequency texture sub-bands using wavelet transform, and utilizes the Mamba state-space model to model the long-range dependencies of the low-frequency components and enhance the details of the high-frequency components, achieving joint optimization in the frequency and spatial domains. This innovative method significantly improves the color realism, structural clarity, and detail perceptibility of underwater images, while reducing computational complexity and improving the model's real-time inference capability, demonstrating high efficiency and deployability. Furthermore, this invention employs an end-to-end training approach, optimizing the image enhancement effect through a weighted combination of reconstruction loss and frequency loss, further enhancing the structural fidelity and detail representation of the enhanced image. Testing on a standard experimental dataset demonstrates that this method performs exceptionally well in objective metrics such as peak signal-to-noise ratio and structural similarity index, significantly outperforming existing methods and validating its advantages in brightness restoration, texture reconstruction, and structural fidelity. Attached Figure Description

[0156] To more clearly illustrate the technical solutions in the embodiments of the present invention or the prior art, the drawings used in the description of the embodiments or the prior art will be briefly introduced below. Obviously, the drawings described below are only embodiments of the present invention. For those skilled in the art, other drawings can be obtained based on the provided drawings without creative effort.

[0157] Figure 1 This is a flowchart of the underwater image enhancement method based on wavelet Mamba in this embodiment.

[0158] Figure 2 This is a schematic diagram of the overall network structure of the underwater image enhancement method based on wavelet Mamba in this embodiment.

[0159] Figure 3 This is a schematic diagram of the wavelet Mmaba module composition proposed in the underwater image enhancement method of this embodiment.

[0160] Figure 4 This is a schematic diagram of the lightweight gated feedforward network proposed in the underwater image enhancement method of this embodiment. Detailed Implementation

[0161] The present invention will now be described in detail with reference to the accompanying drawings and specific embodiments. The following embodiments are explanations of the present invention, but the present invention is not limited to the following implementation methods.

[0162] Example 1:

[0163] like Figure 1 As shown, this invention provides an underwater image enhancement method based on wavelet Mamba, specifically including the following steps:

[0164] Step S1, the step of extracting shallow features, in which:

[0165] Input a low-quality underwater image and perform shallow feature extraction; the shallow features include the image's basic color and texture information;

[0166] Step S2, the multi-scale coding and feature compression step, in which:

[0167] The extracted shallow features are input into a multi-scale coding network. During the coding process, wavelet Mamba units are used to perform wavelet frequency domain decomposition and state space modeling on the image to extract low-frequency color information and high-frequency detail features. Global semantic features are obtained by combining downsampling and layer-by-layer compression.

[0168] Step S3, the deep representation and global modeling of multi-scale fused features, in which:

[0169] Deep representation of multi-scale fusion features is performed in the bottleneck layer of the network to complete the global modeling of underwater image color shift correction and contrast enhancement.

[0170] Step S4, the step of upsampling to restore resolution and reconstructing and enhancing texture details and color information, in which:

[0171] Spatial resolution is gradually restored through upsampling, and texture details and color information are reconstructed and enhanced by combining skip connections and wavelet Mamba units.

[0172] Step S5, mapping features to the image domain and residual fusion to generate an enhanced underwater image, in which:

[0173] Features are mapped to the image domain and residually fused with the input image to generate enhanced underwater images with natural colors and clear textures.

[0174] like Figure 2 The diagram shown illustrates the overall network structure of the underwater image enhancement method based on wavelet Mamba.

[0175] In this embodiment, the network architecture of the underwater image enhancement method based on wavelet Mamba is based on a symmetric U-Net multi-scale encoder-decoder structure. The encoder and decoder consist of three layers of symmetric wavelet Mamba units and a bottleneck layer of one layer of wavelet Mamba units. Each encoder-decoder layer contains multiple wavelet Mamba units.

[0176] Input low-quality underwater images:

[0177]

[0178] in, Represents underwater images, Indicates "belongs to". Represents the real number field. The dimensions of the image are represented by H (height), W (width), and 3 (number of channels).

[0179] First, a 3×3 convolutional layer is used to extract basic texture information and color distribution features to obtain a shallow feature map:

[0180]

[0181] in, This represents the shallow feature map obtained after initial processing, which serves as the input for subsequent convolution operations. This indicates that the value of the feature map belongs to the real number field; These represent the height and width of the feature map, respectively; C represents the number of channels in the feature map.

[0182] Furthermore, the shallow features are input into wavelet Mamba units in the encoder network for layer-by-layer feature extraction. Starting with high-resolution features, the encoder extracts low-frequency global information and high-frequency detail information at each layer through wavelet decomposition and state-space modeling, respectively. It also progressively reduces the spatial size of the features through downsampling while increasing the number of channels to enhance feature expressiveness. After stacking multi-scale wavelet Mamba units, a deep feature representation containing rich semantic information is obtained.

[0183]

[0184] in, This represents the deep feature map obtained after processing with four layers of wavelet Mamba units. This indicates that the numerical value of the deep feature map belongs to the real number field; These represent the height and width of the deep feature map, respectively. This represents the number of channels in the deep feature map; where C is the number of channels in the shallow feature map.

[0185] Furthermore, the aforementioned deep features are input into the decoder network. The decoding process starts with the lowest resolution features and gradually restores the spatial resolution through upsampling operations. In each decoding layer, the features at the current scale are concatenated with the output features of the wavelet Mamba units at the corresponding scale in the encoder using skip connections, and then the number of channels is adjusted through 1×1 convolution to achieve effective fusion of frequency domain information and spatial details. This decoding process restores the image structure and texture information layer by layer and generates a reconstructed feature map. Finally, the reconstructed features output by the decoder are mapped to a residual signal R through a 3×3 convolution, and then added pixel by pixel with the original input image to generate an enhanced underwater image.

[0186] =

[0187] in, To enhance underwater images, The original input image, For residual images;

[0188] Through this residual fusion structure, the present invention not only restores the brightness, color and texture details of underwater images, but also ensures that the enhancement process does not introduce structural distortion, and the output result has a natural and realistic visual effect.

[0189] Figure 3 and Figure 4 These are schematic diagrams of the wavelet Mmaba module and the lightweight gated feedforward network proposed in the underwater image enhancement method of this embodiment.

[0190] In this embodiment, the wavelet Mamba unit consists of three parts: wavelet transform, visual state space model, and lightweight gated feedforward network. This module achieves color restoration and texture detail enhancement of underwater images by combining frequency domain decoupling, global modeling, and nonlinear enhancement. The specific implementation method is as follows:

[0191] First, input features:

[0192]

[0193] in, The input feature map to the wavelet Mamba unit is represented by ∈i; ∈i indicates that the value of the feature map belongs to the real number domain; H×W represents the spatial size of the feature map; C represents the number of channels of the feature map;

[0194] The frequency components are decomposed into one low-frequency component and three high-frequency components using a two-dimensional discrete wavelet transform (2D-DWT). For the low-frequency component, a visual state-space model (VSSB) is used for modeling. This model captures long-range dependencies through a 2D selective scanning mechanism while maintaining linear computational complexity, yielding the enhanced low-frequency output features. For the high-frequency components, 3×3 convolution is used to enhance detail information, resulting in enhanced high-frequency features. These enhanced low-frequency and high-frequency features are then input into an inverse wavelet transform (IWT) module to reconstruct the frequency domain features back into the spatial domain, yielding reconstructed features. These reconstructed features are then input into a lightweight gated feedforward network (GFN) for nonlinear enhancement, ultimately obtaining the output features of the wavelet Mamba unit. The specific expression for this process is as follows:

[0195]

[0196] in, These are the original input features of the wavelet Mamba unit; For two-dimensional discrete wavelet transform, the core operation is to decompose the input features into one low-frequency component and three high-frequency components; Low-frequency components; As a visual state space model, it processes low-frequency components and captures long-range dependencies of low-frequency components in the frequency domain through a 2D selective scanning mechanism, while maintaining linear computational complexity, and outputs low-frequency enhanced features. This is a 3×3 convolution operation; For inverse wavelet transform, there is a dual-input module: the first input is the low-frequency enhanced feature output by VSSB, and the second input is the high-frequency enhanced feature after 3×3 convolution. The core function is to reconstruct the frequency domain features after separation and enhancement back to the spatial domain and output the reconstructed features. To achieve a lightweight gated feedforward network, it receives the reconstructed features output by IWT and uses a gating mechanism to achieve nonlinear enhancement, further optimizing the feature representation capability; For the final output features;

[0197] In this implementation, dynamic propagation and global modeling of input features are achieved by sequentially scanning the input features in both horizontal and vertical directions. Specifically, the input features are first normalized to ensure numerical stability across different scales. Then, the features enter a state recursion unit, where information is propagated sequentially from left to right and from right to left in the horizontal direction to capture semantic relationships across columns. Next, features are propagated vertically from top to bottom and from bottom to top to characterize structural dependencies across rows. By weighted fusion of the recursion results in the four directions, long-distance modeling of low-frequency structural information is achieved. This two-dimensional selective scanning mechanism not only possesses the advantage of linear complexity but also effectively avoids the quadratic computational overhead of self-attention mechanisms. In the wavelet Mamba unit, this selective scanning mechanism is used to process low-frequency features to enhance global brightness consistency and color restoration capabilities, and to provide a high-quality semantic foundation for subsequent wavelet reconstruction and texture enhancement.

[0198] In this embodiment, an end-to-end training method is adopted during the model training process. The loss function includes a weighted combination of reconstruction loss and frequency loss, which is used to jointly constrain the reconstruction quality of the network output in the spatial and frequency domains, thereby improving the structural fidelity and detail representation ability of the enhanced image. The total loss function... Defined as:

[0199]

[0200] in, Represents the reconstruction loss, used to measure the difference between the output image and the reference image in the pixel domain. This is an adjustable non-negative weighting parameter used to balance spatial reconstruction accuracy and frequency domain feature consistency; it is set to = 0.05. Representing frequency loss, it is used to measure the fidelity of the enhanced image in terms of spectral structure, thereby enhancing the model's ability to recover texture details and structural information; where Calculated using the following formula:

[0201]

[0202] in, This indicates the enhanced underwater image. Represents a real image. express Norm;

[0203] Furthermore, by inputting the underwater image into the trained network, a clear and enhanced image can be obtained as the output.

[0204] Furthermore, to verify the effectiveness of the method of this invention in underwater image enhancement tasks, this embodiment selects a standard experimental dataset for testing and uses the luminance component (Y channel) in the YCbCr color space as the main evaluation criterion to perform quantitative comparative analysis between the enhanced image and the real image. In this evaluation process, two widely recognized objective indicators, Peak Signal-to-Noise Ratio (PSNR) and Structural Similarity Index (SSIM), are selected as performance measurement standards to objectively reflect the advantages of the method of this invention in luminance restoration, texture reconstruction, and structural fidelity.

[0205] PSNR is used to evaluate the degree of difference between the enhanced image and the reference image at the pixel level, and can reflect the overall image restoration quality. The higher the value, the less image distortion and the better the enhancement effect. Its formula is as follows:

[0206]

[0207] in, is the peak signal-to-noise ratio; MSE is the mean squared error; n is the number of bits (integer) for the pixel / signal quantization, representing the number of bits per sampling point, which determines the dynamic range of the signal;

[0208] SSIM comprehensively evaluates the enhancement results from multiple aspects, including brightness preservation, contrast restoration, and structural consistency. Its numerical value ranges from 0 to 1, with values ​​closer to 1 indicating that the enhanced image closely resembles the human eye's perception of the real image. The formula is:

[0209]

[0210] in, and These represent the average intensity of image x and y within the current local region, reflecting the brightness distribution of that region; and These represent the variance values ​​of image x and y within this local region, respectively, and are used to measure the degree of contrast change in the image; It is the covariance of images x and y within the same region, reflecting the degree of consistency of structural information between the two images;

[0211] The above calculation process introduces constants. and To enhance numerical stability, the value is typically set to:

[0212]

[0213]

[0214] Where K1 and K2 are empirical parameters, K1 is usually set to 0.01 and K2 to 0.03; L represents the dynamic range of the image pixels;

[0215] The above-disclosed embodiments are merely preferred embodiments of the present invention, but the present invention is not limited thereto. Any non-creative variations that can be conceived by those skilled in the art, as well as any improvements and modifications made without departing from the principles of the present invention, should fall within the protection scope of the present invention.

Claims

1. A wavelet Mamba based underwater image enhancement method, characterized in that, Specifically comprising the following steps: Step S1, a step of extracting shallow features, in which: An underwater low-quality image is input and shallow features are extracted; the shallow features include basic color and texture information of the image; Step S2, a step of multi-scale encoding and feature compression, in which: The extracted shallow features are input into a multi-scale encoding network, and in the encoding process, the image is decomposed in wavelet frequency domain and state space modeling is performed by a wavelet Mamba unit, low-frequency color information and high-frequency detail features are extracted, and global semantic features are obtained by combining layer-by-layer compression of down-sampling; Step S3, a step of deep representation of multi-scale fusion features and global modeling, in which: The multi-scale fusion features are represented in a deep layer at the bottleneck layer of the network, and global modeling of underwater image color offset correction and contrast enhancement is completed; Step S4, a step of recovering resolution by up-sampling and reconstructing enhanced texture details and color information, in which: The spatial resolution is gradually recovered by up-sampling, and the texture details and color information are reconstructed and enhanced by combining the jump connection and the wavelet Mamba unit; Step S5, a step of mapping features to the image domain and residual fusion to generate an enhanced underwater image, in which: The features are mapped to the image domain, and residual fusion is performed with the input image to generate an enhanced underwater image with natural color and clear texture; The multi-scale encoding network of step S2 is constructed based on a U-Net architecture, including an encoder and a decoder; the encoder and the decoder include 3 layers of symmetric structure of wavelet Mamba units and 1 layer of wavelet Mamba units at the bottleneck layer, and each layer of the encoder-decoder includes multiple wavelet Mamba units; the encoding part is composed of 3 wavelet Mamba units connected in turn and combined with down-sampling operations, and the decoding part is composed of wavelet Mamba units corresponding in number to the encoding part and up-sampling operations; in the encoding process, the spatial resolution is compressed layer by layer and multi-scale features are extracted, and in the decoding process, the spatial resolution is recovered layer by layer and the features of each scale are fused; The wavelet Mamba unit includes the following operations: The input feature map is decomposed into low-frequency components and high-frequency components by two-dimensional discrete wavelet transform, and the low-frequency and high-frequency components are recombined by two-dimensional inverse wavelet transform to generate a reconstructed feature map; The low-frequency components are used to extract the structure and color information of the image, and the high-frequency components are used to extract the texture and detail information of the image; A Mamba state space model is introduced in the low-frequency component; in the high-frequency component, a 3x3 convolution operation is used for detail enhancement, and the spatial resolution of the feature map is gradually reduced by down-sampling operation while the number of channels is increased, to extract deeper semantic features; The wavelet Mamba unit also includes a light-weight gated feed-forward network that extracts local context from the IWT-reconstructed features through depthwise separable convolutions, and then splits the feature channels into two paths , , to implement gated activation, with the specific expression as follows: wherein, , are two paths of feature channel split respectively, denotes element-wise multiplication, denotes GeLU activation function; Further, the channel attention and 1x1 convolution are fused, and the output feature expression is: wherein, is an IWT reconstruction feature, denotes a channel attention, denotes a convolution operation, is a feed-forward output representation.

2. The method of claim 1, wherein, The step S1 specifically includes: The shallow feature extraction is realized by inputting the underwater low-quality image and performing feature mapping by 3x3 convolution; the shallow features include extracting the basic color and texture information of the image; the specific implementation manner is as follows: An original underwater low-quality image is input, and the image is represented as: wherein, represents an underwater image, represents a real number field, represents a dimension of an image; wherein H is a height, W is a width, and 3 is a number of channels; The shallow features are extracted by convolution with a kernel size of 3x3, and are represented as: wherein, represents a shallow feature map obtained after initial processing, which is the input of subsequent convolution operations; represents that the numerical value of the feature map belongs to the real number field, that is, the feature value is a real number; respectively represent the height and width of the feature map, which describe the spatial dimensions of the feature map; C represents the number of channels of the feature map.

3. The method of claim 2, wherein, The Mamba state space modeling in the wavelet Mamba unit models the global dependence feature through a visual state space model; The state space model maps the input signal to a latent state, and then generates an output response from the latent state, and the dynamic process thereof can be represented by a first-order linear state equation set as follows: where N denotes the state size, A, B, C and D are learnable weight matrices; x(t) is the input signal and x(t) e R; is the potential state and y(t) is the output response and y(t) e R; is the output signal of the system at time t; is the state vector derivative with respect to time t; The continuous form needs to be discretized, and the discrete expression thereof can be represented as: wherein, is the discretized time step, , are the discrete input, latent state, output at the tth step, respectively; is the state transition matrix and input matrix after zero-order hold discretization of continuous parameters A, B, and C, D are the output matrix and direct transmission coefficient after discretization.

4. The method of claim 3, wherein, The step S4 further includes: The deep feature restores the spatial resolution through the upsampling operation, is fused with the corresponding encoding layer feature, and then is subjected to frequency domain enhancement and structure reconstruction through the wavelet Mamba unit, and is outputted: wherein, denotes an up-sampling operation, denotes a concatenation operation, denotes a wavelet Mamba unit; is an encoded layer feature, the encoded stage being the intermediate feature of the layer; is an up-sampled feature decoded from the first fused feature decoded from the second fused feature decoded from the final fused feature decoded from the output feature of the layer decoded from the layer decoded from the layer decoded from the 5. The method of claim 4, wherein, The output feature mapped to the image domain generates an enhanced image through 3*3 convolution, and is added to the input image pixel by pixel.

6. The method of claim 5, wherein, The step S5 further includes: The enhanced feature outputted in the decoding stage is mapped to the image domain, and a final enhanced underwater image is generated through a residual fusion mechanism, and the specific implementation process is as follows: The feature outputted by the last layer of the decoder is denoted as: where H and W represent the height and width of the image respectively, and C represents the number of channels; for the output feature of the last layer; The feature is mapped to the RGB image domain, and a 3*3 convolution operation is performed on the output feature of the last layer, and the specific representation is as follows: wherein R is a residual image, is the output feature of the last layer, is a 3x3 convolution operation; The residual image is added to the original input image pixel by pixel, and a final output enhanced underwater image is obtained, and the specific expression is as follows: = wherein, to enhance underwater images, is the original input image, is the residual image.

7. The method of claim 6, wherein, The method adopts an end-to-end training manner, and a loss function includes a weighted combination of a reconstruction loss and a frequency loss, and is used for jointly constraining the reconstruction quality of the network output in the spatial domain and the frequency domain, so as to improve the structure fidelity and detail expression capability of the enhanced image, and the total loss function L is defined as: wherein, represents a reconstruction loss for measuring the difference between the output image and the reference image in the pixel domain, is an adjustable non-negative weight parameter for balancing the spatial reconstruction accuracy and the frequency domain feature consistency, and the weight is set to =0.05 according to the empirical value; represents a frequency loss for measuring the fidelity of the enhanced image in the frequency spectrum structure, so as to enhance the recovery ability of the model to the texture details and structural information; wherein the reconstruction loss is calculated by the following formula: wherein, denotes the enhanced underwater image, denotes the real image, denotes norm.

Images