Underwater image enhancement method and network based on space-frequency domain fusion

Abstract

The invention discloses an underwater image enhancement method based on space-frequency domain fusion, and the method comprises the following steps: 1, constructing a basic GAN network, and enabling a generator of the GAN network to comprise a PCM module and an RCDM module; 2, constructing a space-frequency domain feature fusion module, embedding the space-frequency domain feature fusion module between the PCM module and the RCDM module, taking CCI and CDM output by the PCM module as the input of the space-frequency domain feature fusion module, and taking the output of the space-frequency domain feature fusion module as the input of the RCDM module, the space-frequency domain feature fusion module fuses the space domain feature and the frequency domain feature of the input image; step 3, constructing a loss function of the GAN network containing the frequency domain perception enhancement loss; and step 4, the underwater image is enhanced through the constructed GAN network. According to the GAN network designed by the invention, the enhancement capability and robustness of the method are effectively improved by fusing the image frequency domain information in the image enhancement process.

Description

TECHNICAL FIELD

[0001] The application design relates to the technical field of underwater image processing, and in particular to an underwater image enhancement method and network based on space-frequency domain fusion. BACKGROUND

[0002] Underwater images are crucial in marine scientific research, seabed resource exploration, underwater archaeology and other fields, as they provide a valuable perspective on the underwater world. Due to the absorption and scattering of water, underwater images are often affected by severe color distortion, low contrast, blurring and noise, which seriously affects the analysis of subsequent image tasks. Therefore, effective underwater image enhancement (UIE) algorithms are very important for more accurate identification and classification of seabed objects, more accurate measurement of seabed topography, and more in-depth observation of marine ecosystems. Existing image enhancement solutions can be roughly divided into three categories: physical model-based methods, image prior-based methods and deep learning-based methods. Physical model-based methods require accurate estimation of complex physical parameters and have relatively large computational complexity; image prior-based methods rely on priors that may not always hold due to the complex and variable underwater environment, limiting the effectiveness of the method in certain situations; deep learning-based methods mainly rely on spatial features to restore and enhance degraded images, ignoring the effective representation of global information provided by frequency domain features and the influence of abnormal spectral distribution, limiting the ability of the method to identify and correct degradation problems in complex underwater environments. SUMMARY

[0003] To address the deficiencies in the prior art, the present application provides an underwater image enhancement method and network based on space-frequency domain fusion to solve the technical problem of low ability to identify and correct degradation problems in complex underwater environments in the prior art.

[0004] The present application provides an underwater image enhancement method based on space-frequency domain fusion, comprising the following steps:

[0005] Step 1: Construct a basic GAN network, wherein the generator of the GAN network includes a PCM module and a RCDM module;

[0006] Step 2: Construct a space-frequency domain feature fusion module and embed it between the PCM module and the RCDM module, taking the CCI and CDM output by the PCM module as the input of the space-frequency domain feature fusion module, and taking the output of the space-frequency domain feature fusion module as the input of the RCDM module, wherein the space-frequency domain feature fusion module fuses the spatial domain features and frequency domain features of the input image;

[0007] Step 3: Construct a loss function of the GAN network containing a frequency domain perception enhancement loss;

[0008] Step 4: Enhance the underwater image using the constructed GAN network.

[0009] Furthermore, the PCM module calculates a depth map and attenuation coefficients from the original degraded image, combines the depth map and attenuation coefficients to generate a TM, performs color calibration on the original degraded image using the TM to generate a CCI, and calculates the corresponding CDM using the CCI.

[0010] Furthermore, the spatial frequency domain feature fusion module also embeds a frequency domain attention mechanism, which dynamically adjusts the processing weights of different frequency domain components.

[0011] Furthermore, the input to the spatial frequency domain feature fusion module is:

[0012] F input =CCI+CDM

[0013] In the formula, CCI is the color-enhanced image output by the PCM module; CDM is the depth map output by the PCM module, and CDM = d(CCI).

[0014] Furthermore, the specific method for dynamically adjusting the processing weights of different frequency domain components through the frequency domain attention mechanism is as follows:

[0015] Mix mag =FFT(F input )

[0016] att w =γ*mix mag +β

[0017] In the formula, att w The weight matrix; mix mag γ is the fused amplitude after Fourier block processing; β and γ are hyperparameters that control the attention of frequency domain components.

[0018] Furthermore, the output of the spatial frequency domain feature fusion module is:

[0019] F out =α*iFFT(att) w *cos(mix pha ),att w *sin(mix pha ))+(1-α)F input

[0020] In the formula, mix pha The fused amplitude after Fourier block processing, mix pha =FFT(F input ); α hyperparameter.

[0021] Furthermore, the frequency domain sensing enhancement loss is:

[0022] FAEL=α*MSE loss(FEI,GT) +β*BCELoss(att w ,1)

[0023] In the formula, β and β are hyperparameters; MSE is the mean squared error loss; BCELoss is the binary cross-entropy loss; att w This is the weight matrix.

[0024] Furthermore, the loss function of the GAN network is:

[0025]

[0026] In the formula, L gsl For global similarity loss; L pl For perceptual loss; FAEL is the frequency domain perceptual enhancement loss; λ1~λ5 are scaling factors that adjust the proportion of each loss; G is the generator; D V and D A There are two discriminators.

[0027] This invention also provides an underwater image enhancement network based on spatial-frequency domain fusion, applicable to the aforementioned underwater image enhancement method based on spatial-frequency domain fusion, comprising: a GAN network, wherein the GAN network includes: a generator and a discriminator, the generator including: a PCM module, an ASFFM module, and an RCDM module; the original degraded image is used as the input of the PCM module; the output of the PCM module is used as the input of the ASFFM module; the output of the ASFFM module is used as the input of the RCDM module; the output of the RCDM module is used as the input of the discriminator.

[0028] The beneficial effects of this invention are:

[0029] The GAN network designed in this invention effectively improves the enhancement capability and robustness of the method by incorporating image frequency domain information during the image enhancement process. This invention expands the dimension of feature representation and performs spatial-frequency domain feature fusion through the constructed ASFFM module; simultaneously, a frequency domain attention mechanism is embedded in the ASFFM module to dynamically adjust the attention to different frequency components and promote smooth transitions between adjacent frequencies. This invention adds a frequency-aware enhancement loss to the loss function of the GAN network to ensure that the enhanced image has higher sharpness and contrast while maintaining more natural colors. Attached Figure Description

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

[0031] Figure 1 This is a schematic diagram of the network architecture of a specific embodiment of the present invention;

[0032] Figure 2 This is a structural diagram of the PCM module in a specific embodiment of the present invention;

[0033] Figure 3 This is a structural diagram of the ASFFM module in a specific embodiment of the present invention;

[0034] Figure 4 This is a structural diagram of the RCDM module in a specific embodiment of the present invention;

[0035] Figure 5 These are visual effect diagrams of different methods in specific embodiments of the present invention;

[0036] Figure 6 These are visual effect diagrams of ablation experiments using different methods in specific embodiments of the present invention. Detailed Implementation

[0037] To make the objectives, technical solutions, and advantages of the embodiments of the present invention clearer, the technical solutions of the embodiments of the present invention will be clearly and completely described below with reference to the accompanying drawings. Obviously, the described embodiments are only some embodiments of the present invention, not all embodiments. Based on the embodiments of the present invention, all other embodiments obtained by those skilled in the art without creative effort are within the scope of protection of the present invention.

[0038] The present invention will be further illustrated below with reference to specific embodiments. Those skilled in the art should understand that these embodiments are for illustrative purposes only and are not intended to limit the scope of the invention. Modifications to the present invention in various equivalent forms all fall within the scope defined by the appended claims.

[0039] like Figure 1 As shown, this invention provides an underwater image enhancement method based on spatial-frequency domain fusion generative adversarial networks, comprising the following steps:

[0040] Step 1: Construct the basic generative adversarial network, namely the GAN network. The generator of the GAN network includes a parameter estimation and color correction module, namely the PCM module, and a region correction module based on the degree of degradation, namely the RCDM module.

[0041] like Figure 2As shown, the input to the PCM module is the original degraded image, i.e., I. The PCM module calculates the depth map and attenuation coefficient from the original degraded image, and combines the depth map and attenuation coefficient to generate the transmission map of the original image, i.e., TM. The original degraded image is color-calibrated using TM to generate a color correction map, i.e., CCI. The corresponding depth map, i.e., CDM, is calculated using CCI.

[0042] like Figure 4 As shown, the input of the RCDM module is the output of the ASFFM module, namely FEI, I, and TM. A series of convolutional pooling blocks are used to extract and fuse multi-scale features from FEI and I. The decoder uses the fused features, and the RCDM generates the final enhanced image, E, by analyzing the degree of FEI attenuation.

[0043] Step 2: Construct a spatial-frequency domain feature fusion module with embedded frequency domain attention mechanism, namely the ASFFM module, and embed the ASFFM module between the PCM module and the RCDM module. Use the CCI and CDM outputs of the PCM module as inputs to the ASFFM module, and use the output of the ASFFM module as inputs to the RCDM module. The ASFFM module fuses the spatial and frequency domain features of the input image.

[0044] The ASFFM module not only captures local texture details but also considers the global frequency distribution of the image, specifically enhancing key information and achieving a complementary balance between the two domains through an attention fusion strategy.

[0045] The ASFFM module extracts twice the valuable information from limited datasets, addressing the current scarcity of high-quality underwater datasets. The ASFFM module works by using Fourier Transform (FFT) to perform a space-to-frequency domain transformation on the image. The specific formula is as follows:

[0046]

[0047] In the formula, u and v are frequency domain coordinates; X(u,v) is the Fourier transform result of the image in the frequency domain; x(h,w) is the pixel value of the image in the spatial domain; H and W are the height and width of the image; is the kernel function of the Fourier transform, which maps the image from the spatial domain to the frequency domain; i is the imaginary unit.

[0048] Frequency domain information can be represented by amplitude component A(X(u,v)) and phase component P(X(u,v)):

[0049]

[0050] In the formula, R and I are the real and imaginary parts of X(u,v), respectively.

[0051] like Figure 3As shown, the ASFFM module takes CCI and CDM as inputs:

[0052] F input =CCI+CDM

[0053] In the formula, CCI is the color-enhanced image output by the PCM module; CDM is the depth map output by the PCM module, and CDM = d(CCI).

[0054] The addition of depth maps helps the network better understand the relationship between image content and objects, provides three-dimensional geometric information of underwater scenes, makes up for the susceptibility of single RGB images to distortion, and thus improves the robustness of the network.

[0055] The amplitude and phase of the fused components are obtained through FFT, and these two components are enhanced in the frequency domain to more effectively capture the global information and structural features of the image. To enable the network to dynamically adjust the attention given to different frequency components, a specific method for dynamically adjusting the processing weights of different frequency components using a frequency domain attention mechanism is as follows:

[0056] Mix mag =FFT(F input )

[0057] att w =γ*mix mag +β

[0058] In the formula, att w The weight matrix; mix mag γ is the fused amplitude after Fourier block processing; β and γ are hyperparameters that control the attention of frequency domain components.

[0059] The real and imaginary components of the image in the frequency domain are obtained using trigonometric functions. The image is then restored from the frequency domain to the spatial domain using inverse Fourier transform, resulting in a secondary enhanced image (FEI) fused with spatial and frequency domain information.

[0060] F out =α*iFFT(att) w *cos(mix pha ),att w *sin(mix pha ))+(1-α)F input

[0061] In the formula, mix pha The fused amplitude after Fourier block processing, mix pha =FFT(F inputα is a hyperparameter that controls the spatial-frequency domain transformation effect of the image. It controls the inverse Fourier transform operation to ensure that the image is closer to the real and natural image when it is converted from the frequency domain to the spatial domain, that is, from the amplitude and phase to the image itself.

[0062] Step 3: Construct the loss function of the GAN network that includes frequency domain sensing enhancement loss, i.e., FAEL;

[0063] Frequency domain loss is used to further optimize image details and color distribution. Targeting the overall spectral distribution, frequency domain loss can address spectral anomalies typically manifested in the frequency domain, such as global color deviation and uneven illumination. Frequency domain loss enables the network to learn local and global features, resulting in a more balanced representation and improved image enhancement. FAEL can enhance high-frequency components, optimize low-frequency components, and minimize unwanted frequency components, thereby preventing aberrant frequency shifts in the generator and ensuring image quality and visual appeal.

[0064] The frequency domain sensing enhancement loss is:

[0065] FAEL=α*MSE loss(FEI,GT) +β*BCELoss(att w ,1)

[0066] In the formula, α and β are hyperparameters; MSE is the mean squared error loss; BCELoss is the binary cross-entropy loss; att w This is the weight matrix.

[0067] To address physical quality and visual appeal, the loss function of the GAN network also incorporates global similarity loss and perceptual loss. The loss function of the GAN network is as follows:

[0068]

[0069] In the formula, L gsl For global similarity loss; L pl For perceptual loss; FAEL is the frequency domain perceptual enhancement loss; λ1~λ5 are scaling factors that adjust the proportion of each loss; G is the generator; D V and D A There are two discriminators.

[0070] Step 4: Enhance the underwater image using the constructed GAN network.

[0071] like Figure 5 As shown, the following is a comparative analysis of the present invention with 7 traditional methods and 6 learning-based methods:

[0072]

[0073] Table 1

[0074] Table 1 shows the evaluation of different methods on three datasets based on the average PSNR (DB), MSE, and SSIM values.

[0075] Table 1 shows the performance of ACDE, UDCP, UNTV, GDCP, HLRP, MLLE, RetinexBased, U-shapeTransformer, UWCNN, deep-sesr, FUnIE-GAN, and PUGAN on the UIEB, UFO-120, and EUVP datasets based on PSNR, MSE, and SSIM scores. These three metrics are evaluated based on the ground truth (GT). PSNR evaluates the peak signal-to-noise ratio between the enhanced image (E) and GT; a higher value indicates better quality. MSE quantizes the mean squared error; a lower value indicates better enhancement. SSIM evaluates the structural similarity of the image; a higher score reflects better preservation of structural information. Based on these evaluations, the method of this invention achieves the highest PSNR and SSIM scores and the lowest MSE value across all three datasets, demonstrating that the method improves image quality while preserving structural information. Furthermore, the method of this invention achieves good scores on all three datasets, indicating strong adaptability to various underwater environments. In contrast, other methods typically cannot optimize all three metrics simultaneously and usually only perform well on one dataset. For example, the U-shape Transformer achieved the second-highest PSNR and the second-lowest MSE score, performing well on the UIEB and EUVP datasets but poorly on the UFO-120 dataset. PUGAN, ranked third, performed well in both PSNR and MSE metrics but poorly in SSIM. Overall, the method of this invention demonstrates significant advantages in both image quality and visual aesthetics.

[0076] The following is the ablation experiment of this invention, which includes four experimental procedures:

[0077]

[0078] Table 2

[0079] Table 2 shows the quantitative results of the mean PSNR (DB), MSE, and SSIM of the ablation experiments on the UIEB dataset.

[0080] Experiment 1: CDM is removed, leaving only CCI as input, corresponding to NO.1 in Table 2, w / o CDM for input of ASFFM. Experiment 1 is an experiment to test the effect obtained when the ASFFM module does not input CDM.

[0081] Experiment 2: Remove the ASFFM module from the GAN network, corresponding to NO.2 in Table 2, w / o ASFFM. Experiment 2 is an experiment to test the effect obtained when the ASFFM module is not added to the GAN network.

[0082] Experiment 3: Remove the frequency domain attention mechanism from the ASFFM module, corresponding to NO.3 in Table 2, with / o frequency domain attention mechanism. Experiment 3 is an experiment to test the effect obtained when the frequency domain attention mechanism is not added to the ASFFM module.

[0083] Experiment 4: Remove FAEL from the GAN network loss function, corresponding to NO.4 in Table 2, w / o FAEL. Experiment 4 is an experiment to test the effect obtained by not adding FAEL to the GAN network loss function.

[0084] Table 2 provides validated metric scores for the entire network and each component, including PSNR, MSE, and SSIM.

[0085] In Experiment 1, removing CDM reduced the network's PSNR from 21.63 to 20.12, indicating that the enhancement process introduces additional noise without depth information, thus reducing PSNR. MSE increased from 572.6667 to 774.6097, which degraded the image enhancement quality. SSIM decreased from 0.8692 to 0.8376, indicating a reduction in structural accuracy, which is particularly challenging when processing complex textures to achieve optimal results. For example, red artifacts appeared... Figure 6 In the second image of No.1, in terms of overall visual effect, removing CDM reduces color richness, limits detail sharpening, and results in somewhat flat contrast.

[0086] Therefore, the absence of CDM leads to a significant decrease in image quality, highlighting the importance and necessity of designing multimodal inputs. Providing both CCI and CDM as multimodal inputs simultaneously allows for the capture of more comprehensive scene information, thereby improving the network's multimodal task performance. The inclusion of CDM enables the network to more effectively interpret the relationships between image content and objects, thus compensating for the limitations of a single, easily distorted CCI.

[0087] In Experiment 2, PSNR decreased from 21.63 to 19.75, a drop of 8.6%, while MSE increased by 44% and SSIM decreased by 10%. These changes indicate a significant decline in the ability and accuracy of image detail restoration, especially for complex textures and structures. Figure 6Image No. 2 exhibits varying degrees of detail loss, particularly in the high-frequency regions where blurring is pronounced. Removing the ASFFM module leads to a significant degradation in image quality, highlighting the crucial role of the ASFFM module in quality enhancement and providing valuable guidance for subsequent network enhancement processes.

[0088] Similar to the results of experiments 3 and 4 above, after removing the corresponding objects, the quantization results showed a certain degree of reduction in PSNR, MSE, and SSIM. Figure 6 As can be seen from images NO.3 and NO.4, there are also certain flaws in terms of visual effect, such as artifacts and blurred boundaries in some areas of the image.

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

Claims

1. An underwater image enhancement method based on spatial-frequency domain fusion, characterized in that, Includes the following steps: Step 1: Construct the basic GAN network, where the generator of the GAN network includes a PCM module and an RCDM module; Step 2: Construct a spatial-frequency domain feature fusion module and embed it between the PCM module and the RCDM module. Use the CCI and CDM outputs of the PCM module as inputs to the spatial-frequency domain feature fusion module, and use the output of the spatial-frequency domain feature fusion module as inputs to the RCDM module. The spatial-frequency domain feature fusion module fuses the spatial and frequency domain features of the input image. Step 3: Construct the loss function of the GAN network that includes frequency domain sensing enhancement loss; Step 4: Enhance the underwater image using the constructed GAN network.

2. The underwater image enhancement method based on spatial-frequency domain fusion as described in claim 1, characterized in that, The PCM module calculates a depth map and attenuation coefficients from the original degraded image, combines the depth map and attenuation coefficients to generate a TM, performs color calibration on the original degraded image using the TM to generate a CCI, and calculates the corresponding CDM using the CCI.

3. The underwater image enhancement method based on spatial-frequency domain fusion as described in claim 1 or 2, characterized in that, The input to the spatial frequency domain feature fusion module is: F input =CCI+CDM In the formula, CCI is the color-enhanced image output by the PCM module; CDM is the depth map output by the PCM module, and CDM = d(CCI).

4. The underwater image enhancement method based on spatial-frequency domain fusion as described in claim 1, characterized in that, The spatial frequency domain feature fusion module also embeds a frequency domain attention mechanism, which dynamically adjusts the processing weights of different frequency domain components.

5. The underwater image enhancement method based on spatial-frequency domain fusion as described in claim 4, characterized in that, The specific method for dynamically adjusting the processing weights of different frequency components through the frequency domain attention mechanism is as follows: mix mag =FFT(F input ) att w =γ*mix mag +b In the formula, att w The weight matrix; mix mag γ is the fused amplitude after Fourier block processing; β and γ are hyperparameters that control the attention of frequency domain components.

6. The underwater image enhancement method based on spatial-frequency domain fusion as described in any one of claims 1-5, characterized in that, The output of the spatial frequency domain feature fusion module is: F out =α*iFFT(att w *cos(mix pha ),att w *sin(mix pha ))+(1-α)F input In the formula, mix pha The fused amplitude after Fourier block processing, mix pha =FFT(F input ); α hyperparameter.

7. The underwater image enhancement method based on spatial-frequency domain fusion as described in claim 1, characterized in that, The frequency domain sensing enhancement loss is: FAFL=α*MSE loss(FEI,GT) +β*BCELoss(att w ,1) In the formula, α and β are hyperparameters; MSE is the mean squared error loss; BCELoss is the binary cross-entropy loss; att w This is the weight matrix.

8. The underwater image enhancement method based on spatial-frequency domain fusion as described in claim 7, characterized in that, The loss function of the GAN network is: In the formula, L gsl For global similarity loss; L pl For perceptual loss; FAEL is the frequency domain perceptual enhancement loss; λ1~λ5 are scaling factors that adjust the proportion of each loss; G is the generator; D V and D A There are two discriminators.

Images