An underwater image enhancement method based on diffusion model and adversarial learning
Patent Information
- Application Number
- CN202410738392.5
- Authority / Receiving Office
- CN · China
- Patent Type
- Patents(China)
- Current Assignee / Owner
- Filing Date
- 2024-06-07
- Publication Date
- 2026-09-15
- Estimated Expiration
- 2044-06-07
AI Technical Summary
然而,传统的GAN网络大多基于卷积神经网络(CNN)构建,虽然CNN因其固有的局部相关性偏差而对局部特征有着强大的感知能力,但在处理全局上下文信息时却显得力不从心
[0028] 1. This invention not only fully leverages the advantages of diffusion models in feature extraction and image enhancement, but also draws upon the unique capabilities of adversarial learning in image enhancement. Through the organic combination of the two, this invention can effectively remove blur and distortion from underwater images while preserving image details and texture, significantly improving image clarity and contrast.
Smart Images

Figure CN118736178B_ABST
Abstract
Description
Technical Field
[0001] This invention relates to the field of underwater image processing technology, and more specifically, to an underwater image enhancement method based on a diffusion model and adversarial learning. Background Technology
[0002] Due to the unique characteristics of the underwater environment, underwater images are often accompanied by severe quality degradation and distortion. In water, light is selectively absorbed, leading to varying degrees of color distortion. Furthermore, scattering alters the direction of light propagation, further reducing image contrast. This degraded imagery has a significant negative impact on many practical applications, such as underwater biometrics, topographic analysis, and underwater resource exploration. Therefore, to meet the needs of practical applications, improving the degradation of underwater images and restoring their potential high quality has become an urgent problem, prompting in-depth research into underwater image processing technologies.
[0003] In recent years, with the increasing popularity of underwater exploration, researchers have devoted themselves to the study of underwater image enhancement. Based on the different methods employed, underwater image enhancement can currently be divided into traditional methods and data-driven methods.
[0004] Traditional methods include model-free and model-based methods. Model-free methods do not directly consider the physical processes of underwater imaging; they primarily achieve their effect by modifying pixels. For example, histogram stretching improves image contrast, and filters reduce noise. Model-free methods are generally more flexible and intuitive, but may lack a deep understanding of the physical processes of underwater imaging. Model-based methods, on the other hand, focus on building underwater imaging models and obtaining clear images through the inverse operations of these models. For example, scattering models need to consider the scattering and absorption processes of underwater light, improving image quality by constructing corresponding physical models. These models typically consider multiple factors such as water quality, light propagation, and camera parameters, making the image enhancement process more scientifically grounded. While model-based methods are logically more rigorous and interpretable, they usually require more prior knowledge and may involve complex computational processes.
[0005] Data-driven deep learning methods, especially Generative Adversarial Networks (GANs), have garnered significant attention for their outstanding performance in feature learning. These models automatically learn effective image enhancement strategies by training on large amounts of data, without requiring explicit physical models. However, most traditional GANs are built upon Convolutional Neural Networks (CNNs). While CNNs possess strong perceptual capabilities for local features due to their inherent local correlation bias, they struggle with processing global contextual information. Underwater images often exhibit wide variations in brightness and color, and this uneven degradation is difficult for traditional CNNs to overcome. Summary of the Invention
[0006] In view of the shortcomings of existing technologies, this invention provides an underwater image enhancement method based on a diffusion model and adversarial learning. Starting from the characteristics of underwater images, this invention incorporates a cross-scale similarity and non-uniformity capture mechanism into a deep learning method, which helps to improve the effectiveness and applicability of underwater image enhancement.
[0007] The technical means employed in this invention are as follows:
[0008] An underwater image enhancement method based on diffusion models and adversarial learning includes the following steps:
[0009] S1. Obtain paired images as training data, wherein the paired images include degraded underwater images and corresponding enhanced reference images; obtain pure noise images based on the reference images corresponding to the degraded underwater images;
[0010] S2. Adversarial training is performed on the generator based on the diffusion model according to the training data. The generator is used to enhance the spliced degraded image, which is obtained by channel-level concatenation of the degraded underwater image and the pure noise image. The generator adopts a diffusion model based on the U-Net structure, replacing the original attention module in the U-Net structure with a multi-scale dynamic window attention module. The multi-scale dynamic window attention module is used to extract spatial features of the input features at different scales, and obtains a proxy feature token rich in multi-scale information after extracting spatial features of different scales. The multi-scale dynamic window attention module includes a multi-scale convolutional block and a rectangular window self-attention block. The rectangular window self-attention block is used to capture features in different directions in the input image by simulating the dependency relationship between horizontal and vertical rectangles.
[0011] S3. Obtain the underwater image to be enhanced as input to the trained generator, and obtain the generator output as the enhanced image.
[0012] Furthermore, a purely noisy image is obtained based on a reference image corresponding to the degraded underwater image, including:
[0013] The enhanced image corresponding to the degraded underwater image in the training data is forward diffused using a diffusion model forward process to obtain a pure noise image.
[0014] Further, adversarial training is performed on the diffusion-based generator based on the training data, including:
[0015] S201. Obtain the enhanced image output by the generator;
[0016] S202. Input the enhanced image and the reference image together into the PatchGAN-based discriminator, and extract features through the feature extraction network in the discriminator;
[0017] S203. The discriminator compares the features of the enhanced image with those of the reference image, enabling the discriminator to accurately distinguish between the enhanced image and the reference image, and providing effective feedback to the generator.
[0018] Furthermore, adversarial training of the diffusion-based generator based on the training data also includes:
[0019] The generator and discriminator optimize each other: when the discriminator can easily distinguish the enhanced image from the reference image, it provides feedback to the generator, instructing the generator to improve the generation effect. The generator adjusts its parameters and structure based on the discriminator's feedback.
[0020] Furthermore, adversarial training of the diffusion-based generator based on the training data also includes:
[0021] The training loss is set to include: the difference loss between the augmented image and the reference image, the optimization loss of the generator, and the adversarial loss between the generator and the discriminator.
[0022] Furthermore, the operation of the multi-scale dynamic window attention module includes:
[0023] After feature extraction from the convolutional operation in the encoding part of the stitched degraded image, a feature map X is obtained. X is then input into a multi-scale convolutional block layer and a linear layer. The multi-scale convolutional block layer generates surrogate feature tokens, and the linear layer generates key tokens, value tokens, and query tokens. The MCB layer takes the input feature X∈R... C×H×W It is divided into four parts, each with the same number of channels. Feature maps after segmentation For each partitioned feature map, convolutional kernels of different sizes are used to extract its spatial information, thus obtaining a series of feature maps. By concatenating the series of feature maps along the channel dimension, the proxy feature token can be obtained.
[0024] The proxy feature token and key token are processed using a rectangular window self-attention mechanism as follows:
[0025] Divide the feature map X into a series of non-overlapping rectangular windows X j ∈R C×sh×sw Each window has a height of sh and a width of sw, where j represents the number of the j-th window, and there are a total of Each window has one side that matches the width or length of the feature map;
[0026] The feature map X is split into two groups along the channel dimension, and an attention mechanism with different directional windows is applied to each group. Within each window, the dot product of the surrogate feature token and the key token is calculated, resulting in a score matrix of size sh×sw. The Softmax function is then applied to this matrix to convert the scores into normalized weights. Next, these weights are used to perform a weighted summation with the value token. This step generates an attention-weighted feature representation for each group of feature maps. After completing the attention calculation for both groups, these attention-weighted feature representations are concatenated along the channel dimension to restore the number of channels in the original feature map. Finally, this interactively processed information is passed to the query token.
[0027] Compared with the prior art, the present invention has the following advantages:
[0028] 1. This invention not only fully leverages the advantages of diffusion models in feature extraction and image enhancement, but also draws upon the unique capabilities of adversarial learning in image enhancement. Through the organic combination of the two, this invention can effectively remove blur and distortion from underwater images while preserving image details and texture, significantly improving image clarity and contrast.
[0029] 2. This invention also proposes a multi-scale dynamic window attention mechanism. This mechanism can keenly capture the cross-scale similarity features that are common in underwater images, while also finely depicting the texture details of non-uniformity.
[0030] Based on the above reasons, this invention can be widely applied in fields such as underwater photography, underwater robot vision systems, and underwater environmental monitoring. Attached Figure Description
[0031] 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.
[0032] Figure 1 This is a flowchart of an underwater image enhancement method based on diffusion model and adversarial learning in an embodiment of the present invention.
[0033] Figure 2 This is a schematic diagram of the generator training process in an embodiment of the present invention.
[0034] Figure 3 This is the model architecture in the embodiments of the present invention. Detailed Implementation
[0035] 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.
[0036] This invention provides an underwater image enhancement method based on a diffusion model and adversarial learning, using the publicly available underwater UIEB dataset, which includes degraded underwater images and corresponding enhanced reference images. During network training, the reference image y0 is first noise-added, gradually obtaining a pure noise image with an approximate Gaussian distribution. Next, the degraded underwater image and the noise image are input together into the network for inverse denoising. In the denoising process, an attention module is designed, and a diffusion model is used for iterative denoising, ultimately obtaining the enhanced underwater image. The enhanced underwater image is then... The reference image y0 is input into the discriminator, which judges its authenticity and returns the result to the generator based on the diffusion mode model. The two then optimize each other. Figure 1 As shown, the specific steps include:
[0037] S1. Obtain paired images as training data. The paired images include degraded underwater images and corresponding enhanced reference images. Use a forward diffusion model to perform forward diffusion on the enhanced image y0 corresponding to the degraded underwater image x in the UIEB dataset, obtaining a pure noise image y. T .
[0038] S2. Adversarial training is performed on the generator based on the diffusion model according to the training data. The generator is used to enhance the stitched degraded image, which is obtained by channel-level concatenation of the degraded underwater image and the pure noise image. The generator adopts a diffusion model based on the U-Net structure, replacing the original attention module in the U-Net structure with a multi-scale dynamic window attention module. The multi-scale dynamic window attention module is used to extract spatial features of the input features at different scales, and obtains a proxy feature token rich in multi-scale information after extracting spatial features of different scales. The multi-scale dynamic window attention module includes a multi-scale convolutional block and a rectangular window self-attention block. The rectangular window self-attention block is used to capture features in different directions in the input image by simulating the dependency relationship between the horizontal rectangle and the vertical rectangle.
[0039] In this invention, one of the key components of the generator is the attention module. This application abandons the original attention module of U-Net and instead designs a multi-scale dynamic window attention mechanism. Then, the degraded underwater image x and the purely noisy image y are... T Channel-level connections are made, and the input is fed into the generator for enhancement. The generator adopts a U-Net-like overall structure, but with improvements to the attention mechanism. A multi-scale dynamic window attention mechanism is proposed, which can capture cross-scale similarity and non-uniformity in underwater images. In the generator, the degraded underwater image x and the purely noisy image y are first processed. T After downsampling the input to the convolutional layers in the original U-Net, the resulting feature map is then fed into the improved attention mechanism to obtain a noisy feature map, such as... Figure 2 As shown, the training process specifically includes:
[0040] Step 2.1: Design an improved attention module, replacing the original attention module in U-Net with a multi-scale dynamic window attention mechanism, and applying it to the degraded image x∈R. H×W×3 Compared with pure noise image Y T ∈R H×W×3 The feature map X is obtained by inputting the channel-level connection into the convolution module in the generator, and then the feature map X is input into the multi-scale dynamic window attention.
[0041] Step 2.2: In multi-scale dynamic window attention, the obtained feature map X∈R H×W×C Perform different operations, among which the steps for performing multi-scale convolutional blocks are: input features X∈R C×H×W It is divided into four parts, each with the same number of channels. Feature maps after segmentation For each partitioned feature map, convolutional kernels of different sizes are used to extract its spatial information, resulting in a series of feature maps. This allows for the full capture of spatial information at different scales. Subsequently, by concatenating these feature maps along the channel dimension, a proxy feature token rich in multi-scale information can be obtained.
[0042] A = Cat([X'1,X'2,X'3,X'4])
[0043] Where A∈R C×H×W On the other hand, for feature maps X∈R H×W×C Perform linear operations to obtain keys, values, and queries, providing the necessary input for subsequent attention mechanisms.
[0044] Step 2.3: After obtaining A∈R C×H×W Key∈R C×H×W Value∈R C×H×W ,Query∈R C×H×W Subsequently, a rectangular window self-attention mechanism is introduced to process surrogate feature tokens and key tokens. This mechanism effectively captures features in different directions in the image by simulating the long-term dependency between horizontal and vertical rectangles. Specifically, the feature map X is first divided into a series of non-overlapping rectangular windows X. j ∈R C×sh×sw Each window has a height of sh and a width of sw, where j represents the number of the j-th window, and there are a total of There are 3 windows, and one side of each window is consistent with the width or length of the feature map, i.e., X. j ∈R C×H×sw or X j ∈R C×sh×W This provides a larger interaction space and stronger feature interaction capabilities. Subsequently, we split the feature map into two groups along the channel dimension. The attention mechanism employs different directional windows for each application group. Specifically, within each window, the dot product of the proxy feature token and the key token is calculated to obtain a score matrix of size sh×sw. Then, a Softmax function is applied to this matrix to convert the scores into normalized weights. Next, these weights are used to perform a weighted summation operation with the value tokens. This step generates an attention-weighted feature representation for each feature map group. After completing the attention calculation for both sets of feature maps, these two attention-weighted feature representations are concatenated along the channel dimension to restore the number of channels X'∈R of the original feature map. C×sh×sw This approach allows us to establish in-window dependencies over a larger scope within a specific dimension. Ultimately, this interactively processed information is passed to the query token to further enrich its informative content.
[0045]
[0046] in The constant factor is used to prevent weight concentration and attention collapse.
[0047] Through a series of nested operations, including dimensionality transformation, linear transformation, attention computation, and shape rearrangement, we obtain the final feature representation. Subsequently, convolution, group normalization, and SiLU activation are applied to the feature representation to enhance its expressive power. Simultaneously, Dropout is randomly applied along the channel dimension to prevent overfitting. Next, upsampling is performed to restore the spatial resolution of the image. At each upsampling layer, a residual connection and a multi-scale dynamic window attention module are used to further enhance the feature representation. Finally, a convolutional layer output yields the noisy feature map.
[0048] In the generator, the aforementioned diffusion model based on multi-scale dynamic window attention is used to process the pure noise image y. T Iterative denoising was performed, with a total of 2000 iterations, to obtain a significantly enhanced image. The enhanced image was then input into a discriminator for further processing and analysis.
[0049] After obtaining the enhanced image, it is input together with the reference image into a PatchGAN-based discriminator for feature extraction. During feature extraction, the discriminator progressively extracts high-level semantic information from the image through multiple convolutional and pooling operations. Next, the discriminator compares the features of the enhanced image with those of the reference image, aiming to accurately distinguish between them and provide effective feedback to the generator. Simultaneously, the generator and discriminator optimize each other. When the discriminator can easily distinguish the enhanced image from the reference image, it provides feedback to the generator, indicating that the generator needs to improve its generation. The generator adjusts its parameters and structure based on the discriminator's feedback to generate an enhanced image that more closely resembles the reference image. As training progresses, the generator's enhancement effect gradually improves, and the discriminator's discrimination ability also continuously strengthens. Ultimately, a high-quality enhanced image is obtained. Furthermore, an L2 loss function is introduced during the optimization process to further reduce the difference between the enhanced image and the reference image, providing more explicit guidance for the training of the generator and discriminator. The network architecture of this invention is as follows: Figure 3 As shown.
[0050] S3. Obtain the underwater image to be enhanced as input to the trained generator, and obtain the generator output as the enhanced image.
[0051] The following specific application examples will further illustrate the solution and effects of the present invention.
[0052] This embodiment provides an underwater image processing method based on diffusion model and adversarial learning. This method combines the advantages of diffusion model and adversarial learning to generate enhanced images consistent with the features of degraded underwater images.
[0053] During training, paired images (degraded underwater images and their corresponding enhanced reference images) are used as training data. In the diffusion-based generator, the reference image is first subjected to a forward diffusion operation. The forward chain perturbs the data by gradually adding Gaussian noise to the data according to a pre-designed noise schedule until the data distribution tends towards a standard Gaussian distribution. Notably, the noise addition process is driven by a mathematical law and requires no additional learning. However, the denoising process relies on the training of the neural network. By providing pure noise as input, the neural network gradually learns and recovers the distribution of the original data. In the generator design of this invention, we focus on learning a specific inverse process p. θ (y (0:T) |x), while maintaining the forward diffusion process q(y) (1:t) The invariance of |y0). To ensure that the generated enhanced image and the degraded image are consistent in high-level semantic information, the degraded underwater image x needs to be used to guide the inverse process of the diffusion model. Specifically, given the input (y0), T ,x,t), where y T The image is a noisy image obtained by progressively adding Gaussian noise to a clean image y0, where x is the degraded underwater image and t is the time step. At time t, the relationship between the noisy state and the previous time t-1 can be described by a specific mathematical expression:
[0054]
[0055] in Let be the hyperparameters of the noise table, t∈{0,1,.....T}, and I be an identity matrix with the same dimensions as the initial state y0. Given y0, y1, y2...y can be derived. T The joint distribution can be expressed as:
[0056]
[0057] Furthermore, by utilizing the properties of Markov chains, we can derive the noisy state at any time t after inputting y0.
[0058]
[0059] Finally, the noise-enhanced state y TThe noise distribution will approximate a Gaussian distribution. For each image in the training set, we can generate a series of noisy images with varying levels of noise. These images, together with the degraded underwater image, constitute the actual training samples. In the reverse process, using the existing noisy state, we learn to estimate the noise distribution and further obtain the state from the previous time step. This process gradually reconstructs the real data from the Gaussian distribution. Since the goal of underwater image enhancement is to generate the desired enhanced image, we introduce the degraded underwater image x as a condition into the noise estimation network to estimate the conditional noise distribution. Specifically, we concatenate the degraded image x with the noisy state as the input to the noise estimation network. Therefore, the joint probability distribution of the reverse diffusion process can be expressed as:
[0060]
[0061] This invention uses a neural network model δ θ To learn the inverse process, predict the noise vector ∈, and the enhanced image. Represented as:
[0062]
[0063] Where γ represents the current noise level. Each iteration of the inverse process can be represented as: (The image at step t-1 is obtained by subtracting noise from the image at step t, but this noise is fitted by the network. In addition, each sampling step also adds pure noise sampled from a normal distribution.)
[0064]
[0065] Where ∈ t ~N(0,I). That is, in each iteration of the inverse process, we update the image state using noise fitted by the network and pure noise sampled from a normal distribution. This process is repeated iteratively until the final enhanced image is generated.
[0066] To ensure that the model gradually reduces the difference between the enhanced image and the reference image during the training phase, thereby optimizing the image enhancement effect, this invention uses L2 loss as the key optimization metric. The specific formula is as follows:
[0067]
[0068] Where ∈ represents the noise vector. The purpose of this loss function is to minimize the error in the image enhancement process. By minimizing this error, the model can more accurately capture and learn the structure and characteristics in the reference image, thus more effectively preserving the detailed features and structure of the image when enhancing underwater images. The core objective of this invention is to enhance degraded underwater images x. Based on variational inference theory, the optimization objective of the generator is ultimately to minimize q(y). t-1 |y t ,y0) and p θ (y t-1 |y t The KL divergence of two distributions, x and y, is:
[0069]
[0070] Furthermore, to enhance the robustness of the model, this invention introduces a discriminator D based on PatchGAN. The discriminator is adversarially trained against a generator based on a diffusion model, enabling the generator to produce samples sufficient to deceive the discriminator. During GAN training, the discriminator strives to accurately distinguish between real and generated input samples, while the generator aims to generate samples sufficient to confuse the discriminator. This adversarial training mechanism significantly improves the generator's generation capability and enhances the model's robustness. The specific expression of the adversarial loss function is as follows:
[0071]
[0072] Considering all the above factors, the overall optimization objective of the final model can be expressed as:
[0073] L = L m +L θ (y0|x)+L adv (G,D) (10)
[0074] In diffusion-based generators, U-Net is primarily used as the underlying architecture. Building upon this, this invention proposes a multi-scale dynamic window attention mechanism. Underwater images often exhibit significant cross-scale similarity and non-uniformity, features crucial for image enhancement tasks. Cross-scale similarity means that certain features and textures in underwater images exhibit similar characteristics at different scales. To more accurately capture these similarities, feature maps are scaled, and multi-scale convolutional blocks are used to extract local features, achieving fine-grained capture of information at different scales. On the other hand, underwater images, influenced by factors such as light diffusion and water flow, often exhibit features with different directions and structures, requiring special attention to non-uniformity during processing. To effectively address this issue, a rectangular window self-attention mechanism is introduced. The core idea of this mechanism is to divide the input features into two independent parts and apply horizontal and vertical window self-attention mechanisms respectively, thereby simulating long-term dependencies in different directions. In the specific implementation, the feature map X is first input to the multi-scale convolutional block (MCB) layer and the Linear layer, respectively. The MCB layer is designed to generate surrogate feature tokens, while the Linear layer is used to generate the key Q (query), K (key), and V (value) components. The MCB layer is designed to fully consider the multi-scale information of the input data; specifically, it takes the input features X∈R... C×H×W It is divided into four parts, each with the same number of channels. Feature maps after segmentation For each partitioned feature map, convolutional kernels of different sizes (e.g., 3×3, 5×5, 7×7, 9×9) are used to extract their respective spatial information, thus obtaining a series of feature maps. This allows for the full capture of spatial information at different scales. Subsequently, by concatenating these feature maps along the channel dimension, a proxy feature token rich in multi-scale information can be obtained.
[0075] A=Cat([X'1,X'2,X'3,X'4]) (11)
[0076] Where A∈R C×H×W This token not only preserves the local details of the original feature map but also incorporates contextual information at different scales, thus providing rich feature representations for subsequent image processing. Next, we utilize a rectangular window self-attention mechanism to process the surrogate feature token and the key token. This mechanism effectively captures features in different directions in the image by simulating the long-term dependency between horizontal and vertical rectangles. Specifically, the feature map X is divided into a series of non-overlapping rectangular windows X. j ∈R C×sh×sw Each window has a height of sh and a width of sw, where j represents the number of the j-th window, and there are a total of Each window is a separate window. In the design, it's intentionally ensured that one side of each window aligns with the width or length of the feature map. This design, compared to traditional rectangular windows, provides more interaction space and enhances the interactivity between features. Subsequently, the feature map is split into two groups along the channel dimension, and an attention mechanism with windows of different orientations is applied to each group. This approach allows us to establish intra-window dependencies over a larger scope within a specific dimension. After window attention processing, the two groups of features are reconnected along the channel dimension to aggregate key information. Finally, this interactively processed information is passed to a query token to further enrich its information content.
[0077] This interaction process is achieved through the following formula:
[0078]
[0079] By combining multi-scale convolutional blocks and rectangular window self-attention mechanisms, not only can the unique characteristics of underwater images be accurately captured and processed, but also indirect interaction is achieved through the introduction of surrogate feature tokens. This not only reduces the computational burden but also improves the efficiency of information integration and processing performance. This optimized information aggregation helps to enhance the perception of complex structures and changes in underwater scenes.
[0080] 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 a diffusion model and adversarial learning, characterized in that, Includes the following steps: S1. Obtain paired images as training data, wherein the paired images include degraded underwater images and corresponding enhanced reference images; obtain pure noise images based on the reference images corresponding to the degraded underwater images; S2. Adversarial training is performed on the generator based on the diffusion model according to the training data. The generator is used to enhance the spliced degraded image, which is obtained by channel-level concatenation of the degraded underwater image and the pure noise image. The generator adopts a diffusion model based on the U-Net structure, replacing the original attention module in the U-Net structure with a multi-scale dynamic window attention module. The multi-scale dynamic window attention module is used to extract spatial features at different scales from the input features, and obtains a proxy feature token rich in multi-scale information after extracting the spatial features at different scales. The multi-scale dynamic window attention module includes a multi-scale convolutional block and a rectangular window self-attention block. The rectangular window self-attention block is used to capture features in different directions of the input image by simulating the dependency relationship between horizontal and vertical rectangles. S3. Input the underwater image to be enhanced into the trained generator, and obtain the generator output as the enhanced image. The working process of the multi-scale dynamic window attention module includes: After the stitching degraded image is subjected to feature extraction through convolution operation of the encoding part in the U-Net, a feature map X is obtained, and then X is input into a multi-scale convolution block layer and a Linear layer respectively, the multi-scale convolution block layer is used to generate proxy feature tokens, and the Linear layer is used to generate key tokens, value tokens and query tokens; the MCB layer generates proxy feature tokens by inputting the feature X is divided into four parts, each part has the same number of channels , the feature map after cutting ; for each divided feature map, a convolution kernel of different size is used to extract respective spatial information, thereby obtaining a series of feature maps ; the series of feature maps are spliced in the channel dimension, thereby obtaining proxy feature tokens; The proxy feature token and key token are processed using a rectangular window self-attention mechanism as follows: Divide the feature map X into a series of non-overlapping rectangular windows X j Each window has a height of sh and a width of sw, where j represents the number of the j-th window, and there are a total of One edge of each window is consistent with the width or length of the feature map; The feature map X is split into two groups along the channel dimension, and an attention mechanism with different directional windows is applied to each group. This includes calculating the dot product of the proxy feature token and the key token within each window to obtain a score matrix of size sh×sw. Then, the Softmax function is applied to this matrix to convert the scores into normalized weights. Next, these weights are used to perform a weighted summation operation with the value token. This step generates an attention-weighted feature representation for each group of feature maps. After completing the attention calculation for the two groups of feature maps, the two groups of attention-weighted feature representations are concatenated along the channel dimension to restore the number of channels of the original feature map. Finally, this interactively processed information is passed to the query token. This interaction process is achieved through the following formula: in, To query the token, For proxy feature tokens, For key tokens, For value tokens, It is a constant factor.
2. The underwater image enhancement method based on diffusion model and adversarial learning according to claim 1, characterized in that, Obtaining a pure noise image based on a reference image corresponding to a degraded underwater image includes: The enhanced image corresponding to the degraded underwater image in the training data is forward diffused using a diffusion model forward process to obtain a pure noise image.
3. The underwater image enhancement method based on diffusion model and adversarial learning according to claim 1, characterized in that, Adversarial training of the diffusion-based generator based on the training data includes: S201. Obtain the enhanced image output by the generator; S202. Input the enhanced image and the reference image together into the PatchGAN-based discriminator, and extract features through the feature extraction network in the discriminator; S203. The discriminator compares the features of the enhanced image with those of the reference image, enabling the discriminator to accurately distinguish between the enhanced image and the reference image, and providing effective feedback to the generator.
4. The underwater image enhancement method based on diffusion model and adversarial learning according to claim 3, characterized in that, The adversarial training of the diffusion-based generator based on the training data also includes: The generator and discriminator optimize each other: when the discriminator can easily distinguish the enhanced image from the reference image, it provides feedback to the generator, instructing the generator to improve the generation effect. The generator adjusts its parameters and structure based on the discriminator's feedback.
5. The underwater image enhancement method based on diffusion model and adversarial learning according to claim 3, characterized in that, The adversarial training of the diffusion-based generator based on the training data also includes: The training loss is set to include: the difference loss between the augmented image and the reference image, the optimization loss of the generator, and the adversarial loss between the generator and the discriminator.
Citation Information
Patent Citations
Underwater image enhancement method based on de-noising diffusion probability model
CN116883259A
Underwater image restoration method based on improved generative adversarial network model
CN117058037A