Underwater image enhancement method and device based on diffusion model, and medium

Abstract

The invention relates to an underwater image enhancement method and device based on a diffusion model, and a medium. The method comprises the following steps: obtaining a degraded image; respectively inputting the degraded images into a teacher model and a student model; calculating tracks of the teacher model and the student model; inputting the degraded image to a reflection estimation module, and calculating a reflection image; through a forward diffusion process of the diffusion model, mapping the reflection graph into a track correction amount of the teacher model, adding the track correction amount and the track of the teacher model bit by bit according to a preset weight, and correcting the track of the teacher model; calculating the distillation loss of the corrected teacher model track and student model track, freezing the parameters of the teacher model, and optimizing the parameters of the student model based on gradient descent; and performing image enhancement on the degraded image by using the optimized student model to obtain an enhanced underwater image. Compared with the prior art, the method has the advantages that the calculation amount is reduced, and meanwhile, the performance of the model in the aspects of detail and texture recovery is improved.

Description

Technical Field

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

[0002] Image enhancement technology is a processing method that makes images visually clearer and more detailed by improving image parameters such as brightness, contrast, and clarity. This technology has a wide range of applications in computer vision, image processing, autonomous driving, medical imaging, surveillance, and other fields. Existing image enhancement technologies can be divided into two categories: traditional methods and deep learning-based methods. Traditional methods include histogram equalization, Laplace transform, filtering, etc. These methods usually enhance visual effects by directly adjusting the brightness and contrast of the image. Deep learning-based methods automatically learn features in images through models such as convolutional neural networks and perform adaptive enhancement processing. The diffusion model is an iterative deep learning method that has shown strong performance in the field of image generation in recent years and has some applications in the field of image enhancement.

[0003] Diffusion model-based methods usually use the diffusion model as the backbone network, take the original image as the conditional input, and output the enhanced image. In order to improve the enhancement effect, the existing diffusion model-based methods are inefficient in the image enhancement process and require a large number of sampling steps for iteration. Usually, 25 steps or more iterations are required to obtain good image quality in terms of brightness, contrast, color restoration, etc., which consumes a lot of computing resources and limits practical applications.

[0004] However, traditional distillation techniques often suffer from significant performance degradation when compressing the number of sampling steps, resulting in poor enhanced image quality and poor restoration of details such as texture, structure and color.

[0005] Therefore, there is currently a lack of an image enhancement method that can balance computing resources and restoration effects. Summary of the invention

[0006] The purpose of the present invention is to provide an underwater image enhancement method, device and medium based on a diffusion model. By introducing trajectory distillation technology to reduce the number of reasoning steps of the diffusion model, and combining the reflection estimation module to optimize the model trajectory, the performance of the model in detail and texture restoration is significantly improved while reducing the amount of calculation.

[0007] The purpose of the present invention can be achieved by the following technical solutions:

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

[0009] Step 1) obtaining a degraded image, wherein the degraded image is a low-quality underwater image to be enhanced;

[0010] Step 2) inputting the degraded image into the teacher model and the student model respectively, wherein the teacher model is a pre-trained conditional diffusion model, and the student model has the same architecture as the teacher model;

[0011] Step 3) Calculate the trajectories of the teacher model and the student model;

[0012] Step 4) inputting the degraded image into the reflection estimation module to calculate the reflection map;

[0013] Step 5) Mapping the reflection map into the trajectory correction amount of the teacher model through the forward diffusion process of the diffusion model, and adding the trajectory correction amount and the trajectory of the teacher model bit by bit with a preset weight to correct the trajectory of the teacher model;

[0014] Step 6) Calculate the distillation loss of the corrected teacher model trajectory and the student model trajectory, freeze the parameters of the teacher model, and optimize the student model parameters based on gradient descent;

[0015] Step 7) Use the optimized student model to perform image enhancement on the degraded image.

[0016] The teacher model and the student model adopt a diffusion model architecture based on SR3, and the diffusion process is expressed as follows;

[0017]

[0018] Among them, x t represents the intermediate noise result of the diffusion model, x0 represents the true value of the normal image, ∈ represents the Gaussian noise of the standard distribution, represents the value of the noise schedule at step t, σ t represents the variance of the noise at the tth step.

[0019] In the step 3), a first-order ordinary differential equation sampler is used as a trajectory decoder to calculate the trajectories of the teacher model and the student model.

[0020] The trajectory decoder G of the teacher model from step t to step s η It is expressed as:

[0021]

[0022] The trajectory of the teacher model decoded by the trajectory decoder from step t to step u and then to step s is expressed as:

[0023]

[0024] Among them, x tdenotes the intermediate noisy result of the diffusion model, y denotes the degraded image, and σ t , σ s respectively represent the variances of the noise at the t-th and s-th steps, respectively represent the values of the noise schedule at the t-th and s-th steps, ∈ η denotes the output of the teacher model. t, u, and s are all intermediate steps of the diffusion model, and 0 < t < u < s < T, where T is the total number of steps of the diffusion model, and η represents the model parameters of the teacher model.

[0025] The reflection estimation module performs the following steps on the input degraded image:

[0026] Step 41) Estimate the noise map through Gaussian filtering: z′ = |y - ψ(y)|, where ψ represents the Gaussian blur operation and z′ represents the estimated noise map;

[0027] Step 42) Estimate the illumination distribution h′ based on the maximum channel value in the RGB three channels of the degraded image;

[0028] Step 43) Calculate the reflection component of the degraded image based on the Retinex model to obtain the reflection map:

[0029]

[0030] where, denotes the reflection map.

[0031] Step 5) is specifically:

[0032] Through the forward diffusion process of the diffusion model, map the reflection map to the trajectory correction amount of the teacher model:

[0033]

[0034] With a preset weight w refine Add the trajectory correction amount and the trajectory of the teacher model bit by bit to correct the trajectory of the teacher model:

[0035]

[0036]

[0037] where, x target is the corrected trajectory output by the teacher model.

[0038] The trajectory decoder G of the student model from the t-th step to the s-th step θ is expressed as:

[0039]

[0040] The trajectory of the student model decoded by the trajectory decoder from step t to step s is expressed as:

[0041]

[0042] Then, the trajectory x output by the student model eps is:

[0043]

[0044] where x t represents the intermediate noisy result of the diffusion model, y represents the degraded image, and σ t , σ s represent the variances of the noise at steps t and s respectively, represent the values of the noise schedule at steps t and s respectively, ∈ θ represents the output of the student model. Both t and s are intermediate steps of the diffusion model, and 0 < t < s < T, where T is the total number of steps of the diffusion model, and θ represents the model parameters of the student model.

[0045] The distillation loss is expressed as:

[0046] L distill = λ(t)||x target - x est ||

[0047] where x target is the corrected trajectory output by the teacher model, and x eps is the trajectory output by the student model,

[0048]

[0049] σ t represents the variance of the noise at step t, represents the value of the noise schedule at step t.

[0050] An underwater image enhancement device based on a diffusion model, comprising a memory, a processor, and a program stored in the memory, wherein the processor implements the method described above when executing the program.

[0051] A storage medium, on which a program is stored, and the program implements the method described above when executed.

[0052] Compared with the prior art, the present invention has the following beneficial effects:

[0053] (1) The present invention uses trajectory distillation technology to compress the model that originally requires multiple steps of iteration into only two steps to achieve high-quality image enhancement, which significantly improves the efficiency of underwater image enhancement tasks while ensuring a better image enhancement effect.

[0054] (2) The present invention introduces a reflection estimation module and uses the reflection information of the image to optimize the trajectory of the diffusion model, thereby reducing the performance loss in the distillation process and greatly improving the performance of the model after distillation, so that it can achieve excellent results in detail and texture restoration of underwater images while reducing the amount of calculation. BRIEF DESCRIPTION OF THE DRAWINGS

[0055] Figure 1 is a flow chart of the method of the present invention;

[0056] Figure 2 Schematic diagram of the trajectory distillation training process of the present invention. DETAILED DESCRIPTION

[0057] The present invention is described in detail below in conjunction with the accompanying drawings and specific embodiments. This embodiment is implemented based on the technical solution of the present invention, and provides a detailed implementation method and specific operation process, but the protection scope of the present invention is not limited to the following embodiments.

[0058] This embodiment provides an underwater image enhancement method based on a diffusion model. Figure 1 As shown, the following steps are included:

[0059] Step 1) obtaining a degraded image. In this embodiment, the degraded image is a low-quality underwater image to be enhanced.

[0060] Step 2) Input the degraded image into the teacher model and the student model respectively.

[0061] The teacher model is a pre-trained diffusion model-based image enhancement model, which usually requires more iterations to enhance the image and provides guidance for the student model during training. The student model is an untrained diffusion model-based image enhancement model. After distillation training guided by the teacher model, it can achieve image enhancement with very few iterations.

[0062] In this embodiment, the teacher model and the student model adopt a diffusion model architecture based on SR3, and the U-Nets convolutional neural network is used as the backbone network of the diffusion model. The diffusion process is expressed as follows:

[0063]

[0064] Among them, x t represents the intermediate noise result of the diffusion model, x0 represents the true value of the normal image, ∈ represents the Gaussian noise of the standard distribution, denotes the value of the noise schedule at the \(t\)-th step, \(\sigma\) t represents the variance of the noise at the \(t\)-th step.

[0065] In other embodiments, different diffusion model architectures may also be employed.

[0066] Step 3) Calculate the trajectories of the teacher model and the student model.

[0067] In this embodiment, a first-order ODE (ordinary differential equation) sampler is used as the trajectory decoder to calculate the trajectories of the teacher model and the student model, where the trajectory represents the set of all intermediate results generated by the diffusion model during the step-by-step iteration process. In other embodiments, different trajectory decoders may also be used, and the use of different trajectory decoders does not affect the achievement of the object of the present invention.

[0068] In this embodiment, the trajectory decoder \(G\) of the teacher model from the \(t\)-th step to the \(s\)-th step η is expressed as:

[0069]

[0070] The trajectory of the teacher model decoded by the trajectory decoder from the \(t\)-th step to the \(u\)-th step and then to the \(s\)-th step is expressed as:

[0071]

[0072] where \(x\) t represents the intermediate noisy result of the diffusion model, \(y\) represents the degraded image, \(\sigma\) t , \(\sigma\) s respectively represent the variances of the noise at the \(t\)-th and \(s\)-th steps, respectively represent the values of the noise schedule at the \(t\)-th and \(s\)-th steps, \(\epsilon\) η represents the output of the teacher model, \(t, u, s\) are all intermediate steps of the diffusion model, and \(0 \lt t \lt u \lt s \lt T\), \(T\) is the total number of steps of the diffusion model, and \(\eta\) represents the model parameters of the teacher model.

[0073] The trajectory decoder \(G\) of the student model from the \(t\)-th step to the \(s\)-th step θ is expressed as:

[0074]

[0075] The trajectory of the student model decoded by the trajectory decoder from the \(t\)-th step to the \(s\)-th step is expressed as:

[0076]

[0077] Then, the trajectory \(x\) output by the student model eps is:

[0078]

[0079] Among them, ∈ θ represents the output of the student model, and θ represents the model parameters of the student model.

[0080] Step 4) Input the degraded image to the reflection estimation module to calculate the reflection map, which is a distribution map used to describe the reflection characteristics of the image.

[0081] The reflectance estimation module is a computational unit for estimating the reflectance map of an image. It performs the following steps on the input degraded image:

[0082] Step 41) Estimate the noise map by Gaussian filtering: z′=|y-ψ(y)|, where ψ represents the Gaussian blur operation and z′ represents the estimated noise map;

[0083] Step 42) Estimate the illumination distribution h′ based on the maximum channel value in the three RGB channels of the degraded image:

[0084]

[0085] Step 43) Calculate the reflection component of the degraded image based on the Retinex model to obtain a reflection map:

[0086]

[0087] in, Represents a reflection map.

[0088] In other embodiments, different denoising algorithms may be used to estimate the noise map and different sampling methods may be used to estimate the illumination distribution. The difference in the adopted methods does not affect the realization of the purpose of the invention.

[0089] Step 5) Through the forward diffusion process of the diffusion model, the reflection map is mapped into the trajectory correction amount of the teacher model, and the trajectory correction amount and the trajectory of the teacher model are added bit by bit with a preset weight to correct the trajectory of the teacher model.

[0090] Step 51) Through the forward diffusion process of the diffusion model (the computational process of gradually adding specific noise to the image), the reflection map is mapped to the trajectory correction amount of the teacher model:

[0091]

[0092] Step 52) with the preset weight w refine The trajectory correction amount and the trajectory of the teacher model Add bit by bit to correct the trajectory of the teacher model:

[0093]

[0094]

[0095] Among them, x target is the corrected trajectory output by the teacher model; w refine The value range is [0, 1] and is determined by experience. For the 8-step model, w refine The best value of is 0.2; for the 4-step model, w refine The best value of is 0.6; for the 2-step model, w refine The optimal value of is 1.

[0096] Step 6) Calculate the distillation loss of the modified teacher model trajectory and student model trajectory, freeze the parameters of the teacher model, and optimize the student model parameters based on gradient descent.

[0097] In this embodiment, the distillation loss adopts L2 loss, which is expressed as:

[0098] L distill =λ(t)||x iarget -x est ||

[0099] Among them, x target is the corrected trajectory output by the teacher model, x eps is the trajectory output by the student model,

[0100]

[0101] Step 7) Use the optimized student model to perform image enhancement on the degraded image.

[0102] The present invention aims to enhance the degraded underwater image into a high-quality clear image. Figure 2 As shown in the figure, for a given degraded image, the diffusion model is used to gradually enhance the image, and finally generate a high-quality output image. In the training stage, a conditional diffusion model is first trained as the teacher model, which can effectively enhance the degraded image through multi-step iterative sampling; then, a student model with the same architecture as the teacher model is constructed, and the trajectories of the two are decoded; then, the reflectance map of the degraded image is estimated through the reflection estimation module, and it is forward diffused to obtain the trajectory correction of the teacher model, and the trajectory of the teacher model is corrected. Finally, the distillation loss is calculated for the corrected teacher model trajectory and student model trajectory, and the gradient descent method is used to optimize the network, and finally a diffusion model that can complete high-quality image enhancement with very few steps is obtained.

[0103] The invention designs a trajectory distillation framework for image enhancement tasks, and for the first time realizes an acceleration method for the diffusion model in image enhancement tasks, so that the diffusion model can achieve image enhancement with fewer iterations. In addition, the present invention proposes a trajectory optimization method based on reflection estimation, which optimizes the trajectory of the teacher model, while reducing the amount of calculation of the diffusion model, ensures the restoration effect on details such as texture and color, improves the naturalness and realism of the enhanced image, and retains more details and information. Through experimental verification on multiple benchmark data sets, the invention verifies its ability to enhance complex degraded images, and successfully solves the problem of poor enhancement effect of traditional methods when dealing with complex illumination distribution and noise.

[0104] This embodiment also provides an underwater image enhancement device based on a diffusion model, including a memory, a processor, and a program stored in the memory, and the processor implements the above method when executing the program.

[0105] In a preferred embodiment, the device comprises:

[0106] Image acquisition module: used to acquire degraded images;

[0107] Model training module: used to perform the following steps: input the degraded image into the teacher model and the student model respectively; calculate the trajectories of the teacher model and the student model; input the degraded image into the reflection estimation module to calculate the reflection map; through the forward diffusion process of the diffusion model, map the reflection map to the trajectory correction amount of the teacher model, and add the trajectory correction amount and the trajectory of the teacher model bit by bit with the preset weight to correct the trajectory of the teacher model; calculate the distillation loss of the corrected teacher model trajectory and the student model trajectory, freeze the parameters of the teacher model, and optimize the student model parameters based on gradient descent;

[0108] Image enhancement module: used to enhance the degraded image using the optimized student model.

[0109] Those skilled in the art can clearly understand that, for the convenience and brevity of description, the specific working process of the described module can refer to the corresponding process in the aforementioned method embodiment, and will not be repeated here.

[0110] If the above functions are implemented in the form of software functional units and sold or used as independent products, they can be stored in a computer-readable storage medium. Based on this understanding, the technical solution of the present invention, or the part that contributes to the prior art, or the part of the technical solution, can be embodied in the form of a software product. The computer software product is stored in a storage medium, including several instructions for a computer device (which can be a personal computer, a server, or a network device, etc.) to perform all or part of the steps of the method described in each embodiment of the present invention. The aforementioned storage medium includes: U disk, mobile hard disk, read-only memory (ROM, Read-Only Memory), random access memory (RAM, Random Access Memory), disk or optical disk, etc., which can store program code.

[0111] The preferred specific embodiments of the present invention are described in detail above. It should be understood that a person skilled in the art can make many modifications and changes based on the concept of the present invention without creative work. Therefore, any technical solution that can be obtained by a person skilled in the art through logical analysis, reasoning, or limited experiments based on the concept of the present invention on the basis of the prior art should be within the scope of protection determined by the claims.

Claims

1. An underwater image enhancement method based on a diffusion model, characterized in that: The following steps are involved: Step 1) obtaining a degraded image, wherein the degraded image is a low-quality underwater image to be enhanced; Step 2) inputting the degraded image into the teacher model and the student model respectively, wherein the teacher model is a pre-trained conditional diffusion model, and the student model has the same architecture as the teacher model; Step 3) Calculate the trajectories of the teacher model and the student model; Step 4) inputting the degraded image into the reflection estimation module to calculate the reflection map; Step 5) Mapping the reflection map into the trajectory correction amount of the teacher model through the forward diffusion process of the diffusion model, and adding the trajectory correction amount and the trajectory of the teacher model bit by bit with a preset weight to correct the trajectory of the teacher model; Step 6) Calculate the distillation loss of the corrected teacher model trajectory and the student model trajectory, freeze the parameters of the teacher model, and optimize the student model parameters based on gradient descent; Step 7) Use the optimized student model to perform image enhancement on the degraded image.

2. The underwater image enhancement method based on diffusion model according to claim 1, characterized in that: The teacher model and the student model adopt a diffusion model architecture based on SR3, and the diffusion process is expressed as follows; Among them, x t represents the intermediate noise result of the diffusion model, x0 represents the true value of the normal image, ∈ represents the Gaussian noise of the standard distribution, represents the value of the noise schedule at step t, σ t represents the variance of the noise at the tth step.

3. The underwater image enhancement method based on diffusion model according to claim 1, characterized in that: In the step 3), a first-order ordinary differential equation sampler is used as a trajectory decoder to calculate the trajectories of the teacher model and the student model.

4. The underwater image enhancement method based on diffusion model according to claim 3, characterized in that: The trajectory decoder G of the teacher model from step t to step s η It is expressed as: The trajectory of the teacher model decoded by the trajectory decoder from step t to step u and then to step s is expressed as: Among them, x t represents the intermediate noisy result of the diffusion model, y represents the degraded image, and σ t , σ s represent the variances of the noises at the t-th and s-th steps respectively, represent the values of the noise schedule at the t-th and s-th steps respectively, ∈ η represents the output of the teacher model. t, u, and s are all intermediate steps of the diffusion model, and 0 < t < u < s < T, where T is the total number of steps of the diffusion model, and η represents the model parameters of the teacher model.

5. The underwater image enhancement method based on diffusion model according to claim 4, characterized in that: The reflection estimation module performs the following steps on the input degraded image: Step 41) Estimate the noise map by Gaussian filtering: z′=|y-ψ(y)|, where ψ represents the Gaussian blur operation, z ′ represents the estimated noise map; Step 42) Estimate the illumination distribution h based on the maximum channel value of the three RGB channels of the degraded image ′ ; Step 43) Calculate the reflection component of the degraded image based on the Retinex model to obtain a reflection map: in, Represents a reflection map.

6. The underwater image enhancement method based on diffusion model according to claim 5, characterized in that: The step 5) is specifically as follows: Through the forward diffusion process of the diffusion model, the reflection map is mapped to the trajectory correction of the teacher model: With the preset weight w refine The trajectory correction amount and the trajectory of the teacher model Add bit by bit to correct the trajectory of the teacher model: Among them, x target is the corrected trajectory output by the teacher model.

7. The underwater image enhancement method based on diffusion model according to claim 3, characterized in that: The trajectory decoder G of the student model from step t to step s θ It is expressed as: The trajectory of the student model from step t to step s obtained by decoding with the trajectory decoder It is expressed as: Then, the trajectory x output by the student model is eps for: Among them, x t represents the intermediate noisy result of the diffusion model, y represents the degraded image, and σ t , σ s represent the variances of the noise at the t-th and s-th steps respectively, represent the values of the noise schedule at the t-th and s-th steps respectively, ∈ θ represents the output of the student model. Both t and s are intermediate steps of the diffusion model, and 0 < t < s < T, where T is the total number of steps of the diffusion model, and θ represents the model parameters of the student model.

8. The underwater image enhancement method based on diffusion model according to claim 1, characterized in that: The distillation loss is expressed as: L distill =λ(t)||x target -x est || Among them, x target is the corrected trajectory output by the teacher model, x eps is the trajectory output by the student model, σ t represents the variance of the noise at step t, Represents the value of the noise schedule at step t.

9. An underwater image enhancement device based on a diffusion model, comprising a memory, a processor, and a program stored in the memory, characterized in that: When the processor executes the program, the method according to any one of claims 1 to 8 is implemented.

10. A storage medium having a program stored thereon, characterized in that: When the program is executed, the method according to any one of claims 1 to 8 is implemented.

Images