An underwater image enhancement method and apparatus
By using the DDPM model, conditional input module, and content compensation module, the mode collapse problem in underwater image enhancement methods during training is solved, achieving efficient image enhancement in variable underwater environments and improving the model's generalization ability and visual effect.
Patent Information
- Application Number
- CN202410024732.8
- Authority / Receiving Office
- CN · China
- Patent Type
- Patents(China)
- Current Assignee / Owner
- Filing Date
- 2024-01-08
- Publication Date
- 2025-12-05
- Estimated Expiration
- 2044-01-08
AI Technical Summary
Existing underwater image enhancement methods are prone to mode collapse during training, have insufficient generalization ability, and are difficult to adapt to the changing underwater imaging environment, resulting in poor performance on different datasets.
The Denoising Diffusion Probability Model (DDPM) is used as the base model. Combined with the conditional input module and the content compensation module, the model is trained by inputting the difference between the original underwater image and the noisy image and low-level features to improve stability and generalization ability.
It improves the model's generalization ability in different underwater imaging environments, achieving better visual effects and generalization, and is applicable to a variety of different datasets.
Smart Images

Figure CN117952849B_ABST
Abstract
Description
Technical Field
[0001] This application relates to the field of underwater image processing technology, and in particular to an underwater image enhancement method and apparatus. Background Technology
[0002] Currently, with increasing human activity in the ocean, underwater image enhancement (UIE) has received considerable attention. However, the complex and variable underwater environment and harsh lighting conditions pose significant challenges to improving the quality of degraded underwater images. The degradation of underwater images is primarily caused by the selective absorption and scattering of visible light wavelengths in the underwater environment, resulting in images exhibiting low contrast, low brightness, significant color deviation, blurred details, uneven bright spots, and other defects. These limitations severely hinder practical applications in fields such as marine ecology, marine biology and archaeology, remotely operated vehicles (ROVs), and autonomous underwater vehicles (AUVs). Therefore, UIE research is of great significance for advancing practical research in these fields.
[0003] Chinese patent publication CN116167941A discloses an underwater image enhancement method based on RGB channel restoration, comprising: acquiring an image set consisting of a degraded underwater image and a clear underwater image; constructing an underwater image enhancement model using a generative adversarial network and designing a multinomial loss function; training the underwater image enhancement model based on the acquired image set and the designed multinomial loss function; and inputting the degraded underwater image to be enhanced into the trained underwater image enhancement model to obtain the clear image corresponding to the degraded underwater image to be enhanced. However, its use of a generative adversarial network as its basic architecture makes training difficult, and the presence of multiple discriminators during training makes it difficult for the generator and discriminators to reach Nash equilibrium, resulting in insufficient generalization ability of the model.
[0004] Chinese Patent Publication No. CN115034985A discloses an underwater image enhancement method, which includes inverting an original underwater image to obtain a dark channel image; calculating global background light and underwater ambient light based on the dark channel image; calculating the light transmittance of the water medium; calculating a preliminary dehazing underwater image based on the global background light and the light transmittance of the water medium; enhancing the preliminary dehazing underwater image using a multi-scale Retinex algorithm; compensating for image distortion using a color restoration function; using a histogram to extract the proportion of multi-scale Retinex processing values of different color channels in the final processed image to obtain a secondary processed underwater image; and using a bilateral filter consisting of Gaussian filtering and guided filtering to process noise in the secondary processed underwater image to obtain the final image. However, this method requires secondary image processing, which is cumbersome, and the image enhancement effect is highly dependent on the imaging environment, resulting in insufficient model generalization ability.
[0005] Chinese patent publication CN114972546A discloses an underwater image enhancement method based on adaptive color restoration of the blue-green channels. By analyzing the characteristics of underwater imaging models, and based on the proportion of the blue and green channel mean values in the sum of the three channel mean values, the underwater scene depth is divided into three levels. Then, color is adaptively compensated using light attenuation rate characteristics to achieve multi-scene color correction. The color-compensated image is then divided into four regions: dark tones, mid-dark tones, mid-bright tones, and bright tones. A dark region mapping function maps the dark regions to the bright regions, improving contrast while suppressing noise. Finally, bilinear interpolation is used to solve the block effect caused by the segmented processing. However, its hierarchical adaptive color compensation method is difficult to adapt to the unpredictable underwater imaging environment, and the model's generalization ability is insufficient. Summary of the Invention
[0006] To improve the generalization ability of the model, this application provides an underwater image enhancement method and apparatus.
[0007] In one aspect, this application provides an underwater image enhancement method.
[0008] This application is achieved through the following technical solution:
[0009] An underwater image enhancement method includes the following steps:
[0010] Obtain the original underwater images and their corresponding clear real images, construct paired data in different batches, and form a training set;
[0011] Based on the DDPM model, the original underwater image, the noise image at each time step, and the difference between the original underwater image and the noise image at the current time step are randomly selected from any batch of paired data in the training set. The DDPM model is trained, and the noise added at the current time step is output as the output value. A loss function for the input parameters is designed.
[0012] When the output value of the DDPM model meets the preset condition compared with the noise at the current time step of uniform sampling, the network parameters of the DDPM model are optimized to complete one training cycle, and it is determined whether the loss function has converged.
[0013] The DDPM model is trained iteratively until the loss function converges, at which point the current DDPM model is used as the target model.
[0014] Randomly sample any original underwater image and any Gaussian noise, input them into the target model, combine them with a preset denoising formula, predict the enhanced underwater image of the previous time step, and use the enhanced underwater image of the previous time step as the input result of the denoising formula in the next round, iterating until a real and clear image corresponding to the sampled original underwater image is obtained.
[0015] In a preferred embodiment, this application can be further configured such that the main network architecture of the DDPM model is based on a Unet network, and the input of the Unet network is also connected to a content compensation module for seamlessly integrating the low-level information extracted from the original underwater image into each layer of the Unet network.
[0016] In a preferred embodiment, the Unet network may be further configured as follows: the Unet network includes a first downsampling layer, a second downsampling layer, a third downsampling layer and a fourth downsampling layer that sample different feature dimensions in series, a first upsampling layer, a second upsampling layer, a third upsampling layer and a fourth upsampling layer that sample different feature dimensions in series, and a first connection layer, a second connection layer and a third connection layer located between the downsampling layer and the upsampling layer.
[0017] The first connection layer concatenates the sampling results of the first downsampling layer and the second downsampling layer to obtain a first concatenation result;
[0018] The second connection layer concatenates the first splicing result and the sampling result of the third downsampling layer to obtain a second splicing result;
[0019] The third connection layer splices the second splicing result and the sampling result of the fourth downsampling layer to obtain the third splicing result;
[0020] The sampling results of the fourth downsampling layer are input into the fourth upsampling layer;
[0021] The third splicing result is input into the third upsampling layer;
[0022] Input the second splicing result into the second upsampling layer;
[0023] The first splicing result is input into the first upsampling layer.
[0024] In a preferred embodiment, this application may be further configured such that the content compensation module includes a first convolutional layer, a second convolutional layer, a third convolutional layer, and a fourth convolutional layer;
[0025] The first convolutional layer extracts low-level information with the same feature dimensions as the first downsampling layer from the original underwater image and inputs it into the first downsampling layer;
[0026] The second convolutional layer extracts low-level information with the same feature dimensions as the second downsampling layer from the original underwater image and inputs it into the second downsampling layer;
[0027] The third convolutional layer extracts low-level information with the same feature dimensions as the third downsampling layer from the original underwater image and inputs it into the third downsampling layer;
[0028] The fourth convolutional layer extracts low-level information with the same feature dimensions as the fourth downsampling layer from the original underwater image and inputs it into the fourth downsampling layer.
[0029] In a preferred example, this application can be further configured such that the expression for the loss function is as follows:
[0030]
[0031] In the formula, L simple This represents the loss value of the loss function; This represents the output value of the noise prediction network, used to predict the noise ∈ added during each time step of the noise addition process; This represents the mathematical expectation of the noise prediction network.
[0032] In a preferred embodiment, this application can be further configured to employ gradient descent when optimizing the network parameters of the DDPM model.
[0033] In a preferred example, this application can be further configured such that: the expression for the probability distribution of the noise image at the previous time step output by the DDPM model at a unit time step includes,
[0034]
[0035] In the formula, This represents the model at each time step t in the noisy image x given the current time step. t Under the premise that the previous time step noisy image x t-1 The probability distribution satisfies the following condition: y0 represents the original underwater image, and N represents the Gaussian distribution function. This represents the mathematical expectation of the Gaussian distribution. This represents the standard deviation of the Gaussian distribution.
[0036] In a preferred example, this application can be further configured to: the mathematical expectation of the Gaussian distribution The expressions include,
[0037]
[0038] α t =1―βt
[0039]
[0040] In the formula, This represents the output value of the noise prediction network, used to predict the noise ∈ β added during each time step of the noise addition process. t The β value gradually increases from 0.0001 to 0.02, and t gradually increases from 1 to 1000, for a total of 1000 β values. t (β1, β2, ..., β) 1000 ), Represents α1 to α t The product of two digits.
[0041] In a preferred embodiment, this application can be further configured such that: the standard deviation of the Gaussian distribution... The expressions include,
[0042]
[0043] α t =1―β t
[0044]
[0045] In the formula, β t The β value gradually increases from 0.0001 to 0.02, and t gradually increases from 1 to 1000, for a total of 1000 β values. t (β1, β2, ..., β) 1000 ), Represents α1 to α t The product of two digits.
[0046] Secondly, this application provides an underwater image enhancement device.
[0047] This application is achieved through the following technical solution:
[0048] An underwater image enhancement device, comprising,
[0049] The data module is used to acquire raw underwater images and their corresponding real, clear images, construct paired data in different batches, and form a training set.
[0050] The training module is used to train the DDPM model based on the input of the original underwater image, the noise image at each time step, and the difference between the original underwater image and the noise image at the current time step from any batch of paired data in the training set. The module outputs the noise added at the current time step as the output value and designs the loss function of the input parameters.
[0051] The judgment module is used to optimize the network parameters of the DDPM model, complete one training cycle, and determine whether the loss function has converged when the output value of the DDPM model meets the preset conditions compared with the noise at the current time step of uniform sampling.
[0052] The target model module is used to iteratively train the DDPM model until the loss function converges, at which point the current DDPM model is used as the target model.
[0053] An enhancement module is used to randomly sample any original underwater image and any Gaussian noise, input them into the target model, combine them with a preset denoising formula, predict the enhanced underwater image of the previous time step, and use the enhanced underwater image of the previous time step as the input result of the denoising formula for the next round, iterating until a real and clear image corresponding to the sampled original underwater image is obtained.
[0054] Thirdly, this application provides a computer device.
[0055] This application is achieved through the following technical solution:
[0056] A computer device includes a memory, a processor, and a computer program stored in the memory and executable on the processor, wherein the processor executes the computer program to implement the steps of any of the above-described underwater image enhancement methods.
[0057] Fourthly, this application provides a computer-readable storage medium.
[0058] This application is achieved through the following technical solution:
[0059] A computer-readable storage medium storing a computer program that, when executed by a processor, implements the steps of any of the above-described underwater image enhancement methods.
[0060] In summary, compared with the prior art, the beneficial effects of the technical solution provided in this application include at least the following:
[0061] Raw underwater images and their corresponding clear real-world images are acquired, and paired data from different batches are constructed to form a rich training set for model training. The DDPM model is employed to avoid the pattern collapse phenomenon that is common in generative adversarial networks, resulting in more stable model training. When training the DDPM model, the input includes a randomly selected batch of paired data from the training set containing the raw underwater image, the noise image at each time step, and the difference between the raw underwater image and the noise image at the current time step. This facilitates the extraction of image difference features at different time steps, and allows the model to better adapt to the underwater imaging environment. When the output value of the DDPM model meets a preset condition compared to the uniformly sampled noise at the current time step, indicating relatively accurate noise prediction, the DDPM model is optimized. The network parameters are set to complete one training iteration, and the convergence of the loss function is determined. The DDPM model is iteratively trained until the loss function converges, ensuring that the final accuracy of the model meets the requirements. At this point, the current DDPM model is used as the target model. Any original underwater image and any Gaussian noise are randomly sampled and input into the target model. Combined with a preset denoising formula, the enhanced underwater image of the previous time step is predicted. The enhanced underwater image of the previous time step is used as the input result of the denoising formula in the next round. This process is iterated until a clear image corresponding to the sampled original underwater image is obtained. The target model trained is thus well-suited for the UIE task, adapts to different underwater imaging environments, greatly improves the model's generalization ability, and enables the model to achieve better generalization and visual effects on various datasets. Attached Figure Description
[0062] Figure 1 The main flowchart of an underwater image enhancement method provided for an exemplary embodiment of this application is shown.
[0063] Figure 2 This is a block diagram of the main model structure of an underwater image enhancement method provided as another exemplary embodiment of this application.
[0064] Figure 3 The diagram illustrates the visual effects of the model outputs using the Test_L400 (LSUI) dataset, with the following input conditions: inputting only the original underwater image, inputting both the original underwater image and the difference between the original underwater image and the noisy image at the current time step, inputting both the original underwater image and the underlying information extracted by the content compensation module, and inputting both the original underwater image, the difference between the original underwater image and the noisy image at the current time step, and the underlying information extracted by the content compensation module.
[0065] Figure 4The underwater image enhancement method of this application is used to sequentially perform WaterNet, FUnIE, Ucolor, Ushape, and other algorithms on the input raw underwater image using the Test_L400 (LSUI) dataset to output an enhanced underwater image. Detailed Implementation
[0066] This specific embodiment is merely an explanation of this application and is not intended to limit it. After reading this specification, those skilled in the art can make modifications to this embodiment without contributing any inventive step, but such modifications are protected by patent law as long as they fall within the scope of the claims of this application.
[0067] To make the objectives, technical solutions, and advantages of the embodiments of this application clearer, the technical solutions of the embodiments of this application will be clearly and completely described below with reference to the accompanying drawings. Obviously, the described embodiments are only some embodiments of this application, not all embodiments. Based on the embodiments of this application, all other embodiments obtained by those skilled in the art without creative effort are within the scope of protection of this application.
[0068] Furthermore, the term "and / or" in this article is merely a description of the relationship between related objects, indicating that three relationships can exist. For example, A and / or B can represent: A existing alone, A and B existing simultaneously, or B existing alone. Additionally, the character " / " in this article, unless otherwise specified, generally indicates that the preceding and following related objects have an "or" relationship.
[0069] Existing methods for improving underwater image quality using UIE (Underwater Image Enhancement) techniques can generally be categorized into two types: physical modeling and data-driven approaches. Physical modeling methods aim to model the physical processes of light propagation in water by considering the absorption, scattering, and other optical properties of the underwater environment. These methods typically involve complex mathematical models and algorithms to simulate degradation. However, because the aquatic environment changes over time, methods developed for a particular physical scenario cannot adapt to different scenarios, resulting in poor generalization ability of the established physical models. Inspired by the success of deep learning in a wide range of fields, data-driven methods have been proposed by learning the mapping between degraded underwater images and their corresponding high-quality reference images. Data-driven methods rely on large-scale datasets for model training and, based on the learned patterns and features, effectively improve the quality of underwater images. However, the UIE datasets currently established are usually collected in specific underwater environments, such as low light, various turbidities, and particulate matter of different densities. Therefore, models trained on specific datasets perform poorly when crossing datasets. This leads to the phenomenon of mode collapse in most Generative Adversarial Networks (GANs) based on data-driven methods, making the models difficult to train.
[0070] To address this, this application leverages the training stability of the Denoising Diffusion Probabilistic Model (DDPM) as a foundational model, making the training of the model more stable. During model training, the difference between the original image, the original image at each time step, and the noisy image is introduced as a merging condition input to facilitate the extraction of image difference features at different time steps. Consequently, the trained target model is well-suited for the UIE task, adapting to different underwater imaging environments and significantly improving the model's generalization ability, achieving better generalization and visual effects on various datasets.
[0071] The embodiments of this application will now be described in further detail with reference to the accompanying drawings.
[0072] Reference Figure 1 This application provides an underwater image enhancement method, the main steps of which are described below.
[0073] S1: Obtain the original underwater images and their corresponding real clear images, construct paired data in different batches, and form a training set;
[0074] S2: Based on the DDPM model, input the original underwater image, the noise image at each time step, and the difference between the original underwater image and the noise image at the current time step from any batch of paired data in the training set, train the DDPM model, output the noise added at the current time step as the output value, and design the loss function of the input parameters.
[0075] S3: When the output value of the DDPM model meets the preset condition compared with the noise at the current time step of uniform sampling, optimize the network parameters of the DDPM model, complete one training, and determine whether the loss function has converged.
[0076] S4: Iteratively train the DDPM model until the loss function converges, at which point the current DDPM model is used as the target model;
[0077] S5: Randomly sample any original underwater image and any Gaussian noise, input them into the target model, combine them with the preset denoising formula, predict the enhanced underwater image of the previous time step, and use the enhanced underwater image of the previous time step as the input result of the denoising formula in the next round, iterating until a real and clear image corresponding to the sampled original underwater image is obtained.
[0078] Specifically, a dataset is constructed for model training. Pairing dataset. Where S is the size of the dataset, and Let x0 and y0 represent the i-th original underwater image and its corresponding sharp image, respectively. For simplicity, we use sample (x0, y0) to represent any pair of training samples.
[0079] Based on the training stability of the Denoising Diffusion Probabilistic Model (DDPM), this application selects the DDPM model as the base model. The structural features of DDPM make its training process stable, which effectively solves the mode collapse problem that may occur in GAN-based models, making the model training more stable.
[0080] The training of the DDPM model includes a forward noise addition process and a backward noise reduction process.
[0081] During the forward noise addition process, noise is gradually added to the real, clear image x0 of the sample according to the gradually increasing diffusion rate (β∈[0.0001,0.02]). As the time step t gradually increases, the noise image x0... t It exhibits a distribution that increasingly approximates the standard Gaussian distribution. The forward noise-adding process is a Markov chain process.
[0082] The backward denoising process is the inverse of the forward process. The backward denoising process uses a random noise from a standard Gaussian distribution to remove the noise step by step, restoring the original underwater image. It is also a Markov chain process.
[0083] To apply DDPM to the UIE task, during model training, in addition to the input noisy image (x... t In addition, at each time step t, the original underwater image (y0) is input as a supervision condition to guide the diffusion model in generating enhanced underwater images. Furthermore, to address the issue that diffusion models based on the Unet network structure cannot extract sufficient image information from the original underwater image when applied to the UIE task, resulting in poor generalization performance, this application introduces an additional condition during model training: the original underwater image (y0) and the noisy image (x) at the current time step. t The difference between (y0-x) is calculated by taking this difference. t The input is combined into another conditional input, which allows the trained DDPM model network to extract more useful information and cues about the original underwater image and the noisy image. The extracted information also helps the diffusion model produce more accurate and visually appealing enhanced images.
[0084] By innovatively introducing the original underwater image and the difference between the original underwater image and the noisy image as conditional inputs for model training, the model can be well applied to different underwater imaging environments, thus improving its generalization ability.
[0085] The main training and sampling process of the model is shown in Table 1 below.
[0086] Table 1
[0087]
[0088]
[0089] The CPDM model for underwater image enhancement proposed in this application differs from existing diffusion models in its training and sampling processes.
[0090] During training, as shown in the steps of Algorithm 1, a batch of data pairs (x0, y0) is first randomly selected. Then, a uniform sampling method is used from {1, 2, ..., T} to sample a time interval t. A noise ∈ is sampled from a standard Gaussian distribution. This noise is considered to be added to x0 at time t. The network is then constructed. Used to predict this noise, to expect ∈ and When the value is close enough, the network parameters of the model are optimized to complete a training process.
[0091] When the output value of the DDPM model meets a preset condition compared with the noise at the current time step of uniform sampling, such as the difference between the output value of the DDPM model and the noise at the current time step of uniform sampling being less than a preset threshold, that is, the output value of the DDPM model is sufficiently close to the noise at the current time step of uniform sampling.
[0092] In one embodiment, gradient descent is used to optimize the network parameters of the DDPM model, which is more efficient.
[0093] During training, a batch of data pairs (x0, y0) is randomly selected each time to train the model, and this process is repeated until the designed loss function converges.
[0094] Once the model is trained, sampling is performed using the steps in Algorithm 2, that is, recovering the corresponding enhanced underwater image x0 from an underwater image y0. During the sampling process, a y0 and a Gaussian noise z are randomly sampled first, and then x0 is calculated according to the expression in line 4 of Algorithm 2. t―1 The value of x, then put x t―1 As input, we get x t―2 …, until x0 is obtained, which is the enhanced underwater image obtained from y0. According to the original model of the diffusion model, when x0 is obtained, the mean value at x1 is directly obtained, that is, z = 0. At this time, the plus sign at the end of the denoising formula in the 4th row adds a value of 0.
[0095] During the sampling process, the time step T can be set to 1000.
[0096] In one embodiment, the expression for the probability distribution of the noise image at the previous time step output by the DDPM model at a unit time step includes,
[0097]
[0098] In the formula, This represents the model at each time step t in the noisy image x given the current time step. t Under the premise that the previous time step noisy image x t-1 The probability distribution satisfies the following condition: y0 represents the original underwater image, and N represents the Gaussian distribution function. This represents the mathematical expectation of the Gaussian distribution. This represents the standard deviation of the Gaussian distribution.
[0099] Specifically, θ cpdm That is, the posterior probability of the diffusion model, using the expression for the probability distribution of the noise image at the previous time step output by the DDPM model at a unit time step, given the noise image x. t This allows us to gradually obtain the true, clear image x0, which is the enhanced underwater image.
[0100] In one embodiment, the mathematical expectation of the Gaussian distribution The expressions include,
[0101]
[0102] α t =1―β t
[0103]
[0104] In the formula, This represents the output value of the noise prediction network, used to predict the noise ∈ β added during each time step of the noise addition process. t The β value gradually increases from 0.0001 to 0.02, and t gradually increases from 1 to 1000, for a total of 1000 β values. t (β1, β2, ..., β) 1000 ), Represents α1 to α t The product of two digits.
[0105] Specifically, the mathematical expectation of the aforementioned Gaussian distribution is adopted. The average value of the predicted noise is then used to obtain the noise ∈ added during the noise addition process at each time step, making the model's prediction results more accurate.
[0106] In one embodiment, the standard deviation of the Gaussian distribution The expressions include,
[0107]
[0108] α t =1―β t
[0109]
[0110] In the formula, β t The β value gradually increases from 0.0001 to 0.02, and t gradually increases from 1 to 1000, for a total of 1000 β values. t (β1, β2, ..., β) 1000 ), Represents α1 to α t The product of the products. The standard deviation of the Gaussian distribution in this application. Not making predictions makes it easier to obtain prediction results faster.
[0111] In one embodiment, the expression for the loss function of the model input parameters is as follows:
[0112]
[0113] In the formula, L simple This represents the loss value of the loss function; This represents the output value of the noise prediction network, used to predict the noise ∈ added during each time step of the noise addition process; This represents the mathematical expectation of the noise prediction network.
[0114] By designing a loss function, the noise predicted by the trained model can be made sufficiently close to the actual noise.
[0115] In one embodiment, the main network architecture of the DDPM model is based on a Unet network, and the input of the Unet network is also connected to a content compensation module for seamlessly integrating the low-level information extracted from the original underwater image into each layer of the Unet network.
[0116] Specifically, to ensure the stability of model training, this application uses the DDPM model as the base model. Furthermore, since the main network architecture of DDPM is based on UNet and has a simple network structure, while the UIE task requires the preservation of important low-level features, including color, contour, edge, texture, and shape, this application designs a content compensation module to extract low-level features of the original image for content-aware training in order to ensure that the model retains the basic content features of the original underwater image. This improves the information aggregation ability within the noise prediction network, thus forming a new UIE framework called the Content Preserving Diffusion Model (CPDM). This model can improve the visual effect of the images generated by the model, making the model well applicable to the UIE task and achieving better generalization and visual effects on various datasets, greatly improving the quality of underwater images.
[0117] Reference Figure 2 In one embodiment, the Unet network includes a first downsampling layer, a second downsampling layer, a third downsampling layer and a fourth downsampling layer that sample different feature dimensions in series, a first upsampling layer, a second upsampling layer, a third upsampling layer and a fourth upsampling layer that sample different feature dimensions in series, and a first connection layer, a second connection layer and a third connection layer located between the downsampling layer and the upsampling layer.
[0118] The first connection layer concatenates the sampling results of the first downsampling layer and the second downsampling layer to obtain a first concatenation result;
[0119] The second connection layer concatenates the first splicing result and the sampling result of the third downsampling layer to obtain a second splicing result;
[0120] The third connection layer splices the second splicing result and the sampling result of the fourth downsampling layer to obtain the third splicing result;
[0121] The sampling results of the fourth downsampling layer are input into the fourth upsampling layer;
[0122] The third splicing result is input into the third upsampling layer;
[0123] Input the second splicing result into the second upsampling layer;
[0124] The first splicing result is input into the first upsampling layer.
[0125] By inputting a 3x64x64 image into the model network of the new UIE framework, four downsampling blocks are used, each block including four feature layers of different dimensions (128,64,64), (256,32,32), (384,16,16), and (512,8,8); three connection layers; and four upsampling blocks, each block including four feature layers of different dimensions corresponding to the feature dimensions of the downsampling blocks. This allows the model to have a higher number of feature layers, increasing the model's feature extraction capability. The number of downsampling layers is also appropriate, avoiding low sampling efficiency due to an excessive number of layers. This solves the problems of model collapse and difficulty in training that are prone to occur in data-driven methods.
[0126] In one embodiment, the content compensation module includes a first convolutional layer, a second convolutional layer, a third convolutional layer, and a fourth convolutional layer;
[0127] The first convolutional layer extracts low-level information with the same feature dimensions as the first downsampling layer from the original underwater image and inputs it into the first downsampling layer;
[0128] The second convolutional layer extracts low-level information with the same feature dimensions as the second downsampling layer from the original underwater image and inputs it into the second downsampling layer;
[0129] The third convolutional layer extracts low-level information with the same feature dimensions as the third downsampling layer from the original underwater image and inputs it into the third downsampling layer;
[0130] The fourth convolutional layer extracts low-level information with the same feature dimensions as the fourth downsampling layer from the original underwater image and inputs it into the fourth downsampling layer.
[0131] Specifically, such as Figure 2As shown, the content compensation module includes four convolutional layers. Before each downsampling, the original underwater image y0 is first processed by a 1x1 convolution to extract low-level information and obtain low-level features. Then, the image size of the extracted low-level features is transformed to be consistent with the feature dimension size of the feature layer corresponding to the downsampling block. Finally, the transformed low-level features of the original underwater image are directly added to the feature layer of the downsampling block, so as to achieve the goal of seamlessly integrating the low-level information extracted from the original underwater image y0 into each layer of the Unet network.
[0132] The content compensation module inputs low-level features from the original underwater image into the last block (i.e., the downsampling part) of each layer of the encoder. This allows the low-level features to control the network to preserve image content, which is beneficial for recovering a high-quality target image corresponding to the original input underwater image. The content compensation module effectively preserves low-level features during the sampling process, thereby comprehensively improving the quality of the recovered underwater image. It also enhances the information aggregation capability within the noise prediction network, making the trained target model better suited for underwater image enhancement tasks, adaptable to diverse underwater environments, and enhancing the model's generalization ability.
[0133] Furthermore, considering the need to preserve as many low-level features as possible, such as structure, texture, and contour, in the original underwater image, the first, second, third, and fourth convolutional layers all use 1x1 convolutions as low-level feature extraction networks. The network structure is relatively simple, and the efficiency of seamlessly integrating the low-level information extracted from the original underwater image into each layer of the Unet network is higher.
[0134] The CPDM model employs an iterative refinement paradigm during training, embedding conditional input modules and content compensation modules at each time step during training and sampling. This allows for better preservation of the original underwater image content in each denoising step, thereby improving the quality of the recovered underwater image.
[0135] Based on the existing DDPM model framework and content compensation module, a new content-preserving diffusion model CPDM based on the UIE framework is formed, which can improve the visual effect of the generated images, making the model well applicable to UIE tasks and achieving better generalization and visual effects on various datasets. It greatly improves the quality of underwater images, while the model training is stable.
[0136] Extensive experiments were conducted on the LSUI underwater dataset. 3789 image pairs were used for training, 400 image pairs (Test_L400) were used for testing, and 200 image pairs (Test_E200) were randomly selected from the EUVP dataset for testing on data outside the training dataset. The experimental results are shown in Tables 2 and 3 below.
[0137] Table 2
[0138]
[0139] Table 2 shows the quantitative metrics (PSNR, SSIM, and MSE) of the Content Preservation Diffusion Model (CPDM) proposed in this application compared with several other methods, illustrating the quantitative comparison of different UIE methods on the LSUI and EUVP datasets. The best results are highlighted in bold. The experimental results demonstrate the superior capabilities of CPDM in terms of both numerical evaluation metrics and visual effects. The CPDM method proposed in this application achieves the best image visual results on the quantitative metrics and also achieves the best performance on the test data (Test_E200) outside the training dataset, indicating that the CPDM model has excellent generalization performance.
[0140] Table 3
[0141]
[0142]
[0143] Table 3 presents ablation studies on the Test_L400 dataset. Model A represents the model with only y0 as input, Model B represents the model with both y0 and y0-xt as input, Model C represents the model with both y0 and the content compensation module as input, and Model D represents the complete model with y0, y0-xt, and the content compensation module as input. The best results are highlighted in bold.
[0144] Reference Figure 3 The image enhancement effect of the CPDM model in this application can be seen from the ablation experiment. The input image is the original underwater image. Model A represents inputting only the original underwater image y0 as the input condition, while Model B represents inputting both the original underwater image y0 and the difference y0-x between the original underwater image and the noisy image. t As input conditions, Model C represents the simultaneous input of the original underwater image y0 and the low-level features extracted by the content compensation module. Model D represents the CPDM model of this application, which simultaneously inputs the original underwater image y0 and the difference y0-x between the original underwater image and the noisy image. t The underlying features extracted by the content compensation module are used as input conditions. It can be seen that Model B (which only contains the conditional input module) or Model C (which only contains the content compensation module) outperforms Model A (the base model) in terms of both numerical indicators and visual image quality. Furthermore, Model D (which contains both the conditional input module and the content compensation module) achieves the best visual image quality, greatly improving the quality of underwater images.
[0145] Reference Figure 4 The paper compares the visual effects of WaterNet, FUnIE, Ucolor, and Ushape methods with the CPDM model of this application. The input is the original underwater image, and GT is the reference image. It can be seen that the application has better visual effects in color restoration and maintains better color consistency.
[0146] The CPDM method designed in this application can also be easily extended to other conditional generation tasks, including but not limited to image deraining, dehazing, deblurring, and image inpainting.
[0147] In summary, an underwater image enhancement method acquires the original underwater image and its corresponding clear real-world image, constructs paired data in different batches to form a relatively rich training set for model training. Employing the DDPM model avoids the pattern collapse phenomenon that easily occurs with generative adversarial networks, making the model training more stable. When training the DDPM model, the input includes the original underwater image from any batch of paired data randomly selected from the training set, the noise image at each time step, and the difference between the original underwater image and the noise image at the current time step. This facilitates the extraction of image difference features at different time steps, and allows the model to better adapt to the underwater imaging environment. When the output value of the DDPM model meets a preset condition compared to the uniformly sampled noise at the current time step, indicating that the model's noise prediction is relatively accurate, then... Optimize the network parameters of the DDPM model, complete one training iteration, and determine whether the loss function has converged. Iteratively train the DDPM model until the loss function converges, ensuring the final accuracy of the model meets the requirements. At this point, the current DDPM model is used as the target model. Randomly sample any original underwater image and any Gaussian noise, input them into the target model, and combine them with a preset denoising formula to predict the enhanced underwater image of the previous time step. Use the enhanced underwater image of the previous time step as the input result of the denoising formula in the next round, iterating until a clear image corresponding to the sampled original underwater image is obtained. The trained target model is thus well-suited for the UIE task, adapts to different underwater imaging environments, greatly improves the model's generalization ability, and enables the model to achieve better generalization and visual effects on various datasets.
[0148] This application uses a diffusion model as its foundation, overcoming the training instability problem of traditional GANs. Furthermore, by designing a conditional input module and a content preservation module, it uses the original underwater image as the conditional input, making the diffusion model applicable to UIE tasks. It uses the difference between the original underwater image and the noisy image at the current time step as the conditional input, improving the model's feature extraction capability. During the model's backdiffusion process, a content compensation module is used to extract low-level features from the original underwater image for target image generation. This enables the model to adapt to the ever-changing underwater environment and enhances its generalization ability.
[0149] This application does not require secondary processing of underwater images, and can directly obtain enhanced images from the original underwater images. It is independent of the underwater imaging environment and the model has strong generalization ability.
[0150] The CPDM model proposed in this application is better suited for underwater image enhancement (UIE) tasks, adapts to various underwater environments, and ensures the preservation of the original underwater image content to generate high-quality visual image results.
[0151] It should be understood that the sequence number of each step in the above embodiments does not imply the order of execution. The execution order of each process should be determined by its function and internal logic, and should not constitute any limitation on the implementation process of the embodiments of this application.
[0152] This application also provides an underwater image enhancement device, which corresponds one-to-one with the underwater image enhancement method described in the above embodiments. The underwater image enhancement device includes...
[0153] The data module is used to acquire raw underwater images and their corresponding real, clear images, construct paired data in different batches, and form a training set.
[0154] The training module is used to train the DDPM model based on the input of the original underwater image, the noise image at each time step, and the difference between the original underwater image and the noise image at the current time step from any batch of paired data in the training set. The module outputs the noise added at the current time step as the output value and designs the loss function of the input parameters.
[0155] The judgment module is used to optimize the network parameters of the DDPM model, complete one training cycle, and determine whether the loss function has converged when the output value of the DDPM model meets the preset conditions compared with the noise at the current time step of uniform sampling.
[0156] The target model module is used to iteratively train the DDPM model until the loss function converges, at which point the current DDPM model is used as the target model.
[0157] An enhancement module is used to randomly sample any original underwater image and any Gaussian noise, input them into the target model, combine them with a preset denoising formula, predict the enhanced underwater image of the previous time step, and use the enhanced underwater image of the previous time step as the input result of the denoising formula for the next round, iterating until a real and clear image corresponding to the sampled original underwater image is obtained.
[0158] For specific limitations regarding an underwater image enhancement device, please refer to the limitations regarding an underwater image enhancement method mentioned above, which will not be repeated here.
[0159] The various modules in the aforementioned underwater image enhancement device can be implemented entirely or partially through software, hardware, or a combination thereof. These modules can be embedded in or independent of the processor in a computer device, or stored in the memory of a computer device as software, so that the processor can call and execute the operations corresponding to each module.
[0160] In one embodiment, a computer device is provided, which may be a server. The computer device includes a processor, memory, a network interface, and a database connected via a system bus. The processor provides computing and control capabilities. The memory includes a non-volatile storage medium and internal memory. The non-volatile storage medium stores an operating system, computer programs, and a database. The internal memory provides an environment for the operation of the operating system and computer programs in the non-volatile storage medium. The network interface is used to communicate with external terminals via a network connection. When the computer program is executed by the processor, it implements any of the aforementioned underwater image enhancement methods.
[0161] In one embodiment, a computer-readable storage medium is provided, including a memory, a processor, and a computer program stored in the memory and executable on the processor, wherein the processor executes the computer program to implement any of the above-described underwater image enhancement methods.
[0162] Those skilled in the art will understand that all or part of the processes in the methods of the above embodiments can be implemented by a computer program instructing related hardware. The computer program can be stored in a non-volatile computer-readable storage medium and includes several instructions to cause a computer device (which may be a personal computer, server, or network device, etc.) to execute all or part of the steps of the methods described in the various embodiments of this application. When executed, the computer program may include the processes of the embodiments of the above methods. Any references to memory, storage, databases, or other media used in the embodiments provided in this application may include non-volatile and / or volatile memory. Non-volatile memory may include read-only memory (ROM), programmable ROM (PROM), electrically programmable ROM (EPROM), electrically erasable programmable ROM (EEPROM), or flash memory. Volatile memory may include random access memory (RAM) or external cache memory. By way of illustration and not limitation, RAM is available in a variety of forms, such as static RAM (SRAM), dynamic RAM (DRAM), synchronous DRAM (SDRAM), dual data rate SDRAM (DDRSDRAM), enhanced SDRAM (ESDRAM), synchronous link DRAM (SLDRAM), RAMbus direct RAM (RDRAM), direct memory bus dynamic RAM (DRDRAM), and memory bus dynamic RAM (RDRAM).
[0163] Those skilled in the art will clearly understand that, for the sake of convenience and brevity, the above-described division of functional units and modules is used as an example. In practical applications, the above functions can be assigned to different functional units and modules as needed, that is, the internal structure of the system can be divided into different functional units or modules to complete all or part of the functions described above.
Claims
1. An underwater image enhancement method, characterized by, The method comprises the following steps, obtaining original underwater images and their corresponding real clear images, constructing paired data in different batches to form a training set; based on the DDPM model, inputting the original underwater image in any batch of paired data in the training set, the noise image at each time step, and the difference between the original underwater image and the noise image at the current time step, training the DDPM model, outputting the added noise at the current time step as the output value, and designing a loss function for the input parameters; the main network architecture of the DDPM model is based on the Unet network, and the input end of the Unet network is also connected with a content compensation module for seamlessly integrating the bottom layer information extracted from the original underwater image into each layer of the Unet network; when the output value of the DDPM model meets the preset condition compared with the noise uniformly sampled at the current time step, the network parameters of the DDPM model are optimized, one training is completed, and whether the loss function converges is judged; iteratively training the DDPM model until the loss function converges, at which time the current DDPM model is taken as a target model; randomly sampling any original underwater image and any Gaussian noise, inputting them into the target model, combining a preset denoising formula, predicting an enhanced underwater image at the previous time step, and taking the enhanced underwater image at the previous time step as the next round of input result of the denoising formula, and iterating until the real clear image corresponding to the sampled original underwater image is obtained.
2. The underwater image enhancement method according to claim 1, characterized in that, The Unet network comprises first, second, third and fourth down-sampling layers for sampling different feature dimensions in series, first, second, third and fourth up-sampling layers for sampling different feature dimensions in series, and first, second and third connection layers between the down-sampling layers and the up-sampling layers; the first connection layer splices the sampling results of the first and second down-sampling layers to obtain a first splicing result; the second connection layer splices the first splicing result and the sampling result of the third down-sampling layer to obtain a second splicing result; the third connection layer splices the second splicing result and the sampling result of the fourth down-sampling layer to obtain a third splicing result; the sampling result of the fourth down-sampling layer is input into the fourth up-sampling layer; the third splicing result is input into the third up-sampling layer; the second splicing result is input into the second up-sampling layer; the first splicing result is input into the first up-sampling layer.
3. The underwater image enhancement method of claim 2, wherein, The content compensation module comprises first, second, third and fourth convolutional layers; the first convolutional layer extracts bottom layer information of the same feature dimension from the original underwater image and inputs it into the first down-sampling layer; the second convolutional layer extracts bottom layer information of the same feature dimension from the original underwater image and inputs it into the second down-sampling layer; the third convolutional layer extracts bottom layer information of the same feature dimension from the original underwater image and inputs it into the third down-sampling layer; and the fourth convolutional layer extracts bottom layer information of the same feature dimension from the original underwater image and inputs it into the fourth down-sampling layer. The third convolutional layer extracts bottom layer information with the same feature dimension as the third down-sampling layer from the original underwater image and inputs the bottom layer information into the third down-sampling layer; The fourth convolutional layer extracts bottom layer information with the same feature dimension as the fourth down-sampling layer from the original underwater image and inputs the bottom layer information into the fourth down-sampling layer.
4. The underwater image enhancement method of claim 1, wherein, The expression of the loss function is as follows: In the formula, a loss value of the loss function; represents an output value of the noise prediction network, used to predict the noise added by the noise adding process at each time step; represents the mathematical expectation of the noise prediction network; represents the added noise; represents the noise image; represents the sampling time; represents the original underwater image; represents the difference between the original underwater image and the noise image at the current time step.
5. The underwater image enhancement method of claim 1, wherein, The gradient descent method is used to optimize the network parameters of the DDPM model.
6. The underwater image enhancement method of claim 1, wherein, The expression of the probability distribution of the noise image output by the DDPM model at a unit time step includes, wherein denotes the model at each time step t given the noisy image at the current time step denotes the noisy image at the previous time step under the premise that the model at each time step t given the noisy image at the current time step satisfies the probability distribution, denotes the original underwater image, N denotes the distribution function of the Gaussian distribution, denotes the mathematical expectation of the Gaussian distribution, denotes the standard deviation of the Gaussian distribution, denotes the difference between the original underwater image and the noisy image at the current time step.
7. The underwater image enhancement method of claim 6, wherein, the mathematical expectation of the gaussian distribution the expression includes, In the formula, denotes the output value of the noise prediction network, used to predict the noise added by the noise adding process at each time step , from 0.0001 to 0.02, t from 1 to 1000, a total of 1000 ( ), multiplication.
8. The underwater image enhancement method of claim 6, wherein, a standard deviation of the Gaussian distribution the expression includes, In the formula, Step up from 0.0001 to 0.02, t from 1 to 1000, a total of 1000 Multiplication of 9. An underwater image enhancement device, characterized by including, The data module is configured to obtain original underwater images and corresponding real clear images, construct paired data in different batches, and form a training set. The training module is configured to input the original underwater image in the paired data in any batch in the training set, the noise image at each time step, and the difference between the original underwater image and the noise image at the current time step, train the DDPM model based on the DDPM model, output the noise added at the current time step as an output value, and design a loss function for input parameters. The judgment module is configured to optimize the network parameters of the DDPM model when the output value of the DDPM model meets a preset condition compared with the noise at the current time step uniformly sampled, complete one training, and determine whether the loss function converges. The target model module is configured to iteratively train the DDPM model until the loss function converges, and the current DDPM model is used as a target model. The enhancement module is configured to randomly sample any original underwater image and any Gaussian noise, input the original underwater image and the Gaussian noise into the target model, combine a preset denoising formula, predict an enhanced underwater image at a previous time step, and use the enhanced underwater image at the previous time step as an input result of the next round of the denoising formula, and iterate until a real clear image corresponding to the sampled original underwater image is obtained.
10. A computer device, comprising: The computer readable storage medium stores a computer program, and the computer program is executed by the processor to realize the steps of the method of any one of claims 1 to 8.
11. A computer readable storage medium, characterized in that, The computer readable storage medium stores a computer program, and the computer program is executed by the processor to realize the steps of the method of any one of claims 1 to 8.
Citation Information
Patent Citations
Underwater image enhancement method based on blue-green channel adaptive color recovery
CN114972546A
Underwater image enhancement method based on RGB channel recovery
CN116167941A