A damaged image restoration method based on diffusion model

Abstract

The invention discloses a damaged image restoration method based on a diffusion model. Firstly, the damaged image to be restored is compressed into a low-dimensional space through a feature extraction module to extract a low-dimensional feature map; then the low-dimensional feature map is input into an adaptive noise enhancement module to calculate a noise enhancement coefficient; at the same time, the low-dimensional feature is linearly combined with Gaussian noise by obtaining the noise enhancement coefficient to obtain the image feature after noise enhancement; then, a trained noise predictor is applied to the image feature after noise enhancement to gradually remove the noise in the noise enhancement feature through iteration; finally, the image feature after iterative denoising is input into a pixel restoration module to obtain a restored image.

Description

Technical Field

[0001] The present invention belongs to the technical field of image processing, and more specifically, relates to a damaged image repair method based on a diffusion model. Background Art

[0002] Image restoration refers to the use of certain algorithms and techniques to repair and restore damaged or degraded images in order to restore the integrity and accuracy of the original image as much as possible.

[0003] With the advent of the digital age, image restoration technology has been widely used in various fields, such as cultural heritage protection, historical photo restoration, medical image processing, digital art creation, etc. Therefore, the research on image restoration technology is of great significance. In recent years, researchers have proposed many image restoration algorithms, including algorithms based on partial differential equations, algorithms based on statistical learning, algorithms based on deep learning, etc. These algorithms have achieved certain results in image restoration, but there are also some problems, such as unsatisfactory restoration effects, large amount of calculations, and poor robustness. Recently, the diffusion model, as a generative model, has achieved outstanding performance in image generation, and the generated effect is better than the previous convolutional neural network. The image restoration problem can be regarded as a problem of generating the content of the damaged area of ​​the image. Therefore, it is a reasonable idea to apply the diffusion model to the field of image restoration. Summary of the invention

[0004] The purpose of the present invention is to overcome the shortcomings of the prior art and provide a damaged image repair method based on a diffusion model to repair or restore damaged or missing images to restore the original quality and content of the image.

[0005] To achieve the above-mentioned object of the invention, the present invention provides a damaged image restoration method based on a diffusion model, characterized in that it comprises the following steps:

[0006] (1) Obtain a damaged image X to be repaired;

[0007] (2) Input the damaged image X into the image compression module and extract the low-dimensional feature map Z;

[0008] (3) Input the low-dimensional feature map Z into the trained adaptive noise enhancer φ θ In the above formula, we get the enhancement coefficient γ, and then enhance the low-dimensional feature map Z by the following formula to get the enhanced feature map

[0009]

[0010] Among them, ε is random Gaussian noise;

[0011] (4) Through the trained noise predictor ε θ For feature maps Perform denoising;

[0012] (4.1) Set the maximum number of iteration steps T, and the current number of iteration steps t = T, T-1, ..., 1;

[0013] (4.2), the feature map Input to the noise predictor ε θ The noise amount predicted for the Tth iteration is obtained

[0014] (4.3) Calculate the mean after the Tth iteration and variance All obey Gaussian distribution;

[0015]

[0016]

[0017] in, β T =1-α T , α t is the preset parameter at the tth iteration;

[0018] (4.4), for the mean and variance Sampling is performed to obtain the input of the noise predictor at the T-1th iteration Where I is the identity matrix;

[0019] (4.5) By iterating steps (4.2)-(4.4) T times, the denoised feature map can be obtained.

[0020] (5) The denoised feature map Z is input into the pixel restoration module, and decoded by the pixel restoration module to obtain a restored image Z with the same size as Z. * .

[0021] The object of the invention of the present invention is achieved in this way:

[0022] The damaged image restoration method based on the diffusion model of the present invention first compresses the damaged image to be restored into a low-dimensional space through a feature extraction module to extract a low-dimensional feature map; then the low-dimensional feature map is input into an adaptive noise enhancement module to calculate the noise enhancement coefficient; at the same time, the low-dimensional feature is linearly combined with Gaussian noise by obtaining the noise enhancement coefficient to obtain the image feature after noise enhancement; then, the trained noise predictor is applied to the image feature after noise enhancement, and the noise in the noise enhancement feature is gradually removed through iteration; finally, the image feature after iterative denoising is input into a pixel restoration module to obtain a restored image.

[0023] At the same time, the damaged image restoration method based on the diffusion model of the present invention also has the following beneficial effects:

[0024] (1) By using the diffusion model training method, a noise predictor that can be used for iterative denoising is obtained. Compared with the one-time denoising of traditional image denoising methods, this method iteratively removes noise and retains more detailed information of the image.

[0025] (2) An adaptive noise enhancement module is used to allow image features to adaptively diffuse to a reasonable denoising range, and iterative denoising is performed by setting a smaller denoising step number. This greatly reduces the amount of computation compared to the traditional diffusion model that requires iterative denoising from Gaussian noise.

[0026] (3) By compressing the damaged image into a low-dimensional space for denoising, it not only speeds up the calculation and reduces memory consumption, but also reduces the interference of redundant information in the image, thus improving the final denoising effect. BRIEF DESCRIPTION OF THE DRAWINGS

[0027] Figure 1 It is a flow chart of the damaged image repair method based on the diffusion model of the present invention;

[0028] Figure 2 is a schematic diagram of an adaptive noise enhancer;

[0029] Figure 3 is a schematic diagram of the noise predictor. DETAILED DESCRIPTION

[0030] The specific implementation of the present invention is described below in conjunction with the accompanying drawings so that those skilled in the art can better understand the present invention. It should be noted that in the following description, when the detailed description of known functions and designs may dilute the main content of the present invention, these descriptions will be omitted here.

[0031] Example

[0032] Figure 1 It is a flow chart of the damaged image repair method based on the diffusion model of the present invention.

[0033] In this embodiment, an adaptive noise enhancer φ is used when repairing a damaged image based on a diffusion model. θ and the noise predictor ε θ , whose structure is as follows Figure 2 , 3 As shown; therefore we need to first use a large number of pictures in the picture library to train the adaptive noise enhancer φ θ and the noise predictor ε θ , the specific training process is as follows:

[0034] First, train the adaptive noise enhancer φ θ :

[0035] (1) Download N complete images and input each complete image into the mask image generation module to generate the corresponding mask image;

[0036] In this embodiment, the mask image generation module mainly achieves the effect of blocking the image by generating a random mask. The specific generation process is as follows: assuming that the pixels of the existing complete image Z are w*h*c, where w is the width of the image, h is the height of the image, and c is the number of channels of the image; then, a mask matrix M∈R with all 1s equal to the image size is generated. w*h ; Then, replace the values ​​of some random areas in the mask matrix with 0; set random numbers a and b, and 0<a<w, 0<b<h. Randomly select the i-th row and j-th column in the mask matrix as the mask starting point, and set all values ​​in the area from the i-th row to the i+b row and the j-th column to the j+a column in the mask matrix to 0. Repeat the above operations to generate 1000 mask matrices with different mask shapes. Randomly select a mask matrix and multiply it by the image pixel value to get the mask image MZ=M*Z;

[0037] (2) Each mask image is input into the image compression module, and features are extracted through the image compression module to extract the low-dimensional feature map MZ i , where MZ i Represents the low-dimensional feature map extracted from the i-th mask image, i = 1, 2, ..., N;

[0038] In this embodiment, in order to better extract the semantic information in the image and speed up the processing speed of the image, this embodiment uses an image compression module to extract features from the original size image. The specific details are as follows: each complete image Z and its corresponding mask image MZ are input into the image compression module as a set of training sample pairs. The image compression module consists of two encoders E1 and E2, and the two encoders are used to compress the complete image and the mask image respectively. The two encoders have the same network structure, each encoder consists of 4 layers of two-dimensional convolutional layers and activation layers, the convolution kernel size of each two-dimensional convolutional layer is 4, the step size is 2, and the padding is 1. Each complete image Z and its corresponding mask image MZ are input into an encoder respectively, so as to extract the corresponding low-dimensional feature map.

[0039] (3) In this embodiment, considering that the noise in the MZ cannot be predicted in advance, the noise predictor ε is directly used. θ It is not possible to obtain accurate denoising results by denoising it. Therefore, an adaptive noise enhancer φ is proposed. θ Adjust the noise of MZ. The details are as follows:

[0040] The i-th low-dimensional feature map MZ i Input to the adaptive noise enhancer φ θ In the output, a noise enhancement coefficient γ is output;

[0041] γ=softmax(Con(Relu(Con(Relu(Con(MZ i ))))))

[0042] Where γ is a real number between 0 and 1;

[0043] (4) For the i-th dimension feature map MZ i Perform noise enhancement processing:

[0044]

[0045] Among them, ε is random Gaussian noise;

[0046] (5) Through the trained noise predictor ε θ Noise Enhancement Feature Map Perform denoising;

[0047] (5.1) Set the maximum number of iteration steps T = 1000, and the current number of iteration steps t = T, T-1, ..., 1;

[0048] (5.2) The noise enhancement feature map Input to the trained noise predictor ε θ The noise amount predicted for the Tth iteration is obtained

[0049] (5.3) Calculate the mean after the Tth iteration and variance All obey Gaussian distribution;

[0050]

[0051]

[0052] in, β T =1-α T , α t is the preset parameter for the tth iteration, α t It is generally set to an arithmetic progression from 0.001 to 0.02. In this embodiment, a total of 1000 items are set, namely, α t ∈{0.001,...,0.02};

[0053] (5.4), for the mean and variance Sampling is performed to obtain the input of the noise predictor at the T-1th iteration Where I is the identity matrix;

[0054] (5.5) By iterating steps (5.2)-(5.4) T times, the denoised feature map can be obtained.

[0055] (6) The feature map Input to the pixel restoration module, decode it through the pixel restoration module, and get the same as MZ i The restored images have the same size

[0056] (7) Calculate the loss function value after this round of iteration;

[0057]

[0058] (8) Update the adaptive noise enhancer φ through back propagation of the loss value θ Then the next round of training is carried out until the adaptive noise enhancer φ θ convergence.

[0059] Next, train the noise predictor ε θ :

[0060] 1) Download N complete images and input them into the image compression module, perform feature extraction through the image compression module, and extract low-dimensional feature maps Z i , where Zi Represents the low-dimensional feature map extracted from the i-th complete image, i = 1, 2, ..., N;

[0061] 2) Use the noise perturbation module to perturb the i-th low-dimensional feature map Z i Add random noise T times continuously; the formula is as follows:

[0062]

[0063] Where, t=1,2,…,T, parameter α k is the preset parameter for the kth iteration, ε t is the Gaussian noise added for the tth time;

[0064] 3) The noise feature map Z i (t) Input to the noise predictor ε θ , predict each time to add to the i-th low-dimensional feature map Z i The amount of noise

[0065] 4) Calculate the loss function value after this round of iteration;

[0066]

[0067] 5) Update the noise predictor ε by back propagating the loss value θ Then the next round of training is carried out until the noise predictor ε θ convergence.

[0068] After training the adaptive noise enhancer φ θ and the noise predictor ε θ Next, we will explain in detail the damaged image restoration method based on the diffusion model of the present invention. Figure 1 As shown, the following steps are included:

[0069] S1. Obtain a damaged image X to be repaired;

[0070] S2, input the damaged image X into the image compression module and extract the low-dimensional feature map Z;

[0071] S3, input the low-dimensional feature map Z into the trained adaptive noise enhancer φ θ In the above formula, we get the enhancement coefficient γ, and then enhance the low-dimensional feature map Z by the following formula to get the enhanced feature map

[0072]

[0073] Among them, ε is random Gaussian noise;

[0074] S4, through the trained noise predictor ε θ For feature maps Perform denoising;

[0075] S4.1, set the maximum number of iteration steps T, the current number of iteration steps t = T, T-1, ..., 1;

[0076] S4.2. Feature map Input to the noise predictor ε θ The noise amount predicted for the Tth iteration is obtained

[0077]

[0078] S4.3. Calculate the mean after the Tth iteration and variance All obey Gaussian distribution;

[0079]

[0080]

[0081] S4.4. Mean and variance Sampling is performed to obtain the input of the noise predictor at the T-1th iteration

[0082] S4.5. By iterating steps S4.2-S4.4 T times, the denoised feature map can be obtained.

[0083] S5. De-noised feature map Input to the pixel restoration module, decode it through the pixel restoration module, and obtain the repaired image Z with the same size as Z * .

[0084] Although the above describes the illustrative specific embodiments of the present invention to facilitate those skilled in the art to understand the present invention, it should be clear that the present invention is not limited to the scope of the specific embodiments. For those of ordinary skill in the art, as long as various changes are within the spirit and scope of the present invention as defined and determined by the attached claims, these changes are obvious, and all inventions and creations using the concept of the present invention are protected.

Claims

1. A damaged image restoration method based on a diffusion model, characterized in that: The following steps are involved: (1) Obtain a damaged image X to be repaired; (2) Input the damaged image X into the image compression module and extract the low-dimensional feature map Z; (3) Input the low-dimensional feature map Z into the trained adaptive noise enhancer φ θ In the above formula, we get the enhancement coefficient γ, and then enhance the low-dimensional feature map Z by the following formula to get the enhanced feature map Among them, ε is random Gaussian noise; (4) Through the trained noise predictor ε θ For feature maps Perform denoising; (4.1) Set the maximum number of iteration steps T, and the current number of iteration steps t = T, T-1, ..., 1; (4.2), the feature map Input to the noise predictor ε θ The noise amount predicted for the Tth iteration is obtained (4.3), calculate the mean μ after the Tth iteration T and variance β T , μ T , β T All obey Gaussian distribution; in, β T =1-α T , α t is the preset parameter at the tth iteration; (4.4), for the mean μ T and variance β T Sampling is performed to obtain the input of the noise predictor at the T-1th iteration Where I is the identity matrix; (4.5) By iterating steps (4.2)-(4.4) T times, the denoised feature map can be obtained. Among them, the noise predictor ε θ The pre-training process is: S1), download N complete images and input them into the image compression module, perform feature extraction through the image compression module, and extract low-dimensional feature maps Z i , where Z i Represents the low-dimensional feature map extracted from the i-th complete image, i = 1, 2, ..., N; S2), use the noise perturbation module to the i-th low-dimensional feature map Z i Add random noise T times continuously; the formula is as follows: Where, t=1,2,…,T, parameter α k is the preset parameter for the kth iteration, ε t is the Gaussian noise added for the tth time; S3), the noise feature map Z i (t) Input to the noise predictor ε θ , predict each time to add to the i-th low-dimensional feature map Z i The amount of noise S4), calculate the loss function value after this round of iteration; S5) Update the noise predictor ε by back propagating the loss value θ Then the next round of training is carried out until the noise predictor ε θ convergence; (5) The denoised feature map Input to the pixel restoration module, decode it through the pixel restoration module, and obtain the repaired image Z with the same size as Z * .

2. The damaged image restoration method based on the diffusion model according to claim 1 is characterized in that: The adaptive noise enhancer φ θ The pre-training process is: (2.1) Download N complete images, input each complete image into the mask image generation module, and generate the corresponding mask image; (2.2) Each mask image is input into the image compression module, and features are extracted through the image compression module to extract the low-dimensional feature map MZ i , where MZ i Represents the low-dimensional feature map extracted from the i-th mask image, i = 1, 2, ..., N; (2.3) Take the i-th low-dimensional feature map MZ i Input to the adaptive noise enhancer φ θ In the output, a noise enhancement coefficient γ is output; γ=softmax(Con(Relu(Con(Relu(Con(MZ i )))))) Where γ is a real number between 0 and 1; (2.4), for the i-th dimension feature map MZ i Perform noise enhancement processing: Among them, ε is random Gaussian noise; (2.5) through the trained noise predictor ε θ Noise Enhancement Feature Map Perform denoising; (2.5.1) Set the maximum number of iteration steps T, the current number of iteration steps t = T, T-1, ..., 1; (2.5.2) The noise enhancement feature map Input to the trained noise predictor ε θ The noise amount predicted for the Tth iteration is obtained (2.5.3) Calculate the mean μ after the Tth iteration T and variance β T , μ T , β T All obey Gaussian distribution; in, β T =1-α T , α t is the preset parameter at the tth iteration; (2.5.4), for the mean μ T and variance β T Sampling is performed to obtain the input of the noise predictor at the T-1th iteration Where I is the identity matrix; (2.5.5) By iterating steps (2.5.2)-(2.5.4) T times, the denoised feature map Z can be obtained; (2.6) The feature map Input to the pixel restoration module, decode it through the pixel restoration module, and get the same as MZ i The restored images of the same size M·MZ i ; (2.7) Calculate the loss function value after this round of iteration; (2.8) The loss value is updated through back propagation to the adaptive noise enhancer φ θ Then the next round of training is carried out until the adaptive noise enhancer φ θ convergence.

Images