Underwater image enhancement method based on unsupervised diffusion model

Abstract

The invention discloses an underwater image enhancement method based on an unsupervised diffusion model. The method comprises the following steps: acquiring a non-paired underwater image training set; a to-be-processed underwater image is input into a trained diffusion model based on unsupervised comparative learning, the diffusion model comprises a pre-enhancement network based on antagonistic comparative learning and a conditional denoising diffusion probability model based on discrete wavelet transform, and the pre-enhancement network comprises a generator and a discriminator. The conditional denoising diffusion probability model comprises a K-time discrete wavelet transform process, a low-frequency sub-band diffusion process and a high-frequency sub-band enhancement module; and obtaining the output of the wavelet transform-based conditional denoising diffusion probability model as an enhanced underwater image. According to the method, the generative adversarial network and the diffusion model are combined, training can be carried out in a real underwater scene without a reference image, the problems of color cast, low contrast and the like existing in a real underwater image are effectively relieved, and the enhancement effect comparable with supervised training is achieved.

Description

Technical Field

[0001] The present invention relates to the field of computer vision technology, and in particular to an underwater image enhancement method based on an unsupervised diffusion model. Background Art

[0002] Underwater operations such as underwater biological detection and marine ranching have high requirements for the environmental perception accuracy of underwater robots. As a carrier of marine environmental perception, the clarity of underwater images determines the strength of underwater robots' environmental perception capabilities. Therefore, underwater image enhancement technology has always been one of the key research directions for researchers to explore the mysteries of the ocean.

[0003] However, the abundance of marine life causes turbid water, the water properties cause the absorption and attenuation of light in water to be extremely severe and wavelength-dependent, and the refraction of light caused by suspended particles in the water makes the underwater imaging process complicated. Therefore, it is difficult to obtain sufficient underwater information directly from underwater images taken by underwater robots. To solve this problem, more and more scholars have joined the field of underwater image enhancement, committed to designing an efficient and robust underwater image enhancement method to improve the quality of underwater images and enhance the environmental perception ability of underwater robots.

[0004] Underwater image enhancement technology can be divided into three categories, including methods based on non-physical models, methods based on physical models, and methods based on data-driven. Among them, the non-physical model method mainly enhances the image pixel by pixel, including methods based on fusion, Retinex theory, and color constancy. Such methods ignore the underwater degradation process, and their enhancement results are often prone to oversaturation or undersaturation. The physical model-based method mainly analyzes the underwater imaging process, constructs an underwater imaging model, and infers a clear underwater image by solving related parameters. Such methods are generally based on specific assumptions and priors. Due to the complexity of the underwater imaging process, specific assumptions and priors are not always valid, resulting in new color deviations in the enhancement results. Data-driven methods mainly include supervised methods, semi-supervised methods, and unsupervised methods. Among them, supervised underwater image enhancement methods mainly establish a mapping relationship from underwater images to their clear images, and train on a rich and diverse paired data set. However, high-quality images corresponding to underwater images are often unavailable, so such methods often use synthetic data sets or pseudo labels, and their generalization ability is usually poor. Semi-supervised underwater image enhancement methods usually use a teacher-student model, guide the training of the student model through unsupervised teacher-student consistency loss and contrast loss, and update the weight of the student model by minimizing supervised loss and unsupervised loss. Since a small amount of labeled data is still unavailable, such methods cannot completely eliminate the need for real high-quality underwater images. Unsupervised learning-based methods mainly improve image quality by automatically discovering and utilizing the intrinsic structure and patterns in underwater images. Common technologies include autoencoders, generative adversarial networks, etc. Since such methods do not require labels, the stability of the model and the quality of the generated images are challenging.

[0005] In recent years, the denoising diffusion probability model has emerged as a new generative framework. It consists of two processes, namely the forward diffusion process and the reverse diffusion process. The forward process refers to gradually adding Gaussian noise to the image, making it more and more blurred and random, until the image is close to an isotropic Gaussian distribution; the reverse process is to remove noise from the random Gaussian distribution and reconstruct the original data. Based on the strong theoretical foundation of the diffusion model, more and more scholars have used the diffusion model to enhance underwater images, such as combining the diffusion model with wavelet transform, preprocessing the image as the conditional input of the diffusion model, etc. Compared with the generative adversarial network, this method is more stable in training, but still requires a large paired data set for training, and the generalization performance of the network still needs to be improved. Summary of the invention

[0006] In view of the limitation that existing underwater image enhancement methods require paired high-quality data for model training, the present invention provides an underwater image enhancement method based on an unsupervised diffusion model. By combining a generative adversarial network with a diffusion model, the present invention can be trained in real underwater scenes without reference images, effectively alleviating the problems of color cast and low contrast in real underwater images.

[0007] The technical means adopted by the present invention are as follows:

[0008] An underwater image enhancement method based on an unsupervised diffusion model comprises the following steps:

[0009] Step 1: Obtain an underwater image training set, wherein the underwater image training set includes underwater degraded images and unpaired underwater clear images;

[0010] Step 2: construct a pre-enhancement network based on adversarial contrastive learning, wherein the pre-enhancement network based on adversarial contrastive learning is used to convert the degraded underwater image domain into the clear underwater image domain; the pre-enhancement network based on adversarial contrastive learning includes a generator and a discriminator, wherein the generator is used to process the underwater degraded image to generate a pre-enhancement result on the one hand, and to process the underwater clear image to generate a consistency result on the other hand; the discriminator is used to determine whether a given image is a real clear image or a pseudo clear image generated by the generator;

[0011] Step 3: Use the underwater image training set to train a pre-enhancement network based on adversarial contrastive learning;

[0012] Step 4: construct a wavelet transform-based conditional denoising diffusion probability model, wherein the input data of the wavelet transform-based conditional denoising diffusion probability model are the degraded underwater image and the pre-enhanced underwater image; the wavelet transform-based conditional denoising diffusion probability model is configured to perform the following steps:

[0013] Firstly, the degraded underwater image and the pre-enhanced underwater image are subjected to K-times two-dimensional discrete wavelet transform to obtain the low-frequency subband and high-frequency subband in the wavelet domain.

[0014] Secondly, the low-frequency subband in the wavelet domain is subjected to diffusion processing, and the diffusion processing includes a forward diffusion process and a reverse diffusion process. The forward diffusion process is set to gradually add Gaussian noise to the real data sample until it becomes a Markov chain of random noise, and the reverse diffusion process is set to gradually denoise the random noise until a real sample is generated.

[0015] At the same time, the high-frequency sub-band after the K-th wavelet transform is subjected to detail enhancement processing through the high-frequency enhancement module. The high-frequency enhancement module extracts the features of the high-frequency sub-band through a deep separable convolution operation, calculates the cross attention between the features and performs concatenation to obtain complementary features, performs a hole convolution on the feature map of the sub-band to perform better detail enhancement, and performs a deep separable convolution operation on the feature map after detail enhancement to obtain a reconstructed high-frequency sub-band.

[0016] Finally, the processing results in the wavelet domain are subjected to K-times two-dimensional inverse discrete wavelet transform to generate underwater image enhancement results;

[0017] Step 5: Using the degraded underwater images in the underwater image training set and the pre-enhanced images output by the pre-enhancement network based on adversarial contrastive learning, a conditional denoising diffusion probability model based on wavelet transform is trained;

[0018] Step 6: Input the real underwater image to be processed into the trained pre-enhancement network based on adversarial contrastive learning to obtain the pre-enhanced image, and then input the pre-enhanced image into the trained conditional denoising diffusion probability model based on wavelet transform to obtain the final underwater enhanced image.

[0019] Furthermore, the loss function of the pre-enhancement network based on adversarial contrastive learning is obtained according to the following calculation:

[0020]

[0021] Among them, λ1, λ2 and λ3 represent the weight coefficients of balancing various loss functions, represents the least squares GAN loss of the generator G, represents the block-wise multi-layer contrast loss of the constructed pre-enhancement network, represents the consistency loss of the generator G.

[0022] Furthermore, the least squares GAN loss representing the generator G is obtained according to the following calculation:

[0023]

[0024] in, is the objective function used to optimize the generator, is the objective function used to optimize the discriminator. G(x) represents the result produced by the generator. D(G(x)) represents the probability that the discriminator D determines whether the result generated by the generator G is true. represents the expected value of the distribution function, P fake represents the distribution of the results of the degraded underwater image generated by the generator, D(y) represents the probability that the discriminator D identifies the clear underwater image as real data, P realrepresents the distribution of the results of clear underwater images generated by the generator, k2=1 represents the label of real data, k3=0 represents the label of fake data, and k1=1 represents the value that the generator hopes the discriminator believes in the fake data.

[0025] Furthermore, the consistency loss of the generator G Obtained according to the following calculation:

[0026]

[0027] in, represents the expected value of the distribution function, P y represents the distribution of clear underwater images, G(y) represents the result of clear underwater images generated by generator G, and y represents a clear image.

[0028] Furthermore, the block-wise multi-layer contrast loss of the constructed pre-enhancement network Obtained according to the following calculation:

[0029]

[0030] Among them, G represents the generator, M represents the two-layer MLP network, X represents the input degraded underwater image, and L is the ec The number of layers selected in , l represents the specific lth layer, S l refers to the number of spatial locations in each layer, s∈{1,...S l} represents the spatial position in the selected layer, Represents the cross entropy loss, which is calculated as follows:

[0031]

[0032] in, Indicates the use of encoder G ec The feature stack obtained by extracting features from the degraded image domain X with a two-layer MLP network M, For each selected query code, the corresponding positive sample feature is recorded as The negative sample feature is recorded as C l is the number of channels in each layer, sim(u,v)=u T v / ||u||||v|| represents the cosine similarity between u and v, and τ represents the temperature parameter used to scale the distance between the query and other examples.

[0033] Furthermore, the loss function of the wavelet transform-based conditional denoising diffusion probability model is obtained according to the following calculation:

[0034]

[0035] Among them, θ represents the denoising network model parameters of the reverse diffusion process, G is the generator, M is a two-layer MLP network, κ1, κ2, κ3 represent the weight coefficients of balancing various loss functions, represents the pre-enhancement result, represents the final prediction result, represents the diffusion loss function, represents the content loss function, Represents the block-wise multi-layer contrastive loss of the diffusion probabilistic model for conditional denoising based on wavelet transform.

[0036] Furthermore, the diffusion loss function Obtained according to the following calculation:

[0037]

[0038] in, represents the expectation of the distribution function, ∈ t represents the Gaussian noise added during the forward diffusion process, Represents the noise map predicted by the reverse diffusion process, x t represents the image after adding Gaussian noise at the tth time step, It represents the low-frequency subband of the input degraded underwater image after discrete wavelet transform, and t represents the time step of the diffusion process.

[0039] Furthermore, the content loss function Obtained according to the following calculation:

[0040]

[0041] in, represents the pre-enhancement result, represents the final prediction result, express and The structural similarity measure.

[0042] Furthermore, the block-wise multi-layer contrast loss of the wavelet transform-based conditional denoising diffusion probability model Obtained according to the following calculation:

[0043]

[0044] Among them, G represents the generator, M represents the two-layer MLP network, Represents the result of the pre-enhancement network based on adversarial contrast learning, L is from the encoder G ec The number of layers selected in , l represents the specific lth layer, S l refers to the number of spatial locations in each layer, s∈{1,...Sl} represents the spatial position in the selected layer, Represents the cross entropy loss, which is calculated as follows:

[0045]

[0046] in, Indicates the use of encoder G ec and a two-layer MLP network M from the pre-enhanced image domain The feature stack obtained by extracting features from is a query code selected each time, then its corresponding positive sample feature is recorded as The negative sample feature is recorded as C l is the number of channels in each layer, sim(u,v)=u T v / ||u||||v|| represents the cosine similarity between u and v, and τ represents the temperature parameter used to scale the distance between the query and other examples.

[0047] Compared with the prior art, the present invention has the following advantages:

[0048] The present invention proposes a real-world underwater image enhancement method based on an unsupervised diffusion model, which can effectively correct color, improve contrast, and enhance details. Specifically:

[0049] The present invention proposes a bidirectional optimization method combining a generative adversarial network and a conditional diffusion probability model for training without underwater real reference images, wherein the generative adversarial network provides a preliminary optimization of the structural image and unknown degradation factors for the diffusion process, and the reverse diffusion process further bridges the gap between the degraded image and the normal image;

[0050] At the same time, in order to guide better enhancement, a block contrast loss is introduced between the pre-enhancement result and the input image, and between the enhancement result and the pre-enhancement result, making full use of the advantages of the learned domain-related feature distribution and learning image restoration in the deep feature space through contrast constraints;

[0051] In addition, considering that wavelet transform can decompose images into sub-bands of different frequencies, a conditional denoising diffusion model is designed for low-frequency sub-bands to enhance the contrast of the image, and a high-frequency enhancement module is designed for high-frequency sub-bands to enrich the details of the image. BRIEF DESCRIPTION OF THE DRAWINGS

[0052] In order to more clearly illustrate the embodiments of the present invention or the technical solutions in the prior art, the drawings required for use in the embodiments or the description of the prior art will be briefly introduced below. Obviously, the drawings described below are some embodiments of the present invention. For ordinary technicians in this field, other drawings can be obtained based on these drawings without paying creative labor.

[0053] Figure 1 A network architecture diagram for implementing an underwater image enhancement method based on an unsupervised diffusion model in an embodiment of the present invention.

[0054] Figure 2 is a residual block structure diagram in an embodiment of the present invention.

[0055] Figure 3 FIG. 4 is a structural diagram of a high frequency enhancement module in an embodiment of the present invention.

[0056] Figure 4 This is a graph of comparative experimental results of the present invention in the embodiment with 12 existing enhancement methods on the UIEB-challenging60 dataset, mainly including 2 methods based on non-physical models (ACDC method, HLRP method), 2 methods based on physical models (GDCP method, UDCP method), two methods based on supervised convolutional neural networks (Ucolor method and PUIE method), two methods based on contrastive learning (TACL method, CWR method), two methods based on Transformer (U-shape method, URSCT method), and two methods based on diffusion models (DM_Underwater method, DiffWater method). DETAILED DESCRIPTION

[0057] In order to enable those skilled in the art to better understand the scheme of the present invention, the technical scheme in the embodiments of the present invention will be clearly and completely described below in conjunction with the drawings in the embodiments of the present invention. Obviously, the described embodiments are only part of the embodiments of the present invention, not all of the embodiments. Based on the embodiments of the present invention, all other embodiments obtained by ordinary technicians in this field without creative work should fall within the scope of protection of the present invention.

[0058] The present invention provides 1. an underwater image enhancement method based on an unsupervised diffusion model, characterized in that it comprises the following steps:

[0059] Step 1: Obtain an underwater image training set, where the underwater image training set includes underwater degraded images and unpaired underwater clear images.

[0060] Specifically, the present invention first defines two data fields: and Construct underwater image training set, validation set and test set, where the training set consists of instance set X containing 1000 real degraded underwater images x and instance set Y containing 1000 unpaired high-definition images y. The validation set contains 90 real underwater degraded images, and the test set is all underwater datasets that need to be tested.

[0061] Step 2: construct a pre-enhancement network based on adversarial contrastive learning, wherein the pre-enhancement network based on adversarial contrastive learning is used to convert the underwater degraded image domain into the underwater clear image domain; the pre-enhancement network based on adversarial contrastive learning includes a generator and a discriminator, wherein the generator is used to process the underwater degraded image to generate a pre-enhancement result on the one hand, and to process the underwater clear image to generate a consistency result on the other hand; the discriminator is used to determine whether a given image is a real clear image or a pseudo clear image generated by the generator.

[0062] Specifically, constructing a pre-enhancement network based on adversarial contrastive learning is to regard image pre-enhancement as an image-to-image conversion task, and then convert it into a mapping function for learning underwater degraded images to underwater high-definition images, such as Figure 1 As shown:

[0063] The overall architecture of the pre-enhancement network consists of a generator and a discriminator. The generator adopts a U-shaped architecture, that is, the spatial dimension of the extracted feature map is first reduced by downsampling technology, then the deep features are enhanced, and finally the spatial dimension of the feature map is restored by upsampling technology. It can capture global dependencies and local context dependencies at the same time. The first half is represented by the encoder G ec , the second half is represented by the decoder G dc The encoder consists of 3 convolutional layers with a kernel size of 3×3 and a stride of 2, and 5 residual blocks, which are used to downsample the degraded underwater image and encode it into a low-resolution feature map. The decoder consists of 4 residual blocks and 3 transposed convolutions with a kernel size of 3×3 and a stride of 2, which upsample the image and restore the final pre-enhanced image. The 9 residual blocks in the middle have the same structure, as shown in Figure 2 As shown in Figure 1, it is used to extract more complex and deeper features in the low-resolution feature space and perform feature enhancement at the same time. The discriminator uses a Markov discriminator, which maps the input image into a matrix. Each pixel in the matrix represents the probability that each patch is a true sample. The average of the matrix is ​​the final output of the discriminator.

[0064] The unpaired real degraded underwater image set X and high-definition image set Y are input to the generator G, and the mapping relationship between the degraded underwater image domain and the high-definition image domain is learned. That is, G:X→Y, a clearer pre-enhancement result is obtained. and consistency results The process is expressed as:

[0065]

[0066] At the same time, the discriminator D is introduced to determine whether the given image is a real high-definition image or a pseudo high-definition image generated by the generator, thereby guiding the generator to generate more realistic high-definition underwater images.

[0067] Step 3: Train the pre-enhancement network based on adversarial contrastive learning using the degraded underwater images and clear images in the underwater image training set.

[0068] Specifically, the loss function of the pre-enhancement network based on adversarial contrastive learning is constructed, including adversarial loss, block contrast loss and consistency loss, as follows:

[0069] Step 3-1: In order to perform more stable training, a least squares GAN loss function is designed during the training process, which is defined as follows:

[0070]

[0071] Among them, k2=1 and k3=0 are the labels of real data and fake data respectively, and k1=1 means that G hopes D to believe the value of fake data.

[0072] Step 3-2: In order to maximize the mutual information between the input and output corresponding blocks, a block contrast loss is designed. This application uses the encoder G ec and a two-layer MLP network M extracts features from the X domain and embeds them into the feature stack Where L is from G ec The number of layers selected in , l is the specific lth layer. These feature stacks represent different patches in the image, and the spatial positions in the selected layer are represented as s∈{1,...S l}, where S l Refers to the number of spatial positions in each layer. Each time a query code is sampled from the X domain, its feature is represented as The corresponding positive sample feature is expressed as The negative sample feature is expressed as Among them C l is the number of channels in each layer. Finally, the block-wise multi-layer contrast loss of the constructed pre-enhancement network is defined as:

[0073]

[0074] in, Represents the cross entropy loss, which is calculated as follows:

[0075]

[0076] Where sim(u,v)=u T v / ||u||||v|| represents the cosine similarity between u and v, τ represents the temperature parameter, which is used to scale the distance between the query and other examples, and is set to 0.07 in the present invention. At the same time, the number of negative samples N is set to 255.

[0077] Step 3-3: In order to prevent the generator from keeping the structure and color composition the same after domain conversion, a consistency loss is added, which is defined as follows:

[0078]

[0079] Step 3-4: The overall loss function of the pre-enhancement network based on adversarial contrastive learning is:

[0080]

[0081] Among them, λ1, λ2, λ3 are set to 1.0, 1.0 and 5.0 respectively in the present invention.

[0082] Step 4: construct a wavelet transform-based conditional denoising diffusion probability model, wherein the input data of the wavelet transform-based conditional denoising diffusion probability model are the degraded underwater image and the pre-enhanced underwater image; the wavelet transform-based conditional denoising diffusion probability model is configured to perform the following steps:

[0083] Step 4-1: For the input degraded underwater image And the pre-enhanced underwater image Perform K two-dimensional discrete wavelet transforms simultaneously and convert them into wavelet domain. The process is as follows:

[0084]

[0085] Among them, the input image can be expressed as Represents the low-frequency subband of the image converted to the wavelet domain, Represents the high frequency sub-band.

[0086] Step 4-2: Diffusion processing is performed on the low-frequency subband in the wavelet domain. It mainly includes the forward diffusion process and the reverse diffusion process, as follows:

[0087] The low-frequency subbands of the input image and the pre-enhanced image are represented as and x0, where the forward diffusion process can be viewed as a Markov chain that starts from a real data sample and gradually adds Gaussian noise to it until it becomes random noise, as shown below:

[0088]

[0089] in, is a Gaussian distribution, I is the identity covariance matrix with the same dimension as the input image x0, β t is a variable that controls the noise variance, and let α t =1-β t , then:

[0090]

[0091] Then, by merging and deriving the Gaussian distribution of formula (8), it can be further expressed as:

[0092]

[0093] in, represents the noise added at t time steps.

[0094] For the reverse diffusion process, the model starts with random noise and gradually removes the noise until a real sample is generated. Each step of the reverse process can be described by a conditional probability distribution:

[0095]

[0096] in, refers to the conditional image, which is the low-frequency subband obtained after wavelet transforming the input image. and They refer to the mean and variance of the time step t estimate respectively. By deriving formula 11, we can get the mean and variance as follows:

[0097]

[0098] in, Represents the noise estimated by U-net during the reverse diffusion process.

[0099] Step 4-3: Construct a high frequency enhancement module (HPEM) to perform detail enhancement on the high frequency subband after the Kth wavelet transform, such as Figure 3 First, the three high-frequency sub-bands are subjected to depth-separable convolution operations to extract the corresponding features. The features of as well as The cross attention is connected in series to obtain the complementary features; then, the feature maps of the three sub-bands are respectively subjected to the dilated convolution with a kernel size of 3×3. In the present invention, the dilation rate is set to {1, 2, 3, 2, 1} to increase the receptive field while keeping the feature scale unchanged, so as to perform better detail enhancement; finally, the feature maps after detail enhancement are respectively subjected to the depthwise separable convolution operation to obtain the reconstruction result and

[0100] Step 4-4: Perform K-times two-dimensional inverse discrete wavelet transform on the processing result in the wavelet domain. The specific process is as follows:

[0101]

[0102] in, This is the final prediction result.

[0103] Step 5: Train a wavelet transform-based conditional denoising diffusion probability model using the degraded underwater images in the underwater image training set and the pre-enhanced images output by the pre-enhancement network based on adversarial contrastive learning.

[0104] Specifically, the loss function of the conditional denoising diffusion probability model based on wavelet transform is constructed, including diffusion loss, content loss and block contrast loss, as follows:

[0105] The diffusion loss function is designed for the noise estimated by the network and the Gaussian noise added in the forward diffusion process as follows:

[0106]

[0107] In addition, the present invention also designs a content loss function combining L1 loss and SSIM loss to minimize the difference between the prediction result of the diffusion model and the pre-enhancement result, as follows:

[0108]

[0109] in, is the pre-enhancement result, is the final prediction result.

[0110] Finally, in order to maximize the mutual information between the pre-enhancement result and the final prediction result, this application also uses block contrast loss for the two. We use the encoder G ec And the two-layer MLP network M from domain and embed them into the feature stack Where L is from G ec The number of layers selected in , l is the specific lth layer. These feature stacks represent different patches in the image, and we denote the spatial position in the selected layer as s∈{1,...S l}, where S l Refers to the number of spatial locations in each layer. A query code is sampled from the domain and its features are represented as Among them C l is the number of channels in each layer, then the corresponding positive sample feature is expressed as The negative sample feature is expressed as Finally, the block-by-block multi-layer contrast loss of the constructed wavelet transform-based conditional denoising diffusion probability model is defined as:

[0111]

[0112] In summary, the total loss of the conditional denoising diffusion probability model based on wavelet transform is as follows:

[0113]

[0114] Among them, κ1, κ2, and κ3 are set to 1.0, 1.0, and 0.1 respectively in the present invention.

[0115] Step 6: Input the real underwater image to be processed into the trained pre-enhancement network based on adversarial contrastive learning to obtain the pre-enhanced image, and then input the pre-enhanced image into the trained conditional denoising diffusion probability model based on wavelet transform to obtain the final underwater enhanced image.

[0116] like Figure 4 As shown in the figure, the application effect comparison between the method of the present invention and the existing method is given. As can be seen from the figure, the GDCP method and the DiffWater method cannot correct the color of the image, the ACDC method and the CWR method will reduce the color richness in the enhancement result, the brightness of the enhancement result of the UDCP method is reduced, the enhancement result of the HLRP method is reddish, and the enhancement results of the Ucolor, PUIE-mc, TACL and U-shape methods still have the problem of low contrast. In comparison, the enhancement results of the URSCT and DM_underwater methods are better, but in terms of backscatter removal, they are still not as good as the method of the present invention.

[0117] As shown in Table 1, the objective comparison experimental results of the present invention in the embodiment with 12 existing enhancement methods on the UIEB-challenging60 data set are shown, which mainly include 8 evaluation indicators, namely: e (indicates the ratio of the number of newly added visible edges to the number of original visible edges in the enhanced image, the larger the value, the richer the details), r (indicates the enhancement of the gradient value of edge pixels in the image, the larger the value, the higher the contrast), blur (indicates the blur degree of the enhancement result, the smaller the value, the clearer the details), ENTROPY (reflects the amount of information in the image, the larger the value, the more information), UIQM (used to evaluate the color, clarity and contrast of underwater images, the larger the value, the clearer the underwater image), UCIQE (used to evaluate the color, brightness and saturation of underwater images, the larger the value, the better the enhancement effect), NIQE (uses distance to measure the quality of the image, the smaller the value, the higher the image quality), FADE (used to perceive the fog density in the image, the smaller the value, the cleaner the image and the higher the clarity). Among them, the optimal result of each indicator is bolded, and the suboptimal result is underlined. It can be seen from the table that the method proposed in the present invention achieves the best performance in 5 evaluation indicators and performs second best in the other 3 indicators, which fully proves the superiority of the proposed method.

[0118] Table 1 Comparison of application effects of the method of the present invention and the existing method

[0119]

[0120]

[0121] Finally, 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 aforementioned embodiments, those skilled in the art should understand that they can still modify the technical solutions described in the aforementioned embodiments, or replace some or all of the technical features therein with equivalents. However, these modifications or replacements do not cause the essence of the corresponding technical solutions to deviate from the scope of the technical solutions of the embodiments of the present invention.

Claims

1. An underwater image enhancement method based on an unsupervised diffusion model, characterized in that: The following steps are involved: Step 1: Obtain an underwater image training set, wherein the underwater image training set includes underwater degraded images and unpaired underwater clear images; Step 2: construct a pre-enhancement network based on adversarial contrastive learning, wherein the pre-enhancement network based on adversarial contrastive learning is used to convert the degraded underwater image domain into the clear underwater image domain; the pre-enhancement network based on adversarial contrastive learning includes a generator and a discriminator, wherein the generator is used to process the underwater degraded image to generate a pre-enhancement result on the one hand, and to process the underwater clear image to generate a consistency result on the other hand; the discriminator is used to determine whether a given image is a real clear image or a pseudo clear image generated by the generator; Step 3: Use the underwater image training set to train a pre-enhancement network based on adversarial contrastive learning; Step 4: construct a wavelet transform-based conditional denoising diffusion probability model, wherein the input data of the wavelet transform-based conditional denoising diffusion probability model are the degraded underwater image and the pre-enhanced underwater image; the wavelet transform-based conditional denoising diffusion probability model is configured to perform the following steps: Firstly, the degraded underwater image and the pre-enhanced underwater image are subjected to K-times two-dimensional discrete wavelet transform to obtain the low-frequency subband and high-frequency subband in the wavelet domain. Secondly, the low-frequency subband in the wavelet domain is subjected to diffusion processing, and the diffusion processing includes a forward diffusion process and a reverse diffusion process. The forward diffusion process is set to gradually add Gaussian noise to the real data sample until it becomes a Markov chain of random noise, and the reverse diffusion process is set to gradually denoise the random noise until a real sample is generated. At the same time, the high-frequency sub-band after the K-th wavelet transform is subjected to detail enhancement processing through the high-frequency enhancement module. The high-frequency enhancement module extracts the features of the high-frequency sub-band through a deep separable convolution operation, calculates the cross attention between the features and performs concatenation to obtain complementary features, performs a hole convolution on the feature map of the sub-band to perform better detail enhancement, and performs a deep separable convolution operation on the feature map after detail enhancement to obtain a reconstructed high-frequency sub-band. Finally, the processing results in the wavelet domain are subjected to K-times two-dimensional inverse discrete wavelet transform to generate underwater image enhancement results; Step 5: Using the degraded underwater images in the underwater image training set and the pre-enhanced images output by the pre-enhancement network based on adversarial contrastive learning, a conditional denoising diffusion probability model based on wavelet transform is trained; Step 6: Input the real underwater image to be processed into the trained pre-enhancement network based on adversarial contrastive learning to obtain the pre-enhanced image, and then input the pre-enhanced image into the trained conditional denoising diffusion probability model based on wavelet transform to obtain the final underwater enhanced image.

2. The underwater image enhancement method based on the unsupervised diffusion model according to claim 1, characterized in that: The loss function of the pre-enhancement network based on adversarial contrastive learning is obtained according to the following calculation: Among them, λ1, λ2 and λ3 represent the weight coefficients of balancing various loss functions, represents the least squares GAN loss of the generator G, represents the block-wise multi-layer contrast loss of the constructed pre-enhancement network, represents the consistency loss of the generator G.

3. The underwater image enhancement method based on the unsupervised diffusion model according to claim 2, characterized in that: The least squares GAN loss representing the generator G is obtained according to the following calculation: in, is the objective function used to optimize the generator, is the objective function used to optimize the discriminator. G(x) represents the result produced by the generator. D(G(x)) represents the probability that the discriminator D determines whether the result generated by the generator G is true. represents the expected value of the distribution function, P fake represents the distribution of the results of the degraded underwater image generated by the generator, D(y) represents the probability that the discriminator D identifies the clear underwater image as real data, P real represents the distribution of the results of clear underwater images generated by the generator, k2=1 represents the label of real data, k3=0 represents the label of fake data, and k1=1 represents the value that the generator hopes the discriminator believes in the fake data.

4. The underwater image enhancement method based on the unsupervised diffusion model according to claim 2, characterized in that: Consistency loss of generator G Obtained according to the following calculation: in, represents the expected value of the distribution function, P y represents the distribution of clear underwater images, G(y) represents the result of clear underwater images generated by generator G, and y represents a clear image.

5. The underwater image enhancement method based on unsupervised diffusion model according to claim 2, characterized in that: Block-wise multi-layer contrast loss of the constructed pre-enhancement network Obtained according to the following calculation: Among them, G represents the generator, M represents the two-layer MLP network, X represents the input degraded underwater image, and L is the ec The number of layers selected in , l represents the specific lth layer, S l refers to the number of spatial locations in each layer, s∈{1,...S l } represents the spatial position in the selected layer, Represents the cross entropy loss, which is calculated as follows: in, Indicates the use of encoder G ec The feature stack obtained by extracting features from the degraded image domain X with a two-layer MLP network M, For each selected query code, the corresponding positive sample feature is recorded as The negative sample feature is recorded as C l is the number of channels in each layer, sim(u,v)=u T v / ||u||||v|| represents the cosine similarity between u and v, and τ represents the temperature parameter used to scale the distance between the query and other examples.

6. The underwater image enhancement method based on unsupervised diffusion model according to claim 1, characterized in that: The loss function of the wavelet transform-based conditional denoising diffusion probability model is obtained according to the following calculation: Among them, θ represents the denoising network model parameters of the reverse diffusion process, G is the generator, M is a two-layer MLP network, κ1, κ2, κ3 represent the weight coefficients of balancing various loss functions, represents the pre-enhancement result, represents the final prediction result, represents the diffusion loss function, represents the content loss function, Represents the block-wise multi-layer contrastive loss of the diffusion probabilistic model for conditional denoising based on wavelet transform.

7. The underwater image enhancement method based on unsupervised diffusion model according to claim 6, characterized in that: The diffusion loss function Obtained according to the following calculation: in, represents the expectation of the distribution function, ∈ t represents the Gaussian noise added during the forward diffusion process, Represents the noise map predicted by the reverse diffusion process, x t represents the image after adding Gaussian noise at the tth time step, It represents the low-frequency subband of the input degraded underwater image after discrete wavelet transform, and t represents the time step of the diffusion process.

8. The underwater image enhancement method based on unsupervised diffusion model according to claim 6, characterized in that: The content loss function Obtained according to the following calculation: in, represents the pre-enhancement result, represents the final prediction result, express and The structural similarity measure.

9. The underwater image enhancement method based on unsupervised diffusion model according to claim 6, characterized in that: Block-by-block multi-layer contrast loss for wavelet transform-based conditional denoising diffusion probability model Obtained according to the following calculation: Among them, G represents the generator, M represents the two-layer MLP network, Represents the result of the pre-enhancement network based on adversarial contrast learning, L is from the encoder G ec The number of layers selected in , l represents the specific lth layer, S l refers to the number of spatial locations in each layer, s∈{1,...S l } represents the spatial position in the selected layer, Represents the cross entropy loss, which is calculated as follows: in, Indicates the use of encoder G ec and a two-layer MLP network M from the pre-enhanced image domain The feature stack obtained by extracting features from is a query code selected each time, then its corresponding positive sample feature is recorded as The negative sample feature is recorded as C l is the number of channels in each layer, sim(u,v)=u T v / ||u||||v|| represents the cosine similarity between u and v, and τ represents the temperature parameter used to scale the distance between the query and other examples.

Images