Image restoration network training method based on double attribute constraints of diffusion model

Abstract

The invention relates to the technical field of image restoration, and discloses a diffusion model double attribute constraint-based image restoration network training method, which comprises the following steps of: enabling a degraded image to pass through an image restoration network to obtain a restored image; the clear image and the restored image are input into a diffusion model after noise adding, predicted noise is obtained respectively, and the predicted noise difference between the clear image and the restored image is used as a natural attribute constraint; similarly, the clear image and the restored image are input into the diffusion model after noise addition, semantic features at the intermediate layer are networked respectively, and semantic attribute constraints are constructed by using the semantic features between the clear image and the restored image; constructing a loss function by using the natural attribute constraint and the semantic attribute constraint in combination with the original reconstruction constraint, and training an existing image restoration network; therefore, the visual quality and semantic attributes of the image output by the image restoration network are improved, and an optimal restoration result is obtained.

Description

Technical Field

[0001] The present invention relates to the technical field of image restoration, and in particular to an image restoration network training method based on dual attribute constraints of a diffusion model. Background Art

[0002] In complex imaging environments, image quality is often adversely affected by unpredictable low-light environments and non-uniform media, resulting in missing content information and color distortion, which is also not conducive to downstream tasks such as target detection and scene understanding. In order to obtain high-quality images, image restoration technology has made great progress. Traditional image restoration technology uses the statistical characteristics of degraded images as prior knowledge to restore images, but this type of method derived from specific observations is difficult to robustly simulate the intrinsic characteristics of degraded images. In recent years, with the rise of deep learning, many neural network-based methods have emerged that can effectively learn the mapping from degraded images to clear images, opening up new paths in the field of image restoration.

[0003] Although the above methods are quite effective, due to limited model capabilities and data, the restored images still have color and texture distortion and semantic loss. Some studies have proposed using multi-stage training methods or complex image restoration networks to solve the above problems, but these methods are too cumbersome to be implemented in general scenarios. The diffusion model proposed in recent years has powerful natural image generation and semantic representation capabilities. Some studies have proposed using it as a backbone network for image restoration, but it is difficult to use due to its long inference time and high memory usage.

[0004] To this end, the present invention proposes an image restoration network training method based on dual attribute constraints of a diffusion model. Summary of the invention

[0005] In order to solve the above technical problems, the present invention provides an image restoration network training method based on the dual attribute constraints of the diffusion model, in order to train an image restoration network model with good visual quality and rich semantic content output results, improve the visual effect and semantic attributes of the image, and do not need to increase additional calculations during reasoning. The proposed method for training an image restoration network based on the dual attribute constraints of the diffusion model is simple and effective. By extracting the natural attribute features and semantic attribute features of the clear image and the restored image in the diffusion model and using them as constraints for training the image restoration network, the image restoration network can achieve a more natural and more semantically informative restoration effect. After training, only the image restoration network needs to be retained. The present invention has significant advantages in practical applications and can be used as a general training paradigm in multiple image restoration tasks.

[0006] In order to solve the above technical problems, the present invention adopts the following technical solutions:

[0007] A method for training an image restoration network based on dual attribute constraints of a diffusion model, comprising:

[0008] Design the natural attribute constraint extraction process of the diffusion model: pass the degraded image through the image restoration network to obtain the restored image; input the clear image and the restored image into the diffusion model after adding noise, and obtain the predicted noise of the clear image and the restored image respectively, and use the predicted noise difference between the clear image and the restored image as the natural attribute constraint L nat ;

[0009] Design the diffusion model semantic attribute constraint extraction process: Similarly, the clear image and the restored image are input into the diffusion model after adding noise, and the semantic features of the image restoration network middle layer corresponding to the clear image and the restored image are obtained respectively. The semantic attribute constraint L is constructed using the semantic features between the clear image and the restored image. sem ;

[0010] Design a dual-attribute-constrained image restoration network training process: use the natural attribute constraints and semantic attribute constraints, combined with the original reconstruction constraints to construct a loss function, and train the existing image restoration network; thereby improving the visual quality and semantic attributes of the restored image output by the image restoration network, and obtaining the best restoration result.

[0011] Furthermore, the natural attribute constraint extraction process of the diffusion model specifically includes:

[0012] First, the degraded image y is passed through the image restoration network to obtain the restored image z. z is then denoised together with the clear image x according to the forward denoising process of the diffusion model with a time step of t. In the forward denoising process of the diffusion model, Gaussian noise is repeatedly added to the clear image x0 at t=0. t-1 Add noise to x t The process is expressed as:

[0013]

[0014] In the formula, represents Gaussian distribution, α t is a fixed hyperparameter, I is the identity matrix;

[0015] Then we get the marginal distribution q(x t |x0), that is, adding noise from x0 to x t :

[0016]

[0017] In the formula, the reverse process of the diffusion model is based on x t And through the neural network f θ (·) Infer xt-1 , predict the noise at each step:

[0018] p(x t-1 |x t )=q(x t-1 |x t ,f θ (x t ,t));

[0019] Therefore, the restored image z and the clear image x are denoised with a time step t:

[0020]

[0021]

[0022] Where x0 and z0 are the initial clear image and restored image at t = 0 during the diffusion process, i.e., no noise is added; is a hyperparameter of the diffusion model, reflecting the intensity of the noise addition; ε is the noise sampled from the Gaussian distribution, t is the time step sampled from the uniform distribution; x t and z t The clear image and restored image are obtained by adding noise to x0 and z0 at time step t;

[0023] The obtained x t and z t Input them into the pre-trained diffusion model respectively, and let the diffusion model output the results under the condition of time step t:

[0024]

[0025]

[0026] In the formula, and Respectively represent x t and z t The predicted noise output at time step t after being input into the diffusion model;

[0027] based on and The natural property constraint L constructed nat for:

[0028]

[0029] In the formula, ‖·‖2 means taking the 2-norm.

[0030] Furthermore, the diffusion model semantic attribute constraint extraction process specifically includes:

[0031] Disassemble the U-net network in the diffusion model into the encoder ε θ (·) and decoder The noisy clear image and the noisy restored image are used as input to the diffusion model to extract the encoder ε θ The intermediate features of the output of (·) are used as features of semantic attributes:

[0032]

[0033]

[0034] In the formula, x t and z t The clear image and the restored image are obtained by adding noise to the initial clear image and the initial restored image z0 at the time step t. and Respectively represent x t and z t After inputting into the diffusion model, the encoder ε at time step t θ (·) The output intermediate features;

[0035] Will and Construct semantic attribute constraints L as semantic features sem :

[0036]

[0037] In the formula, ‖·‖2 means taking the 2-norm.

[0038] Furthermore, the dual attribute constrained image restoration network training process specifically includes:

[0039] Constraints based on natural attributes L nat , semantic attribute constraint L sem As well as the existing reconstruction constraints used to train the image restoration network, the overall loss function L of the image restoration network is constructed total :

[0040] L total =‖xz‖2+γL sem +λL nat ;

[0041] Where ‖xz‖2 is the existing reconstruction constraint used to train the image restoration network, ‖·‖2 means taking the 2-norm; λ and γ are weight factors.

[0042] Generally, λ is set to 10. -3 , set γ to 10 -5After the training is completed, the diffusion model can be removed and only the trained image restoration network is retained, so no additional calculation is introduced.

[0043] Compared with the prior art, the beneficial technical effects of the present invention are:

[0044] 1. The present invention designs a new training loss function based on deep learning technology to improve the performance of the image restoration network, and proposes a method for training the image restoration network using the dual attribute constraints of the diffusion model, so as to achieve better visual quality. The method can be used as a general training paradigm for various image restoration tasks. First, a diffusion model natural attribute constraint extraction process is designed, and the clear image and the restored image are subjected to the diffusion model after being denoised to obtain the predicted noise as the natural attribute constraint. Then, a diffusion model semantic attribute constraint extraction process is designed, and the intermediate features of the diffusion model are used as the semantic attribute constraint. Finally, a dual attribute constraint image restoration network training process is designed, and the natural attribute constraint and the semantic attribute constraint are combined with the original reconstruction loss to construct the loss function, so as to realize multi-angle loss constraints, thereby training an image restoration network model with better performance.

[0045] 2. The present invention can implicitly use the diffusion model in the process of training the network by extracting natural attribute constraints and semantic attribute constraints in the diffusion model, thus avoiding the time-consuming diffusion process and reasoning process. The present invention has shown excellent performance and generalization ability when applied to several data sets of image restoration, including image deraining, image defogging and low-light enhancement tasks. Due to its semantic restoration ability, the present invention also improves the performance of downstream image recognition tasks. BRIEF DESCRIPTION OF THE DRAWINGS

[0046] Figure 1 This is a flow chart of the image restoration network training method of the present invention;

[0047] Figure 2 A schematic diagram showing the comparison between the distribution-level constraints of the diffusion model introduced in the present invention and the original pixel-level constraints;

[0048] Figure 3 A schematic diagram of training an image restoration network under the dual constraints of a diffusion model according to the present invention;

[0049] Figure 4 Detailed schematic diagram of the dual constraints of the diffusion model of the present invention. DETAILED DESCRIPTION

[0050] A preferred embodiment of the present invention is described in detail below with reference to the accompanying drawings.

[0051] like Figure 1As shown, the present invention is applied to the field of image restoration technology and proposes an image restoration network training method based on dual attribute constraints of a diffusion model, comprising the following steps:

[0052] S1: First, a diffusion model natural attribute constraint extraction process is designed. The clear image and the restored image are input into the diffusion model after adding noise, and the obtained noise prediction is used as the constraint on the natural attribute;

[0053] S2: Then, a diffusion model semantic attribute constraint extraction process is designed. The clear image and the restored image are input into the diffusion model after adding noise, and the obtained intermediate features are used as constraints on semantic attributes.

[0054] S3: Finally, a dual-attribute-constrained image restoration network training process is designed. The above two attribute constraints are combined with the original reconstruction constraints to construct a loss function and train the existing image restoration network.

[0055] The overall framework of the present invention is as follows Figure 2 shown.

[0056] In step S1, the natural attribute constraint extraction process of the designed diffusion model is as follows: Figure 3 As shown in the figure, firstly, the degraded image y is passed through the image restoration network to obtain the restored image z, which is then denoised together with the clear image x according to the forward denoising process of the diffusion model with a time step of t. In the forward denoising process of the diffusion model, Gaussian noise is generally repeatedly added to a clear image x0 at t=0, and x is denoised. t-1 Add noise to x t The process should be carried out according to the following formula:

[0057]

[0058] In the formula, represents Gaussian distribution, α t is a fixed hyperparameter and I is the identity matrix.

[0059] From the above formula, we can directly get its marginal distribution q(x t |x0), that is, adding noise from x0 to x t , which can be expressed as:

[0060]

[0061] In the formula, is also a fixed hyperparameter.

[0062] The reverse process of the diffusion model is to use x t Infer x t-1 , this process is to use the neural network f θ (·) to predict the noise at each step, which can be expressed as:

[0063] p(x t-1 |x t )=q(x t-1 |x t ,f θ (x t ,t));

[0064] When adding noise to the restored image z and the clear image x with a time step t, the following formula should be used:

[0065]

[0066]

[0067] Where x0 and z0 are the initial states at t = 0 during the diffusion process, i.e., no noise is added. is a hyperparameter of the diffusion model, reflecting the strength of the noise addition. ε is the noise sampled from a Gaussian distribution, and t is the time step sampled from a uniform distribution. t and z t It is the state after adding noise to x0 and z0 at time step t.

[0068] The x obtained by the above process t and z t They are respectively used as input into the diffusion model trained on a large-scale data set, and the diffusion model is allowed to output the results under the condition of time step t. This process can be expressed as:

[0069]

[0070]

[0071] Where f θ (·) represents the neural network used in the diffusion model, and Respectively represent x t and z t The noise prediction result output at time step t after being input into the diffusion model.

[0072] Get the above output results and After that, it is used as the constraint design loss function of natural attributes, which can be expressed as:

[0073]

[0074] In the formula, ‖·‖2 means taking the 2-norm.

[0075] In step S2, the designed diffusion model semantic attribute constraint extraction process is as follows: Figure 3 As shown, the U-net network in the diffusion model is disassembled into the encoder ε θ (·) and decoder The above clear image after adding noise x t and the restored image z after adding noise t As input to the diffusion model, extract the encoder ε θ The output feature of (·) is used as the feature of semantic attribute, which can be expressed as:

[0076]

[0077]

[0078] In the formula, and Respectively represent x t and z t After inputting into the diffusion model, the encoder ε at time step t θ After obtaining the above output results, it is used as the constraint design loss function of the semantic attribute, which can be expressed as:

[0079]

[0080] In step S3, when designing the dual attribute constrained image restoration network training process, the natural attribute constraint L obtained above is nat and semantic attribute constraints L sem Combined with the existing reconstruction constraints used to train the image restoration network model, the network is trained as the overall loss function, which can be expressed as:

[0081] L total =‖xz‖2+γL sem +λL nat ;

[0082] Where ‖xz‖2 is the reconstruction constraint of the original training image restoration network model, and λ and γ are weight factors. Generally, λ is set to 10 -3 , set γ to 10 -5 It is appropriate.

[0083] The image restoration network training method based on the dual attribute constraints of the diffusion model proposed in the present invention can restore the natural attributes of the image, and has improved performance in three common image restoration tasks: image deraining, dehazing, and low-light enhancement. In the image deraining task, the Rain13k dataset is used for training, including 13,712 pairs of data, and the Rain100H dataset is used for testing, including 100 pairs of data. In the image dehazing task, the Dense-Haze dataset is used for training and testing, which includes 50 pairs of training data and 5 pairs of test data. When the LOL dataset is used for training and testing, it includes 450 pairs of training data and 50 pairs of test data.

[0084] The image restoration network training method based on dual attribute constraints of the diffusion model proposed in the present invention can restore the semantic attributes of the image and is tested on the CUB dataset, which includes 5994 training images and 5794 test images.

[0085] In the training phase, the Adam optimizer was used for parameter optimization, the learning rate was set to 0.001, and the number of training rounds was 500. The entire experiment was conducted on an NVIDIA GeForce RTX 3090 graphics card, and PSNR, SSIM, LPIPS, and recognition accuracy were used as evaluation indicators.

[0086] After the training is completed, the diffusion model can be removed and only the trained image restoration network is retained, so no additional calculations are introduced.

[0087] In summary, the present invention proposes a method for training an image restoration network based on dual attribute constraints of a diffusion model. The clear image and the restored image are noisy and then input into the diffusion model to obtain the predicted noise representing the natural attribute and the intermediate feature representing the semantic attribute. The features of the two attributes are combined as constraints with the original reconstruction constraints to train the existing image restoration network to achieve a more natural and semantically informative restoration effect. This method can be used as a general training paradigm in multiple image restoration tasks and is helpful for downstream recognition tasks.

[0088] It is obvious to those skilled in the art that the present invention is not limited to the details of the exemplary embodiments described above, and that the present invention can be implemented in other specific forms without departing from the spirit or essential features of the present invention. Therefore, from any point of view, the embodiments should be regarded as exemplary and non-limiting, and the scope of the present invention is defined by the appended claims rather than the above description, and it is intended that all changes falling within the meaning and scope of the equivalent elements of the claims are included in the present invention, and any reference numerals in the claims should not be regarded as limiting the claims involved.

[0089] In addition, it should be understood that although this specification is described in terms of implementation methods, not every implementation method contains only one independent technical solution. This narrative method of the specification is only for the sake of clarity. Those skilled in the art should regard the specification as a whole. The technical solutions in each embodiment may also be appropriately combined to form other implementation methods that can be understood by those skilled in the art.

Claims

1. A method for training an image restoration network based on dual attribute constraints of a diffusion model, characterized in that: include: Design the natural attribute constraint extraction process of the diffusion model: pass the degraded image through the image restoration network to obtain the restored image; The clear image and the restored image are input into the diffusion model after adding noise, and the predicted noise of the clear image and the restored image are obtained respectively. The predicted noise difference between the clear image and the restored image is used as the natural attribute constraint L nat ; Design the diffusion model semantic attribute constraint extraction process: Similarly, the clear image and the restored image are input into the diffusion model after adding noise, and the semantic features of the image restoration network middle layer corresponding to the clear image and the restored image are obtained respectively. The semantic attribute constraint L is constructed using the semantic features between the clear image and the restored image. sem ; Design a dual-attribute-constrained image restoration network training process: use the natural attribute constraints and semantic attribute constraints, combined with the original reconstruction constraints to construct a loss function, and train the existing image restoration network; thereby improving the visual quality and semantic attributes of the restored image output by the image restoration network, and obtaining the best restoration result.

2. The image restoration network training method based on dual attribute constraints of diffusion model according to claim 1 is characterized in that: The natural attribute constraint extraction process of the diffusion model specifically includes: First, the degraded image y is passed through the image restoration network to obtain the restored image z. z is then denoised together with the clear image x according to the forward denoising process of the diffusion model with a time step of t. In the forward denoising process of the diffusion model, Gaussian noise is repeatedly added to the clear image x0 at t=0. t-1 Add noise to x t The process is expressed as: In the formula, represents Gaussian distribution, α t is a fixed hyperparameter, I is the identity matrix; Then we get the marginal distribution q(x t |x0), that is, adding noise from x0 to x t : In the formula, the reverse process of the diffusion model is based on x t And through the neural network f θ (·) Infer x t-1 , predict the noise at each step: p(x t-1 |x t )=q(x t-1 |x t ,f β (x t ,t)); Therefore, the restored image z and the clear image x are denoised with a time step t: Where x0 and z0 are the initial clear image and restored image at t = 0 during the diffusion process, i.e., no noise is added; is a hyperparameter of the diffusion model, reflecting the intensity of the noise addition; ε is the noise sampled from the Gaussian distribution, t is the time step sampled from the uniform distribution; x t and z t The clear image and restored image are obtained by adding noise to x0 and z0 at time step t; The obtained x t and z t Input them into the pre-trained diffusion model respectively, and let the diffusion model output the results under the condition of time step t: In the formula, and Respectively represent x t and z t The predicted noise output at time step t after being input into the diffusion model; based on and The natural property constraint L constructed nat for: In the formula, ‖·‖2 means taking the 2-norm.

3. The image restoration network training method based on dual attribute constraints of diffusion model according to claim 1 is characterized in that: The diffusion model semantic attribute constraint extraction process specifically includes: Disassemble the U-net network in the diffusion model into the encoder ε θ (·) and decoder The noisy clear image and the noisy restored image are used as input to the diffusion model to extract the encoder ε θ The intermediate features of the output of (·) are used as features of semantic attributes: In the formula, x t and z t The clear image and the restored image are obtained by adding noise to the initial clear image and the initial restored image z0 at the time step t. and Respectively represent x t and z t After inputting into the diffusion model, the encoder ε at time step t θ (·) The output intermediate features; Will and Construct semantic attribute constraints L as semantic features sem : In the formula, ‖·‖2 means taking the 2-norm.

4. The image restoration network training method based on dual attribute constraints of diffusion model according to claim 1 is characterized in that: The dual attribute constrained image restoration network training process specifically includes: Constraints based on natural attributes nat , semantic attribute constraint L sem As well as the existing reconstruction constraints used to train the image restoration network, the overall loss function L of the image restoration network is constructed total : THE total =‖xz‖2+γL sem +λL nat ; Where ‖xz‖2 is the existing reconstruction constraint used to train the image restoration network, ‖·‖2 means taking the 2-norm; λ and γ are weight factors.

Images