Brain Tumor Segmentation via SAM-based fine-tuning on structural MRI images

Project definition

Background

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.

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)

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.

1. 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).
2. 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.

Deliverables

At the end of this project, we will have:

• The current markdown document, completed and revised.
• 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
• The introduction slides of this project.

Results

Progress overview

The SAM foundation model shows an amazing generalization ability to adapt downstream tasks even though the fine-tuning dataset is relatively small in scale.

Tools I learned during this project

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

Results

Deliverable 1: Introduction Slides

You can find the introduction slides of this project here

Deliverable 2: Jupyter Notebooks

You can find jupyter notebooks of this project at my github repository

Deliverable 3: Instructions

1. For fine-tuning on your own datasets, please follow the tutorial provided by MedSAM
2. For reproducing my training results, please follow this jupyter notebook

If you have any questions or suggestions, feel free to contact me via email alexpeng517@gmail.com

Conclusion and acknowledgement

Thanks to BoWang Team from University of Toronto for open sourcing their recent research, so I can learn how to fine-tune the model from their work.

Thanks to SNMG of Dept. CSIE, National Central University, for providing powerful computation resources to make this project practical.

Thanks to Professor Kevin Chun-Hsien Hsu, Institute of Cognitive Neuroscience, National Central University, giving me the oppotunity to attend BrainHackSchool 2023.

Thanks to all the Professors, Moderators, TAs and Classmates that make this BrainHackSchool happend!