Underwater image enhancement method

Abstract

The invention discloses an underwater image enhancement method, and the method comprises the following steps: 1, constructing a generator network, and training the generator network through adversarial training; step 2, acquiring underwater images in RGB and HSV color spaces; 3, obtaining a plurality of scale features corresponding to the underwater images in the two color spaces through an encoder respectively; 4, performing feature fusion on the scale features corresponding to the two color spaces under the same scale through an FFM module to obtain fused scale features under the corresponding scales; and 5, taking the fusion scale feature under the last scale as the input of a decoder, carrying out decoding reconstruction through the decoder, and obtaining an enhanced underwater image. According to the method, the problems of global structure distortion and partial detail loss possibly caused by a convolutional neural network are solved, and meanwhile, the problem that most methods can only relieve color cast to a certain extent but are often insufficient in the aspect of enhancing color saturation is solved.

Description

Technical Field

[0001] The invention relates to the technical field of image enhancement, and in particular to an underwater image enhancement method based on self-attention and dual color space. Background Art

[0002] Autonomous underwater vehicles (AUVs) have been widely used in various underwater operations, such as hydrographic surveys, environmental assessments, and underwater salvage operations. During their missions, AUVs often need to use underwater images to analyze environmental data related to the underwater environment. However, due to the harsh underwater environment, the acquired images usually exhibit severe color cast and blur issues. These issues can severely limit high-level visual tasks such as image segmentation, object classification, and detection. Therefore, it is crucial to effectively process the acquired underwater images, which makes underwater image enhancement one of the hottest research issues in the field of computer vision.

[0003] In recent years, deep learning-based methods have achieved very good results in the field of underwater image enhancement. Deep learning-based methods use the power of deep neural networks to learn complex mappings from degraded underwater images to their reference images. Li et al. proposed a convolutional neural network model for underwater image enhancement, which directly reconstructs clear underwater images by learning the mapping relationship between degraded images and high-definition images. Wang et al. proposed Uice^2Net, which cleverly merged RGB and HSV color spaces in a unified network architecture. To overcome the challenge of scarce datasets, Li et al. designed WaterGAN to synthesize underwater images using atmospheric and depth images from existing datasets. After that, Li et al. proposed the gated fusion network WaterNet and established the UIEB dataset at the same time to facilitate training. Fabbri et al. proposed using CycleGAN to generate underwater images and then train with the corresponding image pairs. Islam et al. established the synthetic underwater dataset EUVP and designed FUnIE-GAN, a lightweight network architecture characterized by simplicity and concise parameters. Current deep learning algorithms widely adopt a strategy of performing end-to-end learning in the RGB color space using a series of convolutional layers. Although this approach alleviates color casts to a certain extent, it is often insufficient in enhancing color saturation. More importantly, each neuron of a convolutional neural network can only perceive a local area of ​​the input data, which may lead to distortion of the global structure and loss of some details. Summary of the invention

[0004] In view of the deficiencies in the prior art, the present invention provides an underwater image enhancement method to solve the technical problems of low color saturation, global structure distortion and detail loss in the prior art.

[0005] The present invention provides an underwater image enhancement method, comprising the following steps:

[0006] Step 1: Build a generator network and use adversarial training to train the generator network. The generator network has a dual encoder-single decoder backbone structure, and the dual encoders have the same structure.

[0007] Step 2: Obtain underwater images in RGB and HSV color spaces;

[0008] Step 3: Pass the underwater images in the two color spaces through an encoder to obtain several corresponding scale features;

[0009] Step 4: The scale features corresponding to the two color spaces at the same scale are fused through the FFM module to obtain the fused scale features at the corresponding scale;

[0010] Step 5: Use the fused scale features at the last scale as the input of the decoder, perform decoding and reconstruction through the decoder, and obtain the enhanced underwater image. During the decoding and reconstruction process of the decoder, the output of the previous decoding block is cascaded with the fused scale features at the corresponding scale as the input of the next decoding block.

[0011] Furthermore, in the step 1, the discriminator structure used in the training process of the generator network is: a dual discriminant network, and the original size image can be used as the input of the discriminator, wherein the receptive fields of the two discriminant networks are different, and the results of the two discriminant networks on reducing the scale of image features are the same.

[0012] Furthermore, one discriminator network structure in the discriminator is: four 4×4 convolution operations connected in series and one 1×1 convolution operation; another discriminator network structure is: four 2×2 convolution operations connected in series and one 1×1 convolution operation.

[0013] Furthermore, in step 1, during the training process of the generator network, a weighted sum function of a WGAN-GP loss function, an L1 loss function, a content loss function, and an SSIM loss function with a gradient penalty term is used as the loss function.

[0014] Furthermore, in step 3, the specific process of obtaining multiple scale features is as follows:

[0015] First, the underwater image is subjected to a 3×3 convolution operation to obtain the first scale feature of the underwater image, which is used as the input of the first decoding block;

[0016] Each subsequent encoding block takes the previous scale feature as input, downsamples the previous scale feature through a 4×4 convolution operation, and passes the downsampled image features through the MDTA module for feature extraction and 3×3 convolution operation, and then concatenates the output of the MDTA module and the result of the 3×3 convolution operation as the next scale feature.

[0017] Furthermore, the specific process of step 4 is as follows:

[0018] Step 41: performing an element-by-element addition operation on the scale features of the two color spaces at the same scale;

[0019] Step 42: converting the scale features after the addition operation from the spatial domain to the frequency domain using fast Fourier transform;

[0020] Step 43: Extract favorable frequency domain features using complex convolution operations;

[0021] Step 44: performing inverse Fourier transform on the extracted frequency domain features;

[0022] Step 45: multiply the scale features of the two color spaces input in step 41 by the corresponding inverse transformation results respectively;

[0023] Step 46: Add the multiplication results element by element to obtain the fused scale feature at the current scale.

[0024] Beneficial effects of the present invention:

[0025] The present invention combines convolution operations with MDTA, using the ability of self-attention to understand long-range dependencies and the complementarity of the convolution mechanism in local feature perception to ensure the global consistency and integrity of the generated image, and solves the problem that convolutional neural networks may cause distortion of the global structure and loss of some details.

[0026] The present invention adopts the FFM module to fuse the information flows from the frequency domain and the spatial domain of the RGB and HSV color spaces to form a fused scale feature, and realizes the innovative integration of multiple color space image features to enhance the robustness of the algorithm and improve the color saturation of the enhanced image, solving the problem that most methods can only alleviate the color cast to a certain extent, but are often insufficient in enhancing the color saturation.

[0027] The present invention designs a dual-discriminant network with different receptive fields to guide the generation network to synthesize more realistic underwater images, solving the problem that the details of the enhanced images are not delicate enough. BRIEF DESCRIPTION OF THE DRAWINGS

[0028] The features and advantages of the present invention will be more clearly understood by referring to the accompanying drawings, which are schematic and should not be construed as limiting the present invention in any way. In the accompanying drawings:

[0029] Figure 1 This is a schematic diagram of a generator network structure in a specific embodiment of the present invention;

[0030] Figure 2 It is a schematic diagram of the structure of the coding block in a specific embodiment of the present invention;

[0031] Figure 3 It is a structural schematic diagram of a decoding block in a specific embodiment of the present invention;

[0032] Figure 4 This is a schematic diagram of the FFM module structure in a specific embodiment of the present invention;

[0033] Figure 5 It is a schematic diagram of the discriminator structure in a specific embodiment of the present invention;

[0034] Figure 6 This is a comparison chart of the enhancement results of each model on the EUVP dataset;

[0035] Figure 7 The following is a comparison chart of the enhanced results of each model on the UIEB dataset. DETAILED DESCRIPTION

[0036] In order to make the purpose, technical solution and advantages of the embodiments of the present invention clearer, the technical solution in the embodiments of the present invention will be clearly and completely described below in conjunction with the drawings in the embodiments of the present invention. Obviously, the described embodiments are part of the embodiments of the present invention, not all of the embodiments. Based on the embodiments of the present invention, all other embodiments obtained by those skilled in the art without creative work are within the scope of protection of the present invention.

[0037] The present invention is further illustrated below in conjunction with specific embodiments. Those skilled in the art should understand that these embodiments are only used to illustrate the present invention and are not used to limit the scope of the present invention, and modifications to various equivalent forms of the present invention fall within the scope defined by the appended claims of this application.

[0038] The present invention provides an underwater image enhancement method, comprising the following steps:

[0039] Step 1: Build Figure 1-3The generator network shown in the figure has a dual encoder-single decoder backbone structure, and the dual encoder structure is the same. Each encoder includes a 3×3 convolution operation and 4 encoding blocks Enc_Block, and the five parts are connected in series. Each encoding block Enc_Block has a 4×4 convolution operation, a 3×3 convolution operation, and an MDTA module. MDTA is a self-attention module, which applies the self-attention mechanism to the channel dimension and can effectively avoid the problem of exponential growth of computational complexity with the increase of spatial resolution. The results of the 4×4 convolution operation are used as the input of the 3×3 convolution operation and the MDTA module respectively. The output of the 3×3 convolution operation is cascaded with the output of the MDTA module as the output of the encoding block. The output of the encoding blocks at the same position in the two encoders is used as the input of the FFM module for scale feature fusion. The scale feature output by the last encoding block is used as the input of the decoder after scale feature fusion.

[0040] The decoder consists of 4 decoding blocks Dec_Block, 1 3×3 convolution operation, and the five parts are connected in series. Each decoding block Dec_Block has 1 4×4 deconvolution operation, 1 3×3 convolution operation and 1 MDTA module. The results of the 4×4 deconvolution operation are used as the input of the 3×3 convolution operation and the MDTA module respectively. The output of the 3×3 convolution operation is cascaded with the output of the MDTA module as the output of the encoding block. The output of the previous decoding block is cascaded with the fused scale feature output by the corresponding FFM module as the input of the next decoding block.

[0041] In the adversarial training process of the generator network, such as Figure 5 As shown, the discriminator structure used is: a dual discriminant network, the original size image can be used as the input of the discriminator, wherein the receptive fields of the two discriminant networks are different, and the results of the two discriminant networks for reducing the scale of image features are the same. In a specific embodiment of the present invention, the main discriminant network: uses 4 4×4 convolution operations and 1 1×1 convolution operation in succession; the auxiliary discriminant network: uses 4 2×2 convolution operations and 1 1×1 convolution operation in succession. The receptive field sizes of the two discriminant networks are different, but both must satisfy the consistency of the results of reducing the scale of image features. The auxiliary discriminator is designed with a smaller receptive field to specifically identify subtle differences between the generated image and the reference image.

[0042] Loss function design: In order to improve the quality of enhanced images, the present invention selects the WGAN-GP loss function with gradient penalty, L1 loss function, content loss function and SSIM loss function for weighted summation to obtain the combined loss function. Let G be the generator, D p is the main discriminator, D a is the auxiliary discriminator, I Dis the original underwater image, I R is the reference image.

[0043] WGAN-GP loss function: Use Wasserstein distance to measure the difference between the generated image and the reference image, and use gradient penalty to strengthen the Lipschitz constraint. The loss function can be expressed as:

[0044]

[0045] L1 loss function: Using L1 loss can enable the model to better learn the global learning of real images and generate samples with higher global similarity to real images. The L1 loss function can be expressed as:

[0046] L1(G)=E X,Y [||YG(X)||1]

[0047] Content loss function: Adding a content loss term to the objective function encourages the generator to generate enhanced images with similar content to the real image. The content loss function is expressed as follows:

[0048] L con (G) = E X,Y [||φ(Y)-φ(G(X))||1]

[0049] Where φ is the content function, and the feature map of the second convolutional layer of the fifth convolutional block of VGG19 is selected as the output of the content function.

[0050] SSIM loss function: measures the similarity between the generated image and the corresponding real image in terms of brightness, contrast and structure. The loss function is expressed as follows:

[0051] L SSIM (G) = E X,Y [1-SSIM(Y,G(X))]

[0052] Where SSIM(Y,G(X)) represents the structural similarity score between the generated image and the real image.

[0053] Therefore, the total loss function is described as:

[0054]

[0055] Among them, λ l ,λ c and λ s Represents the loss weight, set λ l =60,λ c =20,λ s =20.

[0056] Step 2: Obtain underwater images in RGB and HSV color spaces;

[0057] Step 3: Pass the underwater images in the two color spaces through an encoder to obtain several corresponding scale features;

[0058] Take the underwater image in two color spaces scaled to 256×256 as an example, where f i rgb Represents the i-th scale feature corresponding to the RGB color space; Represents the i-th scale feature corresponding to the HSV color space;

[0059] Step 31: Use a 3×3 convolution operation to extract the first scale feature of the underwater image or At this time, the feature scale is 256×256 and is used as the input of the first encoding block;

[0060] Step 32: In the first encoding block, a 4×4 convolution operation is performed on or Down-sample, and input the down-sampled image features into the MDTA module and perform a 3×3 convolution operation respectively; then the output cascade of the MDTA module and the 3×3 convolution operation is used as the second scale feature or At this time, the feature scale is 128×128 and is used as the input of the second encoding block;

[0061] Step 33: In the second encoding block, a 4×4 convolution operation is performed on or Down-sample, and input the down-sampled image features into the MDTA module and perform a 3×3 convolution operation respectively; then the output cascade of the MDTA module and the 3×3 convolution operation is used as the third scale feature or At this time, the feature scale is 64×64 and is used as the input of the third encoding block;

[0062] Step 24, in the third encoding block, use a 4×4 convolution operation on or Down-sample, and input the down-sampled image features into the MDTA module and perform a 3×3 convolution operation respectively; then the output cascade of the MDTA module and the 3×3 convolution operation is used as the fourth scale feature of the encoder or At this time, the feature scale is 32×32 and is used as the input of the fourth encoding block;

[0063] Step 25, in the fourth encoding block, use a 4×4 convolution operation on or Down-sample, and input the down-sampled image features into the MDTA module and perform a 3×3 convolution operation respectively; then the output cascade of the MDTA module and the 3×3 convolution operation is used as the fifth scale feature of the encoder or At this time, the feature scale is 16×16.

[0064] At this time, five scale features corresponding to the RGB color space are formed in one encoder, and five scale features corresponding to the HSV color space are formed in the other encoder. The two scale features at the same scale in the two encoders are paired as the input of the FFM module described in step 4.

[0065] Step 4: The scale features corresponding to the two color spaces at the same scale are fused through the FFM module to obtain the fused scale features at the corresponding scale;

[0066] like Figure 4 As shown in Figure 1, it is a schematic diagram of the structure of the FFM module. The process of the FFM module fusing two scale features at the same scale is as follows:

[0067] Step 41: The two scale features f of the two color spaces at the same scale i rgb 、f i hsv , perform element-by-element addition operation;

[0068] Step 42: converting the scale features after the addition operation from the spatial domain to the frequency domain using fast Fourier transform;

[0069] Step 43: Extract favorable frequency domain features using complex convolution operations;

[0070] Step 44: performing inverse Fourier transform on the extracted frequency domain features;

[0071] Step 45: multiply the two scale features input in step 41 by the corresponding inverse transformation results respectively;

[0072] Step 46: Add the multiplication results element by element to obtain the fused scale feature f at the current scale i , a total of five fused scale features are generated: f1 generated after the first downsampling, f2 generated after the second downsampling, f3 generated after the third downsampling, f4 generated after the fourth downsampling, and f5 generated after the fifth downsampling.

[0073] Step 5: Take the fused scale features at the last scale as the input of the decoder, perform decoding and reconstruction through the decoder, and obtain the enhanced underwater image. During the decoding and reconstruction process of the decoder, the output of the previous decoding block is cascaded with the fused scale features at the corresponding scale as the input of the next decoding block. The specific process is as follows:

[0074] Step 51: Use a 4×4 deconvolution operation on and The fused scale feature f5 is upsampled, and the upsampled image features are input into the MDTA module and the 3×3 convolution operation respectively; then the output of the MDTA module is cascaded with the result of the 3×3 convolution operation, and then and The fusion scale feature f4 is cascaded, and the feature scale is 32×32 at this time;

[0075] Step 52: Use a 4×4 deconvolution operation to upsample the cascade feature input of the previous step, and input the upsampled image features into the MDTA module and perform a 3×3 convolution operation respectively; then cascade the output of the MDTA module with the result of the 3×3 convolution operation, and then and The fusion scale feature f3 is cascaded, and the feature scale is 64×64 at this time;

[0076] Step 53: Use a 4×4 deconvolution operation to upsample the cascade feature input of the previous step, and input the upsampled image features into the MDTA module and perform a 3×3 convolution operation respectively; then cascade the output of the MDTA module with the result of the 3×3 convolution operation, and then and The fusion scale feature f2 is cascaded, and the feature scale is 128×128 at this time;

[0077] Step 54: Use a 4×4 deconvolution operation to upsample the cascade feature input of the previous step, and input the upsampled image features into the MDTA module and perform a 3×3 convolution operation respectively; then cascade the output of the MDTA module with the result of the 3×3 convolution operation, and then and The fusion scale feature f1 is cascaded, and the feature scale is 256×256 at this time;

[0078] Step 55: Use a 3×3 convolution operation to extract the features of the previous cascade operation to complete image reconstruction.

[0079] like Figure 6 , 7 The following Tables 1 and 2 show the comparison results between the present invention and other methods on the EUVP dataset and the UIEB dataset.

[0080]

[0081] Table 1

[0082]

[0083] Table 2

[0084] like Figure 6 As shown in Table 1, the UDCP algorithm has obvious deficiencies in color correction. It not only fails to effectively correct the color cast of the image, but also causes the green tones in the image to be oversaturated. In contrast, the ULAP algorithm shows a relatively ideal image restoration effect in areas with good lighting conditions, but its performance is not satisfactory in environments with insufficient lighting. Retinex-based and UCM algorithms can effectively enhance images, but they may introduce additional noise during the enhancement process, affecting image quality. UGAN-P, FUnIE-GAN, and UWCNN algorithms show good results in color restoration of synthetic images in the EUVP dataset, but when processing real-world images in the UIEB dataset.

[0085] like Figure 7 As shown in Table 2, UGAN-P introduces checkerboard artifacts, color noise appears in the enhanced image by FUnIE-GAN, and the UWCNN model reduces the color saturation of the image. In addition, Shallow-UWnet shows a tendency to blur image details on both EUVP and UIEB datasets. Uice^2Net tends to amplify the blue tones in the image background, resulting in an obvious blue color cast in the enhanced image, especially when processing shallow water blurred images in the UIEB dataset. This problem is particularly significant. In contrast, this patent not only effectively eliminates color deviation, effectively improves contrast and corrects color deviation, but also performs well in reconstructing image details, reflecting excellent generalization performance.

[0086] Although the embodiments of the present invention have been described in conjunction with the accompanying drawings, those skilled in the art may make various modifications and variations without departing from the spirit and scope of the present invention, and such modifications and variations are all within the scope defined by the appended claims.

Claims

1. A method for underwater image enhancement, characterized in that: The steps include: Step 1: Build a generator network and use adversarial training to train the generator network. The generator network has a dual encoder-single decoder backbone structure, and the dual encoders have the same structure. Step 2: Obtain underwater images in RGB and HSV color spaces; Step 3: Pass the underwater images in the two color spaces through an encoder to obtain several corresponding scale features; Step 4: The scale features corresponding to the two color spaces at the same scale are fused through the FFM module to obtain the fused scale features at the corresponding scale; Step 5: Use the fused scale features at the last scale as the input of the decoder, perform decoding and reconstruction through the decoder, and obtain the enhanced underwater image. During the decoding and reconstruction process of the decoder, the output of the previous decoding block is cascaded with the fused scale features at the corresponding scale as the input of the next decoding block.

2. The underwater image enhancement method according to claim 1, characterized in that: In the step 1, the discriminator structure used in the training process of the generator network is: a dual discriminant network, and the original size image can be used as the input of the discriminator, wherein the receptive fields of the two discriminant networks are different, and the results of the two discriminant networks on reducing the scale of image features are the same.

3. The underwater image enhancement method according to claim 2, characterized in that: One discriminator network structure in the discriminator is: four 4×4 convolution operations connected in series and one 1×1 convolution operation; the other discriminator network structure is: four 2×2 convolution operations connected in series and one 1×1 convolution operation.

4. The underwater image enhancement method according to any one of claims 1 to 3, characterized in that: In the step 1, during the training process of the generator network, a weighted sum function of the WGAN-GP loss function, the L1 loss function, the content loss function and the SSIM loss function with a gradient penalty term is used as the loss function.

5. The underwater image enhancement method according to claim 1, characterized in that: In step 3, the specific process of obtaining multiple scale features is as follows: First, the underwater image is subjected to a 3×3 convolution operation to obtain the first scale feature of the underwater image; Each subsequent encoding block takes the previous scale feature as input, downsamples the previous scale feature through a 4×4 convolution operation, and passes the downsampled image features through the MDTA module for feature extraction and 3×3 convolution operation, and then concatenates the output of the MDTA module and the result of the 3×3 convolution operation as the next scale feature.

6. The underwater image enhancement method according to claim 1, characterized in that: The specific process of step 4 is as follows: Step 41: performing an element-by-element addition operation on the scale features of the two color spaces at the same scale; Step 42: Convert the scale features after the addition operation from the spatial domain to the frequency domain using fast Fourier transform; Step 43: Extract favorable frequency domain features using complex convolution operations; Step 44: performing inverse Fourier transform on the extracted frequency domain features; Step 45: multiply the scale features of the two color spaces input in step 41 by the corresponding inverse transformation results respectively; Step 46: Add the multiplication results element by element to obtain the fused scale feature at the current scale.

Images