4807
Probabilistic Brain Tumor Segmentation for Everyone1Massachusetts General Hospital, Harvard Medical School, Boston, MA, United States
Synopsis
Keywords: MR-Guided Interventions, Tumor, Foundation model, Segmentation
Motivation: Probabilistic segmentation offers notable advantages in the context of brain tumor delineation on MRI. However, existing methods are not accessible to most medical institutes, rendering their clinical relevance questionable.
Goal(s): we present a foundation model-based probabilistic brain tumor segmentation approach designed for straightforward integration into clinical applications.
Approach: We employ a parameter-efficient few-shot learning strategy to fine-tune the foundation model, thereby enabling it to output the tumor mask and uncertainty for the brain tumor segmentation task.
Results: The proposed method achieves a competitive performance with training on only five cases. It has the potential for a substantial impact on clinical practice.
Impact: The proposed foundation model-based probabilistic brain tumor segmentation method is open-source and achieves a competitive performance with only five training cases. Such characteristics make it a valuable asset for healthcare institutions that have difficulties developing their proprietary probabilistic segmentation models.
Introduction
Probabilistic segmentation offers notable advantages in the context of brain tumor delineation within MRI images. It provides a more nuanced depiction of the tumor by using soft segmentation masks that allocate a probability value to each voxel. This probability value serves as a quantification of uncertainty, which is particularly valuable in the identification of false positives. While deep learning-based segmentation techniques have exhibited considerable promise [1][2][3], their practical applicability is constrained by a number of challenges. Firstly, the absence of publicly accessible source code or pre-trained models impedes the integration of these methodologies into clinical settings. Secondly, the inadequacy of substantial training data hinders the realization of the advertised performance level, rendering their clinical relevance questionable.In this investigation, we present a foundation model-based probabilistic brain tumor segmentation approach designed for straightforward integration into clinical applications. Our method is openly available as an open-source solution and demonstrates a competitive level of performance even when trained on a limited dataset consisting of only five cases. Such characteristics make it a valuable asset for healthcare institutions that have difficulties developing their proprietary probabilistic segmentation models.
Methods
The emergence of foundational models, which are trained on extensive and diverse datasets, has led to a paradigm shift in the field of intelligent model development. Motivated by their remarkable aptitude for generalization and few-shot learning, the practice of fine-tuning pre-trained large models for a range of downstream tasks has garnered growing interest. The Segment Anything Model (SAM) is a recently developed visual foundation model for image segmentation. SAM is pre-trained on a dataset of 11 million natural images. We employ a parameter-efficient few-shot learning strategy [4] to fine-tune SAM, thereby enabling it to not only achieve competitive performance but also provide segmentation uncertainty for the brain tumor segmentation task.Key components in our fine-tuning strategy include: (1) To address the disparity between 2D natural images and volumetric MR data, we introduce a set of 3D adapters within each transformer block of the image encoder to extract the third-dimensional information. (2) To optimize the parameter tuning efficiency, we employ the FacT technique [5] which retains a substantial portion of the pre-trained SAM weights and only updates lightweight weight increments using as few as five cases. (3) Upon the model's convergence, we preserve a total of 20 parameter checkpoints, which are subsequently employed in constructing a model ensemble for the purpose of quantifying segmentation uncertainty. A summary of the fine-tuning framework is presented in Figure 1.
The validation of the proposed methodology was conducted using the T1-weighted whole tumor segmentation task from the BraTS challenge [6]. In this validation process, a set of five cases was randomly selected for fine-tuning the SAM model, while an additional five cases were designated for testing. This procedure was iteratively repeated three times. The experiments were carried out using PyTorch 2.0, and all parameter configurations adhered to the specifications outlined in the reference [4].
Results
For model evaluation, we employ the common Dice score to assess pixel-wise segmentation accuracy. Initially, our five-shot fine-tuning achieves a Dice score of 0.773. Subsequently, we harness the power of SAM's prompt design by incorporating a tight 3D bounding box per volume into the model. This enhancement results in an improved Dice score of 0.889, placing it in the league of top performers in the BraTS challenge. Notably, our method demonstrated the remarkable practicality of achieving competitive results with just 5 training cases, in contrast to existing networks that necessitate a full ground-up training process. The segmentation uncertainty is calculated by averaging the results of 20 ensemble predictions, and it is informative to gauge the voxel-wise segmentation accuracy. Qualitative visualization of segmentation mask and uncertainty are presented in Fig.2. We also show the volumetric rendering of segmented tumors in Fig.3.Discussion and Conclusion
We introduced a pragmatic foundation model-based probabilistic segmentation method for brain tumor segmentation within MRI images. Our open-source method achieves a competitive performance with training on only five cases. We anticipate that this method holds the potential to significantly impact clinical practices in the field of brain tumor segmentation.Acknowledgements
n/aReferences
[1] Z. Xu, et al. Category-Level Regularized Unlabeled-to-Labeled Learning for Semi-supervised Prostate Segmentation with Multi-site Unlabeled Data. MICCAI 2023
[2] Y. Wang, et al., Cross-domain few-shot learning for rare-disease skin lesion segmentation. ICASSP 2022
[3] Z. Xu, et al., Anti-interference from noisy labels: Mean-teacher-assisted confident learning for medical image segmentation. IEEE Transaction of Medical Imaging 2022
[4] C. Chen, et al., MA-SAM: Modality-agnostic SAM Adaptation for 3D Medical Image Segmentation. arXiv 2023
[5] S. Jie, Z.H. Deng, Fact: Factor-tuning for lightweight adaptation on vision transformer. AAAI 2023
[6] S. Bakas, et al., Advancing the cancer genome atlas glioma mri collections with expert segmentation labels and radiomic features. Scientific data 2017
Figures
