An underwater image enhancement method based on a stable diffusion model reverse iteration process
By introducing noise prediction and step size prediction networks into the underwater image enhancement method, the problem of error accumulation in the inverse iteration process is solved, and a more efficient and accurate underwater image enhancement effect is achieved.
Patent Information
- Authority / Receiving Office
- CN · China
- Patent Type
- Patents(China)
- Current Assignee / Owner
- BEIJING INST OF TECH
- Filing Date
- 2024-05-11
- Publication Date
- 2026-05-29
AI Technical Summary
Existing underwater image enhancement methods suffer from large error accumulation during the inverse iteration process, resulting in uncontrollable generated images and long inference times. Deep learning-based methods are highly dependent on data and lack generalization ability.
A prediction correction strategy and a step size prediction network are adopted. The prediction correction of the iterative process is performed through noise prediction network and step size prediction network to reduce error accumulation, reduce the number of iteration steps, and improve the accuracy of image generation and inference efficiency.
It reduces the accumulation of errors in the image generation process, reduces the number of iterations, shortens the inference time, and improves the accuracy and controllability of image generation.
Smart Images
Figure CN118570087B_ABST
Abstract
Description
Technical Field
[0001] This invention belongs to the field of image processing technology, and in particular relates to an underwater image enhancement method based on a stable diffusion model with a reverse iterative process. Background Technology
[0002] Underwater image enhancement is a key technology for acquiring underwater images and studying the underwater environment. With the rapid increase in demand for underwater exploration, the application of underwater image enhancement technology is becoming increasingly widespread, such as underwater surveillance, underwater scene analysis, seabed exploration, underwater military operations, marine archaeology, biology, and underwater robotics.
[0003] However, underwater imaging environments are harsh and complex. On the one hand, reflected light from underwater scenes is often scattered and absorbed by underwater particles, resulting in blurry, low-contrast, and dark underwater images. On the other hand, the longer the wavelength of light, the weaker its penetrating power, causing blue-green color casts in underwater images. These problems are disastrous for high-quality underwater imaging and pose a significant challenge to underwater image enhancement.
[0004] Currently, many underwater image enhancement algorithms have been proposed, which can be broadly classified into two categories: traditional image processing methods and deep learning-based methods. Traditional image processing methods first analyze the influencing factors of underwater imaging, model the imaging process using prior knowledge, and then adjust the image in terms of contrast and chromaticity to achieve good visual effects. However, traditional image processing methods mostly set certain assumptions or parameters under specific conditions, making them difficult to apply to a wide range of real-world waters. Deep learning is widely used in various image processing fields. CNN and GAN networks have shown excellent performance in underwater image enhancement. However, deep learning-based methods are highly dependent on the original data, requiring a large amount of data, and overfitting may limit their generalization ability to new real-world underwater images. Fortunately, diffusion models can be introduced to address this issue.
[0005] Generally, diffusion models consist of two stages: forward noise addition and backward denoising. The forward process gradually adds noise to the image, making it a completely realistic noisy image, while the backward process simulates the gradual denoising process of the noisy image to obtain a clean image. However, the backward process of the diffusion model requires a large number of iterations, and each iteration causes the error to accumulate, eventually amplifying the error and making the generated image uncontrollable. At the same time, a large number of iterations greatly increases the inference time of the model. Summary of the Invention
[0006] To address the aforementioned issues, this invention provides an underwater image enhancement method based on the inverse iterative process of a stable diffusion model. For the first time, a prediction correction strategy is applied to the inverse iterative process of the diffusion model, and the iteration step size is predicted by the network for underwater image enhancement, which greatly improves the accuracy of the generated images and reduces the inference time.
[0007] An underwater image enhancement method based on a stable diffusion model with inverse iterative process includes the following steps:
[0008] S1: Combine the image to be enhanced (raw), the current step number (t), and the current noise image (x) corresponding to the image to be enhanced (raw). t Input the trained noise prediction network z θ1 The current predicted noise z is obtained. θ1 (x t The current step number t is initialized with T, and the current noise image x corresponds to the image raw to be enhanced. t The initial value is noise ∈; where ∈ is noise that conforms to a standard normal distribution, and T is the upper limit of the set step number.
[0009] S2: Adjust the current prediction noise z θ1 (x t Input the denoising probability diffusion model with (t, raw) to obtain intermediate noise images k1 and k2;
[0010] S3: Use the mean of the intermediate noise images k1 and k2 as the next current noise image x. t-1 And according to x t-1 Update the value of noise ∈ to obtain the updated noise ∈ t-1 ;
[0011] S4: Transfer the noisy image x t-1 and noise ∈ t-1 Input the trained step size prediction network t θ2 This gives the next current step number t′;
[0012] S5: Combine the image to be enhanced (raw), the current step number (t′), and the current noisy image (x). t-1 Repeat steps S1 to S4 until the current step number t′ = 1, and then retrieve the corresponding current noise image x. t-1 As an enhanced image.
[0013] Furthermore, the current prediction noise z θ1 (x t The input denoising probability diffusion model (t, raw) yields intermediate noise images k1 and k2 as follows:
[0014]
[0015]
[0016] Where, α t Let α be the t-th sampling weight obtained by sampling from the weight interval [a,b] corresponding to noise ∈ according to a set step size s, and α t = a + (t-1) × s, where the step number t = 1, 2, ..., T. Indicates rounding down; α t-1 Let σ be the (t-1)th sampled weight obtained by sampling from the corresponding weight interval [a, b] of noise ∈ according to a set step size s, and z be another noise that conforms to a standard normal distribution. t Let α be the noise variance at the current step number t. t-1 The (t-1)th sampled weight is obtained by sampling from the weight interval [a,b] corresponding to the noise ∈ according to a set step size s. θ1 (x t +k1,t-1,raw) represents the image to be enhanced (raw), the step number (t-1), and the current noisy image (x). t The superimposed image x of the intermediate noise image k1 t +k1 is the input to the trained noise prediction network z θ1 The current predicted noise is obtained.
[0017] Furthermore, the noise variance σ at the current step number t t The calculation method is as follows:
[0018]
[0019] Where, α t+1 Let α be the (t+1)th sampled weight obtained by sampling from the weight interval [a,b] corresponding to noise ∈ according to a set step size s, and α T It has a set initial value.
[0020] Furthermore, according to x t-1 Update the value of noise ∈ to obtain the updated noise ∈ t-1 Specifically:
[0021]
[0022] α t Let α be the t-th sampling weight obtained by sampling from the weight interval [a,b] corresponding to noise ∈ according to a set step size s, and α t = a + (t-1) × s, where the step number t = 1, 2, ..., T. This indicates rounding down to the nearest integer.
[0023] Furthermore, the noise prediction network zθ1 and step size prediction network t θ2 The training method is as follows:
[0024] Step 1: Select multiple data pairs [x0, raw] from the SUID dataset, where raw is the image to be enhanced and x0 is the label image after enhancement;
[0025] Step 2: Obtain a noise ∈ that conforms to a standard normal distribution, and set the weight interval [a,b] corresponding to the noise ∈;
[0026] Step 3: Sample T sampling weights α from the weight interval according to the set step size s. t = a + (t-1) × s, where the step number t = 1, 2, ..., T. Indicates rounding down;
[0027] Step 4: Randomly select a step number t, and obtain the current noise image corresponding to the current step number t. And the cumulative weight
[0028] Step 5: Combine the image to be enhanced (raw), the current step number (t), and the current noisy image (x). t Input noise prediction network z θ1 The current predicted noise z is obtained. θ1 (x t ,t,raw);
[0029] Step 6: Convert the current noisy image x t Current prediction noise z θ1 (x t Input step size prediction network t θ2 To obtain the current prediction step size t θ2 (x t ,z θ1 );
[0030] Step 7: Adjust the current prediction noise z θ1 (x t (,t,raw) and the current prediction step size t θ2 (x t ,z θ1 Substitute into the following loss function:
[0031]
[0032] in, This represents the gradient operation, where λ1 and λ2 are the set loss weight parameters, |||| 2 The distance is Euclidean.
[0033] Step 8: Determine whether the loss function is less than a set value or whether the number of iterations has reached a set upper limit. If either is true, then the current noise prediction network z... θ1 and step size prediction network t θ2 For the trained noise prediction network z θ1 and step size prediction network t θ2 If all are negative, proceed to step 9;
[0034] Step 9: Adjust the current noise prediction network z θ1 and step size prediction network t θ2 The network parameters are determined, and the adjusted noise prediction network z is used. θ1 and step size prediction network t θ2 Repeat steps 4 through 9 until a well-trained noise prediction network z is obtained. θ1 and step size prediction network t θ2 .
[0035] Furthermore, λ1 = 0.6, λ2 = 0.4.
[0036] Furthermore, the weight interval [a,b] = [0.02,0.9].
[0037] Beneficial effects:
[0038] This invention provides an underwater image enhancement method based on the inverse iterative process of a stable diffusion model. In the inverse reconstruction process of diffusion, the noise output of the noise prediction network in each iteration is predicted and corrected, instead of directly obtaining the final predicted noise through the noise prediction network. Then, a dedicated step size prediction network is used to predict the corresponding step size for the next step, instead of directly decrementing the current step size. This enhances the robustness of the image generation process, reduces the accumulated error in each iteration, reduces the number of iterations, and reduces the time of the entire model inference process. Attached Figure Description
[0039] Figure 1 A flowchart of an underwater image enhancement method based on a stable diffusion model inverse iterative process is provided by the present invention;
[0040] Figure 2 The overall framework diagram of network training for an underwater image enhancement method based on the inverse iterative process of a stable diffusion model provided by the present invention is shown below.
[0041] Figure 3 The diagram shows the specific framework of the noise prediction network and step size prediction network for the diffusion model.
[0042] Figure 4 This is a diagram of the cross-attention structure in the noise prediction network and the step size prediction network.
[0043] Figure 5 This is a diagram of the fully connected layer structure at the output of the step-size prediction network.
[0044] Figure 6 A visual comparison diagram showing the results of various underwater image enhancement methods and the present invention; Detailed Implementation
[0045] To enable those skilled in the art to better understand the present application, the technical solutions in the embodiments of the present application will be clearly and completely described below with reference to the accompanying drawings.
[0046] like Figure 1 As shown, an underwater image enhancement method based on a stable diffusion model with inverse iterative process includes the following steps:
[0047] S1: Combine the image to be enhanced (raw), the current step number (t), and the current noise image (x) corresponding to the image to be enhanced (raw). t Input the trained noise prediction network z θ1 The current predicted noise z is obtained. θ1 (x t The current step number t is initialized with T, and the current noise image x corresponds to the image raw to be enhanced. t Let noise ∈ be the initial value; where, To ensure that the noise conforms to a standard normal distribution, T is the upper limit of the set step number; for example, T = 1000.
[0048] S2: Adjust the current prediction noise z θ1 (x t Inputting the denoising probability diffusion model (t, raw) yields intermediate noisy images k1 and k2:
[0049]
[0050]
[0051] Where, α t Let α be the t-th sampling weight obtained by sampling from the weight interval [a,b] corresponding to noise ∈ according to a set step size s, and α t = a + (t-1) × s, where the step number t = 1, 2, ..., T. Indicates rounding down; α t-1 Let σ be the (t-1)th sampled weight obtained by sampling from the corresponding weight interval [a, b] of noise ∈ according to a set step size s, and z be another noise that conforms to a standard normal distribution. t Let α be the noise variance at the current step number t. t-1The (t-1)th sampled weight is obtained by sampling from the weight interval [a,b] corresponding to the noise ∈ according to a set step size s. θ1 (x t +k1,t-1,raw) represents the image to be enhanced (raw), the step number (t-1), and the current noisy image (x). t The superimposed image x of the intermediate noise image k1 t +k1 is the input to the trained noise prediction network z θ1 The current predicted noise is obtained.
[0052] The noise variance σ at the current step number t t The calculation method is as follows:
[0053]
[0054] Where, α t+1 Let α be the (t+1)th sampled weight obtained by sampling from the weight interval [a,b] corresponding to noise ∈ according to a set step size s, and α T It has a set initial value.
[0055] S3: Mean the intermediate noise images k1 and k2 As the next current noise image x t-1 And according to x t-1 Update the value of noise ∈ to obtain the updated noise ∈ t-1 :
[0056]
[0057] S4: Transfer the noisy image x t-1 and noise ∈ t-1 Input the trained step size prediction network t θ2 To obtain the next current step number t′ = t θ2 (x t-1 ,∈ t-1 );
[0058] S5: Combine the image to be enhanced (raw), the current step number (t′), and the current noisy image (x). t-1 Repeat steps S1 to S4 until the current step number t′ = 1, and then retrieve the corresponding current noise image x. t-1 As an enhanced image.
[0059] The noise prediction network z is described in detail below. θ1 and step size prediction network t θ2 Training methods, such as Figure 2 As shown, the specific steps include:
[0060] Step 1: Select 890 data pairs [x0, raw] from the SUID dataset, where raw is the image to be enhanced and x0 is the label image after enhancement;
[0061] Step 2: Obtain a noise ∈ that conforms to a standard normal distribution, and set the weight interval [a,b] = [0.02,0.9] corresponding to the noise ∈;
[0062] Step 3: Sample T sampling weights α from the weight interval according to the set step size s. t = a + (t-1) × s, where the step number t = 1, 2, ..., T. Indicates rounding down;
[0063] Step 4: Randomly select a step number t, and obtain the current noise image corresponding to the current step number t. And the cumulative weight That is to say For the first t α t The cumulative product of values;
[0064] It should be noted that α t It is a value that increases linearly between 0.02 and 0.9, and the interval between two adjacent values can be 8.8 × 10^(-4).
[0065] Step 5: Combine the image to be enhanced (raw), the current step number (t), and the current noisy image (x). t Input noise prediction network z θ1 The current predicted noise z is obtained. θ1 (x t ,t,raw);
[0066] Step 6: Convert the current noisy image x t Current prediction noise z θ1 (x t Input step size prediction network t θ2 To obtain the current prediction step size t θ2 (x t ,z θ1 );
[0067] Step 7: Adjust the current prediction noise z θ1 (x t (,t,raw) and the current prediction step size t θ2 (x t ,z θ1 Substitute into the following loss function:
[0068]
[0069] in, This represents the gradient operation, where λ1 and λ2 are the set loss weight parameters, λ1 = 0.6, λ2 = 0.4, |||| 2 The distance is Euclidean.
[0070] Step 8: Determine whether the loss function is less than a set value or whether the number of iterations has reached a set upper limit. If either is true, then the current noise prediction network z... θ1 and step size prediction network t θ2 For the trained noise prediction network z θ1 and step size prediction network t θ2 If all are negative, proceed to step 9;
[0071] Step 9: Adjust the current noise prediction network z θ1 and step size prediction network t θ2 The network parameters are determined, and the adjusted noise prediction network z is used. θ1 and step size prediction network t θ2 Repeat steps 4 through 9 until a well-trained noise prediction network z is obtained. θ1 and step size prediction network t θ2 That is, the diffusion model noise network z was obtained. θ1 and step size prediction network t θ2 The optimal parameters.
[0072] Among them, the noise prediction network z θ1 and step size prediction network t θ2 The specific structure is as follows Figure 3 As shown, the overall architecture adopts the u-net structure, performing four downsampling and four upsampling operations. Residual connections are used between corresponding layers. Each layer performs convolution, normalization, activation, and cross-attention operations. The cross-attention structure is as follows: Figure 4 As shown, the Q and K values of the attention are obtained from the raw image to be enhanced, and the V value is obtained from the feature map after each layer of activation. Then, the cross-attention output is obtained using the following formula:
[0073]
[0074] It is important to note the step size prediction network t θ2 Two fully connected layers were added to the output of u-net to predict the next step size, as shown in the structure below. Figure 5 As shown.
[0075] Therefore, the training process of the network in this invention is simply a continuous repetition of the following process:
[0076] 1. First, randomly select a pair [x0, raw] from the SUID dataset;
[0077] 2. Next, sample a normally distributed noise ∈;
[0078] 3. Next, randomly select a value t from t ~ [1, 1000];
[0079] 4. Then through Get the corresponding x t ;
[0080] 5. Finally, put x t raw, t are used as noise prediction network z θ1 The input is ||∈-z||, and the output is predicted noise. (This invention aims to make the predicted noise as similar as possible to ||∈-z||, so the loss value of the noise prediction network is ||∈-z||.) θ1 (x t ,t,raw)|| 2 ).
[0081] It should be noted that when t = 1000, At this time x 1000 ≈∈, x0 is almost completely masked, x t Basically independent of x0, ∈ is noise that follows a normal distribution, so x 1000 It is also normally distributed noise.
[0082] The following is a theoretical explanation of the error reduction effect of the underwater image enhancement method based on the inverse iterative process of a stable diffusion model provided by this invention:
[0083] For the prediction correction of this invention:
[0084]
[0085]
[0086]
[0087] Error R t-1 :
[0088]
[0089] in, This is the theoretical value;
[0090]
[0091] make
[0092]
[0093]
[0094] For no prediction correction:
[0095]
[0096] Error R' t-1 :
[0097]
[0098]
[0099] The difference in error between the two:
[0100]
[0101]
[0102]
[0103]
[0104] α t From 0.9 to 0.02
[0105]
[0106] z θ1 The output is the predicted noise, and k1 is the derived image, where the noise intensity is much smaller than the image intensity.
[0107]
[0108] Therefore, it can be seen that the method of the present invention can reduce errors.
[0109] Furthermore, Figure 6 A visual comparison chart of underwater image enhancement using seven methods—ULAP, IBLA, GB-dehaze, UDCP, Five-A+, Shallow-Unet, and FunIE-GAN—is presented, showing that the method of this invention has the best visual enhancement effect, good controllability of the generated image, and high stability. Table 1 shows a comparison of the above seven methods with the method of this invention in terms of underwater image enhancement-related indicators, demonstrating that the method of this invention is also among the top in specific indicators.
[0110] Table 1
[0111]
[0112] Table 2 compares the average inference time with and without a step-size prediction network structure during the iteration process. Using a step-size prediction network reduces the average inference time by 45.9198%, significantly shortening the inference time.
[0113] Table 2
[0114]
[0115] In summary, this invention is a stable and fast image enhancement method based on a diffusion model. By leveraging the inverse iteration characteristic of the diffusion model, it can reduce the accumulated error in each iteration, reduce the number of iteration steps, and shorten the inference time.
[0116] Of course, the present invention may have other various embodiments. Without departing from the spirit and essence of the present invention, those skilled in the art can make various corresponding changes and modifications according to the present invention, but these corresponding changes and modifications should all fall within the protection scope of the appended claims.
Claims
1. An underwater image enhancement method based on a stable diffusion model with inverse iterative process, characterized in that, Includes the following steps: S1: Image to be enhanced Current step number t Image to be enhanced Corresponding current noise image Input the trained noise prediction network The current predicted noise is obtained. ; where the current step number t Using T as the initial value, the image to be enhanced Corresponding current noise image With noise As the initial value; where, To ensure that the noise conforms to the standard normal distribution, T is the upper limit of the set step number; S2: Adjust the current prediction noise Input the denoising probability diffusion model to obtain the intermediate noise image. and ; S3: Transform the intermediate noise image and The mean of the current noise image is used as the next noise image. and according to Update noise The value of is used to obtain the updated noise. ; S4: Transfer the noisy image and noise Input the trained stride prediction network Get the next current step number ; S5: Image to be enhanced Current step number Current noisy image Repeat steps S1 to S4 until the current step number. =1, then the corresponding current noise image. As the enhanced image; Noise prediction network and step size prediction network The training method is as follows: Step 1: Select multiple data pairs from the SUID dataset ,in, For the image to be enhanced, To complete the enhanced label image; Step 2: Obtain a noise source that conforms to a standard normal distribution. and set the noise Corresponding weight range ; Step 3: Follow the set step size T sampling weights are sampled from the weight interval. Step number t =1,2,…,T, , Indicates rounding down; Step 4: Randomly select a step number t Get the current step number t Corresponding current noise image And the cumulative weight ; Step 5: Image to be enhanced Current step number t Current noisy image Input noise prediction network The current predicted noise is obtained. ; Step 6: Convert the current noisy image Current prediction noise Input step size prediction network To obtain the current prediction step size ; Step 7: Adjust the current prediction noise and current prediction step size Substitute into the following loss function : in, This represents gradient operation. and The set loss weight parameters, The distance is Euclidean. Step 8: Determine the loss function Whether it is less than a set value or whether the number of iterations has reached a set upper limit, if either is yes, then the current noise prediction network... and step size prediction network For a well-trained noise prediction network and step size prediction network If all are negative, proceed to step 9; Step 9: Adjust the current noise prediction network and step size prediction network The network parameters were determined, and the adjusted noise prediction network was used. and step size prediction network Repeat steps 4 through 9 until a well-trained noise prediction network is obtained. and step size prediction network .
2. The underwater image enhancement method based on the inverse iterative process of a stable diffusion model as described in claim 1, characterized in that, Current prediction noise Input the denoising probability diffusion model to obtain the intermediate noise image. and Specifically: in, To follow the set step size From noise Corresponding weight range The first sample obtained from the middle Each sampling weight, and Step number t =1,2,…,T, , Indicates rounding down; To follow the set step size From noise Corresponding weight range The first sample obtained from the middle Each sampling weight, Another noise that conforms to the standard normal distribution, The current step number t The noise variance under the following conditions To follow the set step size From noise Corresponding weight range The first sample obtained from the middle Each sampling weight, To enhance the image Step number Current noisy image and intermediate noise image Overlay image Input the trained noise prediction network The current predicted noise is obtained.
3. The underwater image enhancement method based on the inverse iterative process of a stable diffusion model as described in claim 2, characterized in that, noise variance The calculation method is as follows: in, To follow the set step size From noise Corresponding weight range The first sample obtained from the middle +1 sampling weight, and It has a set initial value.
4. The underwater image enhancement method based on the inverse iterative process of a stable diffusion model as described in claim 1, characterized in that, according to Update noise The value of is used to obtain the updated noise. Specifically: To follow the set step size From noise Corresponding weight range The first sample obtained from the middle Each sampling weight, and Step number t =1,2,…,T, , This indicates rounding down to the nearest integer.
5. The underwater image enhancement method based on the inverse iterative process of a stable diffusion model as described in claim 1, characterized in that, ， 。 6. The underwater image enhancement method based on the inverse iterative process of a stable diffusion model as described in any one of claims 1 to 5, characterized in that, Weighting range .