Diffusion model image completion method based on boundary perception

Abstract

The invention discloses a diffusion model image completion method based on boundary preserving perception, and belongs to the technical field of artificial intelligence and image generation. The method comprises the following steps: firstly, preprocessing an input image and a corresponding defect mask, and extracting gradient attenuation boundary features and a structure edge graph; a boundary sensing convolution module is introduced into the completion model and is used for boundary information modeling of a coding end; meanwhile, a gating mechanism is adopted in U-Net jump connection, the weight of a gradient attenuation mask is combined, and the boundary information transmission intensity is dynamically adjusted; furthermore, in the diffusion generation process, a region self-adaptive scheduling strategy is adopted, only high-step diffusion iteration is carried out on a defect region and a boundary neighborhood, and low-step or freezing operation is carried out on a non-defect region, so that balance between boundary consistency and calculation efficiency is realized. Experimental results show that the method is superior to an existing diffusion model image completion method in the aspects of structure retentivity, edge continuity and visual quality, and has wide application prospects.

Description

TECHNICAL FIELD

[0001] The present application relates to the field of image processing and generation, in particular to an image completion method based on a diffusion model, and more particularly to a boundary-preserved perceptual local diffusion image completion model, belonging to the cross-application category of artificial intelligence and image generation technology. BACKGROUND

[0002] The image completion task aims to reconstruct the partially missing or occluded image region with high quality and context consistency. Traditional methods mostly rely on convolutional neural networks or generative adversarial networks (GAN), but have limitations in handling large missing areas or maintaining semantic consistency. In recent years, diffusion probability models (Diffusion Models) have become the representative of the new generation of image generation technology due to their high generation quality and strong stability.

[0003] The image completion method using diffusion probability model significantly improves the authenticity of the completion. However, there are still the following shortcomings: boundary discontinuity (obvious edge breakage or fuzzy transition often occurs at the junction of damaged and intact regions), structural distortion (the internal structure of the missing region often lacks consistency with the boundary region, leading to distortion of the completion result), and insufficient use of local information (most existing diffusion models generate in a global range, failing to fully utilize the boundary gradient and structural information of the damaged region).

[0004] To solve the above problems, the present application proposes a diffusion model image completion method based on boundary perception. Specifically, first, the input image and the corresponding missing mask are preprocessed to extract gradient attenuation boundary features and structural edge maps; a boundary perception convolution module is introduced in the completion model to model the boundary information at the encoding end; at the same time, a gating mechanism is used in the U-Net jump connection, combined with the weight of the gradient attenuation mask, to dynamically adjust the boundary information transmission strength. Further, in the diffusion generation process, a region adaptive scheduling strategy is adopted, only the missing region and the boundary neighborhood are iterated with high step number diffusion, while the non-missing region is operated with low step number or frozen to balance the boundary consistency and computational efficiency. This method optimizes the structure of the diffusion model, introduces the boundary perception mechanism, achieves both global consistency and natural boundary transition, and thus obtains a higher quality of the completed image. SUMMARY

[0005] The present application aims to overcome the problems of boundary discontinuity and structural distortion in the prior art, and proposes a boundary-preserved perceptual diffusion image completion model, which introduces a distance attenuation map and a boundary gradient map, optimizes the noise prediction model, and strengthens the boundary constraint during the completion process, thereby significantly improving the structural consistency and detail fidelity of the completion.

[0006] This invention provides an image completion method based on a boundary-aware diffusion model, comprising the following steps:

[0007] Step 1: Build a data processing module, which is used to process data from the training dataset. (H is the image height, W is the image width) Training images are collected. And generate the corresponding mask data for the region to be filled in. Known areas of the original image Distance decay map and boundary gradient map ;

[0008] Step 1.1: Obtain the input and perform preprocessing to obtain the image to be completed. (C is the number of channels, H is the image height, and W is the image width), complete the region mask. Known areas of the original image And perform normalization processing; among them, () represents different numbers of channels, H represents the image height, and W represents the image width;

[0009] Step 1.2: Calculate the distance from each damaged pixel to the nearest boundary based on Euclidean distance transform, normalize the distance, and obtain the distance attenuation mask. ;

[0010] Step 1.3: Perform gradient operator on the input image to obtain the edge map. Combine the boundary region of the mask of the completed region, and retain only the edge information near the completed region to obtain the boundary-aware map. ;

[0011] Step 2: Training the prediction generator network involves a forward diffusion process and a backward denoising process: the forward diffusion process consists of T steps, progressively feeding training images... By adding Gaussian noise, noisy data corresponding to different diffusion time steps are constructed. The above data is then fed into the noise prediction network. The noise prediction network will predict the distribution of Gaussian noise injected in the corresponding diffusion time step based on the input noisy data. The noise prediction network is trained by the loss function to accurately fit the noise distribution (the purpose is to train the inverse denoising capability of the noise prediction network).

[0012] Step 2.1: According to the forward diffusion formula The training preprocessing module obtained Noise-adding processing is performed to obtain noisy data corresponding to different diffusion time steps. Where T is the total number of forward diffusion steps, Indicates the diffusion time step. This represents the noisy data corresponding to the diffusion time step t. represents the noise intensity added at each step, is the multiplication result of the noise corresponding to time step 1 to time step t, , is a noise obeying Gaussian distribution, is a unit variance matrix;

[0013] Step 2.2: Training the Edge U-Net noise prediction network (the U-Net network optimized by the present patent), specifically comprising the following sub-steps:

[0014] Step 2.2.1: Concatenate the completed region mask map, the known region of the original image, the distance attenuation map, and the boundary gradient map in the channel dimension, and input them into the boundary-aware convolution module to obtain the output , after global pooling is concatenated with the sampling map in the channel dimension to obtain the input of the Edge U-Net ; wherein, ) are different channel numbers;

[0015] Step 2.2.2: Input into the first layer of the encoder, and perform convolution on to obtain , input and time step into the residual block to obtain the output of the current encoding layer ; wherein, is the channel number of the time step;

[0016] Step 2.2.3: Downsample by two-dimensional convolution to obtain the input of the next encoding layer;

[0017] Step 2.2.4: Input into the second layer of the encoder, and input and time step into the residual block to obtain the output of the current encoding layer ;

[0018] Step 2.2.5: Downsample by two-dimensional convolution to obtain the input of the next encoding layer;

[0019] Step 2.2.6: Input into the third layer of the encoder, and input and time step into the residual block to obtain the output of the current encoding layer ;

[0020] Step 2.2.7: Downsample the output of the previous encoding layer by a two-dimensional convolution to get the input of the next encoding layer ; ;

[0021] Step 2.2.8: Send the input into the fourth layer of the encoder, merge it with the time step , and send the result into the residual block to get the output of the current encoding layer ; ; ;

[0022] Step 2.2.9: Send the output of the previous decoding layer into the intermediate block (composed of the residual block and the self-attention mechanism) to get the input of the current decoding layer ; ;

[0023] Step 2.2.10: Send the input into the fourth layer of the decoder, merge it with the time step , and send the result into the residual block to get the output of the current decoding layer ; where is the output of the corresponding encoding layer, and is the input of the current decoding layer. ; ; ; Step 2.2.11: Upsample the output of the previous decoding layer by an inverse two-dimensional convolution to get the input of the next decoding layer ;

[0024] ; Step 2.2.12: Send the input into the third layer of the decoder, merge it with the time step , and send the result into the residual block to get the output of the current decoding layer ; where

[0025] is the output of the corresponding encoding layer, and is the input of the current decoding layer. ; ; ; ; ; ; Step 2.2.13: Upsample the output of the previous decoding layer by an inverse two-dimensional convolution to get the input of the next decoding layer ;

[0026] ; Step 2.2.14: Send the input into the second layer of the decoder, merge it with the time step , and send the result into the residual block to get the output of the current decoding layer ; where

[0027] is the output of the corresponding encoding layer, and is the input of the current decoding layer. ; ;After merging, the formula is used to perform skip-connection with the output of the corresponding encoding layer , and then sent into the residual block to obtain the output of the current decoding layer ; wherein, is the output of the corresponding encoding layer, is the input of the current decoding layer;

[0028] Step 2.2.15: upsample by inverse two-dimensional convolution to obtain the input of the next decoding layer ;

[0029] Step 2.2.16: send into the first layer of the decoder, and after merging with the time step , the formula is used to perform skip-connection with the output of the corresponding encoding layer , and then sent into the residual block to obtain the output of the current decoding layer ; wherein, is the output of the corresponding encoding layer, is the input of the current decoding layer;

[0030] Step 2.2.17: perform convolution on to obtain the model prediction output ;

[0031] Step 2.3: the prediction output is the noise estimate of the corresponding time step, denoted as ; the training of each time step will perform additional local iterations on the to-be-completed region; during the training process, the optimization goal is to minimize the loss function ; wherein the diffusion loss , the pixel-level loss , the perception loss , and the boundary gradient loss ; wherein, is the mathematical expectation of the time step t, is the normalization operation, is the time step weight, is the model prediction distribution, is the real distribution, is the weight coefficient of the to-be-completed region, is the to-be-completed region mask, is the predicted image, is the real image, is the feature map of the i-th layer in the network, ​ for In the network Feature mapping of layers, This indicates that the corresponding results for all network layers are summed. To predict the edge gradient map of the image, This is the edge gradient map of the real image. For the weights corresponding to different losses, all learnable parameters in the noise estimation network are updated by performing gradient descent calculation on the loss function until a fit is achieved.

[0032] Step 3: Take the image to be completed as input (the starting point of the reverse denoising process), and use the trained noise prediction network to denoise it step by step. The denoising process takes T steps in total (which can be reduced to T / 2, T / 4, etc. through acceleration methods), and finally obtain a completed image with high structural consistency and high guarantee.

[0033] Step 3.1: Obtain the image to be completed The noise prediction result is obtained by feeding it into the trained noise prediction network. ;

[0034] Step 3.2: According to the reverse diffusion formula Calculations yielded Predicted image at time step According to the formula get Image to be completed at time step ;in, This indicates the noise intensity added at each step. The cumulative product of the noise corresponding to time step 1 to time step t ( ), For noise that follows a Gaussian distribution, These are the weighting coefficients. The known regions of the image to be completed;

[0035] Step 3.3: Repeat the above steps until t=1, which gives you the final completed image.

[0036] The beneficial effects of adopting the above technical solution are as follows:

[0037] The application effectively captures and preserves the gradient features of the damaged area boundary in the completion process by introducing a boundary-aware convolution module and a corresponding skip connection mechanism, thereby avoiding the common edge blur, misplacement or fracture problems in traditional methods. The local diffusion scheduling strategy proposed in the application limits more diffusion iterations to the damaged area and the adjacent boundary, reducing unnecessary computational load, and significantly improving inference efficiency on the premise of ensuring completion quality. The application dynamically weights the boundary area with the help of boundary gradient, making the transition of the completed area to the original image more natural, ensuring the smooth connection and structural consistency of the boundary. The method of the application can adapt to various types of image datasets (such as natural scenes, faces, medical images), and has strong universality and practical value. BRIEF DESCRIPTION OF DRAWINGS

[0038] The technical solutions in the examples of the application will be clearly and completely described below in conjunction with the drawings in the examples of the application. Obviously, the described examples are only some of the embodiments of the application, not all the embodiments; based on the embodiments in the application, all other embodiments obtained by those skilled in the art without creative labor fall within the scope of protection of the application.

[0039] Figure 1 A flowchart of the image completion method based on the boundary-aware diffusion model provided by the embodiments of the application is shown.

[0040] Figure 2 An Edge U-Net noise prediction network diagram provided by the embodiments of the application is shown.

[0041] Figure 3 A boundary-aware convolution module diagram provided by the embodiments of the application is shown.

[0042] Figure 4 An encoder module diagram provided by the embodiments of the application is shown.

[0043] Figure 5 A residual attention module diagram provided by the embodiments of the application is shown.

[0044] Figure 6 An intermediate block module diagram provided by the embodiments of the application is shown.

[0045] Figure 7 A decoder module diagram provided by the embodiments of the application is shown.

[0046] Figure 8 An input image and a damaged mask diagram is shown.

[0047] Figure 9 A distance attenuation and boundary awareness diagram is shown.

[0048] Figure 10This paper compares the image completion method with other mainstream diffusion completion models.

[0049] Figure 11 This paper compares the metrics of our method with those of other mainstream diffusion completion models. Detailed Implementation

[0050] To make the objectives, technical solutions, and advantages of the present invention clearer, the embodiments of the present invention will be further described in detail below with reference to the accompanying drawings.

[0051] like Figure 1 As shown, this embodiment of the invention provides an image completion method based on a boundary-aware diffusion model, including:

[0052] Step 1: Build a data processing module, which is used to process data from the training dataset. In this example (using a cropped 256×256 image dataset), training images are collected. And generate the corresponding mask data for the region to be filled in. Known areas of the original image Distance decay map and boundary gradient map ;

[0053] Step 1.1: Obtain the input and perform preprocessing to obtain the image to be completed. (Number of channels: 3; Image height: 256; Image width: 256) Complete the region mask. Known areas of the original image And perform normalization processing;

[0054] Step 1.2: Calculate the distance from each damaged pixel to the nearest boundary based on Euclidean distance transform, normalize the distance, and obtain the distance attenuation mask. ;

[0055] Step 1.3: Perform gradient operator on the input image to obtain the edge map. Combine the boundary region of the mask of the completed region, and retain only the edge information near the completed region to obtain the boundary-aware map. ;

[0056] Step 2: Training the prediction generator network involves a forward diffusion process and a backward denoising process: the forward diffusion process consists of T steps, progressively feeding training images... By adding Gaussian noise, noisy data corresponding to different diffusion time steps are constructed. The above data is then fed into the noise prediction network. The noise prediction network will predict the distribution of Gaussian noise injected in the corresponding diffusion time step based on the input noisy data. The noise prediction network is trained by the loss function to accurately fit the noise distribution (the purpose is to train the inverse denoising capability of the noise prediction network).

[0057] Step 2.1: According to the forward diffusion formula During the implementation of the example As t increases, it gradually decreases from 1 to 0, which affects the training preprocessing module. Noise-adding processing is performed to obtain noisy data corresponding to different diffusion time steps. Where T is the total number of forward diffusion steps, and in this example, T=2000. Indicates the diffusion time step. This represents the noisy data corresponding to the diffusion time step t. This indicates the noise intensity added at each step. The cumulative product of the noise corresponding to time step 1 to time step t ( ), For noise that follows a Gaussian distribution, It is a unit variance matrix;

[0058] Step 2.2: Train the Edge U-Net noise prediction network (the U-Net network optimized in this patent), which specifically includes the following sub-steps:

[0059] Step 2.2.1: Concatenate the completed region mask image, the known region of the original image, the distance decay map, and the boundary gradient map along the channel dimension, and then feed them into the boundary-aware convolutional module to obtain... ,right Image after global pooling and sampled at time t By concatenating along the channel dimension, we obtain the input to the Edge U-Net. ;

[0060] Step 2.2.2: [The text appears to be incomplete and contains several grammatical errors. A more accurate translation Send to the first layer of the encoder, for Convolution is performed to obtain ,Will With time step The output of the current coding layer is fed into the residual block. ;

[0061] Step 2.2.3: Using two-dimensional convolution pairs Downsampling is performed to obtain the input for the next coding layer. ;

[0062] Step 2.2.4: [The text appears to be incomplete and contains several grammatical errors. A more accurate translation Send it to the second layer of the encoder, with time step into the residual block to get the output of the current encoding layer ;

[0063] Step 2.2.5: down-sample by two-dimensional convolution on to get the input of the next encoding layer ;

[0064] Step 2.2.6: put into the third layer of the encoder, put with time step into the residual block to get the output of the current encoding layer ;

[0065] Step 2.2.7: down-sample by two-dimensional convolution on to get the input of the next encoding layer ;

[0066] Step 2.2.8: put into the fourth layer of the encoder, put with time step into the residual block to get the output of the current encoding layer ;

[0067] Step 2.2.9: put into the intermediate block (composed of the residual block and the self-attention mechanism) to get the input of the next decoding layer ;

[0068] Step 2.2.10: put into the fourth layer of the decoder, put with time step after merging, use the formula to perform skip-connection with of the corresponding encoding layer, and then put it into the residual block to get the output of the current decoding layer ; where is the output of the corresponding encoding layer, is the input of the current decoding layer;

[0069] Step 2.2.11: up-sample by inverse two-dimensional convolution on to get the input of the next decoding layer ;

[0070] Step 2.2.12: put into the third layer of the decoder, put with time step after merging, use the formula to perform skip-connection with is fed into the residual block to obtain the output of the current decoding layer ; where, is the output of the corresponding encoding layer, is the input of the current decoding layer;

[0071] Step 2.2.13: up-sample by inverse 2D convolution on to obtain the input of the next decoding layer ;

[0072] Step 2.2.14: feed into the second decoder layer, merge with the time step , and use the formula to perform skip-connection with of the corresponding encoding layer, and then feed into the residual block to obtain the output of the current decoding layer ; where, is the output of the corresponding encoding layer, is the input of the current decoding layer;

[0073] Step 2.2.15: up-sample by inverse 2D convolution on to obtain the input of the next decoding layer ;

[0074] Step 2.2.16: feed into the first decoder layer, merge with the time step , and use the formula to perform skip-connection with of the corresponding encoding layer, and then feed into the residual block to obtain the output of the current decoding layer ; where, is the output of the corresponding encoding layer, is the input of the current decoding layer;

[0075] Step 2.2.17: perform convolution on to obtain the model predicted output ;

[0076] Step 2.3: the predicted output is the noise estimate of the corresponding time step, denoted as ; the training of each time step performs additional local iteration on the to-be-completed region, and in the present example, each time step performs additional twice local iteration on the to-be-completed region; during the training process, the optimization objective is to minimize the loss function ; where the diffusion loss pixel-level loss Perceived loss Boundary gradient loss ;in, Let be the mathematical expectation of time step t. For normalization operations, As time step weights, To predict the distribution for the model, For the true distribution, To complete the weighting coefficients for the region, For the area mask to be completed, To predict the image, For real images, for In the network Feature mapping of layers, for In the network Feature mapping of layers, This indicates that the corresponding results for all network layers are summed. To predict the edge gradient map of the image, This is the edge gradient map of the real image. For the weights corresponding to different losses, this example uses... By performing gradient descent calculation on the loss function, all learnable parameters in the noise estimation network are updated until a fit is achieved.

[0077] Step 3: Take the image to be completed as input (the starting point of the reverse denoising process), and use the trained noise prediction network to denoise it step by step. The denoising process has a total of 2000 steps (which can be reduced to 250 steps in the actual implementation by using ddim), and finally obtain a completed image with high structural consistency and high guarantee.

[0078] Step 3.1: Obtain the image to be completed The noise prediction result is obtained by feeding it into the trained noise prediction network. ;

[0079] Step 3.2: According to the reverse diffusion formula Calculations yielded Predicted image at time step According to the formula get Image to be completed at time step In the example, ;in, This indicates the noise intensity added at each step. This represents the cumulative product of the noise levels corresponding to time steps 1 to t. for obeying Gaussian distribution noise, for weight coefficient, for known region of image to be completed;

[0080] Step 3.3: repeat the above steps until t = 1, i.e. get the final completed picture.

[0081] As Figure 10 (the rightmost is the completion result of this method, because the Celeba dataset involves face data, Figure 10 only show the completion results of Places2 image dataset) and Figure 11 (the last line is the index data of this method) show that this method (named EdgePaint) is compared with other three mainstream image completion methods based on diffusion model and one based on RNN architecture, and the datasets are Celeba face images and Places2 natural images respectively. Through comparison, it is not difficult to conclude that this method can be comparable or even better than other mainstream models in intuitive feeling and index comparison, and has good technical advancement and practical application scene.

[0082] It should be understood that parts not elaborated in the specification are all prior art.

[0083] It should be understood that the above description of the preferred embodiments is more detailed and should not be considered as limiting the scope of patent protection of the present application. Those skilled in the art can make substitutions or modifications without departing from the scope of protection of the claims, which fall within the scope of protection of the present application. The scope of protection of the present application should be subject to the appended claims.

Claims

1. A boundary-aware diffusion model-based image inpainting method, characterized in that, Comprise: Step 1: build data processing module: The data processing module is used to process data from the training dataset. (H is the image height, W is the image width), collect training images. And generate the corresponding mask data for the region to be filled in. Known regions of the image Distance decay map and boundary gradient map ; Step 2: build boundary-aware noise prediction network: The training of the prediction generation network involves a forward diffusion process and a backward denoising process: the forward diffusion process is T steps, which gradually adds Gaussian noise to the training image to construct noisy data corresponding to different diffusion time steps The above data is input into the noise prediction network, which predicts the distribution of the injected Gaussian noise in the corresponding diffusion time step according to the input noisy data. The noise prediction network is trained by a loss function to accurately fit the noise distribution (the purpose is to train the backward denoising capability of the noise prediction network); Step 3: build image completion module: Take the picture to be completed as input (the starting point of the reverse denoising process), use the trained noise prediction network to denoise it step by step, and the denoising process has T steps (can be reduced to T / 2, T / 4, etc. through acceleration means), and finally get a completed picture with high structural consistency and high guarantee.

2. The boundary-aware diffusion model-based image inpainting method of claim 1, wherein, The implementation process of step 1 is as follows: Step 1.1: Obtain training images and the corresponding to-be-completed region mask data ; wherein in the mask data, a value of 255 represents an intact region, and a value of 0 represents a to-be-completed region; Step 1.2: Based on the mask data , get the binary representation of the mask data (1 represents the intact area and 0 represents the area to be completed), do the Euclidean distance transformation on the binary graph, calculate the minimum Euclidean distance of the boundary between the intact area and the damaged area, and get the distance attenuation graph ; Step 1.3: Calculate the training image gradient map and mask data edge feature map, prune the image gradient according to the edge feature map to obtain a boundary-aware map containing gradient information at the junction of the damaged area and the intact area .

3. The boundary-aware diffusion model-based image inpainting method of claim 1, wherein, The implementation process of step 2 is as follows: Step 2.1: According to the forward diffusion formula The training preprocessing module obtained Noise-adding processing is performed to obtain noisy data corresponding to different diffusion time steps. Where T is the total number of forward diffusion steps, Indicates the diffusion time step. This represents the noisy data corresponding to the diffusion time step t. This indicates the noise intensity added at each step. The cumulative product of the noise corresponding to time step 1 to time step t ( ), For noise that follows a Gaussian distribution, It is a unit variance matrix; Step 2.2: Noisy data at different diffusion time steps, along with corresponding mask maps, distance decay maps, and boundary gradient maps, are used as input conditions to train the noise prediction network. The noise prediction network is the Edge U-Net structure proposed in this patent, and its output is the noise estimate for the corresponding time step. At each time step of training, additional local iterations are performed on the region to be filled. During training, the optimization objective is to minimize the loss function. ;including diffusion loss pixel-level loss Perceived loss Boundary gradient loss ;in, Let be the mathematical expectation of time step t. For mean normalization operation, As time step weights, To predict the distribution for the model, For the true distribution, To complete the weighting coefficients for the region, For the area mask to be completed, To predict the image, For real images, for In the network Feature mapping of layers, for In the network Feature mapping of layers, This indicates that the corresponding results for all network layers are summed. To predict the edge gradient map of the image, This is the edge gradient map of the real image. To assign weights to different losses, gradient descent is used to calculate the loss function, updating all learnable parameters in the noise estimation network until a fit is achieved.

4. The boundary-aware diffusion model-based image inpainting method of claim 1, wherein, The implementation process of step 3 is as follows: Step 3.1: Obtain the image to be completed The noise prediction result is obtained by feeding it into the trained noise prediction network. ;in, This indicates that the image to be completed is input into the prediction model. and corresponding time The predicted results obtained; Step 3.2: Calculate the predicted image of the time step according to the reverse diffusion formula , according to the weighted formula ; wherein, represents the noise intensity added at each step, is the cumulative result of the noise corresponding to time step 1 to time step t ( ), is the noise obeying Gaussian distribution, is the weight coefficient, is the to-be-completed region mask, is the known region of the to-be-completed image.​​​ Step 3.3: repeat the above steps until t=1, that is, get the final completed picture.

Images