Li-Xuan Alex Peng - EECS, National Central University

Brain Tumor Segmentation via SAM-based fine-tuning

on structural MRI images

Background

• Brain tumor segmentation is a classic task in medical image analysis and computer vision, which aims to predict pixel-level masks corresponding to tumor coverage on MRI images.
• There are plenty of open challenges dedicated to accessing the state-of-the-art solution to resolve this task, the most iconic one is the BraTS challenge, of which data will be used to evaluate the model of this BHS project.
• Of the most recent research in computer vision, Meta has released a foundation model i.e. Segments Anything Model (SAM), capable of segmenting arbitrary semantics on any image.
• Researchers from the University of Toronto later fine-tune the SAM model on medical images (MedSAM) and reach a significant improvement on such tasks compared to the original SAM.

Objectives

​

​

• Segmenting brain tumor coverage on brain structural MRI images using foundation models fine-tuning techniques with ground-truths indicated bounding boxes as prompts of the model.
• Follow fine-tuning method proposed by MedSAM and perform fine-tuning on other open-access dataset on Kaggle.
• Compare fine-tuned model performance with original model using BraTS2021 data via simple Dice score statistics and qualitative results

​

Data:

The two datasets below that I found in open-access data all provided 2D MRI images with corresponding ground-truth masks.

Empirically, their quantities and qualities are adequate for fine-tuning.

• Brain Tumor Segmentation@kaggle (Original Source and Credit: Jun Cheng) (For fine-tuning)

This brain tumor dataset containing 3064 T1-weighted contrast-enhanced images (.jpg)

from 233 patients with three kinds of brain tumor: meningioma (708 slices),

glioma (1426 slices), and pituitary tumor (930 slices).

• BRaTS 2021 Task 1 Dataset@RSNA-ASNR-MICCAI BraTS Challenge 2021 (For evaluation)

Consist of 1666 scans with T1, T2, and T2-FLAIR MRI images, delivered as NIfTI files (.nii.gz) and their corresponding ground-truth masks. For detailed explanations please refer to BraTS challenge home page.

​

Tools:

Computation Resource and Running Environment:

• Packages: CUDA, Numpy, Nilearn, Pytorch, SAM, MedSAM, Seaborn
• Environment: Custom built JupyterLab server running on NV-Docker based container
• Hardware: Nvidia Tesla A100 x2 (For training), Nvidia GTX 1080Ti (For inference)

Methods:

The main idea is to let the pre-trained SAM model able to update its parameters through the optimizer when training.

Here we only update the parameters of the decoder that produces the segmentation mask. Code adopted from MedSAM Github Repository

Results

Training loss of fine-tuning over epochs

The training looks stable with little oscillations.

Results

Quantitative Results:

Dice score statistics on 1666 Scans(BraTS Task 1 data), random slice of each scan with T2 weighted MRI images

Original SAM

Fine-tuned SAM

Results

Qualitative Results:

Obtained from random scan, random slices in BraTS dataset

Deliverables

• Two Jupyter notebooks, one for fine-tuning and the other for inference and evaluation. All are adopted from MedSAM with my own minor modifications. Available on my BHS github repository.
• The fine-tuned model checkpoint (.pth) also provided via this google drive folder

Conclusion

• The SAM foundation model shows an amazing generalization ability to adapt downstream tasks even though the fine-tuning dataset is relatively small in scale.
• The original project goal was to find out areas of atrophy induced by Alzheimer's disease. Since I could not find adequate open-access data in quantity and quality criteria, I changed my topic to brain tumor segmentation.
• In the first attempt at fine-tuning the model, I encountered an inadequate computational resource problem on the workstation, then, by coincidence, I was invited to test the new GPU servers by Dept. CSIE, NCU, thanks to this test opportunity, I could finish this project in an unexpectedly efficient manner.

Things Learned From This Project

• Building JupyterLab running environment using Docker container
• Understand the basic concept of MRI imaging
• Pre-process MRI image using Nilearn package
• Understand how foundation model works and perform domain specific adaptation on custom dataset
• Visualize the result in quantitative and qualitative perspectives

Future Work

• Try other fine-tuning techniques
• Add additional custom neural network modules upon SAM basis to further increase the learning performance
• Find more high quality data to further increase the generalization ability.
• Compare to other mainstream methods (e.g. UNet, FCN,Mask R-CNN,etc.)

References

Methodology

• Segment Anything (original paper) https://arxiv.org/abs/2304.02643
• Segment Anything in Medical Images (MedSAM) https://arxiv.org/pdf/2304.12306.pdf
• MedSAM github repository https://github.com/bowang-lab/MedSAM

Dataset

• Cheng, Jun, et al. "Retrieval of Brain Tumors by Adaptive Spatial Pooling and Fisher Vector Representation." PloS one 11.6 (2016).
• BraTS Challenge homepage http://braintumorsegmentation.org/

​