Two-stage network underwater image enhancement method based on color correction and multicolor space stretching

Abstract

The invention discloses a two-stage network underwater image enhancement method based on color correction and multicolor space stretching, and the method comprises the steps: compensating a red channel and a blue channel through a green channel, obtaining a red channel correction image and a blue channel correction image, constructing an improved U-Net network, and outputting the weight mapping of the red channel and the blue channel; performing dot product multiplication on the red and blue channel correction images and the corresponding red and blue channel weight mapping, and splicing the dynamic compensation results of the red and blue channels and the green channel into an RGB space color correction image; the RGB space color correction image is converted into an HSI space and a Lab color space respectively, histogram stretching is carried out on the HSI space, the Lab color space and the RGB space respectively, finally, the three kinds of images subjected to histogram stretching are converted into the RGB color space respectively, and the three images converted into the RGB color space are spliced to obtain a comprehensive input image; and inputting the comprehensive input image into the improved U-Net network, training again, and outputting an enhanced image, so that the imaging resolution and the color correction effect of the underwater camera can be remarkably improved.

Description

Technical Field

[0001] The 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 for 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 been suffering from quality degradation. Specifically, due to the rapid attenuation of red light when it propagates in water, underwater images often show blue and green as the main tones, and lose the color balance of the real world. At the same time, the presence of micro-suspended particles in the underwater environment, such as sediment 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 tasks. Underwater image enhancement methods can be roughly divided into two categories: physical model-based and deep learning-based methods. Physical model-based methods can show excellent performance in certain specific scenarios, but they are more sensitive to the diversity and complexity of underwater image types. This means that in a changeable and complex underwater environment, the predictive ability of these physical models may be limited. Existing deep learning image enhancement methods can be divided into methods based on convolutional neural networks (CNNs) and generative adversarial networks (GANs). Compared with GANs, CNNs are widely used in image denoising, image defogging, image enhancement, and image super-resolution. Although the convolution operation of CNN networks is popular for its ability to efficiently process local features, it also has a significant limitation: it only focuses on information processing in local neighborhoods, which may lead to the loss of global information. In order to make up for this deficiency, the attention mechanism has been widely used and explored in the field of visual processing in recent years. The attention mechanism dynamically adjusts the network's attention to different parts of the input data, allowing the model to automatically ignore unimportant information and allocate more computing resources to key information areas.

[0004] At present, most underwater image enhancement algorithms use deep learning-based technologies, and these algorithms mainly operate in the RGB color space. However, the RGB color space itself has its limitations, and 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 the document with patent publication number CN112419210A first performs color correction on the original RGB dense fog image based on the color correction method of sub-interval linear transformation of the source image, and uses the three-interval histogram equalization method to perform color correction on the source image in the R, G, and B channels respectively, and stretch the pixel values ​​of the R, G, and B channel images to the range of 0 to 255, respectively, to obtain the stretched single-channel image; then calculate the average pixel value of the pixel values ​​of the R, G, and B channel images respectively; by subtracting the pixel value of each pixel from the single-channel image, the color of the original RGB dense fog image is corrected. The average pixel value of the channel is used as the error and the error is squared; the pixel with the largest square error is selected according to the value of the square error of each pixel point. The pixel values ​​near this point will also have a large error with the average pixel value, so the two thresholds required for three-interval division are determined by adding or subtracting three times the variance around this point, and the entire single-channel histogram is divided into three intervals; each sub-interval of the R, G, and B channels is equalized to obtain a single-channel equalized image; finally, the R, G, and B channel images and the equalized R, G, and B channel images are linearly weighted fused to obtain the final defogging image. This underwater image enhancement method still uses operations in the RGB color space and can only achieve defogging of dense fog images.

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

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

[0007] Technical solution:

[0008] The present invention is based on a two-stage network underwater image enhancement method of color correction and multi-color space stretching, comprising the following steps:

[0009] Step 1): Extract the R, G, and B channel components of the original color image respectively, use the green channel to compensate the red and blue channels, and obtain a static red channel correction image. And blue channel correction chart

[0010] Step 2): construct an improved U-Net network, wherein the improved U-Net network includes a hybrid convolution-self-attention feature extraction module for parallel operation, an ASPP void space pyramid pooling layer, and an SFFusion selective feature map fusion method, input the color original 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 blue channel weight map

[0011] Step 3): The red channel correction image Blue channel correction chart Respectively mapped with the corresponding red channel weight and blue channel weight map Perform dot product multiplication to obtain the dynamic compensation results of the red and blue channels respectively, and splice the dynamic 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 correction image into HSI space and Lab color space respectively, and then performing histogram stretching on the three color spaces of HSI space, Lab color space and RGB space respectively to obtain images after histogram stretching of the respective color spaces, and finally converting the three images after histogram stretching back to RGB color space respectively, and merging the three images converted back to RGB color space by splicing 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 jump-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 jump-connected feature map, respectively, to obtain a weighted feature map; the two weighted feature maps are element-wise added to generate the final fused feature map.

[0015] Furthermore, the initial input image size of the improved U-Net network is a color image of H*W*3, with a total of eight layers, and the first half and the second half of the network are respectively an encoder and a decoder; 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, and then further down-samples it, performs the Top-K self-attention mechanism, and finally reshapes the feature map through Reshape.

[0016] Furthermore, the ASPP module includes four parallel 3×3 convolutional layers, each convolutional layer has a different hole rate, namely 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 graph 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 Respectively represent the weight and bias of the i-th convolutional layer; R i represents the hole rate of the i-th convolution; the comprehensive feature map is expressed as: x = Concat (x 1 ,x 2 ,x 3 ,x 4 ), after 1×1 convolution operation: 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, and each feature map Q, K, V undergoes a 3×3 depth-separable convolution operation to further extract features; the feature maps Q, K, 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 effect:

[0020] 1. The present invention is an underwater image enhancement algorithm that integrates traditional and deep learning, which 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. Since the present invention adopts an improved U-Net network, especially the hybrid convolution-self-attention module designed by us, the self-attention mechanism can capture global feature relationships, and the mixed convolution enhances the ability to capture spatial information. By combining these two methods, the module can effectively improve the model's learning ability for global and local features while maintaining computational efficiency. This improves the accuracy and richness of feature extraction, enhances the segmentation effect of complex backgrounds or targets, and improves the robustness and generalization ability of the model. An ASPP module is added between the encoder and the decoder. The ASPP module can expand the receptive field without increasing the amount of calculation, thereby improving the model's recognition ability for multi-scale features and improving segmentation performance, especially when the image contains targets of different sizes. The SFFusion selective feature map fusion method is adopted. This method can fuse multiple layers of features, suppress useless information, and enhance key features through a feature selection mechanism. Compared with traditional direct splicing 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. Since the present invention introduces DCM dynamic correction into SCM static correction, the initial correction of SCM provides a more balanced input basis for DCM, and DCM further dynamically adjusts on this basis, making the color correction process more stable and efficient. The combination of static and dynamic correction enables the model to have both initial rapid adjustment capabilities and detailed deep optimization capabilities. The color adjustment effect of the image is significantly improved, and the color cast is reduced while retaining the image details, which enables the subsequent image enhancement stage to be processed on a higher quality basis, thereby improving the visual effect of the overall image.

[0023] 4. The present invention uses histogram stretching in multiple color spaces. Different color spaces (such as RGB, HSI, and Lab) can describe the color and brightness characteristics of images from multiple angles. Histogram stretching in these color spaces can more effectively enhance the contrast and color performance of the image and correct problems such as uneven lighting, low contrast, and color distortion.

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

[0025] Figure 1 It 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 A 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 It is a schematic structural diagram of an ASPP hollow space 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 of 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 is further described below in conjunction with the accompanying drawings.

[0032] like Figure 1As shown, the present invention comprises 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 Represent the green, red, and blue channels respectively, and ⊙ represents the dot product. Traditional color correction methods mainly rely on region-based information extraction, but the choice of region size has an important impact on the final extraction result, which may lead to inconsistency and poor adaptability of the results in different scenes.

[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 stage are used in the decoding stage to predict the dynamic weight map and generate weights for each pixel area. 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 The color correction image is obtained by performing a dot product operation with the result of static correction. For dynamic fusion, the three-space comprehensive image obtained by histogram stretching is input into the improved U-Net network. The network learns the potential features from the three color spaces of RGB, HSI, and Lab in the input image, and dynamically fuses these features to retain the most important information in each color space, and finally outputs the enhanced image.

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

[0040] The improved U-Net network has eight layers in total. The first half of the network is used for backbone feature extraction and is the encoder. The input image is passed to the initial block for preliminary feature extraction. The first four layers of encoders all use hybrid convolution-self-attention modules for feature extraction.

[0041] The hybrid convolution-self-attention feature extraction module used adopts a dual-branch structure. The upper branch structure contains two convolution layers, each with a 3×3 convolution kernel, a step size of 1, two normalization layers, and two ReLU activation functions. The lower branch structure contains a pooling layer, a Top-K self-attention mechanism, a 1×1 convolution kernel, and a Reshape layer to match the output of other branches. First, the input feature map is subjected to a maximum pooling operation and downsampled. 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 convolution feature map is batch normalized to accelerate network training and improve stability. The ReLU activation function is applied to introduce nonlinear features. Repeat the above operations so that the network can better extract useful information from the input data. The lower branch is a Top-K self-attention mechanism block, which performs an average pooling operation on the downsampled feature map, further downsamples, and executes the Top-K self-attention mechanism. It captures global information by calculating the correlation of the input feature map, and uses the Top-K operator to select important features for processing. Perform a 1×1 convolution operation to reduce the channel dimension, aggregate features, and finally reshape the feature map through Reshape. The standard convolution block in the upper branch focuses on local feature extraction, while the Top-K self-attention mechanism in the lower branch can capture global information. Therefore, adding the two can simultaneously utilize local and global features to enhance the diversity and completeness of feature representation. The ASPP module is added between the encoder and the decoder, which extracts multi-scale features, enhances contextual information, improves feature expression ability, and reduces information loss. The decoder uses the same hybrid convolution-self-attention feature extraction module as the encoder. When connecting feature maps between low-level and high-level layers, the selective feature fusion method SFFusion is used. The selective feature fusion method effectively fuses the features of the main path and the skip connection by generating dynamic selective weights. Each time the upsampling is performed, the output features of the same layer in the corresponding first half of the network are fused to continuously restore the feature resolution. This is repeated four times, and finally the training results are output through a 1×1 convolution kernel.

[0042] like Figure 4 As shown in the figure, the ASPP module used is shown. The ASPP module contains four parallel 3×3 convolutional layers, each with different hole rates, namely 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 ratio, which is used to control the spacing 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.

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

[0046] x=Concat(x 1 ,x 2 ,x 3 ,x 4 ),

[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 depth-separable convolution operation to further extract features. The feature maps Q, K and V are reshaped for subsequent matrix operations. The dot product of Q and K is calculated to obtain the self-attention score matrix. The MASK mask operation is applied to the self-attention score matrix to remove irrelevant or invalid parts, which is usually used to ignore the influence of certain specific positions. The Top-K operator is used to select the top K features with the highest self-attention scores to reduce the amount of calculation and focus on the most important features. The selected Top-K features are then redistributed to the corresponding positions. The self-attention scores after Top-K are Softmax normalized to obtain normalized attention weights. The normalized attention weights are multiplied with 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.

[0051]

[0052] where Y i Represents the vector at 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 ith row of . That is, the K-ranked attention score in the 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 initial input to achieve feature fusion.

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

[0056]

[0057] in is the upsampled feature map, X 2 is the feature map of the skip connection, GAP represents the global average pooling layer, and 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. 1 and α 2 Respectively with the upsampled feature maps and the skip-connected feature map X 2 Perform element-by-element multiplication to obtain a weighted feature map. Add the two weighted feature maps to generate the final fused feature map X. Use weights {a 1 ,a 2} to X 2 Perform point multiplication respectively to obtain the fused feature map X:

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

[0060] 800 original color images were extracted from UIEBD for preprocessing, and a training data set was generated for a random set of 800 images. The images in the training data set were resized to 256×256 and put 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 square 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 trained improved U-Net network and blue channel weight map

[0070] (4) The red channel correction image in step (1) Blue channel correction chart Respectively mapped with the corresponding red channel weights obtained in step (3) and blue channel weight map The dot product multiplication is performed to obtain the final motion compensation results of the red and blue channels, 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 Fusion and splicing along the channel dimension to reconstruct the color RGB space color correction image I cc :

[0073]

[0074] Indicates splicing, I cc That 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, namely, HSI space, Lab color space, and RGB space, 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 Respectively represent 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 with the RGB color space, the HSI color space can directly reflect the way human vision perceives 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 is designed based on the human eye's perception of color. It simulates human perception of color more accurately in the color space. Therefore, stretching the histograms of the three spaces separately can obtain 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 max Need to be manually defined, in deep learning, such a simple setting can not achieve good results, the present invention 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 the 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)) respectively represent the results of the global maximum pooling operation of the input feature map. Through this operation, the input feature map S in (x) is stretched to a new range, which helps to enhance the contrast of the features. The use of trainable 1×1 convolution blocks is used to further process the adjusted feature maps, which can be regarded as the interaction and adjustment between the channels of the feature maps. At the same time, the 1×1 convolution layer contains trainable parameters, which can learn a better representation of the input feature map, thereby improving the 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 in the channel direction, and a comprehensive input image is obtained after merging. This comprehensive image contains the enhanced information extracted from different color spaces.

[0097] (8) Input the composite image obtained in step (7) into the improved U-Net network obtained in step (2) and train 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, retains the most important information in each color space, realizes dynamic fusion, and finally outputs 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 network's image enhancement capability.

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

[0099] In the present invention, 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] The comparison algorithms include: global contrast enhancement (GDCP) based on color prior, underwater Retinex algorithm (Water-Retinex), multi-scale multi-level edge preservation (MMLE), underwater convolutional neural network (UWCNN), underwater image restoration network (Water-Net), fast neural network for underwater image enhancement (FUnIE), and our own algorithm CCMCS-Net. Tests were conducted on two types of test sets. The UIEB dataset was divided into training set, validation set and test set, where the test set contains 90 images. The test set is denoted as Test-U90. The official test set of the SUIM-E dataset contains 110 images. It is denoted as the second test set, and this test set is called Test-S110. The specific experimental results are shown in Tables 1 and 2.

[0101] Table 1 shows the test enhancement comparison results for U90 as follows:

[0102]

[0103] Table 2 shows the test enhancement comparison results for S110 as follows:

[0104]

[0105]

[0106] From the data in Table 1 and Table 2, it can be seen that the method of the present invention is ahead of other underwater image enhancement methods in various indicators, including GDCP, underwater Retinex, MMLE, UWCNN, Water-Net, and FUnIE. Specifically, traditional algorithms such as GDCP have a higher MSE index, indicating that their image errors are large, noise cannot be effectively suppressed, and the processing ability of complex scenes is limited. Although convolutional neural network methods such as UWCNN and Water-Net have certain enhancement effects, they lack effective processing of global features and multi-scale information, resulting in insufficient robustness and generalization of the model. Compared with traditional methods and other deep learning algorithms, our improved U-Net network significantly improves the model's global feature capture ability and local spatial information expression ability due to the introduction of a hybrid convolution-self-attention module. Indicators such as PSNR and SSIM in the table reflect the advantages of the model in processing complex backgrounds or target segmentation. The self-attention mechanism effectively captures global feature relationships, and the introduction of convolution enhances the ability to extract spatial information, thereby improving the accuracy and richness of feature extraction, so that the final generated image has higher clarity and visual quality. 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 complex textures. 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 MSE, PSNR, SSIM and other indicators of each method 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 in that The following steps are involved: Step 1): Extract the R, G, and B channel components of the original color image respectively, 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, wherein the improved U-Net network includes a hybrid convolution-self-attention feature extraction module for parallel operation, an ASPP void space pyramid pooling layer, and an SFFusion selective feature fusion method, inputting the color original image into the improved U-Net network, training the improved U-Net network using the Adam optimizer, obtaining the trained improved U-Net network, and outputting a red channel weight map and a blue channel weight map; Step 3): dot-product multiplying 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 correction image into HSI space and Lab color space respectively, then perform histogram stretching on the three color spaces of HSI space, Lab color space and RGB space respectively, and obtain images after histogram stretching of the respective color spaces respectively, and finally convert the three images after histogram stretching back to RGB color space respectively, and merge the three images converted back to RGB color space by splicing 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 two-stage network of color correction and multicolor 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 They represent green, red, and blue channels respectively, and ⊙ represents the dot product.

3. The method for underwater image enhancement based on two-stage network of color correction and multicolor space stretching according to claim 1, characterized in that: In step 2), the initial input image size of the improved U-Net network is a color image of 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 Padding of 1, two normalization layers BN, and two ReLU activation functions; the lower branch structure includes a pooling layer AvgPool, a Top-K self-attention mechanism, a convolution kernel of size 1×1, and a Reshape reshaping layer for matching the outputs 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 convolutional 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, performs a Top-K self-attention mechanism, and finally reshapes the feature map through Reshape.

4. The method for underwater image enhancement based on two-stage network of color correction and multicolor space stretching according to claim 3 is characterized in that: The ASPP module includes four parallel 3×3 convolutional layers, each of which has 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 two-stage network of color correction and multi-color space stretching according to claim 4 is characterized in that: 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, 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 convolutional layer respectively, dilation represents the dilation convolution rate, which is used to control the interval between convolution kernel elements, and 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 = Comcat (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 two-stage network of color correction and multicolor space stretching according to claim 3 is 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 is subjected to a 3×3 depth-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 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 normalized by Softmax 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 two-stage network of color correction and multicolor space stretching according to claim 1, characterized in that: The SFFusion selective feature map 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 processed by the global average pooling layer 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-by-element 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 color correction and multicolor space stretching in two stages of network 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) is the coordinate 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 square 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 two-stage network of color correction and multicolor 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 convert the XYZ space 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 larger color values.

10. The method for underwater image enhancement based on 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 Respectively represent 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