SAMSeg: U-SAM 🧠

Overview

The Segment Anything Model (SAM) is a cutting-edge, promptable deep learning model designed for 2D natural image segmentation. However, medical imaging—particularly brain tumor segmentation 🧠—presents unique challenges, including highly irregular tumor boundaries, variations in intensity, and the complex, multidimensional nature of medical scans.

To address these challenges, this project extends SAM's segmentation capabilities by integrating Parameter-Efficient Fine-Tuning (PEFT) techniques and training on the BraTS Intracranial Meningioma 2023 dataset. Furthermore, we introduce U-SAM, a novel framework that combines SAM’s segmentation power with the precision and structural awareness of U-Net architectures, thereby enhancing segmentation accuracy for complex medical imaging tasks.

Given the difficulties encountered in programmatic data retrieval from Synapse, I have made our curated dataset (including our specific training/test split) and model checkpoints available for easy access: 🔗 U-SAM Project SharePoint

📢 If you experience any accessibility issues with the dataset or checkpoints, please feel free to contact me.

Project Directory Structure 🗂️

├── U-SAM/                               # Scripts for running and evaluating U-SAM
│   ├── U1-SAM2.py                        > Runs U1-SAM2
│   └── U-SAM-Combo.py                    > Runs the U-SAM combination
├── U-Net/                               # U-Net model scripts
│   ├── unet_test.py                      > Tests U-Net
│   ├── unet_train.py                     > Trains U-Net
│   └── unet3d.py                         > 3D U-Net architecture implementation
├── SAM/                                 # SAM model scripts
│   ├── sam_test.py                       > Tests SAM
│   ├── sam_train.py                      > Trains SAM
│   └── sam_tuning.py                     > Tunes SAM hyperparameters
├── dataset/                             # Dataset folders (must be downloaded separately)
│   ├── BraTS-MEN-Test/
│   ├── BraTS-MEN-Train/
│   └── BraTS-MEN-Validate/
├── checkpoints/                         # Pre-trained model checkpoints (must be downloaded separately)
│   ├── peft-sam_epoch_125.pth.tar
│   ├── unet_epoch_100.pth.tar
│   └── sam_vit_b_01ec64.pth
├── utilities/                           # Utility scripts
│   ├── brats_dataset.py                   > Custom dataset loader for BraTS
│   ├── general_utils.py                   > General helper functions
│   ├── model_utils.py                     > Model-related utilities
│   ├── plot.py                            > Visualization functions
│   ├── preprocess.py                      > Preprocessing functions
├── requirements.txt                     # Dependencies list

Setup & Installation ⚙️

To set up the environment for this project:

1. Create a Virtual Environment (recommended):

   # Using venv
   python3 -m venv samseg-env
   source samseg-env/bin/activate
   
   # Using conda
   conda create --name samseg-env python=3.8
   conda activate samseg-env

2. Install Required Libraries:

   pip install -r requirements.txt

Usage 🚀

If you run into module import issues, ensure you're in the root project directory and run:

export PYTHONPATH=$(pwd):$PYTHONPATH

Running Scripts

Hyperparameter Tuning for SAM 🎯

python3 SAM/sam_tuning.py --checkpoint_path checkpoints/sam_vit_b_01ec64.pth \
                          --output_dir SAM/tuning_outputs --n_trials 50

Training SAM 🏋️‍♂️

python3 SAM/sam_train.py --checkpoint_path checkpoints/peft-sam_epoch_125.pth.tar \
                         --output_dir SAM/training_outputs --init_lr 0.00006 \
                         --decay_rate 0.995 --max_epoch 50 --batch_size 1 \
                         --samples_per_epoch 150 --backup_interval 10 \
                         --validate --validation_interval 5

Testing SAM 🧪

• Vanilla SAM:

  python3 SAM/sam_test.py --vanilla --visualise --results_dir SAM/vanilla_test_results --batch_size 1

• PEFT-SAM:

  python3 SAM/sam_test.py --peft_path checkpoints/peft-sam_epoch_125.pth.tar \
                          --visualise --results_dir SAM/peft_test_results --batch_size 1

Training U-Net 🏋️

python3 UNet/unet_train.py --checkpoint_path checkpoints/unet_epoch_100.pth.tar \
                           --output_dir UNet/training_outputs --init_lr 0.00006 \
                           --decay_rate 0.995 --max_epoch 50 --batch_size 1 \
                           --samples_per_epoch 150 --backup_interval 10 \
                           --validate --validation_interval 5

Testing U-Net 🧪

python3 UNet/unet_test.py --checkpoint_path checkpoints/unet_epoch_100.pth.tar \
                          --visualise --results_dir UNet/unet_test_results --batch_size 1

Running U-SAM-Combo 🤖

python3 U-SAM/U-SAM-Combo.py --unet_checkpoint checkpoints/unet_epoch_100.pth.tar \
                             --sam_checkpoint checkpoints/peft-sam_epoch_125.pth.tar \
                             --visualise --output_dir U-SAM/u-sam-combo-results --bbox_emulation

Running U1-SAM2 🧠

python3 U-SAM/U1-SAM2.py --unet_checkpoint checkpoints/unet_epoch_100.pth.tar \
                         --sam_checkpoint checkpoints/peft-sam_epoch_125.pth.tar \
                         --visualise --output_dir U-SAM/u1-sam2-results --bbox_emulation

Notes 📌

• Dataset & Checkpoints: The dataset must be downloaded separately from the SharePoint link above.
• Ensure paths are correct before running any scripts.
• GPU Support: Requires CUDA for optimal performance.
• Tested Environment: A100 Core GPU on UCT's HPC.
• Citations: Full citations can be found in the accompanying paper and code, with special thanks to the Kurtlab BraTS 2023 submission for their adaptable project setup.

---

🚀 Happy segmenting! 🧠 For any questions or comments about this project, please don’t hesitate to reach out!