fastWDM3D: Fast and Accurate 3D Healthy Tissue Inpainting
This work addresses the need for fast and accurate 3D inpainting in medical imaging applications like tumor growth modeling and image registration, representing an incremental improvement over existing methods.
The paper tackled the slow sampling speed of denoising diffusion probabilistic models (DDPMs) for 3D healthy tissue inpainting by adapting a 2D approach with a variance-preserving noise schedule and reconstruction losses, resulting in fastWDM3D achieving a SSIM of 0.8571, MSE of 0.0079, and PSNR of 22.26 in 1.81 seconds per image, which is up to 800 times faster than other DDPMs.
Healthy tissue inpainting has significant applications, including the generation of pseudo-healthy baselines for tumor growth models and the facilitation of image registration. In previous editions of the BraTS Local Synthesis of Healthy Brain Tissue via Inpainting Challenge, denoising diffusion probabilistic models (DDPMs) demonstrated qualitatively convincing results but suffered from low sampling speed. To mitigate this limitation, we adapted a 2D image generation approach, combining DDPMs with generative adversarial networks (GANs) and employing a variance-preserving noise schedule, for the task of 3D inpainting. Our experiments showed that the variance-preserving noise schedule and the selected reconstruction losses can be effectively utilized for high-quality 3D inpainting in a few time steps without requiring adversarial training. We applied our findings to a different architecture, a 3D wavelet diffusion model (WDM3D) that does not include a GAN component. The resulting model, denoted as fastWDM3D, obtained a SSIM of 0.8571, a MSE of 0.0079, and a PSNR of 22.26 on the BraTS inpainting test set. Remarkably, it achieved these scores using only two time steps, completing the 3D inpainting process in 1.81 s per image. When compared to other DDPMs used for healthy brain tissue inpainting, our model is up to 800 x faster while still achieving superior performance metrics. Our proposed method, fastWDM3D, represents a promising approach for fast and accurate healthy tissue inpainting. Our code is available at https://github.com/AliciaDurrer/fastWDM3D.