Brain Tumour Removing and Missing Modality Generation using 3D WDM

TL;DR

The paper addresses brain tumor inpainting and missing MRI modality generation in BraTS 2024 by introducing conditional 3D wavelet diffusion models that enable full‑resolution training and inference on standard GPUs. It presents robust training pipelines (Default, Default with conditioning, Always Known, Known All Time, Known 3 to 1) for both local inpainting and global missing‑modality synthesis, with specialized data augmentation and wavelet‑domain conditioning. The best test results for missing modality generation achieve $\mathrm{MSE}=0.07\pm0.04$, $\mathrm{PSNR}=22.8\pm4.41$, and $\mathrm{SSIM}=0.91\pm0.15$ at $5000$ sampling steps, though qualitative analyses reveal occasional hallucinations and limited convergence in some models. The work highlights a trade‑off between image fidelity and runtime, and points to future enhancements such as adversarial losses and more efficient sampling schedules to reduce computation while improving morphological accuracy.

Abstract

This paper presents the second-placed solution for task 8 and the participation solution for task 7 of BraTS 2024. The adoption of automated brain analysis algorithms to support clinical practice is increasing. However, many of these algorithms struggle with the presence of brain lesions or the absence of certain MRI modalities. The alterations in the brain's morphology leads to high variability and thus poor performance of predictive models that were trained only on healthy brains. The lack of information that is usually provided by some of the missing MRI modalities also reduces the reliability of the prediction models trained with all modalities. In order to improve the performance of these models, we propose the use of conditional 3D wavelet diffusion models. The wavelet transform enabled full-resolution image training and prediction on a GPU with 48 GB VRAM, without patching or downsampling, preserving all information for prediction. The code for these tasks is available at https://github.com/ShadowTwin41/BraTS_2023_2024_solutions.

Figures (7)

• Figure 1: Training pipeline of the Default model
• Figure 2: Conditional sampling pipeline.
• Figure 3: Training pipeline of the Conditional Default model
• Figure 4: Training pipeline of the Conditional Always known model
• Figure 5: Task 8 results. Results of the inpainting model for the training cases 00084-000, 00723-000, 01163-000, 01274-000, 01502-000 in each row. The first column shows a slice of the real case with a tumour, the second the replacement of the tumour with healthy tissue, the third a random slice where the healthy tissue has been removed and the last the respective inpainted slice.