An adaptive underwater image enhancement method and system based on Retinex theory and Mamba

Abstract

The present invention provides an adaptive underwater image enhancement method and system based on Retinex theory and Mamba, which relates to the field of underwater image enhancement technology. The method and system include taking an original underwater image and its corresponding illumination prior map as input, obtaining an illumination factor map through feature splicing, a first convolutional layer, a multi-scale channel aggregation module layer, and a second convolutional layer, performing a Hadamard product operation on the illumination factor map and the original image to obtain a basic enhanced image; the basic enhanced image is subjected to an encoder to extract multi-scale degradation features step by step, a decoder to restore spatial resolution through jump connections, and a WLMamba bottleneck module to selectively model the state space to enhance illumination-texture correlation, ultimately generating a high-quality underwater image enhancement result. This application establishes a mathematical representation system for the optical properties of water bodies, and the model can adaptively adjust the chromaticity compensation intensity to achieve image restoration under the physical property constraints of the underwater scene.

Description

Technical Field

[0001] The present application relates to the technical field of underwater image enhancement, and more specifically, to an adaptive underwater image enhancement method and system based on Retinex theory and Mamba. Background Art

[0002] The scattering and attenuation of light in water leads to widespread degradation of underwater images, such as low contrast and color shift. Furthermore, the noise pollution and artifacts generated by underwater imaging equipment during the exposure process further exacerbate the blurring and distortion of image details. For underwater image processing, the main technical solutions adopted in existing technologies are color correction, which includes the following:

[0003] (1) The Chinese invention patent with publication number CN119648588A discloses an image enhancement method based on high-frequency and low-frequency adaptive fusion, comprising: performing contrast enhancement on the color channels of the original image to obtain an intermediate feature map; calculating the pixel mean of the intermediate feature map under each color channel; performing color compensation on the intermediate feature map according to the pixel mean to obtain a color reconstructed image; decomposing the color reconstructed image into high-frequency image features and low-frequency image features; normalizing and gamma correcting the high-frequency image features to obtain a high-frequency corrected image; sharpening the high-frequency corrected image to obtain a high-frequency restored image; defogging the low-frequency image feature weights to obtain a low-frequency restored image, and adaptively fusing the high-frequency restored image and the low-frequency restored image to obtain an output image.

[0004] (2) Chinese invention patent publication number CN119741245A discloses an image enhancement algorithm based on adaptive filtering. By analyzing the pixel values ​​of the r, g, and b channels of the image, the attenuation degree of each channel in the low-quality image is evaluated, and adaptive adjustment coefficients are designed to balance the chromatic aberration to obtain a color-corrected image. Next, a backscattered light estimation method is introduced to evaluate the scattering distribution in the color-corrected image. Finally, the transmittance is calculated by combining the dark channel prior theory, and the backscattered light is combined with the transmittance using the atmospheric scattering model to generate a final enhanced image with rich contrast and details.

[0005] It can be seen that the image enhancement methods in the existing schemes all first perform color correction (or color compensation) and then perform image attenuation repair. The image enhanced in this way can show better color effects, but it will cause loss and tampering of image information. Summary of the Invention

[0006] To solve the above problems, the technical solution adopted in this application is an adaptive underwater image enhancement method based on Retinex theory and Mamba, including:

[0007] Light Information Estimation Subnetwork: Taking the original underwater image and its corresponding illumination prior map as input, the network sequentially obtains an illumination factor map through feature concatenation, the first convolutional layer, the multi-scale channel aggregation module layer, and the second convolutional layer. The multi-scale channel aggregation module layer extracts the illumination estimation component of the underwater scene and performs a Hadamard product operation on the illumination factor map and the original image to obtain a basic enhanced image.

[0008] Degradation restoration subnetwork: The encoder extracts multi-scale degradation features of the base enhanced image step by step, the decoder uses skip connections to restore the spatial resolution, and the WLMamba bottleneck module selectively models the state space to enhance the illumination-texture correlation, ultimately generating high-quality underwater image enhancement results.

[0009] Optionally, the multi-scale channel aggregation module layer includes adjusting the detail features initially extracted by the first convolutional layer simultaneously through a first upper branch, a first middle branch, and a first lower branch, and then performing element-by-element multiplication on the results of the three branches to obtain a final illumination estimation component of the underwater scene;

[0010] The first upper branch retains the original input features;

[0011] The first intermediate branch is used for multi-scale aggregation on multiple channels;

[0012] The first lower branch performs a pooling operation in the spatial direction to obtain a spatial descriptor representation, then learns the correlation between channels through the lower branch convolution layer, and generates a channel weight representation through the Sigmoid function.

[0013] Optionally, the first intermediate branch splits the input features into four equal parts according to the number of channels, and then implements a multi-scale downsampling strategy on the split input features through the feature channels of Channel 1 to Channel 4. The specific operations are as follows:

[0014] Channel 1: uses 8x downsampling to extract global semantic features;

[0015] Channel 2: Balances local details and contextual information through 4x downsampling;

[0016] Channel 3: Perform a 2x downsampling operation to retain medium-scale features;

[0017] Channel 4: Maintains original underwater resolution to preserve high-frequency details;

[0018] The features after downsampling from Channel 1 to 3 are subjected to 3×3 depth convolution, and the dimension is restored through upsampling operation. The features of the feature channels of Channel 1 to 4 are spliced, and then the aggregation operation is performed through the first intermediate branch convolution layer to finally obtain multi-scale aggregated features.

[0019] Optionally, the WLMamba bottleneck module adjusts the features it receives through the second upper branch, the second middle branch, and the second lower branch simultaneously;

[0020] The second upper branch retains the original input features;

[0021] The second intermediate branch will input the feature channel First, by linearly increasing , is a predefined channel expansion factor, followed by depthwise convolution, followed by SiLU activation, SSM-2D layer, and layer normalization;

[0022] The second lower branch combines the input features first through a channel attention module modulated by the light information estimation component and a spatial attention module in cascade, and then activates them through SiLU.

[0023] Optionally, the input of the channel attention module consists of an illumination estimation component and an intermediate feature map of image processing. The intermediate feature map is first passed through a global average pooling layer to capture the global statistical features of each channel. Subsequently, the compressed channel feature vector is subjected to a nonlinear transformation through a fully connected layer group to learn the complex relationship between channels. The channel weight is then activated and outputted through a Sigmoid function. In the feature recalibration stage, the above-mentioned channel weight is multiplied element-by-element with the channel corresponding to the illumination estimation component to realize adaptive modulation of the feature map, and the final output is a weighted feature map.

[0024] Optionally, the weighted feature map output by the channel attention module is used as the input of the spatial attention module. First, the input feature map is subjected to average pooling and maximum pooling operations respectively. Then, the average feature map and the maximum feature map are concatenated along the channel dimension to generate a composite feature that integrates global and local features. The composite feature is subjected to spatial information fusion and nonlinear modeling through a 3×3 convolution kernel. After activation by a Sigmoid function, a spatial attention weight is generated. Finally, the spatial attention weight is element-wise multiplied with the illumination estimation component to realize illumination-spatial characteristic modulation and obtain the feature map after spatial attention optimization.

[0025] Optionally, long-range dependencies in underwater images are captured through an SSM-2D layer, which includes converting the input feature map into a one-dimensional sequence in four directions through a flattening operation, arranging the two-dimensional features into a one-dimensional sequence according to the scanning sequence of the directions, applying a dynamically parameterized state space model to the sequence in each direction, restoring the output sequences in the four directions into two-dimensional feature maps, and fusing them through attention weighting.

[0026] Optionally, the overall architecture of the degradation repair subnetwork adopts a four-layer encoder-decoder structure, which includes cross-layer skip connections and WLMamba bottleneck modules; the encoders of layers 1-3 are respectively connected to the decoders of their corresponding layers through skip connections to perform multi-scale feature fusion, and the encoder features of the 4th layer are adjusted by the WLMamba bottleneck module and passed to the decoder of the 4th layer.

[0027] Optionally, the jump connection consists of a feature alignment operation and an adaptive fusion operation, wherein the feature alignment operation includes spatial alignment and channel alignment. The spatial alignment includes upsampling the encoder features to the decoder size by bilinear difference, and then adjusting the number of channels by a convolution kernel of size 1×1 to obtain the aligned features. ;

[0028] The input of the adaptive fusion operation is composed of the output features of the feature alignment operation and the decoder output features, which are spliced ​​on the channel. Two learnable parameters are set: ,in , set up the The decoder feature of the layer is , obtain the adaptive fusion features ; The adaptive fusion features Enter The decoder of the layer completes the adaptive fusion operation.

[0029] This application also provides an adaptive underwater image enhancement system based on Retinex theory and Mamba, including:

[0030] Light Information Estimation Subnetwork: This module takes the original underwater image and its corresponding illumination prior map as input, sequentially passes through feature concatenation, the first convolutional layer, the multi-scale channel aggregation module layer, and the second convolutional layer to obtain an illumination factor map. The multi-scale channel aggregation module layer extracts the illumination estimation component of the underwater scene, and performs a Hadamard product operation on the illumination factor map and the original image to obtain a basic enhanced image.

[0031] Degradation restoration sub-network module: It is used to extract multi-scale degradation features of the basic enhanced image step by step through the encoder, restore the spatial resolution through the decoder jump connection, enhance the illumination-texture correlation through the WLMamba bottleneck module selective state space modeling, and finally generate high-quality underwater image enhancement results.

[0032] The beneficial effects of the adaptive underwater image enhancement method and system provided in this application based on Retinex theory and Mamba are:

[0033] (1) This application provides an underwater image enhancement model that integrates Retinex theory, Mamba architecture, and U-shaped network. By establishing a mathematical representation system for the optical properties of water, it innovatively combines Retinex theory with a deep learning framework, effectively addressing the deficiencies of traditional neural networks in terms of interpretability. The model uses an encoder-decoder topology to achieve multi-scale feature interaction. By combining underwater optical attenuation characteristics with Retinex theory, the model can adaptively adjust the chromaticity compensation intensity to achieve image restoration under the constraints of the physical characteristics of underwater scenes.

[0034] (2) This study innovatively improves the traditional Retinex theory. By integrating the underwater light scattering-attenuation characteristics and the equipment exposure attenuation mechanism, an improved Retinex theoretical framework for underwater environments is constructed, and the theoretical mapping is realized through a neural network architecture. It includes two sub-networks: light information estimation and degradation restoration. The light information estimation sub-network is based on the improved Retinex theory. It adopts a multi-scale feature decoupling mechanism through the multi-scale channel aggregation module layer to accurately extract the illumination estimation component and the basic enhanced image of the underwater scene; the degradation restoration sub-network innovatively designs the WLMamba module. The WLMamba module deeply integrates the spatiotemporal modeling capabilities of the Mamba architecture with the interaction mechanism of the illumination estimation component, and combines the jump connection of the U-shaped network to construct a dual-path feature interaction architecture to form an adaptive degradation restoration network. Experiments show that this method has achieved significant performance improvement on the LSUI, UIEB standard datasets and the reference-free SQUID dataset. While maintaining high image quality indicators, it greatly optimizes the model inference efficiency, verifying the effectiveness of this theoretical framework in complex underwater environments. BRIEF DESCRIPTION OF THE DRAWINGS

[0035] In order to more clearly illustrate the technical solutions in the embodiments of the present application, the following briefly introduces the drawings required for use in the embodiments or descriptions of the prior art.

[0036] Figure 1 This is the overall network structure diagram of the adaptive underwater image enhancement method based on Retinex theory and Mamba provided in the embodiment of the present application;

[0037] Figure 2 This is a diagram of the optical information estimation subnetwork structure provided in an embodiment of the present application;

[0038] Figure 3 This is a structural diagram of a multi-scale channel aggregation module provided in an embodiment of the present application;

[0039] Figure 4 This is a diagram of the degradation repair subnetwork structure provided by an embodiment of the present application;

[0040] Figure 5 This is a diagram of the WLMamba module structure provided in an embodiment of the present application;

[0041] Figure 6 This is a schematic diagram of the working principle of SSM-2D provided in the embodiment of the present application;

[0042] Figure 7 is a training-loss graph provided in an embodiment of the present application;

[0043] Figure 8 is a real underwater image provided by the embodiment of the present application;

[0044] Figure 9 yes Figure 8 Three-dimensional RGB chromaticity distribution diagram;

[0045] Figure 10 yes Figure 8 RGB channel distribution histogram;

[0046] Figure 11 yes Figure 8 RGB intensity curve graph (middle row);

[0047] Figure 12 is the enhanced underwater image provided by the embodiment of the present application;

[0048] Figure 13 yes Figure 12 Three-dimensional RGB chromaticity distribution diagram;

[0049] Figure 14 yes Figure 12 RGB channel distribution histogram;

[0050] Figure 15 yes Figure 12 RGB intensity curve graph (middle row);

[0051] Figure 16 It is a label diagram provided in the embodiment of the present application;

[0052] Figure 17 yes Figure 16 Three-dimensional RGB chromaticity distribution diagram;

[0053] Figure 18 yes Figure 16 RGB channel distribution histogram;

[0054] Figure 19 yes Figure 16 RGB intensity graph (middle row). DETAILED DESCRIPTION

[0055] In order to make the technical problems, technical solutions and beneficial effects to be solved by this application more clearly understood, this application is further described in detail below with reference to the accompanying drawings and embodiments. It should be understood that the specific embodiments described herein are only used to explain this application and are not intended to limit this application.

[0056] Example 1

[0057] like Figure 1 As shown, this application provides an adaptive underwater image enhancement method based on Retinex theory and Mamba, including:

[0058] Light Information Estimation Subnetwork: Taking the original underwater image and its corresponding illumination prior map as input, the network sequentially obtains an illumination factor map through feature concatenation, the first convolutional layer, the multi-scale channel aggregation module layer, and the second convolutional layer. The multi-scale channel aggregation module layer extracts the illumination estimation component of the underwater scene and performs a Hadamard product operation on the illumination factor map and the original image to obtain a basic enhanced image.

[0059] Degradation restoration subnetwork: The encoder extracts multi-scale degradation features of the base enhanced image step by step, the decoder uses skip connections to restore the spatial resolution, and the WLMamba bottleneck module selectively models the state space to enhance the illumination-texture correlation, ultimately generating high-quality underwater image enhancement results.

[0060] According to the traditional Retinex theory, any natural image Can be decomposed into reflection components and light component , the expression is as follows:

[0061] ;

[0062] Where ⊙ represents the Hadamard product operation.

[0063] However, this theory is based on the assumption of non-degraded images, which is significantly different from actual underwater imaging scenarios. Underwater image degradation mainly stems from the following mechanisms:

[0064] The medium scattering effect causes the image quality to deteriorate, which includes forward scattering and backscattering processes:

[0065] Forward scatter (reflected light that deviates from its original path due to medium interference and is captured by the sensor) usually has a smaller impact.

[0066] Backscatter (the incident light is reflected by water or suspended particles and enters the imaging system before reaching the target) is the dominant factor in scattering.

[0067] The wavelength-dependent attenuation characteristics cause color shift. Water absorbs the red light band most significantly, while the absorption of the blue and green bands is relatively weak.

[0068] The imaging process significantly amplifies sensor noise and image artifacts, leading to underexposure or color distortion.

[0069] To address the above degradation mechanism, we introduce a disturbance term into the original Retinex model to establish an improved equation:

[0070] ;

[0071] ;

[0072] In the formula represents the illumination component perturbation term, Represents the disturbance term of the reflection component. Considered as the benchmark image under ideal exposure conditions. To obtain a better quality image, an illumination factor mapping is applied to both sides of the improved formula. Processing, where :

[0073] ;

[0074] Where, Characterize the noise artifacts generated during imaging through light mapping Amplified degradation process, This comprehensively reflects the coupling effect of water scattering attenuation and exposure color shift. The simplified equation can be expressed as:

[0075] ;

[0076] Where, represents the pure image after degraded illumination, Characterize the composite degradation function.

[0077] Then build the degradation optimization processing (ROP) framework:

[0078] ;

[0079] Where, As a lighting feature estimation module, it is responsible for extracting underwater lighting prior features . The original input and its corresponding illumination prior map (obtained by calculating the mean of channel-dimensional pixels) Input this module and the output contains the basic enhanced image and illumination estimation component based on underwater scene feature extraction .

[0080] Then it will and The image is input to the degradation restoration module, and the degradation compensation is achieved through joint optimization to finally generate an enhanced image.

[0081] Optical Information Estimation Subnetwork

[0082] The traditional Retinex algorithm has limitations because it does not consider the underwater light propagation characteristics. Figure 2 As shown in Figure 2, this study innovatively designed a light information estimation subnetwork. By mining the semantic relevance of underwater images and constructing a multi-scale feature fusion mechanism, this subnetwork effectively extracts deep features closely related to the optical properties of water. Specifically, the network employs a cross-channel attention mechanism to capture color attenuation patterns, significantly improving the representational capabilities of feature maps. The specific structure is as follows:

[0083] According to the degradation optimization processing (ROP) framework, the input of the entire light information estimation is the original underwater image input and its corresponding illumination prior map (obtained by calculating the mean value of pixels in the channel dimension) are spliced ​​in the channel dimension and enter feature extraction. The specific structure is as follows:

[0084] Conv1x1 fusion: First, use the convolution layer with a convolution kernel size of 1×1 to fuse the spliced ​​results and perform preliminary detail feature extraction.

[0085] In view of the strong absorption of red band and weak absorption of blue-green band by water, this study innovatively proposes a multi-scale channel aggregation module (MSCA). This module dynamically integrates the multi-channel information of underwater images through an adaptive feature aggregation strategy, effectively compensating for wavelength-dependent attenuation, and thus approaching the spectral balance requirement of the gray world hypothesis (the ratio of the three channels R, G, and B is 1:1:1). The specific structure is as follows Figure 3 As shown, there are three branches of MSCA.

[0086] The first intermediate branch is mainly used for multi-scale aggregation on multiple channels. First, the input features According to the number of channels, it is divided into four equal parts, namely . Then, for Figure 3 The multi-scale downsampling strategy is implemented for the feature channels 1 to 4 in the middle. The specific operations are as follows:

[0087] Channel 1: uses 8x downsampling to extract global semantic features;

[0088] Channel 2: Balances local details and contextual information through 4x downsampling;

[0089] Channel 3: Perform a 2x downsampling operation to retain medium-scale features;

[0090] Channel 4: Maintains original underwater resolution to preserve high-frequency details.

[0091] This operation generates multi-scale feature representation through progressive scale compression. Its design is based on the following: the red light band (corresponding to Channel 1) in the underwater environment requires a larger range of context compensation due to strong absorption, while the blue-green band (corresponding to Channel 2 and Channel 3) can retain more original details due to weak absorption. The part based on the prior illumination map needs to be fully retained (corresponding to Channel 4). feature.

[0092] In order to facilitate calculation, the features after three multi-scale downsampling are Perform a 3×3 depth convolution and restore the dimension through upsampling operation. The formula is as follows:

[0093] ;

[0094] Where, upsample (·) represents the upsampling process, DWConv Represents the depthwise convolution process.

[0095] The current features are spliced ​​on the channel, and then aggregated through a 1×1 convolutional layer to obtain multi-scale aggregated features. , and then Apply the GELU activation function to obtain the final output features .

[0096] ;

[0097] The first lower branch adjusts the channel importance of the entire input feature through channel attention. First, a pooling operation is performed in the spatial direction to obtain a spatial descriptor representation. Then, a 1×1 convolutional layer is used to learn the correlation between channels, and a sigmoid function is used to generate a channel weight representation:

[0098] ;

[0099] Where, AvgPlool (·) represents average pooling, Conv represents the convolution operation, Sigmoid (·) represents the operation of generating channel weights through the Sigmoid function, which corresponds to Figure 3 The first lower branch in Figure 5 The lower branch in the Figure 3 The lower branch in is the first lower branch, and the naming rules of the first upper branch and the first middle branch are the same). The ReLU function in the first lower branch is not reflected in the formula. The ReLU function is used between the convolutional layers of two 1×1 convolutional layers to capture the nonlinear relationship.

[0100] The first upper branch represents the original input features. Since multi-scale operations may cause the loss of some detail information in the spatial dimension, retaining the original input can effectively supplement the spatial details, thereby improving the integrity of feature expression.

[0101] Finally, the results of the three branches are multiplied element by element to obtain the final illumination estimation component of the underwater scene. :

[0102] ;

[0103] (3) conv1x1 convolution: The illumination estimation component is convolved through a 1×1 convolution kernel. Integrate to generate a light factor map .we will Set to a three-channel RGB tensor to better achieve color enhancement effects, and finally With the original image Perform Hadamard product operation to obtain the basic enhanced image .

[0104] Degradation repair subnetwork

[0105] This application uses a degradation restoration sub-network based on a U-shaped adaptive network to correct the image after being processed by the light information estimation sub-network. The residual error in the detail texture and channel response is used to generate high-quality underwater scene reconstruction results. The network structure is as follows Figure 4 shown.

[0106] The overall architecture of the degradation restoration subnetwork adopts a four-layer encoder-decoder structure, including cross-layer skip connections and the WLMamba bottleneck module. This architecture, centered on the design of the WLMamba module, achieves accurate restoration of underwater image details by deeply fusing the illumination estimation component with the channel-spatial domain characteristics of the feature map. The encoder extracts multi-scale degradation features step by step, while the decoder restores spatial resolution through skip connections. The WLMamba bottleneck module enhances illumination-texture correlation through selective state-space modeling, ultimately generating high-quality underwater image enhancement results. The WLMamba design scheme is as follows:

[0107] Use WLMamba module, the module architecture is as follows Figure 5 As shown, the input features Entered three branches.

[0108] The second upper branch retains the original input features;

[0109] Second intermediate branch: characteristics aisle First, we increase it to , is a predefined channel expansion factor, which is set to 2 here, followed by depthwise convolution (DWCNN), SiLU activation function, SSM-2D layer and layer normalization (LN: LayerNorm).

[0110] Second lower branch: characteristics First, the components are estimated by optical information The modulated channel attention module is cascaded with the spatial attention module and finally activated by SiLU.

[0111] The overall fusion formula is as follows:

[0112] ;

[0113] Where X1 represents the process of the second middle branch, X2 represents the process of the second lower branch, and X outRepresents the combination and output process of the second upper branch, the second middle branch, and the second lower branch. Linear (·) represents linear transformation, DWConv (·) represents depthwise convolution, SiLU represents activation through SiLU function, SSM2D (·) represents processing through SSM-2D layer, LN (·) represents layer normalization, the second middle branch of the WL Mamba structure in encoder 1 may not include the layer normalization process in the initial stage of feature processing, the WL Mamba structure in encoders 2-4 and the second middle branch of the WLMamba structure between encoder 4 and decoder 4 need to include the layer normalization process, so the structure in the above formula can represent the WL Mamba structure in encoder 1, and Figure 5 The structure shown can represent the WL Mamba structure in encoders 2-4, Figure 5 The SiLU activation function is not shown in .

[0114] The input of the channel attention module (CA) consists of the illumination estimation component and image processing intermediate feature maps For the WL Mamba structure of encoder 1, the image processing intermediate feature map refers to the basic enhanced image , for the WL Mamba structure in encoders 2-4 and the WL Mamba structure between encoder 4 and decoder 4, the intermediate feature map of image processing refers to the feature map output after processing by the previous encoder. First, the feature map passes through the global average pooling layer to compress the spatial dimension and eliminate the dependency of spatial position, thereby capturing the global statistical features of each channel. Subsequently, the compressed channel feature vector is subjected to a nonlinear transformation through a fully connected layer group to learn the complex relationship between channels, and then activated and outputted through the Sigmoid function as a weight vector. In the feature recalibration stage, the weight vector is combined with the illumination estimation component The corresponding channels are multiplied element by element to achieve adaptive modulation of the feature map. Through this mechanism, the channel response with strong pixel correlation is significantly enhanced, redundant channel information is suppressed, and the model focuses on the key feature dimensions. The final output is the weighted feature map .

[0115] The spatial attention module (SA) first performs an Average pooling and maximum pooling operations are performed separately. Average pooling obtains global statistical features by calculating the mean of the feature maps of each channel, while maximum pooling extracts the local maximum of each channel to retain the significant response area. The two pooling operations capture spatial features from two dimensions: global statistics and local significance: average pooling focuses on overall distribution characteristics, and maximum pooling focuses on strong activation areas. Subsequently, the average feature map and the maximum feature map are spliced ​​along the channel dimension to generate a composite representation that fuses global and local features. The composite features are subjected to spatial information fusion and nonlinear modeling through a 3×3 convolution kernel, and the spatial attention weight is generated after activation by the Sigmoid function. Finally, the image in the garden that is related to the illumination component is integrated. Element-by-element multiplication realizes illumination-spatial characteristic modulation, strengthens regional response and suppresses identical information, and obtains a feature map after spatial attention optimization.

[0116] (3) SSM-2D (State Space Model-2D) module: Based on the Mamba principle, SSM-2D is used to capture long-range dependencies in underwater images. Figure 6 As shown, the mathematical principle is as follows:

[0117] ① Input feature serialization: input feature map After flattening, it is converted into a one-dimensional sequence in four directions, S d Represents the feature sequence after flattening along direction d:

[0118] ;

[0119] in, Indicates direction The scanning sequence arranges the two-dimensional features into a one-dimensional sequence, each sequence length is , the element dimension is . Here, among them The four directions represented by the numbers are:

[0120] Main diagonal scan: upper left → lower right direction, traversal in row-priority order;

[0121] Sub-diagonal scan: bottom right → top left, traversal in reverse order;

[0122] Vertical scan: from lower left to upper right, traversing in column priority order;

[0123] Reverse vertical scanning: from upper right to lower left, traversal is prioritized in reverse order.

[0124] ②Multi-directional selective state space model: for each direction sequence , applying a dynamically parameterized state-space model:

[0125] Dynamic parameter generation: Generate direction-dependent parameter matrices from the input sequence via linear projection:

[0126] ;

[0127] in, represents the learnable projection matrix, is the time step related discretization parameter, represents the input features at time step t, represents the input matrix representing time step t, denotes the output matrix at time step t, and softplus(·) denotes the softplus activation function.

[0128] Discretization of state space equations: Discretization of continuous state space parameters Perform zero-order hold (ZOH) discretization:

[0129] ;

[0130] Where, represents the discretized state transfer matrix, represents the discretized input matrix, represents the matrix exponential, which is used to discretize the continuous state space parameters, where Represents the identity matrix.

[0131] The hidden state update equation after discretization is:

[0132] ;

[0133] Where, represents the hidden state at time step t, represents the hidden state at time step t−1.

[0134] The output features are:

[0135] ;

[0136] in is the skip connection parameter matrix.

[0137] ③ Multi-directional feature fusion, the output sequence of the four directions Restore to a two-dimensional feature map and fuse it with attention weights:

[0138] ;

[0139] Where Y dRepresents the two-dimensional feature map restored by direction d, y (d) represents the output sequence of direction d, Y represents the final fused feature map, Indicates direction The inverse operation restores the tensor to Features, is the directional attention weight, generated by lightweight convolution:

[0140] ;

[0141] in Represented as Sigmoid function, Represents element-by-element multiplication, [Y1;Y2;Y3;Y4]: represents concatenating the feature maps of the four directions by channel.

[0142] The U-shaped network area is composed as follows:

[0143] Encoder: The encoder consists of a WLMamba and a downsampling convolutional layer (convolution kernel size is 3×3, stride is 2, and padding is 1);

[0144] Decoder: The decoder consists of a WLMamba with upsampling transposed convolutional layers (kernel size 3×3, stride 2, padding 1, and extra padding 1).

[0145] Skip connection: Skip connection is used to bridge the corresponding layers of the encoder and decoder in the U-shaped network to achieve multi-scale feature fusion. The specific implementation consists of two steps: feature alignment operation and adaptive fusion operation. The principle is as follows:

[0146] Feature alignment operation: To achieve the To effectively fuse the encoder and decoder features of the layer, it is necessary to ensure that the sizes of the two are matched. The main process is spatial alignment and channel alignment. Spatial alignment is to upsample the encoder features to the decoder size by bilinear difference, and then adjust the number of channels by a convolution kernel of size 1×1 to obtain the aligned features. .

[0147] Adaptive fusion operation: The input of adaptive fusion consists of the output features of the feature alignment operation and the decoder output features spliced ​​on the channel. Here we set two learnable parameters ,in We set up The decoder feature of the layer is , obtain the adaptive fusion features The feature results after adaptive fusion Enter The decoder of the layer completes the adaptive operation.

[0148] Neural network training

[0149] For training, we used two reference datasets: UIEB (Underwater Image Enhancement Benchmark) and LSUI (Large Scale Underwater Image Dataset). The UIEB dataset contains 890 labeled real underwater images, randomly partitioned into UIEB-Train (800 images) and UIEB-Test (90 images) for training and testing. The LSUI dataset contains 4279 labeled real underwater images. We used the official LSUI partition, consisting of LSUI-Train (3879 images) for training and LSUI-Test (400 images) for testing. We also used the C60, U45, and UCCS benchmark datasets for no-reference dataset quality evaluation. The C60 dataset is derived from the 60-image challenge set provided by the UIEB dataset official website. The U45 dataset contains 45 no-reference images, covering both underwater foggy and green scenes. UCCS contains 300 images with no-reference hues of cyan, green, and blue.

[0150] Training details

[0151] This method was implemented in PyTorch using an NVIDIA RTX4090 graphics card. The MAE (Mean Absolute Error) loss was used as the training loss. The total training epochs were set to 100, and the batch size was set to 8. The Adam (Adaptive Moment Estimation) optimizer was used, and the initial learning rate was set to , which steadily decreases to .

[0152] In terms of data preprocessing: the input image is resized to a fixed size (224 × 224), the mean and variance of the entire dataset are normalized, and the training data is expanded by random rotation and flipping. The training-loss image is as follows Figure 7 shown.

[0153] Observing the above training process is sufficient to verify the convergence of this method during the training process.

[0154] Model results verification

[0155] On the reference dataset, the peak signal-to-noise ratio (PSNR), structural similarity (SSIM), UCIQE (Underwater Color Image Quality Evaluation), and UIQM (Underwater Image Quality Measure) are used as evaluation metrics. For the non-reference dataset, UCIQE and UIQM are used as evaluation metrics. The above evaluation metrics are defined as follows:

[0156] Peak signal-to-noise ratio: Peak signal-to-noise ratio (PSNR) is a commonly used image quality assessment metric. It calculates the mean square error (MSE) between the original image and the processed image, and then converts it into decibels to express the image quality. The formula is as follows:

[0157] ;

[0158] in, and are the original image and the processed image respectively. and are the height and width of the image, and The images are at positions The pixel value at .

[0159] The higher the PSNR value, the smaller the difference between the processed image and the original image, and the better the image quality.

[0160] Structural similarity (SSIM) is an indicator used to measure the similarity between two images. It not only considers the grayscale difference of the images, but also the structural information of the images. It is more consistent with the human visual system's perception of image quality. The formula is as follows:

[0161] ;

[0162] in They are images and The mean of Represents images respectively and The variance of is the covariance of the image, C1 and C2 are constants used to stabilize the formula and prevent the denominator from being zero.

[0163] The SSIM value range is usually between [-1, 1]. When SSIM approaches 1, it means that the correlation between the two images is greater. When it approaches -1, it means that the correlation between the two images is smaller. When the SSIM value is 0, it means that there is no linear correlation between the two images.

[0164] UCIQE is an objective metric for evaluating underwater color image quality. It comprehensively considers multiple underwater image characteristics, including brightness, contrast, and saturation, to provide a comprehensive image quality assessment. The formula is as follows:

[0165] ;

[0166] Where, represents the image contrast, Indicates the image saturation, represents the image brightness, Represent the weight coefficients, respectively, used to adjust the contribution of contrast, saturation and brightness in the overall evaluation.

[0167] Contrast It can be measured by calculating the standard deviation of the grayscale histogram of the image. The larger the standard deviation, the higher the contrast of the image. The calculation formula is as follows:

[0168] ;

[0169] Where, Indicates the image The gray value of a pixel, represents the average gray value of the image, Represents the total pixel value of the image.

[0170] Saturation It can be measured by calculating the average value of the saturation channel of the image in the HSV color space. The higher the average value, the higher the saturation of the image. The formula is as follows:

[0171] ;

[0172] Where, The first Pixel saturation.

[0173] brightness It is measured by calculating the average grayscale value of the image. The formula is as follows:

[0174] ;

[0175] UIQM is a comprehensive metric used to evaluate underwater image quality. UIQM is designed to more closely mimic the human visual system's perception of image quality, providing an objective and comprehensive image quality assessment. The formula is as follows:

[0176] ;

[0177] Where, represents the weight parameter, They represent image contrast, image saturation, and image brightness respectively. For the specific formula, see UCIQE.

[0178] Experimental results

[0179] Through systematic experimental analysis, we conducted a horizontal comparative study on nine SOTA (State of the Art) methods and the method provided in this application. The nine SOTA methods specifically include: 2 traditional methods, 3 convolutional neural network (CNN)-type methods, 3 Transformer-type methods, and 1 emerging Mamba-type method.

[0180] Table 1 Experimental results of traditional methods and CNN-based methods

[0181]

[0182] ↑: Indicates that an increase in the indicator value represents an improvement in performance.

[0183] GDCP: Gradient Domain Color Prior (Gradient Domain Color Prior);

[0184] MLLE: Minimum Local Learning Enhancement (local minimum learning enhancement);

[0185] CWR: Color Weighted Reconstruction;

[0186] SCNet: Spatial Convolutional Network;

[0187] FiveA+: FiveA+ (a deep learning-based image enhancement method);

[0188] TIP'18: IEEE Transactions on Image Processing, 2018;

[0189] TIP'22: IEEE Transactions on Image Processing, 2022;

[0190] IGASS'22: International Geoscience and Remote Sensing Symposium, 2022;

[0191] ICASSP'22:IEEE International Conference on Acoustics, Speech, and Signal Processing, 2022

[0192] BMVC'23: British Machine Vision Conference, 2023.

[0193] Table 2 Experimental results based on Transformer and Mamba

[0194]

[0195] ↑: Indicates that an increase in the indicator value represents an improvement in performance.

[0196] U-Trans: U-Transformer (U-network based on Transformer);

[0197] NU2Net: Nested U-Net (Nested U-Net);

[0198] X-CAUNET: X-CAU Network (an improved U-type network);

[0199] UWMamba: Underwater Mamba (Underwater Mamba Network);

[0200] TIP'23: IEEE Transactions on Image Processing, 2023;

[0201] AAAI'23: Association for the Advancement of Artificial Intelligence, 2023 (AAAI Annual Conference, 2023);

[0202] ICASSP'24:IEEE International Conference on Acoustics, Speech, and Signal Processing, 2024

[0203] SPL'24: IEEE Signal Processing Letters, 2024.

[0204] The experimental results shown in Tables 1 and 2 demonstrate that the adaptive underwater image enhancement method based on Retinex theory and Mamba provided in this application has achieved significant results in multiple evaluation indicators. Furthermore, this method outperforms other methods in terms of texture detail and contrast enhancement, fully demonstrating its superiority in enhancing underwater scenes.

[0205] Color channel comparison results

[0206] After underwater image enhancement, the performance of RGB channels is significantly improved: before enhancement, the G / B channel response is too strong and the R channel is significantly attenuated due to the selective light attenuation of the water body. Figures 8-11 As shown in , the histogram distribution is concentrated and there is color deviation; the enhanced result is as follows Figure 12-15 As shown, the channel attenuation difference is compensated, the R channel intensity is significantly improved, the G / B channel distribution is more uniform, the overall histogram dynamic range is expanded, and the contrast is improved. Compared with the label map (true value) Figure 16-Figure 19 Compared with the above, the enhanced image is roughly similar to the label image in channel intensity distribution, and the channel situation is closer to the color balance of the real scene.

[0207] Figure 11 、 Figure 15 and Figure 19 The so-called intermediary bank refers to Figure 8 、 Figure 12 and Figure 16 The row in the middle of the horizontal image, and the sampling positions of the three RGB intensity curves are consistent.

[0208] Through systematic experimental verification, the method proposed in this study demonstrated significant effectiveness, robustness and generalization ability in underwater image enhancement tasks, providing a reliable technical path for visual restoration of complex underwater environments.

[0209] Enhanced method performance analysis

[0210] In the computational efficiency evaluation, this study conducted standardized tests based on 100 standard test images with a resolution of 224×224. The quantitative analysis results shown in Table 3 demonstrate that the proposed method significantly outperforms the comparison method with a minimum computational complexity of 15.20G FLOPs and a single-frame processing time of 9ms, providing important technical support for real-time underwater vision systems.

[0211] Table 3 Operation efficiency experimental results

[0212]

[0213] ↓: A decrease in the indicator value indicates an improvement in performance.

[0214] FLOPs(G): floating point operations per second (in billions), indicating computational complexity;

[0215] TIME(S): time (in seconds), indicating the processing time of a single frame image;

[0216] Params(M): Parameter quantity (in millions), indicating the number of parameters of the model.

[0217] UIEC2: Underwater Image Enhancement CNN-based using two Color Spaces (underwater image enhancement based on convolutional neural networks and two color spaces);

[0218] X-CAUNET: A cross-color channel attention mechanism based on transformer for underwater image enhancement.

[0219] The above embodiments are only used to illustrate the technical solutions of the present application, rather than to limit them. Although the present application has been described in detail with reference to the aforementioned embodiments, those skilled in the art should understand that they can still modify the technical solutions described in the aforementioned embodiments, or make equivalent replacements for some of the technical features therein. These modifications or replacements do not deviate the essence of the corresponding technical solutions from the spirit and scope of the technical solutions of the various embodiments of the present application, and should all be included in the scope of protection of the present application.

Claims

1. An adaptive underwater image enhancement method based on Retinex theory and Mamba, characterized by: include: Light Information Estimation Subnetwork: Taking the original underwater image and its corresponding illumination prior map as input, the network sequentially obtains an illumination factor map through feature concatenation, the first convolutional layer, the multi-scale channel aggregation module layer, and the second convolutional layer. The multi-scale channel aggregation module layer extracts the illumination estimation component of the underwater scene and performs a Hadamard product operation on the illumination factor map and the original image to obtain a basic enhanced image. Degradation Restoration Sub-Network: The encoder extracts multi-scale degradation features from the base enhanced image step by step, the decoder uses skip connections to restore the spatial resolution, and the WLMamba bottleneck module selectively models the state space to enhance the illumination-texture correlation, ultimately generating high-quality underwater image enhancement results. The multi-scale channel aggregation module layer includes adjusting the detail features initially extracted by the first convolutional layer through the first upper branch, the first middle branch, and the first lower branch, and then performing element-by-element multiplication on the results of the three branches to obtain the final illumination estimation component of the underwater scene; The first upper branch retains the original input features; The first intermediate branch is used for multi-scale aggregation on multiple channels; The first lower branch performs a pooling operation in the spatial direction to obtain a spatial descriptor representation, then learns the correlation between channels through the lower branch convolution layer, and generates a channel weight representation through the Sigmoid function; The WLMamba bottleneck module adjusts the features it receives through the second upper branch, the second middle branch, and the second lower branch simultaneously; The second upper branch retains the original input features; The second intermediate branch will input the feature channel First, by linearly increasing , is a predefined channel expansion factor, followed by depthwise convolution, followed by SiLU activation, SSM-2D layer, and layer normalization; The second lower branch combines the input features first through a channel attention module modulated by the light information estimation component and a spatial attention module in cascade, and then activates them through SiLU.

2. The adaptive underwater image enhancement method based on Retinex theory and Mamba according to claim 1, characterized in that: The first intermediate branch splits the input features into four equal parts according to the number of channels. Then, a multi-scale downsampling strategy is implemented on the split input features through the feature channels of Channel 1 to Channel 4. The specific operations are as follows: Channel 1: uses 8x downsampling to extract global semantic features; Channel 2: Balances local details and contextual information through 4x downsampling; Channel 3: Perform a 2x downsampling operation to retain medium-scale features; Channel 4: Maintains original underwater resolution to preserve high-frequency details; The features after downsampling from Channel 1 to 3 are subjected to 3×3 depth convolution, and the dimension is restored through upsampling operation. The features of the feature channels of Channel 1 to 4 are spliced, and then the aggregation operation is performed through the first intermediate branch convolution layer to finally obtain multi-scale aggregated features.

3. The adaptive underwater image enhancement method based on Retinex theory and Mamba according to claim 1, characterized in that: The input of the channel attention module consists of an illumination estimation component and an intermediate feature map of image processing. First, the intermediate feature map is passed through a global average pooling layer to capture the global statistical features of each channel. Subsequently, the compressed channel feature vector is subjected to a nonlinear transformation through a fully connected layer group to learn the complex relationship between channels. The channel weight is then activated and outputted through a Sigmoid function to obtain the channel weight. In the feature recalibration stage, the above-mentioned channel weight is multiplied element-by-element with the channel corresponding to the illumination estimation component to achieve adaptive modulation of the feature map, and the final output is a weighted feature map.

4. The adaptive underwater image enhancement method based on Retinex theory and Mamba according to claim 3, characterized in that: The weighted feature map output by the channel attention module is used as the input of the spatial attention module. First, the input feature map is subjected to average pooling and maximum pooling operations respectively. Then, the average feature map and the maximum feature map are spliced ​​along the channel dimension to generate a composite feature that integrates global and local features. The composite feature is subjected to spatial information fusion and nonlinear modeling through a 3×3 convolution kernel. After activation by a Sigmoid function, a spatial attention weight is generated. Finally, the spatial attention weight is element-wise multiplied with the illumination estimation component to realize illumination-spatial characteristic modulation and obtain the feature map after spatial attention optimization.

5. The adaptive underwater image enhancement method based on Retinex theory and Mamba according to claim 1, characterized in that: The SSM-2D layer is used to capture long-range dependencies in underwater images, which involves converting the input feature map into a one-dimensional sequence in four directions through a flattening operation, arranging the two-dimensional features into a one-dimensional sequence according to the scanning sequence of the directions, applying a dynamically parameterized state space model to each direction sequence, restoring the output sequences in the four directions to a two-dimensional feature map, and fusion through attention weighting.

6. The adaptive underwater image enhancement method based on Retinex theory and Mamba according to claim 1, characterized in that: The overall architecture of the degradation repair sub-network adopts a four-layer encoder-decoder structure, including cross-layer skip connections and WLMamba bottleneck modules; The encoders of layers 1-3 perform multi-scale feature fusion with the decoders of their corresponding layers through skip connections. The encoder features of the 4th layer are adjusted by the WLMamba bottleneck module and then passed to the decoder of the 4th layer.

7. The adaptive underwater image enhancement method based on Retinex theory and Mamba according to claim 6, characterized in that: The jump connection consists of a feature alignment operation and an adaptive fusion operation, wherein the feature alignment operation includes spatial alignment and channel alignment. Spatial alignment includes upsampling the encoder features to the decoder size through bilinear difference, and then adjusting the number of channels through a convolution kernel of size 1×1 to obtain the aligned features. ; The input of the adaptive fusion operation is composed of the output features of the feature alignment operation and the decoder output features, which are spliced ​​on the channel. Two learnable parameters are set: ,in , set up the The decoder feature of the layer is , obtain the adaptive fusion features ; The adaptive fusion features Enter The decoder of the layer completes the adaptive fusion operation.

8. An adaptive underwater image enhancement system based on Retinex theory and Mamba, characterized by: include: Light Information Estimation Subnetwork: This module takes the original underwater image and its corresponding illumination prior map as input, sequentially passes through feature concatenation, the first convolutional layer, the multi-scale channel aggregation module layer, and the second convolutional layer to obtain an illumination factor map. The multi-scale channel aggregation module layer extracts the illumination estimation component of the underwater scene, and performs a Hadamard product operation on the illumination factor map and the original image to obtain a basic enhanced image. Degradation restoration sub-network module: This module extracts multi-scale degradation features from the base enhanced image through the encoder, restores spatial resolution through decoder skip connections, and enhances illumination-texture correlation through selective state space modeling of the WLMamba bottleneck module, ultimately generating high-quality underwater image enhancement results. The multi-scale channel aggregation module layer includes adjusting the detail features initially extracted by the first convolutional layer through the first upper branch, the first middle branch, and the first lower branch, and then performing element-by-element multiplication on the results of the three branches to obtain the final illumination estimation component of the underwater scene; The first upper branch retains the original input features; The first intermediate branch is used for multi-scale aggregation on multiple channels; The first lower branch performs a pooling operation in the spatial direction to obtain a spatial descriptor representation, then learns the correlation between channels through the lower branch convolution layer, and generates a channel weight representation through the Sigmoid function; The WLMamba bottleneck module adjusts the features it receives through the second upper branch, the second middle branch, and the second lower branch simultaneously; The second upper branch retains the original input features; The second intermediate branch will input the feature channel First, by linearly increasing , is a predefined channel expansion factor, followed by depthwise convolution, followed by SiLU activation, SSM-2D layer, and layer normalization; The second lower branch combines the input features first through a channel attention module modulated by the light information estimation component and a spatial attention module in cascade, and then activates them through SiLU.

Images