An underwater image enhancement method based on an unsupervised diffusion model

Abstract

The application discloses an underwater image enhancement method based on an unsupervised diffusion model, and comprises the following steps: acquiring an unpaired underwater image training set; inputting an underwater image to be processed into a trained diffusion model based on unsupervised contrast learning, wherein the diffusion model comprises a pre-enhancement network based on adversarial contrast learning and a conditional denoising diffusion probability model based on discrete wavelet transform; 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 subband diffusion process and a high-frequency subband enhancement module; and the output of the conditional denoising diffusion probability model based on wavelet transform is acquired as an enhanced underwater image. The application combines the generative adversarial network and the diffusion model, can train in a real underwater scene without a reference image, effectively alleviates the problems of color deviation and low contrast in the real underwater image, and achieves an enhancement effect comparable to that of supervised training.

Description

Technical Field

[0001] This invention relates to the field of computer vision technology, and more specifically, to an underwater image enhancement method based on an unsupervised diffusion model. Background Technology

[0002] Underwater operations such as underwater biological detection and marine ranching place high demands on the environmental perception accuracy of underwater robots. Underwater images, as the carrier of marine environmental perception, determine the strength of an underwater robot's environmental perception capabilities. Therefore, underwater image enhancement technology has always been one of the key research directions for researchers exploring the mysteries of the ocean.

[0003] However, abundant marine life leads to turbid water, and the properties of water cause severe and wavelength-dependent absorption and attenuation of light. Light refraction caused by suspended particles in the water complicates the underwater imaging process. Therefore, obtaining sufficient underwater information directly from underwater images captured by underwater robots is quite difficult. To address this challenge, an increasing number of researchers have joined the field of underwater image enhancement, dedicated to designing an efficient and robust underwater image enhancement method to improve the quality of underwater images and enhance the environmental perception capabilities of underwater robots.

[0004] Underwater image enhancement techniques can be broadly categorized into three types: non-physical model-based methods, physical model-based methods, and data-driven methods. Non-physical model-based methods primarily enhance images pixel-by-pixel, including fusion-based methods, methods based on Retinex theory, and color constancy. These methods ignore underwater degradation processes, and their enhancement results often exhibit oversaturation or undersaturation. Physical model-based methods analyze the underwater imaging process, construct underwater imaging models, and derive clear underwater images by solving relevant parameters. These methods are generally based on specific assumptions and priors. Due to the complexity of the underwater imaging process, these assumptions and priors are not always valid, often leading to new color casts in the enhancement results. Data-driven methods include supervised, semi-supervised, and unsupervised methods. Supervised underwater image enhancement methods establish mappings from underwater images to their clearer counterparts, training on diverse paired datasets. However, high-quality images corresponding to underwater images are often unavailable, so these methods often use synthetic datasets or pseudo-labels, resulting in generally poor generalization ability. Semi-supervised underwater image enhancement methods typically employ a teacher-student model, guiding the training of the student model through unsupervised teacher-student consistency loss and contrast loss. The student model's weights are updated by minimizing both supervised and unsupervised losses. However, since a limited amount of labeled data remains unavailable, these methods cannot completely eliminate the need for truly high-quality underwater images. Unsupervised learning-based methods primarily improve image quality by automatically discovering and utilizing the inherent structures and patterns within underwater images. Common techniques include autoencoders and generative adversarial networks. These methods, however, face challenges in terms of model stability and the quality of generated images because they do not require labels.

[0005] In recent years, the denoising diffusion probabilistic model has emerged as a novel generative framework. It comprises two processes: forward diffusion and backward diffusion. The forward process involves gradually adding Gaussian noise to the image, making it increasingly blurred and random until the image approximates an isotropic Gaussian distribution. The backward process removes noise from the random Gaussian distribution, reconstructing the original data. Based on the strong theoretical foundation of diffusion models, more and more researchers are using them to enhance underwater images. For example, combining diffusion models with wavelet transforms to preprocess images as conditional inputs to the diffusion model. Compared to generative adversarial networks (GANs), these methods offer more stable training, but still require large paired datasets for training, and the network's generalization performance still needs improvement. Summary of the Invention

[0006] Given the limitation of existing underwater image enhancement methods that require paired high-quality data for model training, this invention provides an underwater image enhancement method based on an unsupervised diffusion model. By combining generative adversarial networks (GANs) and diffusion models, this invention enables training in real underwater scenes without reference images, effectively mitigating problems such as color cast and low contrast in real underwater images.

[0007] The technical means employed in this invention are as follows:

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

[0009] Step 1: Obtain an underwater image training set, which includes underwater degradation images and unpaired clear underwater images;

[0010] Step 2: Construct a pre-enhancement network based on adversarial contrastive learning. This network is used to transform the degraded underwater image domain into a clear underwater image domain. The pre-enhancement network includes a generator and a discriminator. The generator processes the degraded underwater image to generate a pre-enhancement result and processes the clear underwater image to generate a consistent result. The discriminator determines whether a given image is a genuine clear image or a pseudo-clear image generated by the generator.

[0011] Step 3: Train a pre-enhanced network based on adversarial contrastive learning using the underwater image training set;

[0012] Step 4: Construct a conditional denoising and diffusion probability model based on wavelet transform. The input data for the conditional denoising and diffusion probability model based on wavelet transform are degraded underwater images and pre-enhanced underwater images. The conditional denoising and diffusion probability model based on wavelet transform is configured to perform the following steps:

[0013] First, K-order two-dimensional discrete wavelet transforms were performed on the degraded underwater image and the pre-enhanced underwater image to obtain the low-frequency and high-frequency subbands in the wavelet domain.

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

[0015] Simultaneously, a high-frequency enhancement module performs detail enhancement processing on the high-frequency sub-band after the Kth wavelet transform. This module extracts features from the high-frequency sub-band through depthwise separable convolution, calculates cross-attention between features, and concatenates them to obtain complementary features. The feature map of the sub-band is then subjected to dilated convolution for better detail enhancement. Finally, depthwise separable convolution is performed on the detail-enhanced feature map to obtain the reconstructed high-frequency sub-band.

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

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

[0018] Step 6: Input the real underwater image to be processed into the pre-enhancement network based on adversarial contrastive learning after training to obtain the pre-enhancement image. Then input the pre-enhancement image into the conditional denoising diffusion probability model based on wavelet transform after training 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] Where λ1, λ2, and λ3 represent the weighting coefficients of the loss functions balancing the various terms. This represents the least-squares GAN loss of the generator G. This represents the block-based multi-layer contrastive loss of the constructed pre-enhanced network. This represents the consistency loss of generator G.

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

[0023]

[0024] in, It is the objective function used to optimize the generator. This is the objective function used to optimize the discriminator. G(x) represents the result produced by the generator, and D(G(x)) represents the probability that the discriminator D judges whether the result generated by the generator G is true. P represents the expected value of the distribution function. fake Let D(y) represent the distribution of the resulting underwater images generated by the generator, and let D(y) represent the probability that the discriminator D classifies a clear underwater image as real data. realThe distribution of the results generated by the generator for clear underwater images is represented by k2=1, k3=0, and k1=1, where k2=1 represents the label of the real data, k3=0 represents the label of the fake data, and k1=1 represents the value that the generator wants the discriminator to believe in.

[0025] Furthermore, the consistency loss of generator G Obtained based on the following calculations:

[0026]

[0027] in, P represents the expected value of the distribution function. y Let G(y) represent the distribution of clear underwater images, where G(y) represents the result of generating clear underwater images through generator G, and y represents the clear image.

[0028] Furthermore, the constructed pre-enhanced network uses a block-based multi-layer contrastive loss. Obtained based on the following calculations:

[0029]

[0030] Where G represents the generator, M represents a two-layer MLP network, X represents the input degraded underwater image, and L is the encoder G. ec The selected layer number, l represents the specific l-th layer, S l The number of spatial locations in each layer, s∈{1,...S} l} represents the spatial location within the selected layer. The cross-entropy loss is represented by the following formula:

[0031]

[0032] in, Indicates the use of encoder G ec The feature stack obtained by extracting features from the degraded image domain X using a two-layer MLP network M. For each selected query code, the corresponding positive sample feature is denoted as . Negative sample features are denoted as C l It 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] Where θ represents the parameters of the denoising network model for the reverse diffusion process, G is the generator, M is a two-layer MLP network, and κ1, κ2, and κ3 represent the weight coefficients balancing the various loss functions. This indicates the pre-enhancement results. This indicates the final prediction result. Represents the diffusion loss function. Represents the content loss function. This represents the block-based multi-level contrast loss of the conditional denoising diffusion probability model based on wavelet transform.

[0036] Furthermore, the diffusion loss function Obtained based on the following calculations:

[0037]

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

[0039] Furthermore, the content loss function Obtained based on the following calculations:

[0040]

[0041] in, This indicates the pre-enhancement results. This indicates the final prediction result. express and A measure of structural similarity.

[0042] Furthermore, a block-based multi-layer contrastive loss based on a wavelet transform-based conditional denoising diffusion probability model is used. Obtained based on the following calculations:

[0043]

[0044] Where G represents the generator, and M represents a two-layer MLP network. L represents the result after processing a pre-enhanced network based on adversarial contrastive learning, where L is derived from the encoder G. ec The selected layer number, l represents the specific l-th layer, S l The number of spatial locations in each layer, s∈{1,...S}l} indicates the spatial location representation within the selected layer. The cross-entropy loss is represented by the following formula:

[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 it. For each selected query code, its corresponding positive sample feature is denoted as . Negative sample features are denoted as C l It 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] This invention proposes a real-world underwater image enhancement method based on an unsupervised diffusion model, which can effectively correct colors, improve contrast, and enhance details. Specifically:

[0049] This invention proposes a bidirectional optimization method combining generative adversarial networks and conditional diffusion probability models for training in the absence of real underwater reference images. The generative adversarial network provides preliminary optimization of the structured image and unknown degradation factors for the diffusion process, while the reverse diffusion process further bridges the gap between the degraded image and the normal image.

[0050] Meanwhile, in order to guide better enhancement, block contrast loss is introduced between the pre-enhancement result and the input image, and between the enhancement result and the pre-enhancement result. This fully utilizes the advantages of the learned domain-related feature distribution and learns image restoration in the deep feature space through contrast constraints.

[0051] Furthermore, considering that wavelet transform can decompose an image into sub-bands of different frequencies, a conditional denoising diffusion model is designed for the low-frequency sub-band to enhance image contrast, and a high-frequency enhancement module is designed for the high-frequency sub-band to enrich image details. Attached Figure Description

[0052] To more clearly illustrate the technical solutions in the embodiments of the present invention or the prior art, the drawings used in the description of the embodiments or the prior art will be briefly introduced below. Obviously, the drawings described below are some embodiments of the present invention. For those skilled in the art, other drawings can be obtained based on these drawings without creative effort.

[0053] Figure 1 This is 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 This is a structural diagram of the residual block in an embodiment of the present invention.

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

[0056] Figure 4 The figure shows the comparative experimental results of the present invention on the UIEB-challenging60 dataset with 12 existing augmentation methods in the embodiments. These mainly include two methods based on non-physical models (ACDC method and HLRP method), two methods based on physical models (GDCP method and UDCP method), two methods based on supervised convolutional neural networks (Ucolor method and PUIE method), two methods based on contrastive learning (TACL method and CWR method), two methods based on Transformer (U-shape method and URSCT method), and two methods based on diffusion models (DM_Underwater method and DiffWater method). Detailed Implementation

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

[0058] This invention provides 1. an underwater image enhancement method based on an unsupervised diffusion model, characterized by comprising the following steps:

[0059] Step 1: Obtain an underwater image training set, which includes underwater degradation images and unpaired clear underwater images.

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

[0061] Step 2: Construct a pre-enhancement network based on adversarial contrastive learning. This pre-enhancement network is used to transform the underwater degraded image domain into the underwater clear image domain. The pre-enhancement network includes a generator and a discriminator. The generator is used to process the underwater degraded image to generate a pre-enhancement result and to process the underwater clear image to generate a consistent result. 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 involves treating image pre-enhancement as an image-to-image transformation task, which is then transformed into learning a mapping function between underwater degradation images and underwater high-resolution images, such as... Figure 1 As shown:

[0063] The pre-enhanced network's overall architecture consists of a generator and a discriminator. The generator employs a U-shaped architecture: it first reduces the spatial dimension of the extracted feature map using downsampling techniques, then enhances deeper features, and finally restores the spatial dimension of the feature map using upsampling techniques. This approach can simultaneously capture global and local contextual dependencies. The first part is represented by the encoder G. ec The latter half represents the decoder G. dc The encoder consists of three convolutional layers with 3×3 kernels and a stride of 2, and five residual blocks. It downsamples the degraded underwater image and encodes it into a low-resolution feature map. The decoder consists of four residual blocks and three transposed convolutional layers with 3×3 kernels and a stride of 2. It upsamples the image and recovers the final pre-enhanced image. The nine residual blocks in the middle have the same structure, as shown below. Figure 2 As shown, this method is used to extract more complex and deeper features in a low-resolution feature space while simultaneously performing feature enhancement. The discriminator employs a Markov discriminator, which maps the input image to a matrix where each pixel represents the probability that each patch is a true sample. The mean of the matrix is ​​the final output of the discriminator.

[0064] An unpaired set of real underwater images showing degradation, X, and a set of high-resolution images, Y, are input into a generator G. The generator learns the mapping relationship between the underwater image domain and the high-resolution image domain, i.e., G:X→Y, to obtain a clearer pre-enhanced result. and consistency results This process can be represented as:

[0065]

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

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

[0068] Specifically, the loss function for constructing the pre-enhancement network based on adversarial contrastive learning includes adversarial loss, block contrastive loss, and consistency loss, as follows:

[0069] Step 3-1: To achieve more stable training, a least-squares GAN loss function is designed during the training process, defined as follows:

[0070]

[0071] Here, k2=1 and k3=0 are the labels for real data and fake data, respectively, and k1=1 indicates that G wants D to believe the value of the fake data.

[0072] Step 3-2: To maximize the mutual information between corresponding input and output blocks, a block contrast loss is designed. This application uses encoder G. ec A two-layer MLP network M extracts features from the X domain and embeds them into a feature stack. Where L is from G ec The selected layer number, l, is the specific l-th layer. These feature stacks represent different blocks in the image, and the spatial location within the selected layer is represented as s∈{1,...S}. l}, where S l This refers to the number of spatial locations in each layer. Each time, a query code is sampled from the X domain, and its characteristics are represented as... Then its corresponding positive sample features are represented as Negative sample features are represented as Where C l This refers to the number of channels in each layer. Ultimately, the block-based multi-layer contrastive loss of the constructed pre-enhanced network is defined as:

[0073]

[0074] in, The cross-entropy loss is represented by the following formula:

[0075]

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

[0077] Step 3-3: To prevent the generator from maintaining the same structure and color composition after domain transformation, a consistency loss is added, defined as follows:

[0078]

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

[0080]

[0081] In this invention, λ1, λ2, and λ3 are set to 1.0, 1.0, and 5.0, respectively.

[0082] Step 4: Construct a conditional denoising and diffusion probability model based on wavelet transform. The input data for the conditional denoising and diffusion probability model based on wavelet transform are degraded underwater images and pre-enhanced underwater images. The conditional denoising and diffusion probability model based on wavelet transform is configured to perform the following steps:

[0083] Step 4-1: For the input underwater image of deteriorated water quality and pre-enhanced underwater images Perform K simultaneous two-dimensional discrete wavelet transforms to convert both into the wavelet domain, as follows:

[0084]

[0085] The input image can be represented as This represents the low-frequency subband of the image converted to the wavelet domain. Represents the high-frequency subband.

[0086] Step 4-2: Perform diffusion processing on the low-frequency subband in the wavelet domain. This mainly includes forward diffusion and backward diffusion processes, as detailed below:

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

[0088]

[0089] in, It is a Gaussian distribution, I is the identity covariance matrix with the same dimension as the input image x0, and β t Let α be the variable controlling the noise variance. t =1-β t Then we have:

[0090]

[0091] Then, by combining and deriving Gaussian distributions, formula (8) can be further expressed as:

[0092]

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

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

[0095]

[0096] in, This refers to the conditional image, which is the low-frequency sub-band obtained after performing wavelet transform on the input image. and These refer to the mean and variance of the time step t estimate, respectively. Through the derivation of Equation 11, the mean and variance can be obtained as follows:

[0097]

[0098] in, This 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 processing on the high-frequency subband after the Kth wavelet transform, such as... Figure 3 As shown. First, depthwise separable convolution operations are performed on the three high-frequency subbands to extract the corresponding features. The features are calculated respectively. as well as The cross-attention mechanism is used to concatenate the features to obtain their complementary features. Then, the feature maps of the three sub-bands are subjected to dilated convolutions with a kernel size of 3×3. In this invention, the dilation rate is set to {1,2,3,2,1} to increase the receptive field while maintaining the feature scale, thereby achieving better detail enhancement. Finally, depthwise separable convolution operations are performed on the detail-enhanced feature maps to obtain the reconstruction result. and

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

[0101]

[0102] in, This will be the final prediction result.

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

[0104] Specifically, the loss function for constructing the conditional denoising diffusion probability model based on wavelet transform includes diffusion loss, content loss, and block contrast loss, as follows:

[0105] The diffusion loss function is designed to address the noise estimated by the network and the Gaussian noise added during the forward diffusion process, as follows:

[0106]

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

[0108]

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

[0110] Finally, to maximize the mutual information between the pre-enhancement results and the final prediction results, this application also uses block-contrast loss on both, and we use encoder G. ec and two-layer MLP network M from Extract features from the domain and embed them into a feature stack. Where L is from G ec The selected layer number, l, is the specific l-th layer. These feature stacks represent different blocks in the image, and we denote the spatial location within the selected layer as s∈{1,...S}. l}, where S l This refers to the number of spatial locations within each floor. Each time from... Sample a query code in the domain and represent its features as Where C l If the number of channels in each layer is given, then its corresponding positive sample features are represented as follows: Negative sample features are represented as Finally, the block-level multi-layer contrastive 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] In this invention, κ1, κ2, and κ3 are set to 1.0, 1.0, and 0.1, respectively.

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

[0116] like Figure 4 The figure shows a comparison of the application effects of the method of the present invention with existing methods. As can be seen from the figure, the GDCP method and DiffWater method cannot correct the color of the image; the ACCDC method and CWR method reduce the color richness in the enhancement result; the UDCP method reduces the brightness of the enhancement result; the HLRP method results in a reddish tint; and the Ucolor, PUIE-mc, TACL, and U-shape methods still have low contrast issues in their enhancement results. In comparison, the URSCT and DM_underwater methods produce better enhancement results, but they are still inferior to the method of the present invention in terms of backscatter removal.

[0117] Table 1 shows the objective comparative experimental results of the present invention with 12 existing enhancement methods on the UIEB-challenging60 dataset in the embodiments. It mainly includes eight evaluation metrics: e (the ratio of newly added visible edges to the original number of visible edges in the enhanced image; a larger value indicates richer details), r (the enhancement of edge pixel gradient values ​​in the image; a larger value indicates higher contrast), blur (the degree of blur in the enhancement result; a smaller value indicates clearer details), ENTROPY (reflects the amount of information in the image; a larger value indicates more information), UIQM (used to evaluate the color, sharpness, and contrast of underwater images; a larger value indicates a clearer underwater image), UCIQE (used to evaluate the color, brightness, and saturation of underwater images; a larger value indicates a better enhancement effect), NIQE (using distance to measure image quality; a smaller value indicates higher image quality), and FADE (used to perceive fog density in the image; a smaller value indicates a cleaner image with higher sharpness). The best result for each metric is indicated in bold, and the second best result is indicated by underlining. As can be seen from the table, the method proposed in this invention achieves the best performance among the five evaluation indicators and performs second best among the other three indicators, which fully demonstrates the superiority of the proposed method.

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

[0119]

[0120]

[0121] Finally, it should be noted that the above embodiments are only used to illustrate the technical solutions of the present invention, and not to limit them; although the present invention has been described in detail with reference to the foregoing embodiments, those skilled in the art should understand that modifications can still be made to the technical solutions described in the foregoing embodiments, or equivalent substitutions can be made to some or all of the technical features; and these modifications or substitutions 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, Includes the following steps: Step 1: Obtain an underwater image training set, which includes underwater degradation images and unpaired clear underwater images; Step 2: Construct a pre-enhancement network based on adversarial contrastive learning. This network is used to transform the degraded underwater image domain into a clear underwater image domain. The pre-enhancement network includes a generator and a discriminator. The generator processes the degraded underwater image to generate a pre-enhancement result and processes the clear underwater image to generate a consistent result. The discriminator determines whether a given image is a genuine clear image or a pseudo-clear image generated by the generator. Step 3: Train a pre-enhanced network based on adversarial contrastive learning using the underwater image training set; Step 4: Construct a conditional denoising and diffusion probability model based on wavelet transform. The input data for the conditional denoising and diffusion probability model based on wavelet transform are degraded underwater images and pre-enhanced underwater images. The conditional denoising and diffusion probability model based on wavelet transform is configured to perform the following steps: First, K-order two-dimensional discrete wavelet transforms were performed on the degraded underwater image and the pre-enhanced underwater image to obtain the low-frequency and high-frequency subbands in the wavelet domain. Secondly, the low-frequency subband in the wavelet domain undergoes diffusion processing, which includes a forward diffusion process and a backward diffusion process. The forward diffusion process is configured to progressively add Gaussian noise to the real data sample until it becomes a Markov chain of random noise. The backward diffusion process is configured to progressively denoise the random noise until a real sample is generated. Simultaneously, a high-frequency enhancement module performs detail enhancement processing on the high-frequency sub-band after the Kth wavelet transform. This module extracts features from the high-frequency sub-band through depthwise separable convolution, calculates cross-attention between features, and concatenates them to obtain complementary features. The feature map of the sub-band is then subjected to dilated convolution for better detail enhancement. Finally, depthwise separable convolution is performed on the detail-enhanced feature map to obtain the reconstructed high-frequency sub-band. Finally, the wavelet domain processing results are subjected to K-order two-dimensional inverse discrete wavelet transform to generate underwater image enhancement results; Step 5: Train a conditional denoising diffusion probability model based on wavelet transform using degraded underwater images from the underwater image training set and pre-enhanced images output by the pre-enhanced network based on adversarial contrastive learning. Step 6: Input the real underwater image to be processed into the pre-enhancement network based on adversarial contrastive learning after training to obtain the pre-enhancement image. Then input the pre-enhancement image into the conditional denoising diffusion probability model based on wavelet transform after training to obtain the final underwater enhanced image.

2. The underwater image enhancement method based on an 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: Where λ1, λ2, and λ3 represent the weighting coefficients of the loss functions balancing the various terms. This represents the least-squares GAN loss of the generator G. This represents the block-based multi-layer contrastive loss of the constructed pre-enhanced network. This represents the consistency loss of generator G.

3. The underwater image enhancement method based on an 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, It is the objective function used to optimize the generator. This is the objective function used to optimize the discriminator. G(x) represents the result produced by the generator, and D(G(x)) represents the probability that the discriminator D judges whether the result generated by the generator G is true. P represents the expected value of the distribution function. fake Let D(y) represent the distribution of the resulting underwater images generated by the generator, and let D(y) represent the probability that the discriminator D identifies a clear underwater image as real data. real The distribution of the results generated by the generator for clear underwater images is represented by k2=1, k3=0, and k1=1, where k2=1 represents the label of the real data, k3=0 represents the label of the fake data, and k1=1 represents the value that the generator wants the discriminator to believe in.

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

5. The underwater image enhancement method based on an unsupervised diffusion model according to claim 2, characterized in that, Block-based multi-layer contrastive loss of the constructed pre-enhanced network Obtained based on the following calculations: Where G represents the generator, M represents a two-layer MLP network, X represents the input degraded underwater image, and L is the encoder G. ec The selected layer number, l represents the specific l-th layer, S l The number of spatial locations in each layer, s∈{1,...S} l } represents the spatial location within the selected layer. The cross-entropy loss is represented by the following formula: in, Indicates the use of encoder G ec The feature stack obtained by extracting features from the degraded image domain X using a two-layer MLP network M. For each selected query code, the corresponding positive sample feature is denoted as . Negative sample features are denoted as C l It 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 an unsupervised diffusion model according to claim 1, characterized in that, The loss function of the conditional denoising diffusion probability model based on wavelet transform is obtained according to the following calculation: Where θ represents the parameters of the denoising network model for the reverse diffusion process, G is the generator, M is a two-layer MLP network, and κ1, κ2, and κ3 represent the weight coefficients balancing the various loss functions. This indicates the pre-enhancement results. This indicates the final prediction result. Represents the diffusion loss function. Represents the content loss function. This represents the block-based multi-level contrast loss of the conditional denoising diffusion probability model based on wavelet transform.

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

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

9. The underwater image enhancement method based on an unsupervised diffusion model according to claim 6, characterized in that, Block-based multi-layer contrastive loss based on a conditional denoising diffusion probability model using wavelet transform Obtained based on the following calculations: Where G represents the generator, and M represents a two-layer MLP network. L represents the result after processing a pre-enhanced network based on adversarial contrastive learning, where L is derived from the encoder G. ec The selected layer number, l represents the specific l-th layer, S l The number of spatial locations in each layer, s∈{1,...S} l } indicates the spatial location representation within the selected layer. The cross-entropy loss is represented by the following formula: 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 it. For each selected query code, its corresponding positive sample feature is denoted as . Negative sample features are denoted as C l It 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