Underwater image enhancement method based on two-stage network of color correction and multi-color space stretching

Abstract

The application discloses a two-stage network underwater image enhancement method based on color correction and multi-color space stretching, utilizes a green channel to compensate for a red channel and a blue channel to obtain a red channel correction graph and a blue channel correction graph, constructs an improved U-Net network, and outputs red channel weight mapping and blue channel weight mapping; the red channel correction graph and the blue channel correction graph are respectively multiplied by point product of the corresponding red channel weight mapping and the blue channel weight mapping, a dynamic compensation result of the red channel and the blue channel is spliced into an RGB space color correction image together with the green channel; the RGB space color correction image is respectively converted into an HSI space and a Lab color space, HSI space, Lab color space and RGB space are respectively histogram stretched, finally, the three histogram stretched images are respectively converted back to the RGB color space, and three images converted back to the RGB color space are spliced to obtain a comprehensive input image; the comprehensive input image is input into the improved U-Net network to output an enhanced image, and the resolution of underwater camera imaging and the color correction effect can be significantly improved.

Description

Technical Field

[0001] The present invention belongs to the technical field of computer vision and relates to an underwater image enhancement method, in particular to an underwater image enhancement method based on color correction and multi-color space stretching. Background Art

[0002] Underwater image processing is of great significance to the exploration and research of various marine resources, such as marine biology, underwater target detection, and autonomous underwater vehicle (AUV) navigation. However, underwater images have always suffered from poor quality degradation. Specifically, due to the rapid attenuation of red light when propagating in water, underwater images often appear to have blue and green as the main tones, losing the color balance of the real world. At the same time, the presence of micro-suspended particles in the underwater environment, such as mud and plankton, causes light scattering, which significantly reduces the contrast of underwater images and blurs the image details. These problems severely limit the ability to obtain meaningful information and are not conducive to the further processing of underwater images. Therefore, underwater image enhancement (UIE) plays an important role in solving this problem.

[0003] In recent years, many effective methods have been proposed for underwater image enhancement. These methods can be roughly divided into two categories: physics-based and deep learning-based. While physics-based methods can demonstrate excellent performance in certain scenarios, they are sensitive to the diversity and complexity of underwater image types. This means that the predictive power of these physical models may be limited in the dynamic and complex underwater environment. Existing deep learning image enhancement methods can be categorized as those based on convolutional neural networks (CNNs) and generative adversarial networks (GANs). Compared to GANs, CNNs are widely used in image denoising, image dehazing, image enhancement, and image super-resolution. While the convolution operation of CNNs is popular for its ability to efficiently process local features, it also has a significant limitation: it focuses on processing information in a local neighborhood, which can lead to a loss of global information. To address this shortcoming, attention mechanisms have been widely used and explored in the field of visual processing in recent years. Attention mechanisms dynamically adjust the network's focus on different parts of the input data, enabling the model to automatically ignore unimportant information and allocate more computational resources to key information areas.

[0004] Currently, most underwater image enhancement algorithms use deep learning-based techniques, and these algorithms mainly operate in the RGB color space. However, the RGB color space itself has its limitations, as it does not directly reflect some key parameters and characteristics of the image, such as contrast, saturation, and brightness. For example, the underwater image enhancement method based on color correction and three-interval histogram stretching disclosed in patent publication CN112419210A first performs color correction on the original RGB heavy fog image using a sub-interval linear transformation-based color correction method. Then, the three-interval histogram equalization method is used to process the source image in R, G, and B channels respectively. The pixel values of R, G, and B channel images are stretched to the range of 0-255 to obtain the stretched single-channel images. The average pixel value of R, G, and B channel images is calculated respectively. The error is obtained by subtracting the average pixel value of the single channel from the pixel value of each pixel point, and the error is squared. The pixel point with the maximum error square value is selected, and the pixel values near this point also have a large error with the average pixel value. Therefore, the two threshold values required for three-interval division are determined by adding and subtracting three times the variance to the center point, and the entire single-channel histogram is divided into three intervals. The sub-interval of each R, G, and B channel is equalized to obtain the single-channel equalized image. Finally, the R, G, and B channel images and the equalized R, G, and B channel images are linearly weighted and fused to obtain the final de-fogging image. This underwater image enhancement method still operates in the RGB color space and can only achieve de-fogging of heavy fog images.

[0005] Therefore, in order to more comprehensively improve the quality of underwater images, it is necessary to explore enhancement algorithms that combine multiple color spaces and image correction techniques to more effectively optimize the image. SUMMARY

[0006] The purpose of the present application is to overcome the defects of the prior art and provide an underwater image enhancement method based on color correction and multi-color space stretching two-stage network, which combines multiple color spaces and image correction techniques to more comprehensively and effectively improve the quality of underwater images and more accurately capture and optimize important attributes of the image.

[0007] TECHNICAL SOLUTION

[0008] The underwater image enhancement method based on color correction and multi-color space stretching two-stage network of the present application includes the following steps:

[0009] Step 1): Extract R, G, and B channel components from the color original image, and use the green channel to compensate for the red and blue channels to obtain static red channel correction images and blue channel correction images

[0010] Step 2): Construct an improved U-Net network, which includes a parallel hybrid convolution-self-attention feature extraction module, an ASPP void space pyramid pooling layer, and an SFFusion selective feature map fusion method. Input the original color image into the improved U-Net network, train the improved U-Net network using the Adam optimizer, obtain the trained improved U-Net network, and output the red channel weight map. and the blue channel weight map

[0011] Step 3): The red channel correction image Blue channel correction chart Respectively mapped to the corresponding red channel weight and the blue channel weight map Performing dot product multiplication to obtain motion compensation results of the red and blue channels respectively, and splicing the motion compensation results of the red and blue channels and the green channel to form a color RGB space color correction image;

[0012] Step 4): converting the RGB color-corrected image into an HSI space and a Lab color space, respectively, and then performing histogram stretching on the HSI, Lab, and RGB color spaces to obtain histogram-stretched images of the respective color spaces. Finally, converting the three histogram-stretched images back to the RGB color space, and merging the three images converted back to the RGB color space in the channel direction to obtain a comprehensive input image.

[0013] Step 5): Input the comprehensive input image into the improved U-Net network for retraining, and output the enhanced image.

[0014] Furthermore, the SFFusion selective feature fusion method is to splice the upsampled feature map and the skip-connected feature map in the channel dimension, and the spliced ​​feature map is subjected to the global average pooling layer GAP to generate a global feature vector, and the global feature vector is processed by the multi-layer perceptron MLP to generate a selective weight vector; the MLP is composed of one or more fully connected layers and activation functions, and outputs a selective weight vector; finally, the selective weight vector is element-wise multiplied with the upsampled feature map and the skip-connected feature map respectively to obtain a weighted feature map; the two weighted feature maps are added element-by-element to generate the final fused feature map.

[0015] Furthermore, the initial input image of the improved U-Net network is a color image of size H*W*3, with a total of eight layers. The first half and the second half of the network are the encoder and decoder respectively; an ASPP module is added between the encoder and the decoder; the decoder uses the same hybrid convolution-self-attention feature extraction module as the encoder; the hybrid convolution-self-attention feature extraction module adopts a dual-branch structure, the upper branch structure includes two convolution layers, each layer has a convolution kernel Conv of size 3×3, a step size of 1, two normalization layers BN, and two ReLU activation functions; the lower branch structure includes a pooling layer, a To The p-K self-attention mechanism, 1×1 convolution kernel and Reshape reshaping layer are used to match the output of other branches; first, the input feature map is subjected to the maximum pooling operation and down-sampled; then it is passed to the upper branch and lower branch modules respectively; the upper branch is a standard convolution block, which performs a 3×3 convolution operation and extracts features from the down-sampled feature map; the convolved feature map is batch normalized; the lower branch is a Top-K self-attention mechanism block, which performs an average pooling operation on the down-sampled feature map, further down-samples, executes the Top-K self-attention mechanism, and finally reshapes the feature map.

[0016] Furthermore, the ASPP module includes four parallel 3×3 convolutional layers, each with a different dilation rate of 1, 6, 12, and 18. The output feature maps of the four convolutional layers are concatenated in the channel dimension to form a comprehensive feature map, which is then processed by a 1×1 convolutional layer to generate the final output feature map.

[0017] Furthermore, the feature map of each convolution operation and batch normalization operation is: i =BN(Conv(I,W ci ,b ci ,dilation=R i ,padding=R i ), I is the input feature map, W ci and b ci Represent the weight and bias of the i-th convolutional layer respectively; R i Represents the hole rate of the i-th convolution; the comprehensive feature map is represented as: x = Concat (x1, x2, x3, x4), and after the 1 × 1 convolution operation is: F = Conv (x, W c5 ,b c5 ), F is the output feature map of the ASPP module.

[0018] Furthermore, the input of the Top-K self-attention mechanism first passes through three independent 1×1 convolutional layers to generate three feature maps Q, K and V respectively. Each feature map Q, K, and V undergoes a 3×3 depth-wise separable convolution operation to further extract features; the feature maps Q, K, and V are reshaped, and the dot product of Q and K is calculated to obtain a self-attention score matrix; a MASK mask operation is applied to the self-attention score matrix, and the top-K operator is used to select the top K features with the highest self-attention scores; the selected top-K features are then redistributed to the corresponding positions, and the self-attention scores after Top-K are Softmax normalized to obtain normalized attention weights; the normalized attention weights are multiplied by the feature map V to obtain a weighted feature map; the weighted feature map is reshaped to match the shape of the initial input; and finally, the weighted feature map is added to the initial input to achieve feature fusion.

[0019] Technical effects:

[0020] 1. This invention is an underwater image enhancement algorithm that integrates traditional and deep learning. It can significantly improve the resolution and color correction effect of underwater camera imaging, better restore the detailed information of underwater images, and make the generated images closer to the real scene.

[0021] 2. Because the present invention utilizes an improved U-Net network, specifically our designed hybrid convolution-self-attention module, the self-attention mechanism can capture global feature relationships, while the hybrid convolution enhances the ability to capture spatial information. By combining these two approaches, the module effectively improves the model's ability to learn both global and local features while maintaining computational efficiency. This enhances the accuracy and richness of feature extraction, strengthens the segmentation of complex backgrounds or objects, and improves the model's robustness and generalization capabilities. The ASPP module is added between the encoder and decoder. The ASPP module can expand the receptive field without increasing computational effort, thereby enhancing the model's ability to recognize multi-scale features and improving segmentation performance, especially when images contain objects of varying sizes. SFFusion selective feature map fusion is employed. This method fuses multiple layers of features through a feature selection mechanism, suppressing useless information and enhancing key features. Compared to traditional direct concatenation or addition, SFFusion can selectively retain important feature information. This effectively reduces redundant information and improves the efficiency and quality of feature fusion.

[0022] 3. Because the present invention incorporates dynamic DCM correction into static SCM correction, the initial SCM correction provides a more balanced input foundation for the DCM, which then dynamically adjusts based on this foundation, making the color correction process more stable and efficient. The combination of static and dynamic corrections enables the model to achieve both rapid initial adjustments and in-depth optimization of details. This significantly improves the image's color adjustment, reducing color cast while preserving image detail. This enables subsequent image enhancement to be performed at a higher quality level, thereby enhancing the overall visual quality of the image.

[0023] 4. This invention employs histogram stretching in multiple color spaces. Different color spaces (such as RGB, HSI, and Lab) can describe the color and brightness characteristics of an image from multiple perspectives. Histogram stretching in these color spaces can more effectively enhance image contrast and color representation, correcting issues such as uneven lighting, low contrast, and color distortion.

[0024] 5. This invention combines multi-color spatial histogram stretching with deep learning, enabling the model to learn diverse enhancement features. This fusion of preprocessing and deep learning effectively enhances the model's robustness and generalization capabilities, enabling it to better cope with complex underwater environmental changes. BRIEF DESCRIPTION OF THE DRAWINGS

[0025] Figure 1 This is a flow chart of the two-stage network underwater image enhancement method based on color correction and multi-color space stretching of the present invention;

[0026] Figure 2 Schematic diagram of a trainable histogram stretching module according to an embodiment of the present invention.

[0027] Figure 3 This is a schematic diagram of the overall framework of the improved U-Net network;

[0028] Figure 4 Schematic diagram of the structure of the ASPP atrous spatial pyramid pooling layer module according to an embodiment of the present invention.

[0029] Figure 5 Schematic diagram of the Top-K Self-Attention module structure according to an embodiment of the present invention.

[0030] Figure 6 A schematic diagram of the SFFusion selective feature fusion method according to an embodiment of the present invention. DETAILED DESCRIPTION

[0031] The technical solution of the present invention will be further described below with reference to the accompanying drawings.

[0032] like Figure 1As shown, the present invention includes the following steps:

[0033] (1) Construct a static color correction module (SCM). Input the original color image into the static correction module. Based on the underwater IFM theory, the rich information contained in the green channel is used to compensate the red and blue channels. The correction images of the red and blue channels are obtained by compensating the green channel to achieve static correction.

[0034] The original color image can be an image extracted from UIEBD (Underwater Image Enhancement Benchmark Dataset). The original color image is processed to adjust the color saturation, hue and color accuracy of the original color image to improve the image display effect. The R, G, and B channel components of the original color image are extracted respectively. The extracted R, G, and B channels enter the static color correction for processing. Because the red channel and the blue channel will be severely attenuated underwater, while the green channel retains rich information, the green channel is used to compensate the red channel and the blue channel to obtain a static red channel correction image. and blue channel correction chart

[0035]

[0036] Among them I G , I R , I B Denote the green, red, and blue channels respectively, and ⊙ denotes the dot product. Traditional color correction methods mainly rely on region-based information extraction, but the choice of region size has a significant impact on the final extraction result, which may lead to inconsistent results and poor adaptability in different scenarios.

[0037] (2) Construct an improved U-Net network. The U-Net network is improved and the original color image is input into the improved U-Net network. The improved U-Net network analyzes the input image and extracts features related to color correction. The features extracted in the encoding phase are used in the decoding phase to predict the dynamic weight map and generate weights for each pixel region. These weights determine how each pixel should be adjusted during the color correction process. In this way, the improved U-Net network can achieve adaptive color correction across multiple regions.

[0038] The improved U-Net network is applied to the dynamic correction and dynamic fusion modules of underwater image enhancement: For dynamic correction, the improved U-Net network obtains two dynamic weight mappings by training the original image and A dot product operation is performed with the static correction result to produce the color correction image. For dynamic fusion, the three-space composite image obtained by histogram stretching is input into the improved U-Net network. The network learns the latent features of the input image from the three color spaces of RGB, HSI, and Lab, and dynamically fuses these features, retaining the most important information in each color space, and ultimately outputs the enhanced image.

[0039] The specific construction method of the improved U-Net network is as follows: Figure 3 The improved U-Net network shown in the figure has a U-shaped backbone network. The initial input image size is a color image of H*W*3, including a parallel hybrid convolution-self-attention feature extraction module, an ASPP void spatial pyramid pooling layer, and an SFFusion selective feature map fusion method.

[0040] The improved U-Net network consists of eight layers. The first half of the network is used for backbone feature extraction and serves as the encoder. The input image is passed to the inception block for preliminary feature extraction. The first four encoder layers all use a hybrid convolution-self-attention module for feature extraction.

[0041] The hybrid convolution-self-attention feature extraction module uses a dual-branch structure. The upper branch consists of two convolutional layers, each with a 3×3 convolution kernel with a stride of 1, two normalization layers, and two ReLU activation functions. The lower branch includes a pooling layer, a Top-K self-attention mechanism, a 1×1 convolution kernel, and a reshape layer to match the outputs of the other branches. First, the input feature map is downsampled by max pooling. It is then passed to the upper and lower branches, respectively. The upper branch is a standard convolution block that performs a 3×3 convolution operation and extracts features from the downsampled feature map. Batch normalization is performed on the convolved feature map to accelerate network training and improve stability. The ReLU activation function is applied to introduce nonlinear features. This repeated operation enables the network to better extract useful information from the input data. The lower branch consists of a Top-K self-attention block, which performs average pooling on the downsampled feature maps. It then further downsamples and applies a Top-K self-attention mechanism. This mechanism captures global information by calculating the correlation of the input feature maps and selects important features for processing using the Top-K operator. A 1×1 convolution is then performed to reduce the channel dimension, aggregate features, and finally reshape the feature maps using Reshape. While the standard convolution block in the upper branch focuses on local feature extraction, the Top-K self-attention mechanism in the lower branch captures global information. Therefore, combining these two blocks simultaneously leverages local and global features, enhancing the diversity and completeness of feature representation. An ASPP module is added between the encoder and decoder to extract multi-scale features, enhance contextual information, improve feature representation, and reduce information loss. The decoder uses the same hybrid convolution-self-attention feature extraction module as the encoder. SFFusion is used to connect feature maps between lower and higher layers. This selective feature fusion method effectively fuses features from the main path and skip connections by generating dynamic selective weights. Each upsampling is performed by fusing the output features of the same layer in the corresponding first half of the network to continuously restore the feature resolution. This is repeated four times, and the training results are finally output through a 1×1 convolution kernel.

[0042] like Figure 4 As shown in the figure, the ASPP module used is used. The ASPP module contains four parallel 3×3 convolutional layers, each with a different dilation rate of 1, 6, 12, and 18. The output feature maps of the four convolutional layers are then concatenated in the channel dimension to form a comprehensive feature map. The concatenated feature map is processed by a 1×1 convolutional layer to integrate features of different scales and adjust the number of channels. The 1×1 convolution helps to compress the channel dimension of the feature map and generate the final output feature map. The feature map of each convolution operation and batch normalization operation can be expressed as:

[0043] x i =BN(Conv(I,W ci ,b ci ,dilation=R i ,padding=R i )),

[0044] Where I is the input feature map, x i Represents the output feature map after the i-th convolution operation and batch normalization. BN represents batch normalization, which is used to normalize the output of the convolution layer. Conv represents the convolution operation, W ci and b ci Represent the weight and bias of the i-th convolution layer respectively. Dilation represents the dilation convolution rate, which is used to control the spacing between convolution kernel elements. i Represents the hole rate of the i-th convolution. Padding is used to add pixels to the edge of the input feature map to control the size of the output feature map.

[0045] The operation of feature map splicing can be expressed as:

[0046] x=Concat(x1,x2,x3,x4),

[0047] Finally, after a 1×1 convolution operation:

[0048] F=Conv(x,W c5 ,b c5 ),

[0049] F is the output feature map of the ASPP module.

[0050] like Figure 5As shown in the figure, the input of the Top-K self-attention mechanism first passes through three independent 1×1 convolutional layers to generate three feature maps, named Q (Query), K (Key), and V (Value). Each feature map Q, K, and V undergoes a 3×3 depthwise separable convolution operation to further extract features. The feature maps Q, K, and V are reshaped to facilitate subsequent matrix operations. The dot product of Q and K is calculated to obtain a self-attention score matrix. A mask operation is applied to the self-attention score matrix to remove irrelevant or invalid components, which is often used to ignore the influence of specific locations. The Top-K operator is used to select the top K features with the highest self-attention scores to reduce computation and focus on the most important features. The selected top-K features are then redistributed to the corresponding positions. The Top-K self-attention scores are softmax-normalized to obtain normalized attention weights. The normalized attention weights are multiplied by the feature map V to obtain a weighted feature map. The weighted feature map is reshaped to match the shape of the original input. Finally, the weighted feature map is added to the initial input to achieve feature fusion.

[0051]

[0052] where Y i Represents the vector of the i-th position in the feature map output by the self-attention mechanism, SelectAtt(Q i ,K i ,V i ) represents the function representation of the Top-K self-attention mechanism, which accepts three parameters: Q i (Query), K i (key), V i (value), Indicates that the score after Top-K selection is divided by Where d is the dimension of the key vector, and the softmax function is used to convert the score matrix into normalized attention weights. k (·) Select operators for learnable top-k:

[0053]

[0054] Here [T k (S)] ij is the Top-K function T k (·) is applied to the matrix S to obtain the element in the i-th row and j-th column of the output matrix. This element retains the original value S ij Or set to 0. S ij is an element in the matrix S, representing the attention score of the i-th row and j-th column. i yes The top-k values ​​of the i-th row of the dataset are represented by the K-ranked attention scores in that row. The top-k selection function retains elements greater than the threshold and sets elements less than or equal to the threshold to 0. This threshold is used to select the largest K values ​​in each row. Finally, the weighted feature map is added to the original input to achieve feature fusion.

[0055] like Figure 6 , the SFFusion selective feature fusion method. The upsampled feature map The feature map X2 of the skip connection is concatenated in the channel dimension. The concatenated feature map is passed through a global average pooling layer (GAP) to generate a compact global feature vector. The global feature vector is processed by a multi-layer perceptron (MLP) to generate a selective weight vector. The MLP consists of one or more fully connected layers (FC) and an activation function (ReLU). It outputs the selective weight vectors α1 and α2 after a softmax function.

[0056]

[0057] in is the upsampled feature map, X2 is the jump connection feature map, GAP represents the global average pooling layer, F mlp It is a multi-layer perceptron, consisting of one or more fully connected layers (FC) and activation functions (ReLU).

[0058] The attention weights calculated by MLP and softmax function are dynamic, which means that they will adaptively adjust the weights according to the input of the feature map, so that a more important feature map will get a higher weight. Finally, the selective weight vectors α1 and α2 are respectively compared with the upsampled feature map. Perform element-wise multiplication on the feature map X2 of the jump connection to obtain the weighted feature map. Add the two weighted feature maps to generate the final fusion feature map X. Use the weights {a1, a2} to X2 is multiplied by points to obtain the fused feature map X:

[0059] (3) Train the improved U-Net network and output the red channel weight map and the blue channel weight map Realize dynamic correction.

[0060] 800 original color images were extracted from UIEBD and preprocessed. A training dataset was generated from a random set of these 800 images. The images in the training dataset were resized to 256×256 and placed into the improved U-Net network. The improved U-Net network was trained using the Adam optimizer to obtain the trained improved U-Net network. The learning rate was set to 0.0001, and the batch size was set to 12.

[0061] The loss function of the improved U-Net network is defined as the sum of the mean absolute error loss function and the mean square error loss function, where the mean absolute error loss function is expressed as:

[0062]

[0063] Where H and W are the height and width of the input image, and (m,n) are the coordinates of the pixel. is the predicted value of the pixel at (m,n), y(m,n) is the true value of the pixel at position (m,n), is the absolute difference between the predicted value and the true value.

[0064] Mean squared error loss function:

[0065]

[0066] where f(x i ) and y i They represent the predicted value and the corresponding true value of the i-th sample respectively, and n is the number of samples.

[0067] The total loss function L loss It consists of the above two parts:

[0068] L loss =L MAE +L MSE .

[0069] The final red channel weight map is obtained through the improved U-Net network after training and the blue channel weight map

[0070] (4) Correct the red channel image in step (1) Blue channel correction chart Respectively mapped to the corresponding red channel weights obtained in step (3) and the blue channel weight map The final motion compensation results of the red and blue channels are obtained by multiplying the dot products, that is, the red and blue channel correction images:

[0071]

[0072] Afterwards, the red channel dynamic compensation result I′ r , blue channel dynamic compensation result I′ b , I G Fuse and stitch along the channel dimension to reconstruct the color RGB space color correction image I cc :

[0073]

[0074] Indicates splicing, I cc This is the RGB space color corrected image obtained by the color correction subnet.

[0075] (5) The obtained RGB space color correction image I cc Convert to HSI space and Lab color space respectively. The formula for converting RGB to HSI space is:

[0076]

[0077] First convert RGB to XYZ space:

[0078]

[0079] Υ -1 (x) represents the increased gamma gain.

[0080] Then convert the XYZ space to Lab space

[0081] L=116f(Y / Y n )-16,

[0082] a=500[f(X / X n )-f(Y / Y n )],

[0083] b=200[f(Y / Y n )-f(Z / Z n )],

[0084] in:

[0085] X n =95.047, Y n =100.0, Z n =108.883,

[0086]

[0087] The function f(t) is used to linearize smaller color values ​​and process larger color values ​​so that the converted color values ​​are more consistent with the human eye's perception of brightness.

[0088] (6) Histogram stretching is performed on the three color spaces, HSI, Lab, and RGB, to obtain images after histogram stretching of each color space. Histogram stretching is performed on the image in the RGB color space to enhance its contrast and details; histogram stretching is performed on the image in the HSI color space to enhance its hue and saturation; histogram stretching is performed on the image in the Lab color space to enhance its brightness and color information.

[0089] The histogram stretching formula is:

[0090]

[0091] where s out , s in Represents the intensity value of the output and input pixels, I max , I min Represents the maximum and minimum pixel intensity values ​​of the image, o max , o min It represents the maximum and minimum pixel intensity values ​​of the desired output image. min =0,o max =1.

[0092] Compared to the RGB color space, the HSI color space intuitively reflects the way humans perceive color. It is composed of hue, saturation, and brightness. This design makes the HSI color space very convenient for dealing with color-related issues. The Lab color space, on the other hand, is designed based on the human eye's perception of color and more accurately simulates human color perception. Therefore, performing histogram stretching on each of the three spaces can provide more information about an image.

[0093] like Figure 2 As shown, in the traditional histogram stretching method, I max and I min It's easy to get, but o min and o max Need to define manually, in deep learning, such a simple setting can not achieve good results, the present invention is for I max , I min A trainable stretching block is used to obtain the maximum and minimum values ​​of the image. The trainable stretching block obtains the maximum and minimum values ​​of the image through a global pooling layer and obtains the output through a 1×1 convolution block. Point convolution is used to train the histogram stretching, which is specifically expressed as:

[0094]

[0095] S out(x) is the final stretched pixel value obtained in each color space, gmp(S in (-x)) and gmp(S in (x)) represent the results of the global maximum pooling operation of the input feature map. Through this operation, the input feature map S in Stretching (x) to a new range helps enhance feature contrast. Trainable 1×1 convolutional blocks are used to further process the adjusted feature maps, which can be seen as interaction and adjustment between the channels of the feature maps. The 1×1 convolutional layer also contains trainable parameters that can learn a better representation of the input feature map, thereby improving model performance.

[0096] (7) The three images after histogram stretching obtained in step (6) are converted back to RGB color space, and three RGB images containing different information are obtained. The three RGB images containing different information are then merged by splicing them in the channel direction. After merging, a comprehensive input image is obtained. This comprehensive image contains the enhanced information extracted from different color spaces.

[0097] (8) The composite image obtained in step (7) is input into the improved U-Net network obtained in step (2) and trained again. At this time, the improved U-Net network learns the potential features from the three color spaces of RGB, HSI, and Lab in the input image, and dynamically fuses these features, retaining the most important information in each color space, achieving dynamic fusion, and finally outputting the enhanced image. Since the improved U-Net network can adaptively retain the most important information in each image and dynamically fuse these features, an enhanced image can be obtained. During the training process, similar to step (3), the network parameters should be adjusted by optimizing the loss function to gradually improve the image enhancement capability of the network.

[0098] The following provides verification examples of the present invention:

[0099] In this paper, the performance evaluation indicators used are mean square error (MSE), peak signal-to-noise ratio (PSNR), and structural similarity (SSIM) as full reference indicators. The smaller the MSE, the higher the PSNR and SSIM, which means that the enhanced image is closer to the real image in terms of texture, details, and overall structure.

[0100] Comparison algorithms include: Global Contrast Enhancement (GDCP) based on color priors, the Underwater Retinex algorithm (Water-Retinex), Multi-Scale Multi-Level Edge Preserving (MMLE), the Underwater Convolutional Neural Network (UWCNN), the Underwater Image Restoration Network (Water-Net), the Fast Underwater Image Enhancement Neural Network (FUnIE), and our own algorithm, CCMCS-Net. Testing was performed on two test sets: the UIEB dataset was divided into training, validation, and test sets. The test set contains 90 images, denoted as Test-U90. The official test set of the SUIM-E dataset, containing 110 images, was used as the second test set, denoted as Test-S110. The specific experimental results are shown in Tables 1 and 2.

[0101] Table 1 shows the enhanced comparison results for the U90 test:

[0102]

[0103] Table 2 shows the enhanced comparison results for the S110 test:

[0104]

[0105]

[0106] The data in Tables 1 and 2 show that our method outperforms other underwater image enhancement methods, including GDCP, Underwater Retinex, MMLE, UWCNN, Water-Net, and FUnIE, in all metrics. Specifically, traditional algorithms such as GDCP exhibit high mean square error (MSE), indicating large image errors, ineffective noise suppression, and limited processing capabilities for complex scenes. Convolutional neural network methods such as UWCNN and Water-Net, while effective, lack robustness and generalization due to their lack of effective processing of global features and multi-scale information. Compared to traditional methods and other deep learning algorithms, our improved U-Net network significantly enhances its global feature capture and local spatial information representation capabilities by introducing a hybrid convolution-self-attention module. Metrics such as PSNR and SSIM in the table demonstrate the model's advantages in handling complex backgrounds or object segmentation. The self-attention mechanism effectively captures global feature relationships, while the introduction of convolution enhances the ability to extract spatial information, thereby improving the accuracy and richness of feature extraction and resulting in higher clarity and visual quality for the resulting images. ASPP constructs multiple receptive fields through atrous convolution with different sampling rates, which can capture the feature information of underwater images at multiple scales and enhance the model's ability to capture textures of different sizes and complexities. The SFFusion selective feature map fusion method effectively reduces redundant information, improves the efficiency and quality of feature fusion, and further enhances the image's detail expression and visual effects. By comparing the various methods in terms of MSE, PSNR, SSIM and other indicators in the table, the present invention shows significant technical advantages in many aspects. These advantages are derived from the innovative design of the present invention and the organic combination of deep learning and traditional image processing methods, which makes the underwater image enhancement effect significantly better than existing methods and better restores the real underwater scene.

Claims

1. A two-stage network underwater image enhancement method based on color correction and multi-color space stretching, characterized by The following steps are involved: Step 1): Extract the R, G, and B channel components of the original color image, use the green channel to compensate the red and blue channels, and obtain a static red channel correction image. and blue channel correction chart Step 2): Constructing an improved U-Net network, the improved U-Net network includes a parallel hybrid convolution-self-attention feature extraction module, an ASPP void space pyramid pooling layer, and an SFFusion selective feature fusion method, inputting the original color image into the improved U-Net network, training the improved U-Net network using the Adam optimizer, obtaining a trained improved U-Net network, and outputting a red channel weight map and a blue channel weight map; Step 3): performing dot product multiplication of the red channel correction map and the blue channel correction map with the corresponding red channel weight map and the blue channel weight map, respectively, to obtain the dynamic compensation results of the red and blue channels, respectively, and splicing the dynamic compensation results of the red and blue channels and the green channel to form a color RGB space color correction image; Step 4): first convert the RGB color-corrected image into HSI and Lab color spaces, respectively, then perform histogram stretching on the HSI, Lab, and RGB color spaces, respectively, to obtain histogram-stretched images of their respective color spaces, and finally convert the three histogram-stretched images back to RGB color space, and merge the three images converted back to RGB color space in the channel direction to obtain a comprehensive input image; Step 5): Input the comprehensive input image into the improved U-Net network for retraining, and output the enhanced image.

2. The method for underwater image enhancement based on a two-stage network of color correction and multi-color space stretching according to claim 1, characterized in that: In step 1): the red channel correction image The blue channel correction chart I G , I R , I B Represent the green, red, and blue channels respectively, and ⊙ represents the dot product.

3. The method for underwater image enhancement based on a two-stage network of color correction and multi-color space stretching according to claim 1, characterized in that: In step 2), the initial input image of the improved U-Net network is a color image of size H*W*3, with a total of eight layers. The first half of the network is the encoder, and the second half is the decoder. An ASPP module is added between the encoder and decoder. The decoder uses the same hybrid convolution-self-attention feature extraction module as the encoder. The hybrid convolution-self-attention feature extraction module adopts a dual-branch structure. The upper branch structure contains two convolution layers, each layer has a convolution kernel Conv of size 3×3, a step size Padding of 1, two normalization layers BN, and two ReLU activation functions; the lower branch structure contains a pooling layer AvgPool, a Top-K self-attention mechanism, a 1×1 convolution kernel and a Reshape reshaping layer for matching the output of other branches; first, the input feature map is subjected to the maximum pooling MaxPool operation, that is, downsampling; then it is passed to the upper branch and lower branch modules respectively; the upper branch is a standard convolution block, which performs a 3×3 convolution operation and extracts features from the downsampled feature map; the convolved feature map is batch normalized; the lower branch is a Top-K self-attention mechanism block, which performs an average pooling operation on the downsampled feature map, further downsamples, executes the Top-K self-attention mechanism, and finally reshapes the feature map through Reshape.

4. The method for underwater image enhancement based on a two-stage network of color correction and multi-color space stretching according to claim 3, characterized in that: The ASPP module contains four parallel 3×3 convolutional layers, each with a different dilation rate of 1, 6, 12, and 18. The output feature maps of the four convolutional layers are concatenated in the channel dimension to form a comprehensive feature map, which is then processed by a 1×1 convolutional layer to generate the final output feature map.

5. The method for underwater image enhancement based on a two-stage network of color correction and multi-color space stretching according to claim 4, characterized in that: The feature maps of each convolution operation and batch normalization operation are: x i =BN(Conv(I, W ci , b ci , dilation = R i , padding = R i ), I is the input feature map, x i Represents the output feature map after the i-th convolution operation and batch normalization. BN stands for batch normalization, which is used to normalize the output of the convolution layer. Conv represents the convolution operation, W ci and b ci Represent the weight and bias of the i-th convolution layer respectively, dilation represents the hole convolution rate, which is used to control the interval between convolution kernel elements, R i Represents the hole rate of the i-th convolution; padding is used to add pixels to the edge of the input feature map to control the size of the output feature map; The comprehensive feature map is represented as: x = Concat (x1, x2, x3, x4), and then through a 1 × 1 convolution operation: F = Conv (x, W c5 , b c5 ), F is the output feature map of the ASPP module.

6. The method for underwater image enhancement based on a two-stage network of color correction and multi-color space stretching according to claim 3, characterized in that: The input of the Top-K self-attention mechanism first passes through three independent 1×1 convolutional layers to generate three feature maps Q, K, and V respectively. Each feature map undergoes a 3×3 depthwise separable convolution operation to extract features. The feature maps Q, K, and V are reshaped, and the dot product of Q and K is calculated to obtain the self-attention score matrix. The mask operation is applied to the self-attention score matrix, and the Top-K operator is used to select the top K features with the highest self-attention scores. Then, the selected Top-K features are redistributed to the corresponding positions, and the self-attention scores after Top-K are Softmax normalized to obtain the normalized attention weights; the normalized attention weights are multiplied by the feature map V to obtain the weighted feature map; The weighted feature map is reshaped to match the shape of the initial input; finally, the weighted feature map is added to the initial input to achieve feature fusion.

7. The method for underwater image enhancement based on a two-stage network of color correction and multi-color space stretching according to claim 1, characterized in that: The SFFusion selective feature fusion method is to splice the upsampled feature map and the skip-connected feature map in the channel dimension, and the spliced ​​feature map is passed through the global average pooling layer to generate a global feature vector. The global feature vector is processed by the multi-layer perceptron (MLP) to generate a selective weight vector. The MLP consists of one or more fully connected layers and activation functions and outputs a selective weight vector. Finally, the selective weight vector is element-wise multiplied with the upsampled feature map and the skip-connected feature map to obtain a weighted feature map. The two weighted feature maps are added element by element to generate the final fused feature map.

8. The method for underwater image enhancement based on a two-stage network of color correction and multi-color space stretching according to claim 1, characterized in that: The loss function used in the improved U-Net network is the sum of the mean absolute error loss function and the mean square error loss function. The mean absolute error loss function is: H and W are the height and width of the input image, (m,n) are the coordinates of the pixel, is the predicted value of the pixel at (m,n), y(m,n) is the true value of the pixel at position (m,n), is the absolute difference between the predicted value and the true value; Mean squared error loss function: f(x i ) and y i They represent the predicted value and the corresponding true value of the i-th sample, respectively, and n is the number of samples; The total loss function L loss =L MAE +L MSE .

9. The method for underwater image enhancement based on a two-stage network of color correction and multi-color space stretching according to claim 1, characterized in that: In step 4), the formula for converting RGB to HSI space is: The steps to convert RGB to Lab color space are: First convert RGB to XYZ space: γ -1 (x) represents the increased gamma gain, and then the XYZ space is converted to Lab space: L=116f(Y / Y n )-16,a=500[f(X / X n )-f(Y / Y n )],b=200[f(Y / Y n )-f(Z / Z n )], Where: X n =95.047, Y n =100.0, Z n =108.883, The function f(t) is used to linearize small color values ​​and process large color values.

10. The method for underwater image enhancement based on a two-stage network of color correction and multi-color space stretching according to claim 1, characterized in that: In step 4), the histogram stretching formula is: s out , s in Represents the intensity value of the output and input pixels, I max , I min Represents the maximum and minimum pixel intensity values ​​of the image, o max , o min It represents the maximum and minimum pixel intensity values ​​of the desired output image, o min =0,o max =1; Use point convolution to train histogram stretching: S out (x) is the final stretched pixel value obtained in each color space, gmp(S in (-x)) and gmp(S in (x)) represent the results of the global maximum pooling operation of the input feature map.

Images