Introduction

The human brain's complexity and sensitivity to change underscore the importance of accurate brain tumor detection, where even minor cell growth anomalies can be life-altering. Magnetic Resonance Imaging (MRI) stands out for its detailed visualization of brain structures. We propose an advanced U-Net architecture for refining MRI-based tumor segmentation in particularly challenging, low-contrast regions by combining a diffusion model with uncertainty metrics. Our work not only improves upon traditional methods but also exceeds recent 3D and attention-enhanced U-Net models in detecting intricate cerebral abnormalities.

Methods

The Brain Tumor Segmentation (BraTS) challenge provides a multi-institutional dataset comprising multi-parametric magnetic resonance imaging (mpMRI) for brain glioblastoma analysis¹. In this study, we utilized the 2021 release of the BraTS dataset, which encompasses imaging data from 1,250 patients. Each patient's data set includes imaging modalities such as Fluid Attenuated Inversion Recovery (FLAIR), T1-weighted, T1-weighted post-contrast, and T2-weighted scans. Each sample contains labeled annotations from the tumor, peritumoral edema, and the necrotic and non-necrotic tumor core, exemplified in Figure 1.
We demonstrate our method's effectiveness by comparing it with leading techniques. The first technique encompasses the U-Net, an established segmentation architecture that combines context capture and precise localization². The second is an advanced U-Net, the Attention U-Net, which incorporates attention gates. This architecture allows focusing on key image areas, leading to enhanced accuracy without relying on separate models³. The last baseline is an extended version of the original U-Net with advanced 3D data processing and volumetric training capabilities⁴.
Building on these established methods, our Diffusion U-Net, or Diff-UNet, innovates further by integrating a diffusion model within the U-Net architecture. This integration facilitates the distillation of semantic content from volumetric data, yielding enhanced pixel-level detail crucial for medical segmentation tasks. Distinct from its U-Net predecessor, the Diff-UNet embeds a noise process within the input pipeline and employs iterative label map prediction, which heightens the model's predictive stability. An additional feature of Diff-UNet is its unique uncertainty-based fusion module used during inference, which significantly increases the robustness of the segmentation outcomes⁵.
In this study, we utilized the Diff-UNet pipeline. Initially, this architecture was developed and validated for liver segmentation. We adapted its core principles to address multi-class brain tumor segmentation characterized by high variability and complex volumetric structures. Due to the low resolution and intricacy of these structures, the task is non-straightforward. During training, a denoising module takes multi-modal MRI data and systematically noise-altered labels as input. By combining a feature encoding stream with a denoising U-Net, we yield denoised predictions of the segmented regions of interest, which are aligned with their corresponding labels during training.
In the testing phase, the trained denoising module is applied iteratively to generate denoised segmentations. This step is repeated for a predetermined number of iterations, resulting in a series of outputs that progressively exhibit less noise, culminating in the desired multi-class segmentation. For the final inference, each output in the sequence is combined with an uncertainty map to enhance the method's robustness, preserving critical information obtained at each step. This information is then integrated into the final output of the architecture, producing a multi-class segmentation map of the tumor. Figure 2 presents an overview of this process.

Results and Discussion

Table 1 compares the different baseline models and proposed approach performance, highlighting the mean Dice score of the 3D U-Net over the 2D U-Net despite a lower HD95 measure. The 3D U-Net excels in volumetric data processing, crucial for complex brain tumor segmentation. In contrast, the 2D U-Net is limited to two-dimensional slices and may overlook some 3D features. Attention 3D U-Net slightly outperforms the standard 3D U-Net by employing attention blocks that focus on key regions, enhancing segmentation precision. The diffusion-embedded U-Net shows the best overall results by integrating a denoising diffusion model, improving robustness through noise addition during training, and employing a step uncertainty-based fusion module for testing.

Conclusion

In this study, various models were evaluated, with three traditional U-Net architectures serving as benchmarks. The 3D U-Net demonstrated superior Dice coefficient metrics compared to the 2D U-Net and required the least computational time. The 2D U-Net, however, excelled in delineating boundaries more effectively than its 3D counterpart. The Attention 3D U-Net offered marginal improvements over the conventional 3D U-Net but at the cost of increased computational time. The Diffusion U-Net, when used in the brain tumor segmentation context, utilizes the strengths of the diffusion model and uncertainty to improve robustness, yielding impressive results and showcasing promise for tasks involving multi-modal 3D medical image segmentation.

Acknowledgements

No acknowledgement found.