Two-stage potential diffusion model underwater image enhancement method based on frequency domain

Abstract

The invention discloses a two-stage potential diffusion model underwater image enhancement method based on a frequency domain. The method comprises the following steps: inputting a degraded image and ambient light estimated according to the degraded image into a lightweight underwater image enhancement model to obtain a preliminary recovery image; inputting the preliminary recovery image into a high and low frequency feature pyramid encoder to obtain a potential spatial feature map set and an intermediate feature map set, and inputting the potential spatial feature map set and the intermediate feature map set into a high and low frequency prior extraction module to obtain high and low frequency images and high and low frequency image prior; inputting the low-frequency guidance priori and the potential spatial feature map set of the low-frequency image into a low-frequency guidance condition module to generate a priori guidance condition, and sending the priori guidance condition into a potential diffusion model to obtain a reconstructed potential spatial feature map set; and inputting the reconstructed potential spatial feature map set, the intermediate feature map set and the high and low frequency images into a high and low frequency feature pyramid decoder in prior to obtain an enhanced image. And improving the performance of the potential diffusion model and recovering a high-quality underwater image by using high and low frequency information priori knowledge of the underwater image.

Description

Technical Field

[0001] The invention belongs to the technical field of image processing, and in particular relates to an underwater image enhancement method based on a two-stage potential diffusion model in frequency domain. Background Art

[0002] As an important source of information for human exploration of the ocean, the quality of underwater images directly affects the efficiency and accuracy of underwater visual tasks. Due to the characteristics of water media, underwater images often suffer from distortion, blur, low contrast and color shift. In order to improve these problems, researchers have developed a variety of underwater image enhancement (UIE) technologies, which have broad application prospects in underwater target detection, marine life recognition, seabed mapping and other fields.

[0003] Traditional underwater image enhancement methods usually perform overall enhancement in RGB space, but such methods do not fully consider the different degrees of degradation of high-frequency and low-frequency parts of underwater images. In image processing, high-frequency parts represent fast changes in the image, such as edges and details, while low-frequency parts represent slow changes in the image, such as smooth areas. The impact of the underwater environment on the image often causes more serious degradation of the low-frequency parts, such as reduced overall contrast and color distortion; while the high-frequency parts, although also affected, may not degrade as significantly in terms of details. Therefore, traditional underwater image enhancement methods limit the effect of underwater image enhancement to a certain extent.

[0004] To solve this problem, researchers have begun to explore new underwater image enhancement techniques. An emerging approach is to use deep learning, especially convolutional neural networks (CNNs) and generative adversarial networks (GANs), to learn high-level features of underwater images and achieve more effective underwater image enhancement. Latent Diffusion Models (LDM) is a generative model based on deep learning. By training the diffusion model on a low-dimensional latent space, it can reach a near-optimal point between complexity reduction and detail preservation, significantly improving visual fidelity. The use of latent diffusion models can effectively improve the quality of underwater images and is conducive to further promoting the development of underwater high-frequency vision tasks. Therefore, it is necessary to study an underwater image enhancement method based on latent diffusion models that takes into account the different degrees of degradation of the high and low frequency parts of underwater images. Summary of the invention

[0005] The technical problem to be solved by the present invention is to provide a two-stage latent diffusion model underwater image enhancement method based on the frequency domain, which can effectively improve the performance of the latent diffusion model and restore high-quality underwater images by utilizing the prior knowledge of the high and low frequency information of the underwater image, thereby realizing effective enhancement processing of the underwater image.

[0006] The technical solution adopted by the present invention to solve the above technical problems is: a method for underwater image enhancement based on a two-stage potential diffusion model in the frequency domain, characterized by comprising the following steps:

[0007] Step 1: Input the underwater degraded image x into the Gaussian blur estimation module of the first stage to estimate the degraded ambient light A of the underwater degraded image x;

[0008] Step 2: Input the underwater degraded image x and its degraded ambient light A into the lightweight underwater image enhancement model of the first stage to obtain a preliminary restored underwater image

[0009] Step 3: Initially restore the underwater image Input into the high and low frequency feature pyramid encoder of the second stage to obtain the latent space feature map set z0 and the intermediate feature map set F all ;

[0010] Step 4: Initially restore the underwater image Input into the high and low frequency prior extraction module in the second stage to obtain the preliminary restored underwater image High frequency image Low frequency image High frequency image prior hfp, low frequency image prior lfp;

[0011] Step 5: Low frequency images Perform prior compression to obtain low-frequency guided prior lfp d ; Then the low frequency guided prior lfp d The latent space feature map set z0 is input into the low-frequency guidance condition module of the second stage to generate the low-frequency prior guidance condition z' of the latent diffusion model, and the low-frequency prior guidance condition z' is sent to the latent diffusion model, and the reconstructed latent space feature map set is obtained by iterative sampling.

[0012] Step 6: Collect the reconstructed latent space feature maps Intermediate feature map set F all , high-frequency image prior hfp and low-frequency image prior lfp are input into the high- and low-frequency feature pyramid decoder of the second stage to obtain the underwater enhanced image

[0013] In the step 1, the underwater degraded image x is a three-channel RGB image with a size of H×W, and the degraded ambient light A is a three-channel RGB image with a size of H×W, where H is the height and W is the width.

[0014] The specific process of step 2 is as follows:

[0015] Step 2.1: Construct the first-stage lightweight underwater image enhancement model, which includes the first convolution layer, the first encoding layer, the second encoding layer, the bottom attention enhancement layer, the first decoding layer, the second decoding layer, and the second convolution layer;

[0016] Step 2.2: The 6 feature maps of size H×W obtained by splicing the underwater degraded image x and its degraded ambient light A in the channel dimension are input into the first convolutional layer. The first convolutional layer outputs 24 feature maps of size H×W, and these feature maps constitute the set F L1 ;

[0017] Step 2.3: F L1 Input to the first coding layer, the first coding layer outputs 24 feature maps of size H×W, and these feature maps constitute a set F L2 ; Among them, the first coding layer is composed of 8 first depth residual blocks connected in series, each of which is composed of a jump-connected depth convolution block and a jump-connected first MLP connected in series, the depth convolution block is composed of a first two-dimensional convolution layer and a second two-dimensional convolution layer connected in series, and the first MLP is composed of a third two-dimensional convolution layer, a ReLU activation layer and a fourth two-dimensional convolution layer connected in series;

[0018] Step 2.4: F L2 Input to the second encoding layer, the second encoding layer outputs 48 images of size The feature maps of the set F L3 ; Among them, the second encoding layer is composed of the first downsampling layer and the first deep residual attention layer connected in series, the input dimension and output dimension of the first deep residual attention layer are both 48, the first deep residual attention layer is composed of the first 4 second deep residual blocks and the last 4 first Swin Transformer residual blocks connected in series, the structure of the second deep residual block is the same as that of the first deep residual block, each first Swin Transformer residual block is composed of a jump-connected Swin Transformer block and a jump-connected second MLP connected in series, and the structure of the second MLP is the same as that of the first MLP;

[0019] Step 2.5: F L3 Input to the bottom attention enhancement layer, the bottom attention enhancement layer outputs 48 images of size The feature maps of the set F M ; Among them, the bottom attention enhancement layer is composed of the second downsampling layer, 8 second Swin Transformer residual blocks, and the first upsampling layer in series, and the structure of the second Swin Transformer residual block is the same as that of the first Swin Transformer residual block;

[0020] Step 2.6: F L3 After the first skip convolution and F M All feature maps obtained after element-by-element addition are input into the first decoding layer, which outputs 24 feature maps of size H×W, and these feature maps form a set The first decoding layer is composed of a first depth residual layer and a second upsampling layer connected in series, the first depth residual layer is composed of four third depth residual blocks connected in series, and the structure of the third depth residual block is the same as that of the first depth residual block;

[0021] Step 2.7: F L2 After the second skip convolution All feature maps obtained after element-by-element addition are input into the second decoding layer, which outputs 24 feature maps of size H×W, and these feature maps form a set The second decoding layer is composed of four fourth depth residual blocks connected in series, and the structure of the fourth depth residual block is the same as that of the first depth residual block;

[0022] Step 2.8: Input into the second convolutional layer, the second convolutional layer outputs 6 feature maps of size H×W, the first 3 feature maps are used as the three channels of the attenuation map K, and the last 3 feature maps are used as the three channels of the scattering map B;

[0023] Step 2.9: By formula The preliminary restored underwater image is calculated Preliminary restored underwater images is a three-channel RGB image with size H×W.

[0024] The specific process of step 3 is as follows:

[0025] Step 3.1: Construct the high- and low-frequency feature pyramid encoder of the second stage, which includes the coding convolution, the first residual coding layer, the second residual coding layer, the third residual coding layer, the fourth residual coding layer, the attention coding layer, the first attention layer, and the first integrated convolution layer. The structures of the first to fourth residual coding layers are the same;

[0026] Step 3.2: Initially restore the underwater image Input into the coded convolution, the coded convolution outputs 128 feature maps of size H×W, and these feature maps constitute the set F E1 ;

[0027] Step 3.3: F E1 Input to the first residual coding layer, the first residual coding layer outputs 128 images of size The feature maps of the set F E2; Among them, the first residual coding layer is composed of two residual blocks, an attention block and a downsampling convolution in series;

[0028] Step 3.4: F E2 Input to the second residual coding layer, the second residual coding layer outputs 128 images of size The feature maps of the set F E3 ;

[0029] Step 3.5: F E3 Input to the third residual coding layer, the third residual coding layer outputs 256 images of size The feature maps of the set F E4 ;

[0030] Step 3.6: F E4 Input to the fourth residual coding layer, the fourth residual coding layer outputs 256 images of size The feature maps of the set F E5 ;

[0031] Step 3.7: F E5 Input to the attention encoding layer, the attention encoding layer outputs 512 images of size The feature maps of the set F E6 ; Among them, the attention encoding layer is composed of two residual blocks and one attention block in series;

[0032] Step 3.8: F E6 Input to the first attention layer, the first attention layer outputs 512 images of size The feature maps of the set F E7 ; Among them, the first attention layer is composed of residual block, attention block and residual block in series;

[0033] Step 3.9: F E7 Input to the first integrated convolutional layer, the first integrated convolutional layer outputs 3 images of size The feature maps of the network are the set of latent space feature maps z0; wherein the first integrated convolution layer is composed of group normalization, Swish activation layer, and two-dimensional convolution in series;

[0034] Step 3.10: F E2 、F E3 、F E4 、F E5 All feature maps in are arranged in sequence to form a set F all .

[0035] The specific process of step 4 is as follows:

[0036] Step 4.1: Construct the high and low frequency prior extraction module of the second stage, which includes a gradient map calculation block, a Fourier high and low frequency separation block, a high frequency prior compression block and a low frequency prior compression block;

[0037] Step 4.2: Initially restore the underwater image Input to the gradient map calculation block, the gradient map calculation block outputs a gradient map of size H×W Among them, the gradient map calculation block uses the Laplacian gradient operator to calculate the gradient map;

[0038] Step 4.3: Initially restore the underwater image Input to the Fourier high and low frequency separation block, the Fourier high and low frequency separation block outputs a three-channel high frequency image of size H×W and a three-channel low-frequency image of size H×W The implementation process of the Fourier high-low frequency separation block is as follows: Use a two-dimensional Fourier transform to get a frequency map; then use fftshift on the frequency map to shift the zero frequency of the frequency map to the center position of a matrix of size H×W ; then, according to the size H×W of the frequency map, a low-frequency mask matrix with a center window of the same size of 6 width and height and a content of 1, and a content of 0 in other positions except the center window is generated, and a high-frequency mask matrix with a center window of the same size of 6 width and height and a content of 0, and a content of 1 in other positions except the center window is generated; then, the frequency map is multiplied by the low-frequency mask matrix to obtain the low-frequency part, and the frequency map is multiplied by the high-frequency mask matrix to obtain the high-frequency part; then, ifftshift and two-dimensional inverse Fourier transform are used on the low-frequency part to obtain a low-frequency image And use ifftshift and two-dimensional inverse Fourier transform on the high-frequency part to get the high-frequency image

[0039] Step 4.4: Gradient map and high frequency images All feature maps obtained after splicing in the channel dimension are input into the high-frequency prior compression block, which outputs 64 high-frequency image priors hfp with a size of 1×1; the low-frequency images Input into the low-frequency prior compression block, the low-frequency prior compression block outputs 64 low-frequency image priors lfp with a size of 1×1; wherein, the high-frequency prior compression block and the low-frequency prior compression block have the same structure, which is composed of a first compressed residual block, a second compressed residual block, a compression block and an information extraction block connected in series, and the first compressed residual block and the second compressed residual block have the same structure;

[0040] The frequency prior is input into the first compressed residual block, and the first compressed residual block outputs 16 images of size feature maps, and these feature maps constitute a set H1; where the input dimension of the first compressed residual block is D, for the first compressed residual block in the high-frequency prior compressed block, the frequency prior is a gradient map and high frequency images All feature maps obtained after splicing in the channel dimension, and D = 4, for the first compressed residual block in the low-frequency prior compression block, the frequency prior is the low-frequency image And D=3, the first compressed residual block is composed of a compression block and 4 prior residual blocks in series, the compression block is composed of the first two-dimensional convolution, the LeakyReLU activation layer, and the second two-dimensional convolution in series, the input of the prior residual block is sequentially passed through a two-dimensional convolution, a LeakyReLU activation layer and another two-dimensional convolution, and then added to the input of the prior residual block to obtain the output of the prior residual block;

[0041] H1 is input into the second compressed residual block, and the second compressed residual block outputs 32 images of size , and these feature maps constitute the set H2;

[0042] Input H2 into the compression block, and the compression block outputs 64 images with the size of , and these feature maps constitute a set H3;

[0043] H3 is input into the information extraction block, and the information extraction block outputs 64 image priors of size 1×1; among them, for the information extraction block in the high-frequency prior compression block, the image prior is the high-frequency image prior hfp, and for the information extraction block in the low-frequency prior compression block, the image prior is the low-frequency image prior lfp. The information extraction block is composed of a maximum pooling layer, a first linear layer, a LeakyReLU activation layer, a second linear layer and a LeakyReLU activation layer in series.

[0044] The specific process of step 5 is as follows:

[0045] Step 5.1: Use the prior compression block to reduce the low frequency image Perform prior compression to obtain low-frequency guided prior lfp d ; Among them, the structure of the prior compression block is the same as the structure of the low-frequency prior compression block in the high-low frequency prior extraction module of the second stage, and the low-frequency guided prior lfp d There are 64 feature maps of size 1×1;

[0046] Step 5.2: Construct the low-frequency guidance condition module of the second stage, and its implementation process is as follows:

[0047] Step 5.2.1: The latent space feature map set z0 is sequentially subjected to two-dimensional convolution and group normalization to obtain the feature map set F in ; Then the low frequency guide prior lfp d Through a linear layer, we get the feature map set fp + ; and through another linear layer, we get the feature map set fp × ; Then F in With fp × After multiplication, fp + Add together to get a set of feature maps

[0048] Step 5.2.2: Order Then and Input into the cross attention, the cross attention output feature map set F out ;

[0049] Step 5.2.3: F in Skip connection, with F out The results obtained after MLP are added together to obtain the feature map set F' out ; then F' out After a residual block and a two-dimensional convolution, the low-frequency prior guidance condition z' is obtained;

[0050] Step 5.3: The latent diffusion model has a noise estimation network U-Net, The acquisition process is as follows:

[0051] Step 5.3.1: Let i represent the number of iterations, and the initial value of i is N, where N represents the maximum number of iterations;

[0052] Step 5.3.2: Calculate the sampling time t, t = (i-1) × T s / N+1, where T s Indicates the preset number of sampling moments;

[0053] Step 5.3.3: Sample from the standard Gaussian distribution to obtain the feature map at sampling time t in, It follows a Gaussian distribution with a mean of 0 and a unit variance;

[0054] Step 5.3.4: The result obtained by splicing with z' in the channel dimension and the sampling time t are input into the noise estimation network U-Net. The output of the noise estimation network U-Net is recorded as

[0055] Step 5.3.5: Calculate the next sampling time t next, t next =max((i-2)× T s / N+1,0); then calculate the sampling time t next The characteristic diagram of in, represents the noise parameter at sampling time t, α j =1-β j , β0=0.0015 and increase equidistantly to β j , represents the sampling time t next Noise parameters when ;

[0056] Step 5.3.6: Let i = i-1, where the "=" in i = i-1 is an assignment symbol; then determine whether i is equal to 0. If so, the iteration ends and As a collection of reconstructed latent space feature maps Otherwise, t next Assign to t, Assign to Return to step 5.3.4 and continue.

[0057] The specific process of step 6 is as follows:

[0058] Step 6.1: Construct the high and low frequency feature pyramid decoder of the second stage, which includes decoding convolution, second attention layer, high and low frequency attention decoding layer, first residual decoding layer, second residual decoding layer, third residual decoding layer, fourth residual decoding layer, second integrated convolution layer, and the structures of the first to fourth residual decoding layers are the same;

[0059] Step 6.2: Collect the reconstructed latent space feature maps Input into the decoding convolution, the decoding convolution outputs 512 images of size The feature maps of the set F D1 ;

[0060] Step 6.3: F D1 Input to the second attention layer, the second attention layer outputs 512 images of size The feature maps of the set F D2 ;

[0061] Step 6.4: F D2 , high-frequency image prior hfp and low-frequency image prior lfp are input to the high- and low-frequency attention decoding layer, and the high- and low-frequency attention decoding layer outputs 512 images of size The feature maps of the set F D3; Among them, the high-low frequency attention decoding layer is composed of residual blocks, high-low frequency residual attention blocks and residual blocks in series.

[0062] Before the high- and low-frequency residual attention blocks are implemented, the high-frequency image prior hfp is passed through a linear layer to obtain the feature map set fp hfp+ , and pass the high-frequency image prior hfp through another linear layer to obtain the feature map set fp hfp× , where fp hfp+ and fp hfp× The dimension is Similarly, the low-frequency image prior lfp is passed through a linear layer to obtain the feature map set fp lfp+ And pass the low-frequency image prior lfp through another linear layer to obtain the feature map set fp lfp× , where fp lfp+ and fp lfp× The dimension is Then the set of all feature maps output by the second residual block is After group normalization, first compare with fp hfp× Multiply and add fp hfp+ Add together to get a set of high-frequency fusion feature maps Similarly, the set of all feature maps output by the second residual block is After group normalization, first compare with fp lfp× Multiply and add fp lfp+ Add together to get a set of low-frequency fusion feature maps in, The dimension is

[0063] The high- and low-frequency residual attention blocks consist of two sets of self-attention blocks and MLP blocks connected in parallel and internally in sequence, where the MLP block is a two-dimensional convolution; the implementation process of the high- and low-frequency residual attention blocks is: the input of the high- and low-frequency residual attention blocks is right Perform Fourier transform to get After passing through the first group of self-attention blocks and MLP blocks, we get right and The feature map obtained by adding element by element is Fourier transformed to obtain After passing through the second group of self-attention blocks and MLP blocks, we get right and The feature map obtained by adding element by element is Fourier transformed to obtain Then calculate the high frequency residual R high , Similarly, calculate the low-frequency residual R low , Then R high The result obtained after complex convolution is similar to R low The results obtained after complex convolution are added to obtain the high and low frequency residuals R; finally, the high and low frequency residuals R and Add, and pass the addition result through inverse Fourier transform to get the output of high and low frequency residual attention block, where the output dimension of high and low frequency residual attention block is

[0064] Step 6.5: Set the high-frequency image prior hfp, low-frequency image prior lfp, and F D3 With F all F E5 All feature maps obtained after splicing in the channel dimension are input to the first residual decoding layer, which outputs 256 images of size The feature maps of the set F D4 ; Among them, the first residual decoding layer is composed of two residual blocks, a high-low frequency residual attention block and an upsampling convolution in series;

[0065] Step 6.6: Set the high-frequency image prior hfp, low-frequency image prior lfp, and F D4 With F all F E4 All feature maps obtained after splicing in the channel dimension are input to the second residual decoding layer, and the second residual decoding layer outputs 256 images of size The feature maps of the set F D5 ;

[0066] Step 6.7: Set the high-frequency image prior hfp, low-frequency image prior lfp, and F D5 With F all F E3 All feature maps obtained after splicing in the channel dimension are input to the third residual decoding layer, and the third residual decoding layer outputs 128 images of size The feature maps of the set F D6 ;

[0067] Step 6.8: Set the high-frequency image prior hfp, low-frequency image prior lfp, and F D6 With F all F E2 All feature maps obtained after splicing in the channel dimension are input into the fourth residual decoding layer, which outputs 128 feature maps of size H×W, and these feature maps constitute a set F D7 ;

[0068] Step 6.9: F D7Input to the second integrated convolutional layer, the second integrated convolutional layer outputs a three-channel underwater enhanced image of size H×W

[0069] When training the lightweight underwater image enhancement model, the learning rate is set to 2×10 -5 , the training batch size is set to 16, the number of training rounds is set to 50, and the optimizer uses the AdamW optimizer; the underwater degraded image x and its degraded ambient light A are input into the lightweight underwater image enhancement model, and the lightweight underwater image enhancement model outputs a preliminary restored underwater image calculate With clear image x gt The loss between pen , Among them, ||||1 is the 1-norm operator;

[0070] When the high- and low-frequency feature pyramid encoder and the high- and low-frequency feature pyramid decoder are trained, the learning rate is set to 2×10 -5 , the training batch size is set to 16, the number of training rounds is set to 300, and the optimizer uses the AdamW optimizer; clear image x gt Input into the high- and low-frequency feature pyramid encoder, the high- and low-frequency feature pyramid encoder obtains the latent space feature map set z gt and the set of intermediate feature maps F gt ; Latent space feature map set z gt and the set of intermediate feature maps F gt Input into the high- and low-frequency feature pyramid decoder, and the high- and low-frequency feature pyramid decoder outputs the restored image Calculate x gt and The loss between ed , Among them, ||||2 is the 2-norm operation symbol;

[0071] When training the noise estimation network U-Net in the potential diffusion model, the learning rate is set to 2×10 -5 , the training batch size is set to 16, the number of training rounds is set to 300, and the optimizer uses the AdamW optimizer; the random noise χ is sampled, and the random noise χ obeys a Gaussian distribution with a mean of 0 and a variance of the unit matrix; the clear image x gt Input into the high and low frequency feature pyramid encoder with fixed parameters, the high and low frequency feature pyramid encoder with fixed parameters obtains the latent space feature map set z gt ; Calculate the feature map z at sampling time t gt,t , Where t=1,2,…,T s; The underwater degraded image x and its degraded ambient light A are input into a lightweight underwater image enhancement model with fixed parameters, and the lightweight underwater image enhancement model with fixed parameters outputs a preliminary restored underwater image Preliminary restored underwater images Input into the high and low frequency feature pyramid encoder with fixed parameters, the high and low frequency feature pyramid encoder with fixed parameters obtains the latent space feature map set z0; the initial restored underwater image Input into the high and low frequency prior extraction module to obtain the preliminary restored underwater image High frequency image Low frequency image High-frequency image prior hfp, low-frequency image prior lfp; for low-frequency image Perform prior compression to obtain low-frequency guided prior lfp d ; The low frequency guide prior lfp d and the latent space feature map set z0 are input into the low-frequency guidance condition module to generate the low-frequency prior guidance condition z' of the latent diffusion model; z gt,t The result obtained by splicing z' in the channel dimension and the sampling time t are input into the noise estimation network U-Net, and the noise estimation network U-Net outputs χ(z gt,t ,z',t); calculate χ and χ(z gt,t ,z',t) ldm ,

[0072] L ldm =||χ-χ(z gt,t ,z',t)||2.

[0073] Compared with the prior art, the advantages of the present invention are:

[0074] In the field of underwater imaging, due to the scattering and absorption of light, the acquired underwater images are often severely degraded, resulting in a decrease in the quality of the underwater images. The method of the present invention can effectively improve the performance of the potential diffusion model and restore high-quality underwater images by utilizing the prior knowledge of the high and low frequency information of the underwater degraded images (i.e., high frequency image prior and low frequency image prior), thereby enhancing the acquired underwater degraded images. Specifically, the method of the present invention can enhance the underwater degraded images more accurately by combining the high and low frequency information of the underwater degraded images. The high frequency information usually contains the details and edge information of the image, while the low frequency information contains the overall brightness and color information of the image. By optimizing the information of these two frequency bands, the enhanced underwater images can be closer to the real scene in terms of details and colors, that is, high quality underwater images can be restored.

[0075] In practical applications, the high- and low-frequency parts of underwater images are often subject to varying degrees of degradation. For example, the high-frequency part may become blurred due to light scattering, while the low-frequency part may become dim due to absorption. By using high- and low-frequency residual attention blocks, these degraded parts can be compensated for in a targeted manner, thereby achieving a more refined enhancement effect.

[0076] Aligning the latent diffusion model with the low-frequency information and the low-frequency prior guidance condition of the generated latent diffusion model can effectively reduce the improper learning of the color cast of underwater images by the latent diffusion model. This is because the low-frequency information contains the color and brightness information of underwater images. By aligning with this information, it can ensure that the latent diffusion model can more accurately capture the true color of underwater images during the learning and diffusion process.

[0077] By integrating frequency priors, the method of the present invention can not only enhance the visual effect of underwater images and obtain higher full-reference indicators such as peak signal-to-noise ratio (PSNR) and structural similarity index (SSIM), but also make the enhanced underwater images more in line with the perception habits of the human eye.

[0078] High-quality underwater images contain more available key information, which is crucial for subsequent underwater image analysis tasks. For example, in terms of target recognition, target tracking, and underwater scene detection, high-quality underwater images can provide more feature information, thereby improving the accuracy and efficiency of these tasks.

[0079] In general, this underwater image enhancement method based on the fusion of high and low frequency information is not only innovative in theory, but also shows great potential in practical applications. With the continuous advancement and optimization of technology, this method is expected to play a greater role in the field of underwater imaging in the future and provide more reliable technical support for underwater exploration and monitoring. BRIEF DESCRIPTION OF THE DRAWINGS

[0080] Figure 1 The overall framework diagram of the network for realizing the method of the present invention, part (a) shows the framework of the Gaussian blur estimation module and the lightweight underwater image enhancement model in the first stage, and part (b) shows the schematic flow of the high- and low-frequency feature pyramid encoder, high- and low-frequency prior extraction module, low-frequency guidance condition module, potential diffusion model, and high- and low-frequency feature pyramid decoder in the second stage;

[0081] Figure 2 The framework diagram of the high and low frequency feature pyramid encoder (i.e. latent space encoder) and decoder (i.e. latent space decoder) in the second stage;

[0082] Figure 3The framework diagram of the high and low frequency prior extraction module of the second stage, part (a) shows the framework of the gradient map calculation block and the Fourier high and low frequency separation block, part (b) shows the framework of the high frequency prior compression block and the low frequency prior compression block, and part (c) shows the framework of the first compressed residual block and the second compressed residual block;

[0083] Figure 4 This is the framework diagram of the low-frequency guidance condition module in the second stage. Part (a) shows the overall framework, and part (b) shows the implementation process of cross-attention.

[0084] Figure 5 The framework diagram of the high- and low-frequency residual attention blocks in each residual decoding layer in the high- and low-frequency feature pyramid decoder of the second stage. Part (a) shows the overall framework, and part (b) shows the implementation process of the self-attention block therein;

[0085] Figure 6 Visualization results of different underwater image enhancement methods on the test set L504;

[0086] Figure 7 Visualization results of different underwater image enhancement methods on the test set U90. DETAILED DESCRIPTION

[0087] In order to make the above-mentioned objects, features and advantages of the present invention more obvious and easy to understand, the specific implementation methods of the present invention are described in detail below in conjunction with the drawings. It is obvious that the described embodiments are only part of the embodiments of the present invention, but not all of the embodiments.

[0088] The present invention proposes a method for underwater image enhancement based on a two-stage potential diffusion model in the frequency domain. The overall network framework diagram for implementing the method of the present invention is as follows: Figure 1 As shown, it includes the following steps:

[0089] Step 1: Input the underwater degraded image x into the Gaussian blur estimation module of the first stage to estimate the degraded ambient light A of the underwater degraded image x.

[0090] Here, the underwater degraded image x is a three-channel RGB image with a size of H×W, and the degraded ambient light A is a three-channel RGB image with a size of H×W, where H is the height and W is the width.

[0091] Here, the Gaussian blur estimation module adopts the existing technology, which includes three blur blocks with parameters σ=15, 60, and 90 respectively and arranged in parallel. The input of the three blur blocks is the underwater degraded image x. The output results of the three blur blocks are added and divided by 3 as the degraded ambient light A of the underwater degraded image x. The blur blocks are calculated using a two-dimensional Gaussian blur formula.

[0092] Step 2: Input the underwater degraded image x and its degraded ambient light A into the lightweight underwater image enhancement model of the first stage to obtain a preliminary restored underwater image

[0093] In this embodiment, the specific process of step 2 is:

[0094] Step 2.1: Construct the first stage lightweight underwater image enhancement model, such as Figure 1 As shown in part (a) of the figure, it includes a first convolutional layer, a first encoding layer, a second encoding layer, a bottom attention enhancement layer, a first decoding layer, a second decoding layer, and a second convolutional layer.

[0095] Step 2.2: Concatenate the underwater degraded image x and its degraded ambient light A in the channel dimension ( Figure 1 The 6 feature maps of size H×W obtained after the C operation in (C) are input to the first convolutional layer, and the first convolutional layer outputs 24 feature maps of size H×W, and these feature maps constitute the set F L1 ; Among them, splicing in the channel dimension is a common operation in neural networks, and the first convolution layer is a 3×3 two-dimensional convolution with a stride of 1, an input dimension of 6, and an output dimension of 24.

[0096] Step 2.3: F L1 Input to the first coding layer, the first coding layer outputs 24 feature maps of size H×W, and these feature maps constitute a set F L2 ; Among them, the first coding layer is composed of 8 first deep residual blocks in series, each of which is composed of a jump-connected deep convolution block and a jump-connected first MLP in series, and the deep convolution block is composed of a first two-dimensional convolution layer and a second two-dimensional convolution layer in series. The first two-dimensional convolution layer is a 5×5 two-dimensional convolution with an input dimension and an output dimension of 24, a padding length of 2, a padding method of reflection padding, and a number of groups equal to the input dimension. The second two-dimensional convolution layer is a 1×1 two-dimensional convolution with an input dimension and an output dimension of 24 and a stride of 1. The first MLP is composed of a third two-dimensional convolution layer, a ReLU activation layer, and a fourth two-dimensional convolution layer in series. The third two-dimensional convolution layer is a 1×1 two-dimensional convolution with a stride of 1, an input dimension of 24, and an output dimension of 48. The fourth two-dimensional convolution layer is a 1×1 two-dimensional convolution with a stride of 1, an input dimension of 48, and an output dimension of 24.

[0097] Step 2.4: F L2 Input to the second encoding layer, the second encoding layer outputs 48 images of size The feature maps of the set F L3; Among them, the second encoding layer is composed of the first downsampling layer and the first deep residual attention layer connected in series, the first downsampling layer is a 2×2 two-dimensional convolution with a step size of 2, an input dimension of 24, and an output dimension of 48, the input dimension and output dimension of the first deep residual attention layer are both 48, the first deep residual attention layer is composed of the first 4 second deep residual blocks and the last 4 first Swin Transformer residual blocks connected in series, the structure of the second deep residual block is the same as that of the first deep residual block, and each first Swin Transformer residual block is composed of a jump-connected Swin Transformer block (Swin Transformer blocks are disclosed in Liu, Ze, et al. "Swin transformer: Hierarchical vision transformer using shifted windows." Proceedings of the IEEE / CVF international conference on computer vision. 2021. (Swin Transformer: Hierarchical vision transformer using shifted windows)) and a jump-connected second MLP in series, Swin The input dimension and output dimension of the Transformer block are both 48, the window size is 8, and the number of attention heads is 4. The structure of the second MLP is the same as that of the first MLP, except that in the structure of the second MLP in the first Swin Transformer residual block, the third 2D convolution layer is a 1×1 2D convolution with a stride of 1, an input dimension of 48, and an output dimension of 192, and the fourth 2D convolution layer is a 1×1 2D convolution with a stride of 1, an input dimension of 192, and an output dimension of 48.

[0098] Step 2.5: F L3 Input to the bottom attention enhancement layer, the bottom attention enhancement layer outputs 48 images of size The feature maps of the set F M; Among them, the bottom attention enhancement layer is composed of the second downsampling layer, 8 second Swin Transformer residual blocks, and the first upsampling layer in series. The second downsampling layer is a 2×2 two-dimensional convolution with a step size of 2, an input dimension of 48, and an output dimension of 96. The structure of the second Swin Transformer residual block is the same as that of the first Swin Transformer residual block, except that in the second Swin The input dimension and output dimension of the SwinTransformer block in the Transformer residual block are both 96, the window size is 8, and the number of attention heads is 6. In the structure of the second MLP in the second SwinTransformer residual block, the third two-dimensional convolution layer is a 1×1 two-dimensional convolution with a stride of 1, an input dimension of 96, and an output dimension of 384. The fourth two-dimensional convolution layer is a 1×1 two-dimensional convolution with a stride of 1, an input dimension of 384, and an output dimension of 96. The first upsampling layer is a 2×2 two-dimensional convolution with a stride of 1, an input dimension of 96, an output dimension of 192, and a reflection padding size of 1, which is connected in series with a PixelShuffle (pixel cleaning) operation with a parameter of 2. The output dimension of the first upsampling layer is 48.

[0099] Step 2.6: F L3 After the first skip convolution and F M Element-wise addition ( Figure 1 All feature maps obtained after the + operation in the first decoding layer are input to the first decoding layer, which outputs 24 feature maps of size H×W, and these feature maps constitute a set Among them, the first decoding layer is composed of the first deep residual layer and the second upsampling layer in series, the first deep residual layer is composed of 4 third deep residual blocks in series, the structure of the third deep residual block is the same as that of the first deep residual block, except that the input dimension and output dimension of the third deep residual block are both 48, the third two-dimensional convolution layer in the first MLP in the third deep residual block is a 1×1 two-dimensional convolution with a stride of 1, an input dimension of 48, and an output dimension of 96, the fourth two-dimensional convolution layer is a 1×1 two-dimensional convolution with a stride of 1, an input dimension of 96, and an output dimension of 48, the second upsampling layer is a 2×2 two-dimensional convolution with a stride of 1, an input dimension of 48, an output dimension of 96, and a reflection padding size of 1, and a PixelShuffle operation with a parameter of 2 in series, the output dimension of the second upsampling layer is 24, the first jump convolution is a 1×1 two-dimensional convolution with an input dimension and an output dimension of 48, and element-by-element addition is a conventional operation in neural networks.

[0100] Step 2.7: F L2 After the second skip convolution All feature maps obtained after element-by-element addition are input into the second decoding layer, which outputs 24 feature maps of size H×W, and these feature maps form a set Among them, the second decoding layer is composed of 4 fourth depth residual blocks connected in series. The structure of the fourth depth residual block is the same as that of the first depth residual block, except that the input dimension and output dimension of the fourth depth residual block are both 24. The third two-dimensional convolution layer in the first MLP in the fourth depth residual block is a 1×1 two-dimensional convolution with a stride of 1, an input dimension of 24, and an output dimension of 48. The fourth two-dimensional convolution layer is a 1×1 two-dimensional convolution with a stride of 1, an input dimension of 48, and an output dimension of 24. The second jump convolution is a 1×1 two-dimensional convolution with both input dimension and output dimension of 24.

[0101] Step 2.8: Input into the second convolutional layer, the second convolutional layer outputs 6 feature maps of size H×W, the first 3 feature maps are used as the three channels of the attenuation map K, and the last 3 feature maps are used as the three channels of the scattering map B; among them, the second convolutional layer is composed of a 3×3 two-dimensional convolution with an input dimension of 24, an output dimension of 6, and a reflection padding size of 1, and a PixelShuffle of size 1 in series.

[0102] Step 2.9: By formula The preliminary restored underwater image is calculated Preliminary restored underwater images is a three-channel RGB image with size H×W.

[0103] Step 3: Initially restore the underwater image Input into the high and low frequency feature pyramid encoder of the second stage to obtain the latent space feature map set z0 and the intermediate feature map set F all .

[0104] In this embodiment, the specific process of step 3 is:

[0105] Step 3.1: Construct the second stage high and low frequency feature pyramid encoder ( Figure 1 The block represented by E in (e.g. Figure 2 As shown, it includes a coding convolution, a first residual coding layer, a second residual coding layer, a third residual coding layer, a fourth residual coding layer, an attention coding layer, a first attention layer, and a first integrated convolution layer, and the structures of the first to fourth residual coding layers are the same.

[0106] Step 3.2: Initially restore the underwater image Input into the coded convolution, the coded convolution outputs 128 feature maps of size H×W, and these feature maps constitute the set F E1; Among them, the encoded convolution is a 3×3 two-dimensional convolution with an input dimension of 3, an output dimension of 128, a stride of 1, and a padding of 1.

[0107] Step 3.3: F E1 Input to the first residual coding layer, the first residual coding layer outputs 128 images of size The feature maps of the set F E2 ; Among them, the first residual coding layer is composed of two residual blocks, an attention block and a downsampling convolution in series. In the first residual coding layer, the input dimension of the first residual block is 128, and the output dimension is 128. The input dimension of the second residual block is 128, and the output dimension is 128. The input dimension of the attention block is 128, and the output dimension is 128. The input dimension of the downsampling convolution is 128, and the output dimension is 128. The attention block includes group normalization (GroupNorm) with 16 groups, self-attention and attention convolution. The input of the attention block is sequentially normalized, self-attention and attention convolution, and then added to the input of the attention block to obtain the output of the attention block. The attention convolution is a 1×1 two-dimensional convolution with a step size of 1. The input of the self-attention is set to the feature map set F, and the calculation formula of the self-attention is: is the output of self-attention, Softmax is the normalization function, ⊙ represents element-by-element multiplication, represents matrix multiplication, d represents the dimension of F, W q , W k , W v is a learnable parameter, and the superscript "T" indicates transposition; the input dimension and output dimension of the residual block are set to X and Y respectively. The residual block includes a batch normalization layer (BatchNorm2d) with an input dimension and an output dimension of X, a two-dimensional convolution with a kernel size of 3, an input dimension of X, an output dimension of Y, and a padding of 1, a batch normalization layer (BatchNorm2d) with an input dimension and an output dimension of Y, a 3×3 two-dimensional convolution with an input dimension and an output dimension of Y and a padding of 1, and a 1×1 two-dimensional convolution with an input dimension of X and an output dimension of Y of a jump connection. The 1×1 two-dimensional convolution is used as a jump connection convolution. The input of the residual block passes through the first four parts in sequence, and the result is added to the input of the residual block after the jump connection convolution to obtain the output of the residual block; the downsampling convolution is a 3×3 two-dimensional convolution with the same input dimension and output dimension, a stride of 2, and a padding of 0.

[0108] Step 3.4: F E2 Input to the second residual coding layer, the second residual coding layer outputs 128 images of size The feature maps of the set F E3; Among them, in the second residual coding layer, the input dimension of the first residual block is 128, and the output dimension is 128, the input dimension of the second residual block is 128, and the output dimension is 128, the input dimension of the attention block is 128, and the output dimension is 128, and the input dimension of the downsampling convolution is 128, and the output dimension is 128.

[0109] Step 3.5: F E3 Input to the third residual coding layer, the third residual coding layer outputs 256 images of size The feature maps of the set F E4 ; Among them, in the third residual coding layer, the input dimension of the first residual block is 128, and the output dimension is 256, the input dimension of the second residual block is 256, and the output dimension is 256, the input dimension of the attention block is 256, and the output dimension is 256, and the input dimension of the downsampling convolution is 256, and the output dimension is 256.

[0110] Step 3.6: F E4 Input to the fourth residual coding layer, the fourth residual coding layer outputs 256 images of size The feature maps of the set F E5 ; Among them, in the fourth residual coding layer, the input dimension of the first residual block is 256, and the output dimension is 256, the input dimension of the second residual block is 256, and the output dimension is 256, the input dimension of the attention block is 256, and the output dimension is 256, and the input dimension of the downsampling convolution is 256, and the output dimension is 256.

[0111] Step 3.7: F E5 Input to the attention encoding layer, the attention encoding layer outputs 512 images of size The feature maps of the set F E6 ; Among them, the attention coding layer is composed of two residual blocks and one attention block connected in series. The structure of the residual block in the attention coding layer is the same as that of the residual block in the first residual coding layer. The structure of the attention block in the attention coding layer is the same as that of the attention block in the first residual coding layer. In the attention coding layer, the input dimension of the first residual block is 256 and the output dimension is 512, the input dimension of the second residual block is 512, the output dimension is 512, and the input dimension of the attention block is 512, and the output dimension is 512.

[0112] Step 3.8: F E6 Input to the first attention layer, the first attention layer outputs 512 images of size The feature maps of the set F E7; Among them, the first attention layer is composed of a residual block, an attention block, and a residual block in series. The structure of the residual block in the first attention layer is the same as that of the residual block in the first residual coding layer. The structure of the attention block in the first attention layer is the same as that of the attention block in the first residual coding layer. In the first attention layer, the input dimension of the first residual block is 512 and the output dimension is 512. The input dimension of the attention block is 512 and the output dimension is 512. The input dimension of the second residual block is 512 and the output dimension is 512.

[0113] Step 3.9: F E7 Input to the first integrated convolutional layer, the first integrated convolutional layer outputs 3 images of size The feature maps of the first integrated convolution layer are composed of group normalization (GroupNorm), Swish activation layer, and two-dimensional convolution in series. The number of group normalization groups is 32, and eps = 10. -6 , affine is True, the two-dimensional convolution is a 3×3 padded two-dimensional convolution with a stride of 1, an input dimension of 512, an output dimension of 3, and a padding of 1.

[0114] Step 3.10: F E2 、F E3 、F E4 、F E5 All feature maps in are arranged in sequence to form a set F all .

[0115] Step 4: Initially restore the underwater image Input into the high and low frequency prior extraction module in the second stage to obtain the preliminary restored underwater image High frequency image Low frequency image High frequency image prior hfp, low frequency image prior lfp.

[0116] In this embodiment, the specific process of step 4 is:

[0117] Step 4.1: Construct the high and low frequency prior extraction module of the second stage, such as Figure 3 As shown, it includes a gradient map calculation block, a Fourier high-low frequency separation block, a high-frequency prior compression block and a low-frequency prior compression block.

[0118] Step 4.2: If Figure 3 As shown in part (a) of the figure, the initially restored underwater image Input to the gradient map calculation block, the gradient map calculation block outputs a gradient map of size H×W Among them, the gradient map calculation block uses the Laplacian gradient operator to calculate the gradient map.

[0119] Step 4.3: If Figure 3 As shown in part (a) of the figure, the initially restored underwater image Input to the Fourier high and low frequency separation block, the Fourier high and low frequency separation block outputs a three-channel high frequency image of size H×W and a three-channel low-frequency image of size H×W The implementation process of the Fourier high-low frequency separation block is as follows: Use a two-dimensional Fourier transform to get a frequency map; then use fftshift on the frequency map to shift the zero frequency of the frequency map to the center position of a matrix of size H×W ; then, according to the size H×W of the frequency map, a low-frequency mask matrix with a center window of the same size of 6 and a content of 1, and a content of 0 in other positions except the center window, is generated, and a high-frequency mask matrix with a center window of the same size of 6 and a content of 0 in other positions except the center window is generated; then, the frequency map is multiplied by the low-frequency mask matrix ( Figure 3 The × operation in ( ) obtains the low-frequency part and multiplies the frequency map with the high-frequency mask matrix ( Figure 3 The high-frequency part is obtained by using the × operation in the figure; the low-frequency part is then subjected to ifftshift and two-dimensional inverse Fourier transform to obtain the low-frequency image. And use ifftshift and two-dimensional inverse Fourier transform on the high-frequency part to get the high-frequency image

[0120] Step 4.4: Gradient map and high frequency images Concatenate in the channel dimension ( Figure 3 All feature maps obtained after the C operation in ( ) are input into the high-frequency prior compression block, and the high-frequency prior compression block outputs 64 high-frequency image priors hfp with a size of 1×1; the low-frequency image Input into the low-frequency prior compression block, the low-frequency prior compression block outputs 64 low-frequency image priors lfp with a size of 1×1; among them, the structure of the high-frequency prior compression block and the low-frequency prior compression block is the same, such as Figure 3 As shown in part (b) of FIG. 1 , it is composed of a first compressed residual block, a second compressed residual block, a compression block and an information extraction block connected in series, and the first compressed residual block and the second compressed residual block have the same structure.

[0121] The frequency prior is input into the first compressed residual block, and the first compressed residual block outputs 16 images of size feature maps, and these feature maps constitute a set H1; where the input dimension of the first compressed residual block is D, for the first compressed residual block in the high-frequency prior compressed block, the frequency prior is a gradient map and high frequency images All feature maps obtained after splicing in the channel dimension, and D = 4, for the first compressed residual block in the low-frequency prior compression block, the frequency prior is the low-frequency image And D = 3, such as Figure 3 As shown in part (c) of the figure, the first compressed residual block is composed of a compression block and four prior residual blocks in series. The compression block is composed of the first two-dimensional convolution, the LeakyReLU activation layer, and the second two-dimensional convolution in series. The first two-dimensional convolution in the compression block in the first compressed residual block is a 3×3 two-dimensional convolution with an input dimension of D, an output dimension of 16, a step size of 2, and a reflection padding of 1. The second two-dimensional convolution is a 3×3 two-dimensional convolution with an input dimension of 16, an output dimension of 16, a step size of 2, and a reflection padding of 1. The input of the prior residual block passes through a two-dimensional convolution, a LeakyReLU activation layer, and another two-dimensional convolution in sequence, and is added to the input of the prior residual block to obtain the output of the prior residual block. The two two-dimensional convolutions in the prior residual block are both 3×3 two-dimensional convolutions with an input dimension of 16, an output dimension of 16, and a padding of 1.

[0122] H1 is input into the second compressed residual block, and the second compressed residual block outputs 32 images of size , and these feature maps constitute a set H2; wherein, the first two-dimensional convolution in the compressed block in the second compressed residual block is a 3×3 two-dimensional convolution with an input dimension of 16, an output dimension of 32, a stride of 2, and a reflection padding of 1, and the second two-dimensional convolution is a 3×3 two-dimensional convolution with an input dimension of 32, an output dimension of 32, a stride of 2, and a reflection padding of 1.

[0123] Input H2 into the compression block, and the compression block outputs 64 images with the size of feature maps, and these feature maps constitute a set H3; the structure of the compressed blocks in the high-frequency prior compression block and the low-frequency prior compression block is the same as the structure of the compressed blocks in the first compressed residual block and the second compressed residual block. The first two-dimensional convolution in the compression block is a 3×3 two-dimensional convolution with an input dimension of 32, an output dimension of 64, a stride of 2, and a reflection padding of 1, and the second two-dimensional convolution is a 3×3 two-dimensional convolution with an input dimension of 64, an output dimension of 64, a stride of 2, and a reflection padding of 1.

[0124] H3 is input into the information extraction block, and the information extraction block outputs 64 image priors of size 1×1; among them, for the information extraction block in the high-frequency prior compression block, the image prior is the high-frequency image prior hfp, and for the information extraction block in the low-frequency prior compression block, the image prior is the low-frequency image prior lfp. The information extraction block is composed of a maximum pooling layer, a first linear layer, a LeakyReLU activation layer, a second linear layer and a LeakyReLU activation layer in series. The maximum pooling layer is a 4×4 maximum pooling layer with a step size of 1, and the input dimension and output dimension of the first linear layer and the second linear layer are 64.

[0125] Step 5: Low frequency images Perform prior compression to obtain low-frequency guided prior lfp d ; Then the low frequency guided prior lfp d The latent space feature map set z0 is input into the low-frequency guidance condition module of the second stage to generate the low-frequency prior guidance condition z' of the latent diffusion model, and the low-frequency prior guidance condition z' is sent to the latent diffusion model, and the reconstructed latent space feature map set is obtained by iterative sampling.

[0126] In this embodiment, the specific process of step 5 is:

[0127] Step 5.1: Use the prior compression block to reduce the low frequency image Perform prior compression to obtain low-frequency guided prior lfp d ; Among them, the structure of the prior compression block is the same as the structure of the low-frequency prior compression block in the high-low frequency prior extraction module of the second stage, and the low-frequency guided prior lfp d There are 64 feature maps of size 1×1.

[0128] Step 5.2: Construct the low-frequency guidance condition module of the second stage, such as Figure 4 As shown, the implementation process is as follows:

[0129] Step 5.2.1: The latent space feature map set z0 is sequentially passed through a 1×1 two-dimensional convolution with an input dimension of 3 and an output dimension of a hidden dimension of 96 and a group normalization with a dimension of 96 and a group number of 3 to obtain a feature map set F in ; Then the low frequency guide prior lfp d Through a linear layer, we get the feature map set fp + ; and through another linear layer, we get the feature map set fp × , where the input dimension of the linear layer is 64 and the output dimension is 96; then F in With fp × After multiplication, fp + Add together to get a set of feature maps

[0130] Step 5.2.2: Order Then and Input into the cross attention, the cross attention output feature map set F out ; Among them, Figure 4 As shown in part (b), the calculation formula for cross attention is: Softmax is a normalization function, ⊙ represents element-by-element multiplication, represents matrix multiplication, d' represents the input dimension 96, W q , W k , W v are learnable parameters, and the superscript “T” indicates transpose.

[0131] Step 5.2.3: F in Skip connection, with F out The results obtained after MLP are added together to obtain the feature map set F' out , where MLP is a 1×1 two-dimensional convolution with an input dimension and an output dimension of 96; then F' out After a residual block and a two-dimensional convolution in sequence, the low-frequency prior guidance condition z' is obtained, where the structure of the residual block is the same as that of the residual block in the first residual coding layer, the input dimension of the residual block is 96, and the output dimension is 96, the input dimension of the two-dimensional convolution is 96, and the output dimension is 3, and the feature map in the low-frequency prior guidance condition z' has the same size as the feature map in the latent space feature map set z0.

[0132] Step 5.3: The latent diffusion model has a noise estimation network U-Net, The acquisition process is as follows:

[0133] Step 5.3.1: Let i represent the number of iterations, and the initial value of i is N, where N represents the maximum number of iterations, and in this embodiment, N=50.

[0134] Step 5.3.2: Calculate the sampling time t, t = (i-1) × T s / N+1, where T s represents the preset sampling time number, in this embodiment, T s =1000.

[0135] Step 5.3.3: Sample from the standard Gaussian distribution to obtain the feature map at sampling time t in, It follows a Gaussian distribution with mean 0 and identity variance.

[0136] Step 5.3.4: The result obtained by splicing with z' in the channel dimension and the sampling time t are input into the noise estimation network U-Net. The output of the noise estimation network U-Net is recorded as

[0137] Step 5.3.5: Calculate the next sampling time t next , t next =max((i-2)×T s / N+1,0), where max() is the maximum value function; then calculate the sampling time t next The characteristic diagram of in, represents the noise parameter at sampling time t, α j =1-β j , β0=0.0015 and increase equidistantly to β j , represents the sampling time t next The noise parameters at .

[0138] Step 5.3.6: Let i = i-1, where the "=" in i = i-1 is an assignment symbol; then determine whether i is equal to 0. If so, the iteration ends and As a collection of reconstructed latent space feature maps Otherwise, t next Assign to t, Assign to No sampling is required, return to step 5.3.4 and continue.

[0139] Here, the noise estimation network U-Net is a Unet model with time steps released by StableDiffusion, which is described in Rombach, Robin, et al. "High-resolution image synthesis with latent diffusion models." Proceedings of the IEEE / CVF conference on computer vision and pattern recognition. 2022.

[0140] Step 6: Collect the reconstructed latent space feature maps Intermediate feature map set F all , high-frequency image prior hfp and low-frequency image prior lfp are input into the high- and low-frequency feature pyramid decoder of the second stage to obtain the underwater enhanced image

[0141] In this embodiment, the specific process of step 6 is:

[0142] Step 6.1: Construct the second stage high and low frequency feature pyramid decoder ( Figure 1 The block represented by D in (block), such as Figure 2 As shown, it includes a decoding convolution, a second attention layer, a high and low frequency attention decoding layer, a first residual decoding layer, a second residual decoding layer, a third residual decoding layer, a fourth residual decoding layer, and a second integrated convolution layer. The structures of the first to fourth residual decoding layers are the same.

[0143] Step 6.2: Collect the reconstructed latent space feature maps Input into the decoding convolution, the decoding convolution outputs 512 images of size The feature maps of the set F D1 ; Among them, the decoding convolution is a 3×3 two-dimensional convolution with an input dimension of 3, an output dimension of 512, and a padding of 1.

[0144] Step 6.3: F D1 Input to the second attention layer, the second attention layer outputs 512 images of size The feature maps of the set F D2 ; Among them, the structure of the second attention layer is the same as that of the first attention layer, and the input dimension and output dimension of the second attention layer are both 512.

[0145] Step 6.4: F D2 , high-frequency image prior hfp and low-frequency image prior lfp are input to the high- and low-frequency attention decoding layer, and the high- and low-frequency attention decoding layer outputs 512 images of size The feature maps of the set F D3 ; Among them, the high and low frequency attention decoding layer is composed of residual blocks, high and low frequency residual attention blocks and residual blocks in series. The input dimension of the first residual block is 512 and the output dimension is 512. The input dimension of the high and low frequency residual attention blocks is 512 and the output dimension is 512. The input dimension of the second residual block is 512 and the output dimension is 512. The structure of the residual blocks in the high and low frequency attention decoding layer is the same as that of the residual blocks in the first residual encoding layer.

[0146] like Figure 5 As shown in part (a) of the figure, the high-frequency image prior hfp is passed through a linear layer before the high- and low-frequency residual attention blocks are implemented to obtain the feature map set fp hfp+ , and pass the high-frequency image prior hfp through another linear layer to obtain the feature map set fp hfp×, so that the dimension of hfp is mapped from 64 to the input dimension of the high and low frequency residual attention block Among them, fp hfp+ and fp hfp× The dimension is Similarly, the low-frequency image prior lfp is passed through a linear layer to obtain the feature map set fp lfp+ And pass the low-frequency image prior lfp through another linear layer to obtain the feature map set fp lfp× , so that the dimension of lfp is mapped from 64 to the input dimension of the high and low frequency residual attention block Among them, fp lfp+ and fp lfp× The dimension is Then the set of all feature maps output by the second residual block is After group normalization, first compare with fp hfp× Multiply and add fp hfp+ Add together to get a set of high-frequency fusion feature maps Similarly, the set of all feature maps output by the second residual block is After group normalization, first compare with fp lfp× Multiply and add fp lfp+ Add together to get a set of low-frequency fusion feature maps in, The dimension is

[0147] The high- and low-frequency residual attention blocks include two sets of self-attention blocks and MLP blocks connected in parallel and internally in sequence, where the input of the self-attention block is set to the feature map set F in , then the calculation formula of the self-attention block is: F out is the output of the self-attention block, Softmax is the normalization function, ⊙ represents element-by-element multiplication, represents matrix multiplication, d represents the dimension of F, W q , W k , W v are learnable parameters, the superscript “T” indicates transposition, and the MLP block has input dimensions and output dimensions of And the convolution kernel is a 1×1 two-dimensional convolution; Figure 5 As shown in part (a) of , the implementation process of the high and low frequency residual attention block is: the input of the high and low frequency residual attention block is right Perform Fourier transform to get After passing through the first group of self-attention blocks and MLP blocks, we get right and The feature map obtained by adding element by element is Fourier transformed to obtain After passing through the second group of self-attention blocks and MLP blocks, we get right and The feature map obtained by adding element by element is Fourier transformed to obtain Then calculate the high frequency residual R high , Similarly, calculate the low-frequency residual R low , Then R high The result obtained after complex convolution is similar to R low The results obtained after complex convolution are added to obtain the high and low frequency residual R, where the input and output dimensions of the complex convolution are both The convolution kernel is 1×1; finally, the high and low frequency residuals R and Add, and pass the addition result through inverse Fourier transform to get the output of high and low frequency residual attention block, where the output dimension of high and low frequency residual attention block is

[0148] Step 6.5: Set the high-frequency image prior hfp, low-frequency image prior lfp, and F D3 With F all F E5 All feature maps obtained after splicing in the channel dimension are input to the first residual decoding layer, which outputs 256 images of size The feature maps of the set F D4 ; Among them, the first residual decoding layer is composed of two residual blocks, a high- and low-frequency residual attention block and an upsampling convolution in series. In the first residual decoding layer, the input dimension of the first residual block is 1024 and the output dimension is 256, the input dimension of the second residual block is 256 and the output dimension is 256, the input dimension of the high- and low-frequency residual attention blocks is 256 and the output dimension is 256, the input dimension of the upsampling convolution is 256 and the output dimension is 256, the structure of the residual block in the first residual decoding layer is the same as that of the residual block in the first residual encoding layer, the upsampling convolution is a 3×3 two-dimensional convolution with the same input dimension and output dimension and padding of 1, and the structure of the high- and low-frequency residual attention blocks in the first residual decoding layer is the same as that of the high- and low-frequency residual attention blocks in the high- and low-frequency attention decoding layer.

[0149] Step 6.6: Set the high-frequency image prior hfp, low-frequency image prior lfp, and F D4 With F all F E4 All feature maps obtained after splicing in the channel dimension are input to the second residual decoding layer, and the second residual decoding layer outputs 256 images of size The feature maps of the set F D5 ; Among them, in the second residual decoding layer, the input dimension of the first residual block is 512 and the output dimension is 256, the input dimension of the second residual block is 256 and the output dimension is 256, the input dimension of the high and low frequency residual attention blocks is 256 and the output dimension is 256, and the input dimension of the upsampling convolution is 256 and the output dimension is 256.

[0150] Step 6.7: Set the high-frequency image prior hfp, low-frequency image prior lfp, and F D5 With F all F E3 All feature maps obtained after splicing in the channel dimension are input to the third residual decoding layer, and the third residual decoding layer outputs 128 images of size The feature maps of the set F D6 ; Among them, in the third residual decoding layer, the input dimension of the first residual block is 512 and the output dimension is 128, the input dimension of the second residual block is 128 and the output dimension is 128, the input dimension of the high and low frequency residual attention block is 128 and the output dimension is 128, and the input dimension of the upsampling convolution is 128 and the output dimension is 128.

[0151] Step 6.8: Set the high-frequency image prior hfp, low-frequency image prior lfp, and F D6 With F all F E2 All feature maps obtained after splicing in the channel dimension are input into the fourth residual decoding layer, which outputs 128 feature maps of size H×W, and these feature maps constitute a set F D7 ; Among them, in the fourth residual decoding layer, the input dimension of the first residual block is 256 and the output dimension is 128, the input dimension of the second residual block is 128 and the output dimension is 128, the input dimension of the high and low frequency residual attention blocks is 128 and the output dimension is 128, and the input dimension of the upsampling convolution is 128 and the output dimension is 128.

[0152] Step 6.9: F D7 Input to the second integrated convolutional layer, the second integrated convolutional layer outputs a three-channel underwater enhanced image of size H×W The structure of the second integrated convolutional layer is the same as that of the first integrated convolutional layer.

[0153] When training the above lightweight underwater image enhancement model, the learning rate is set to 2×10 -5, the training batch size is set to 16, the number of training rounds is set to 50, and the optimizer uses the AdamW optimizer; the underwater degraded image x and its degraded ambient light A are input into the lightweight underwater image enhancement model, and the lightweight underwater image enhancement model outputs a preliminary restored underwater image calculate With clear image x gt The loss between pen , Among them, ||||1 is the 1-norm operator.

[0154] When the high- and low-frequency feature pyramid encoder and the high- and low-frequency feature pyramid decoder are trained, the learning rate is set to 2×10 -5 , the training batch size is set to 16, the number of training rounds is set to 300, and the optimizer uses the AdamW optimizer; clear image x gt Input into the high- and low-frequency feature pyramid encoder, the high- and low-frequency feature pyramid encoder obtains the latent space feature map set z gt and the set of intermediate feature maps F gt ; Latent space feature map set z gt and the set of intermediate feature maps F gt Input into the high- and low-frequency feature pyramid decoder, and the high- and low-frequency feature pyramid decoder outputs the restored image Calculate x gt and The loss between ed , Among them, ||||2 is the 2-norm operation symbol.

[0155] When training the noise estimation network U-Net in the above potential diffusion model, the learning rate is set to 2×10 -5 , the training batch size is set to 16, the number of training rounds is set to 300, and the optimizer uses the AdamW optimizer; the random noise χ is sampled, and the random noise χ obeys a Gaussian distribution with a mean of 0 and a variance of the unit matrix; the clear image x gt Input into the high and low frequency feature pyramid encoder with fixed parameters, the high and low frequency feature pyramid encoder with fixed parameters obtains the latent space feature map set z gt ; Calculate the feature map z at sampling time t gt,t , Where t=1,2,…,T s ; The underwater degraded image x and its degraded ambient light A are input into a lightweight underwater image enhancement model with fixed parameters, and the lightweight underwater image enhancement model with fixed parameters outputs a preliminary restored underwater image Preliminary restored underwater images Input into the high and low frequency feature pyramid encoder with fixed parameters, the high and low frequency feature pyramid encoder with fixed parameters obtains the latent space feature map set z0; the initial restored underwater image Input into the high and low frequency prior extraction module to obtain the preliminary restored underwater image High frequency image Low frequency image High-frequency image prior hfp, low-frequency image prior lfp; for low-frequency image Perform prior compression to obtain low-frequency guided prior lfp d ; The low frequency guide prior lfp d and the latent space feature map set z0 are input into the low-frequency guidance condition module to generate the low-frequency prior guidance condition z' of the latent diffusion model; z gt,t The result obtained by splicing z' in the channel dimension and the sampling time t are input into the noise estimation network U-Net, and the noise estimation network U-Net outputs χ(z gt,t ,z',t); calculate χ and χ(z gt,t ,z',t) ldm , L ldm =||χ-χ(z gt,t ,z',t)||2.

[0156] In order to further illustrate the feasibility and effectiveness of the method of the present invention, experiments were conducted on the method of the present invention.

[0157] The proposed method (FD-LDM) is trained and tested in the paired public underwater datasets Underwater Image Enhancement Benchmark (UIEB) and Large Scale Underwater Image Dataset (LSUI). Among them, 90 images are randomly selected from the UIEB dataset as the test set and recorded as the test set U90, and the remaining 800 images are used as the training set; 504 images are randomly selected from the LSUI dataset as the test set and recorded as the test set L504, and the remaining 4500 images are used as the training set. Peak signal-to-noise ratio (PSNR) and structural similarity (SSIM) are used as comparison indicators on the two paired datasets. The results are compared with the current underwater image enhancement methods of different categories and published codes, including UDCP (Drews, Paul, et al. "Transmission estimation inunderwater single images." Proceedings of the IEEE international conference on computer vision workshops. 2013. (Transmittance estimation of underwater single images)), WWPE (Zhang, Weidong, et al. "Underwater image enhancement via weighted wavelet visual perception fusion." IEEE Transactions on Circuits and Systems for Video Technology (2023) (Underwater image enhancement by weighted wavelet visual perception fusion)), Water-net (Li, Chongyi, et al. "An underwater image enhancement benchmark dataset and beyond." IEEE transactions on image processing 29 (2019): 4376-4389. (Underwater image enhancement benchmark dataset and its extension)), and Ucolor (Li, Chongyi, et al. "Underwater image enhancement via medium transmission-guided multi-color space embedding." IEEE Transactions on Image Processing 30(2021):4985-5000.(Underwater image enhancement via medium transfer guided multi-color spatial embedding)), FunIE(Islam, Md Jahidul, Youya Xia, and Junaed Sattar. "Fast underwater image enhancement for improved visual perception." IEEE Robotics and Automation Letters 5.2(2020): 3227-3234.(Fast underwater image enhancement for improved visual perception)), Ushape(Peng, Lintao, Chunli Zhu, and Liheng Bian. "U-shape transformer for underwater image enhancement." IEEE Transactions on Image Processing 32(2023): 3066-3079.(U-shape transformer for underwater image enhancement)), DM-water(Tang, Yi, Hiroshi Kawasaki, and Takafumi Iwaguchi. "Underwater image enhancement by transformer-based diffusion model with non-uniform sampling for skip strategy." Proceedings of the 31st ACM International Conference on Multimedia.2023.(Underwater image enhancement through Transformer-based diffusion model and skip strategy for non-uniform sampling)), DiffWater(Guan,Meisheng,etal. "DiffWater:Underwater image enhancement based on conditional denoisingdiffusion probabilistic model."IEEE Journal of Selected Topics in AppliedEarth Observations and Remote Sensing(2023).(DiffWater:Underwater image enhancement based on conditional denoisingdiffusion probabilistic model)). .

[0158] Table 1 gives a quantitative comparison of the method of the present invention and the UDCP method, WWPE method, Water-net method, Ucolor method, FunIE method, Ushape method, DM-water method, and DiffWater method on the test set U90 and the test set L504.

[0159] Table 1 Quantitative comparison of the method of the present invention and the comparative method on the test set U90 and the test set L504

[0160]

[0161] In Table 1, the Flops (Floating Point Operations Per Second) of the latent diffusion model is measured based on a single sampling operation of U-Net. Flops is used to measure the number of floating point operations that a computer or processor can perform per second. The index before the "+" in the method of the present invention corresponds to the first stage, and the index after the "+" corresponds to the second stage. Thanks to the two-stage strategy proposed by the present invention, the dependence of the latent diffusion model on learning samples is alleviated. The method of the present invention achieves the highest PSNR index on the test set U90 of the UIEB dataset with a small number of samples, which is 7.5% higher than the method with the second highest PSNR index. At the same time, on the UIEB dataset, the best SSIM index is achieved together with the Ushape method. In addition, thanks to the latent diffusion model in the second stage and the high and low frequency priors, the method of the present invention achieves the best PSNR index and obtains a slightly lower SSIM index on the test set L504 of the large sample dataset LSUI.

[0162] Figure 6 The visualization results of different underwater image enhancement methods on the test set L504 are shown. Figure 7 The visualization results of different underwater image enhancement methods on the test set U90 are shown. Figure 6 and Figure 7Among them, (a) is the underwater degraded image, (b) is the visualization result of the underwater enhanced image obtained by the UDCP method, (c) is the visualization result of the underwater enhanced image obtained by the WWPE method, (d) is the visualization result of the underwater enhanced image obtained by the Water-net method, (e) is the visualization result of the underwater enhanced image obtained by the Ucolor method, (f) is the visualization result of the underwater enhanced image obtained by the FunIE method, (g) is the visualization result of the underwater enhanced image obtained by the Ushape method, (h) is the visualization result of the underwater enhanced image obtained by the DM-water method, (i) is the visualization result of the underwater enhanced image obtained by the DiffWater method, (j) is the visualization result of the underwater enhanced image obtained by the method of the present invention (FD-LDM), and (k) is a high-quality clear image. By comparison, it can be found that the method of the present invention (FD-LDM) effectively enhances the underwater degraded images in different color (green, blue) scenes. This is due to the strategy of high and low frequency separation enhancement, and the details of the restored image (underwater enhanced image) are enhanced to varying degrees. Figure 6 The second to last row of pictures and Figure 7 In the last row of images, the proposed method (FD-LDM) effectively enhances the degraded images with severe blue or green degradation, and contains details while being more consistent with the true value images in the dataset (i.e., the high-quality clear images shown in the (k) column). In addition, Figure 7 In the first row of images, the method of the present invention (FD-LDM) effectively enhances the degraded images in extremely dark environments, and shows positive performance compared to the method of the non-diffusion model. While containing more details, the brightness is also consistent with human eye perception. For the diffusion model, the two-stage strategy of the method of the present invention (FD-LDM) effectively reduces the impact of color cast. For other contrast images, the method of the present invention (FD-LDM) also shows the ability to restore details. While retaining more image details, it contains less color cast, effectively avoids the generation of color cast in different underwater scenes, and the restored image is closer to the true value image (high-quality clear image) in color. Although it is still inferior to the WWPE method in contrast, there will be no significant color or contrast degradation in more general underwater scenes.

[0163] By using high and low frequency priors, we can solve the problem of different degradation levels of high and low frequencies in underwater images, and at the same time improve the performance of the latent diffusion model in underwater image enhancement to obtain high-quality images. The retention of more image details and the effective elimination of color degradation are beneficial to tasks such as underwater image recognition, underwater target classification, and underwater detection.

[0164] It should be appreciated that embodiments of the present invention may be implemented or enforced by computer hardware, a combination of hardware and software, or by computer instructions stored in a non-transitory computer-readable memory. The method may be implemented in a computer program using standard programming techniques - including a non-transitory computer-readable storage medium configured with a computer program, wherein the storage medium so configured causes the computer to operate in a specific and predefined manner - according to the methods and drawings described in the specific embodiments. Each program may be implemented in a high-level procedural or object-oriented programming language to communicate with a computer system. However, if desired, the program may be implemented in an assembly or machine language. In any case, the language may be a compiled or interpreted language. In addition, the program may be run on a programmed ASIC for this purpose.

[0165] It should be noted that the above embodiments are only used to illustrate the technical solutions of the present invention rather than to limit it. Although the present invention has been described in detail with reference to the preferred embodiments, those skilled in the art should understand that the technical solutions of the present invention may be modified or replaced by equivalents without departing from the spirit and scope of the technical solutions of the present invention, which should all be covered by the scope of the claims of the present invention.

Claims

1. A two-stage potential diffusion model underwater image enhancement method based on frequency domain, characterized in that The following steps are involved: Step 1: Input the underwater degraded image x into the Gaussian blur estimation module of the first stage to estimate the degraded ambient light A of the underwater degraded image x; Step 2: Input the underwater degraded image x and its degraded ambient light A into the lightweight underwater image enhancement model of the first stage to obtain a preliminary restored underwater image Step 3: Initially restore the underwater image Input into the high and low frequency feature pyramid encoder of the second stage to obtain the latent space feature map set z0 and the intermediate feature map set F all ; Step 4: Initially restore the underwater image Input into the high and low frequency prior extraction module in the second stage to obtain the preliminary restored underwater image High frequency image Low frequency image High frequency image prior hfp, low frequency image prior lfp; Step 5: Low frequency images Perform prior compression to obtain low-frequency guided prior lfp d ; Then the low frequency guided prior lfp d The latent space feature map set z0 is input into the low-frequency guidance condition module of the second stage to generate the low-frequency prior guidance condition z' of the latent diffusion model, and the low-frequency prior guidance condition z' is sent to the latent diffusion model, and the reconstructed latent space feature map set is obtained by iterative sampling. Step 6: Collect the reconstructed latent space feature maps Intermediate feature map set F all , high-frequency image prior hfp and low-frequency image prior lfp are input into the high- and low-frequency feature pyramid decoder of the second stage to obtain the underwater enhanced image 2. The underwater image enhancement method based on the two-stage potential diffusion model in the frequency domain according to claim 1 is characterized in that In the step 1, the underwater degraded image x is a three-channel RGB image with a size of H×W, and the degraded ambient light A is a three-channel RGB image with a size of H×W, where H is the height and W is the width.

3. The underwater image enhancement method based on the two-stage potential diffusion model in the frequency domain according to claim 2 is characterized in that The specific process of step 2 is: Step 2.1: Construct the first-stage lightweight underwater image enhancement model, which includes the first convolution layer, the first encoding layer, the second encoding layer, the bottom attention enhancement layer, the first decoding layer, the second decoding layer, and the second convolution layer; Step 2.2: The 6 feature maps of size H×W obtained by splicing the underwater degraded image x and its degraded ambient light A in the channel dimension are input into the first convolutional layer. The first convolutional layer outputs 24 feature maps of size H×W, and these feature maps constitute the set F L1 ; Step 2.3: F L1 Input to the first coding layer, the first coding layer outputs 24 feature maps of size H×W, and these feature maps constitute a set F L2 ; Among them, the first coding layer is composed of 8 first depth residual blocks connected in series, each of which is composed of a jump-connected depth convolution block and a jump-connected first MLP connected in series, the depth convolution block is composed of a first two-dimensional convolution layer and a second two-dimensional convolution layer connected in series, and the first MLP is composed of a third two-dimensional convolution layer, a ReLU activation layer and a fourth two-dimensional convolution layer connected in series; Step 2.4: F L2 Input to the second encoding layer, the second encoding layer outputs 48 images of size The feature maps of the set F L3 ; Among them, the second encoding layer is composed of the first downsampling layer and the first deep residual attention layer connected in series, the input dimension and output dimension of the first deep residual attention layer are both 48, the first deep residual attention layer is composed of the first 4 second deep residual blocks and the last 4 first Swin Transformer residual blocks connected in series, the structure of the second deep residual block is the same as that of the first deep residual block, each first Swin Transformer residual block is composed of a jump-connected Swin Transformer block and a jump-connected second MLP connected in series, and the structure of the second MLP is the same as that of the first MLP; Step 2.5: F L3 Input to the bottom attention enhancement layer, the bottom attention enhancement layer outputs 48 images of size The feature maps of the set F M ; Among them, the bottom attention enhancement layer is composed of the second downsampling layer, 8 second Swin Transformer residual blocks, and the first upsampling layer in series, and the structure of the second Swin Transformer residual block is the same as that of the first Swin Transformer residual block; Step 2.6: F L3 After the first skip convolution and F M All feature maps obtained after element-by-element addition are input into the first decoding layer, which outputs 24 feature maps of size H×W, and these feature maps form a set The first decoding layer is composed of a first depth residual layer and a second upsampling layer connected in series, the first depth residual layer is composed of four third depth residual blocks connected in series, and the structure of the third depth residual block is the same as that of the first depth residual block; Step 2.7: F L2 After the second skip convolution All feature maps obtained after element-by-element addition are input into the second decoding layer, which outputs 24 feature maps of size H×W, and these feature maps form a set The second decoding layer is composed of four fourth depth residual blocks connected in series, and the structure of the fourth depth residual block is the same as that of the first depth residual block; Step 2.8: Input into the second convolutional layer, the second convolutional layer outputs 6 feature maps of size H×W, the first 3 feature maps are used as the three channels of the attenuation map K, and the last 3 feature maps are used as the three channels of the scattering map B; Step 2.9: By formula The preliminary restored underwater image is calculated Preliminary restored underwater images is a three-channel RGB image with size H×W.

4. The underwater image enhancement method based on the two-stage potential diffusion model in the frequency domain according to claim 1 is characterized in that The specific process of step 3 is as follows: Step 3.1: Construct the high- and low-frequency feature pyramid encoder of the second stage, which includes the coding convolution, the first residual coding layer, the second residual coding layer, the third residual coding layer, the fourth residual coding layer, the attention coding layer, the first attention layer, and the first integrated convolution layer. The structures of the first to fourth residual coding layers are the same; Step 3.2: Initially restore the underwater image Input into the coded convolution, the coded convolution outputs 128 feature maps of size H×W, and these feature maps constitute the set F E1 ; Step 3.3: F E1 Input to the first residual coding layer, the first residual coding layer outputs 128 images of size The feature maps of the set F E2 ; Among them, the first residual coding layer is composed of two residual blocks, an attention block and a downsampling convolution in series; Step 3.4: F E2 Input to the second residual coding layer, the second residual coding layer outputs 128 images of size The feature maps of the set F E3 ; Step 3.5: F E3 Input to the third residual coding layer, the third residual coding layer outputs 256 images of size The feature maps of the set F E4 ; Step 3.6: F E4 Input to the fourth residual coding layer, the fourth residual coding layer outputs 256 images of size The feature maps of the set F E5 ; Step 3.7: F E5 Input to the attention encoding layer, the attention encoding layer outputs 512 images of size The feature maps of the set F E6 ; Among them, the attention encoding layer is composed of two residual blocks and one attention block in series; Step 3.8: F E6 Input to the first attention layer, the first attention layer outputs 512 images of size The feature maps of the set F E7 ; Among them, the first attention layer is composed of a residual block, an attention block, and a residual block connected in series; Step 3.9: F E7 Input to the first integrated convolutional layer, the first integrated convolutional layer outputs 3 images of size The feature maps of the network are the set of latent space feature maps z0; wherein the first integrated convolution layer is composed of group normalization, Swish activation layer, and two-dimensional convolution in series; Step 3.10: F E2 、F E3 、F E4 、F E5 All feature maps in are arranged in sequence to form a set F all .

5. The underwater image enhancement method based on the two-stage potential diffusion model in the frequency domain according to claim 1 is characterized in that The specific process of step 4 is as follows: Step 4.1: Construct the high and low frequency prior extraction module of the second stage, which includes a gradient map calculation block, a Fourier high and low frequency separation block, a high frequency prior compression block and a low frequency prior compression block; Step 4.2: Initially restore the underwater image Input to the gradient map calculation block, the gradient map calculation block outputs a gradient map of size H×W Among them, the gradient map calculation block uses the Laplacian gradient operator to calculate the gradient map; Step 4.3: Initially restore the underwater image Input to the Fourier high and low frequency separation block, the Fourier high and low frequency separation block outputs a three-channel high frequency image of size H×W and a three-channel low-frequency image of size H×W The implementation process of the Fourier high-low frequency separation block is as follows: Use a two-dimensional Fourier transform to get a frequency map; then use fftshift on the frequency map to shift the zero frequency of the frequency map to the center position of a matrix of size H×W ; then, according to the size H×W of the frequency map, a low-frequency mask matrix with a center window of the same size of 6 in width and height and a content of 1, and a content of 0 in other positions except the center window is generated, and a high-frequency mask matrix with a center window of the same size of 6 in width and height and a content of 0, and a content of 1 in other positions except the center window is generated; then, the frequency map is multiplied by the low-frequency mask matrix to obtain the low-frequency part, and the frequency map is multiplied by the high-frequency mask matrix to obtain the high-frequency part; then, ifftshift and two-dimensional inverse Fourier transform are used on the low-frequency part to obtain a low-frequency image And use ifftshift and two-dimensional inverse Fourier transform on the high-frequency part to get the high-frequency image Step 4.4: Gradient map and high frequency images All feature maps obtained after splicing in the channel dimension are input into the high-frequency prior compression block, which outputs 64 high-frequency image priors hfp with a size of 1×1; the low-frequency images Input into the low-frequency prior compression block, the low-frequency prior compression block outputs 64 low-frequency image priors lfp with a size of 1×1; wherein, the high-frequency prior compression block and the low-frequency prior compression block have the same structure, which is composed of a first compressed residual block, a second compressed residual block, a compression block and an information extraction block connected in series, and the first compressed residual block and the second compressed residual block have the same structure; The frequency prior is input into the first compressed residual block, and the first compressed residual block outputs 16 images of size feature maps, and these feature maps constitute a set H1; where the input dimension of the first compressed residual block is D, for the first compressed residual block in the high-frequency prior compressed block, the frequency prior is a gradient map and high frequency images All feature maps obtained after splicing in the channel dimension, and D = 4, for the first compressed residual block in the low-frequency prior compression block, the frequency prior is the low-frequency image And D=3, the first compressed residual block is composed of a compression block and 4 prior residual blocks in series, the compression block is composed of the first two-dimensional convolution, the LeakyReLU activation layer, and the second two-dimensional convolution in series, the input of the prior residual block is sequentially passed through a two-dimensional convolution, a LeakyReLU activation layer and another two-dimensional convolution, and then added to the input of the prior residual block to obtain the output of the prior residual block; H1 is input into the second compressed residual block, and the second compressed residual block outputs 32 images of size , and these feature maps constitute the set H2; Input H2 into the compression block, and the compression block outputs 64 images with the size of , and these feature maps constitute a set H3; H3 is input into the information extraction block, and the information extraction block outputs 64 image priors of size 1×1; among them, for the information extraction block in the high-frequency prior compression block, the image prior is the high-frequency image prior hfp, and for the information extraction block in the low-frequency prior compression block, the image prior is the low-frequency image prior lfp. The information extraction block is composed of a maximum pooling layer, a first linear layer, a LeakyReLU activation layer, a second linear layer and a LeakyReLU activation layer in series.

6. The underwater image enhancement method based on the two-stage potential diffusion model in the frequency domain according to claim 5 is characterized in that The specific process of step 5 is as follows: Step 5.1: Use the prior compression block to reduce the low frequency image Perform prior compression to obtain low-frequency guided prior lfp d ; Among them, the structure of the prior compression block is the same as the structure of the low-frequency prior compression block in the high-low frequency prior extraction module of the second stage, and the low-frequency guided prior lfp d There are 64 feature maps of size 1×1; Step 5.2: Construct the low-frequency guidance condition module of the second stage, and its implementation process is as follows: Step 5.2.1: The latent space feature map set z0 is sequentially subjected to two-dimensional convolution and group normalization to obtain the feature map set F in ; Then the low frequency guide prior lfp d Through a linear layer, we get the feature map set fp + ; and through another linear layer, we get the feature map set fp × ; Then F in With fp × After multiplication, fp + Add together to get a set of feature maps Step 5.2.2: Order Then and Input into the cross attention, the cross attention output feature map set F out ; Step 5.2.3: F in Skip connection, with F out The results obtained after MLP are added together to obtain the feature map set F' out ; then F' out After a residual block and a two-dimensional convolution, the low-frequency prior guidance condition z' is obtained; Step 5.3: The latent diffusion model has a noise estimation network U-Net, The acquisition process is as follows: Step 5.3.1: Let i represent the number of iterations, and the initial value of i is N, where N represents the maximum number of iterations; Step 5.3.2: Calculate the sampling time t, t = (i-1) × T s / N+1, where T s Indicates the preset number of sampling moments; Step 5.3.3: Sample from the standard Gaussian distribution to obtain the feature map at sampling time t in, It follows a Gaussian distribution with a mean of 0 and a unit variance; Step 5.3.4: The result obtained by splicing with z' in the channel dimension and the sampling time t are input into the noise estimation network U-Net. The output of the noise estimation network U-Net is recorded as Step 5.3.5: Calculate the next sampling time t next , t next =max((i-2)×T s / N+1,0); then calculate the sampling time t next The characteristic diagram of in, represents the noise parameter at sampling time t, α j =1-β j , β0=0.0015 and increase equidistantly to β j , represents the sampling time t next Noise parameters when ; Step 5.3.6: Let i = i-1, where "=" in i = i-1 is an assignment symbol; then determine whether i is equal to 0. If so, the iteration ends and As a collection of reconstructed latent space feature maps Otherwise, t next Assign to t, Assign to Return to step 5.3.4 and continue.

7. The underwater image enhancement method based on the two-stage potential diffusion model in the frequency domain according to claim 4 is characterized in that The specific process of step 6 is as follows: Step 6.1: Construct the high and low frequency feature pyramid decoder of the second stage, which includes decoding convolution, second attention layer, high and low frequency attention decoding layer, first residual decoding layer, second residual decoding layer, third residual decoding layer, fourth residual decoding layer, second integrated convolution layer, and the structures of the first to fourth residual decoding layers are the same; Step 6.2: Collect the reconstructed latent space feature maps Input into the decoding convolution, the decoding convolution outputs 512 images of size The feature maps of the set F D1 ; Step 6.3: F D1 Input to the second attention layer, the second attention layer outputs 512 images of size The feature maps of the set F D2 ; Step 6.4: F D2 , high-frequency image prior hfp and low-frequency image prior lfp are input to the high- and low-frequency attention decoding layer, and the high- and low-frequency attention decoding layer outputs 512 images of size The feature maps of the set F D3 ; Among them, the high-low frequency attention decoding layer is composed of residual blocks, high-low frequency residual attention blocks and residual blocks in series. Before the high- and low-frequency residual attention blocks are implemented, the high-frequency image prior hfp is passed through a linear layer to obtain the feature map set fp hfp+ , and pass the high-frequency image prior hfp through another linear layer to obtain the feature map set fp hfp× , where fp hfp+ and fp hfp× The dimension is Similarly, the low-frequency image prior lfp is passed through a linear layer to obtain the feature map set fp lfp+ And pass the low-frequency image prior lfp through another linear layer to obtain the feature map set fp lfp× , where fp lfp+ and fp lfp× The dimension is Then the set of all feature maps output by the second residual block is After group normalization, first compare with fp hfp× Multiply and add fp hfp+ Add together to get a set of high-frequency fusion feature maps Similarly, the set of all feature maps output by the second residual block is After group normalization, first compare with fp lfp× Multiply and add fp lfp+ Add together to get a set of low-frequency fusion feature maps in, The dimension is The high- and low-frequency residual attention blocks consist of two sets of self-attention blocks and MLP blocks connected in parallel and internally in sequence, where the MLP block is a two-dimensional convolution; the implementation process of the high- and low-frequency residual attention blocks is: the input of the high- and low-frequency residual attention blocks is right Perform Fourier transform to get After passing through the first group of self-attention blocks and MLP blocks, we get right and The feature map obtained by adding element by element is Fourier transformed to obtain After passing through the second group of self-attention blocks and MLP blocks, we get right and The feature map obtained by adding element by element is Fourier transformed to obtain Then calculate the high frequency residual R high , Similarly, calculate the low-frequency residual R low , Then R high The result obtained after complex convolution is similar to R low The results obtained after complex convolution are added to obtain the high and low frequency residuals R; finally, the high and low frequency residuals R and Add, and pass the addition result through inverse Fourier transform to get the output of high and low frequency residual attention block, where the output dimension of high and low frequency residual attention block is Step 6.5: Set the high-frequency image prior hfp, low-frequency image prior lfp, and F D3 With F all F E5 All feature maps obtained after splicing in the channel dimension are input to the first residual decoding layer, which outputs 256 images of size The feature maps of the set F D4 ; Among them, the first residual decoding layer is composed of two residual blocks, a high-low frequency residual attention block and an upsampling convolution in series; Step 6.6: Set the high-frequency image prior hfp, low-frequency image prior lfp, and F D4 With F all F E4 All feature maps obtained after splicing in the channel dimension are input to the second residual decoding layer, and the second residual decoding layer outputs 256 images of size The feature maps of the set F D5 ; Step 6.7: Set the high-frequency image prior hfp, low-frequency image prior lfp, and F D5 With F all F E3 All feature maps obtained after splicing in the channel dimension are input to the third residual decoding layer, and the third residual decoding layer outputs 128 images of size The feature maps of the set F D6 ; Step 6.8: Set the high-frequency image prior hfp, low-frequency image prior lfp, and F D6 With F all F E2 All feature maps obtained after splicing in the channel dimension are input into the fourth residual decoding layer, which outputs 128 feature maps of size H×W, and these feature maps constitute a set F D7 ; Step 6.9: F D7 Input to the second integrated convolutional layer, the second integrated convolutional layer outputs a three-channel underwater enhanced image of size H×W 8. The underwater image enhancement method based on the two-stage potential diffusion model in the frequency domain according to claim 1 is characterized in that When training the lightweight underwater image enhancement model, the learning rate is set to 2×10 -5 , the training batch size is set to 16, the number of training rounds is set to 50, and the optimizer uses the AdamW optimizer; the underwater degraded image x and its degraded ambient light A are input into the lightweight underwater image enhancement model, and the lightweight underwater image enhancement model outputs a preliminary restored underwater image calculate With clear image x gt The loss between pen , Among them, || ||1 is the 1-norm operator; When the high- and low-frequency feature pyramid encoder and the high- and low-frequency feature pyramid decoder are trained, the learning rate is set to 2×10 -5 , the training batch size is set to 16, the number of training rounds is set to 300, and the optimizer uses the AdamW optimizer; clear image x gt Input into the high- and low-frequency feature pyramid encoder, the high- and low-frequency feature pyramid encoder obtains the latent space feature map set z gt and the set of intermediate feature maps F gt ; Latent space feature map set z gt and the set of intermediate feature maps F gt Input into the high- and low-frequency feature pyramid decoder, and the high- and low-frequency feature pyramid decoder outputs the restored image Calculate x gt and The loss between ed , Among them, || ||2 is the 2-norm operator; When training the noise estimation network U-Net in the potential diffusion model, the learning rate is set to 2×10 -5 , the training batch size is set to 16, the number of training rounds is set to 300, and the optimizer uses the AdamW optimizer; the random noise χ is sampled, and the random noise χ obeys a Gaussian distribution with a mean of 0 and a variance of the unit matrix; the clear image x gt Input into the high and low frequency feature pyramid encoder with fixed parameters, the high and low frequency feature pyramid encoder with fixed parameters obtains the latent space feature map set z gt ; Calculate the feature map z at sampling time t gt,t , Where t=1,2,…,T s ; The underwater degraded image x and its degraded ambient light A are input into a lightweight underwater image enhancement model with fixed parameters, and the lightweight underwater image enhancement model with fixed parameters outputs a preliminary restored underwater image Preliminary restored underwater images Input into the high and low frequency feature pyramid encoder with fixed parameters, the high and low frequency feature pyramid encoder with fixed parameters obtains the latent space feature map set z0; the initial restored underwater image Input into the high and low frequency prior extraction module to obtain the preliminary restored underwater image High frequency image Low frequency image High-frequency image prior hfp, low-frequency image prior lfp; for low-frequency image Perform prior compression to obtain low-frequency guided prior lfp d ; The low frequency guide prior lfp d and the latent space feature map set z0 are input into the low-frequency guidance condition module to generate the low-frequency prior guidance condition z' of the latent diffusion model; z gt,t The result obtained by splicing z' in the channel dimension and the sampling time t are input into the noise estimation network U-Net, and the noise estimation network U-Net outputs χ(z gt,t ,z',t); calculate χ and χ(z gt,t ,z',t) ldm , L ldm =||χ-χ(z gt,t ,z',t)||2.

Images