Denoising Diffusion Models for Inpainting of Healthy Brain Tissue

TL;DR

This approach builds on denoising diffusion probabilistic models and uses a 2D model that is trained using slices in which healthy tissue was cropped out and is learned to be inpainted again to solve the given inpainting task.

Abstract

This paper is a contribution to the "BraTS 2023 Local Synthesis of Healthy Brain Tissue via Inpainting Challenge". The task of this challenge is to transform tumor tissue into healthy tissue in brain magnetic resonance (MR) images. This idea originates from the problem that MR images can be evaluated using automatic processing tools, however, many of these tools are optimized for the analysis of healthy tissue. By solving the given inpainting task, we enable the automatic analysis of images featuring lesions, and further downstream tasks. Our approach builds on denoising diffusion probabilistic models. We use a 2D model that is trained using slices in which healthy tissue was cropped out and is learned to be inpainted again. This allows us to use the ground truth healthy tissue during training. In the sampling stage, we replace the slices containing diseased tissue in the original 3D volume with the slices containing the healthy tissue inpainting. With our approach, we achieve comparable results to the competing methods. On the validation set our model achieves a mean SSIM of 0.7804, a PSNR of 20.3525 and a MSE of 0.0113. In future we plan to extend our 2D model to a 3D model, allowing to inpaint the region of interest as a whole without losing context information of neighboring slices.

Figures (5)

• Figure 1: Illustration of the problem. MR scans not showing pathologies can be fed directly into automatic processing tools. Often, scans containing pathologies need to be modified first. The task of this challenge is to replace pathological brain tissue with healthy brain tissue. The suggested inpainting requires a masked scan, whereby the masked region will be filled by an inpainting process, for which we suggest a DDPM-based model, visualized by the bold blue arrow. The output, denoted as the transformed scan, then shows a healthy scan. The transformed scan can then be fed into automatic processing tools for further analysis, denoted by the dashed blue arrow.
• Figure 2: Overview of the training process. Context information is provided through the concatenation of input $b_i$ and mask $m_i$ with the noisy image $x_{i,t}$, which originated from the ground truth $x_0$ = $x_{i,GT}$ using Equation \ref{['eqn2']}. The concatenated image $X_t$ is used by the diffusion model to predict a slightly denoised image $x_{i,t-1}$ from $x_{i,t}$ using Equation \ref{['eqn5']}. The diffusion model is trained using an MSE loss between the real and the predicted noise, summarized in Equation \ref{['eqn6']}.
• Figure 3: Sampling method. During sampling, we apply the denoising described by Equation \ref{['eqn5']} for each of the $T$ denoising steps for each slice $b_i$ with a non-zero mask $m_i$. We then replace the slices $b_i$ in $b$ with the samples $x_{i,0}$, containing the inpainting, to obtain our transformed 3D scan.
• Figure 4: Axial slices of images from our validation set. The masked slice, the mask, the generated sample, and the ground truth are provided. In general, the 2D generated samples are of good quality. Moreover, they show a high similarity to the corresponding ground truth. The SSIM, MSE and PSNR that are reported refer to the comparison of the whole generated volume (obtained by stacking of the generated slices) and the corresponding ground truth volume.
• Figure 5: Exemplary image of our validation set in the coronal (top) and sagittal (bottom) view. We observe stripe artifacts (blue boxes) due to the stacking of multiple axial slices. These artifacts can be reduced by applying a Gaussian filter (standard deviation $sigma = 1.075$ for Gaussian kernel).