● Medical SAM Adapter

License

Medical SAM Adapter, or say MSA, is a project to fineturn SAM using Adaption for the Medical Imaging. This method is elaborated in the paper Medical SAM Adapter: Adapting Segment Anything Model for Medical Image Segmentation.

A Quick Overview

News

[TOP] Join in our Discord to ask questions and discuss with others.
23-05-10. This project is still quickly updating 🌝. Check TODO list to see what will be released next.
23-05-11. GitHub Dicussion opened. You guys can now talk, code and make friends on the playground 👨‍❤️‍👨.
23-12-22. Released data loader and example case on REFUGE dataset. Credit: @jiayuanz3
24-01-04. Released the Efficient Med-SAM-Adapter❗️ A new, faster, and more lightweight version incorporates Meta EfficientSAM🏇. Full credit goes to @shinning0821.
24-01-07. The image resolution now can be resized by -image_size. Credit: @shinning0821
24-01-11. Added a detailed guide on utilizing the Efficient Med-SAM-Adapter, complete with a comparison of performance and speed. You can find this resource in guidance/efficient_sam.ipynb. Credit: @shinning0821
24-01-14. We've just launched our first official version, v0.1.0-alpha 🥳. This release includes support for MobileSAM, which can be activated by setting -net mobile_sam. Additionally, you now have the flexibility to use ViT, Tiny ViT, and Efficient ViT as encoders. Check the details here. Credit: @shinning0821

Requirement

Install the environment:

conda env create -f environment.yml

conda activate sam_adapt

Then download SAM checkpoint, and put it at ./checkpoint/sam/

You can run:

wget https://dl.fbaipublicfiles.com/segment_anything/sam_vit_b_01ec64.pth

mv sam_vit_b_01ec64.pth ./checkpoint/sam creat the folder if it does not exist

Example Cases

Melanoma Segmentation from Skin Images (2D)

Download ISIC dataset part 1 from https://challenge.isic-archive.com/data/. Then put the csv files in "./data/isic" under your data path. Your dataset folder under "your_data_path" should be like: ISIC/ ISBI2016_ISIC_Part1_Test_Data/...

ISBI2016_ISIC_Part1_Training_Data/...

ISBI2016_ISIC_Part1_Test_GroundTruth.csv

ISBI2016_ISIC_Part1_Training_GroundTruth.csv
Begin Adapting! run: python train.py -net sam -mod sam_adpt -exp_name *msa_test_isic* -sam_ckpt ./checkpoint/sam/sam_vit_b_01ec64.pth -image_size 1024 -b 32 -dataset isic -data_path *../data* change "data_path" and "exp_name" for your own useage. you can change "exp_name" to anything you want.

You can descrease the image size or batch size b if out of memory.

Evaluation: The code can automatically evaluate the model on the test set during traing, set "--val_freq" to control how many epoches you want to evaluate once. You can also run val.py for the independent evaluation.
Result Visualization: You can set "--vis" parameter to control how many epoches you want to see the results in the training or evaluation process.

In default, everything will be saved at ./logs/

REFUGE: Optic-disc Segmentation from Fundus Images (2D)

REFUGE dataset contains 1200 fundus images with optic disc/cup segmentations and clinical glaucoma labels.

Dowaload the dataset manually from here, or using command lines:

git lfs install

git clone git@hf.co:datasets/realslimman/REFUGE-MultiRater

unzip and put the dataset to the target folder

unzip ./REFUGE-MultiRater.zip

mv REFUGE-MultiRater ./data

For training the adapter, run: python train.py -net sam -mod sam_adpt -exp_name REFUGE-MSAdapt -sam_ckpt ./checkpoint/sam/sam_vit_b_01ec64.pth -image_size 1024 -b 32 -dataset REFUGE -data_path ./data/REFUGE-MultiRater you can change "exp_name" to anything you want.

You can descrease the image size or batch size b if out of memory.

Abdominal Multiple Organs Segmentation (3D)

This tutorial demonstrates how MSA can adapt SAM to 3D multi-organ segmentation task using the BTCV challenge dataset. For BTCV dataset, under Institutional Review Board (IRB) supervision, 50 abdomen CT scans of were randomly selected from a combination of an ongoing colorectal cancer chemotherapy trial, and a retrospective ventral hernia study. The 50 scans were captured during portal venous contrast phase with variable volume sizes (512 x 512 x 85 - 512 x 512 x 198) and field of views (approx. 280 x 280 x 280 mm3 - 500 x 500 x 650 mm3). The in-plane resolution varies from 0.54 x 0.54 mm2 to 0.98 x 0.98 mm2, while the slice thickness ranges from 2.5 mm to 5.0 mm. Target: 13 abdominal organs including Spleen Right Kidney Left Kidney Gallbladder Esophagus Liver Stomach Aorta IVC Portal and Splenic Veins Pancreas Right adrenal gland Left adrenal gland. Modality: CT Size: 30 3D volumes (24 Training + 6 Testing) Challenge: BTCV MICCAI Challenge The following figure shows image patches with the organ sub-regions that are annotated in the CT (top left) and the final labels for the whole dataset (right).

Prepare BTCV dataset following MONAI instruction: Download BTCV dataset from: https://www.synapse.org/#!Synapse:syn3193805/wiki/217752. After you open the link, navigate to the "Files" tab, then download Abdomen/RawData.zip. After downloading the zip file, unzip. Then put images from RawData/Training/img in ../data/imagesTr, and put labels from RawData/Training/label in ../data/labelsTr. Download the json file for data splits from this link. Place the JSON file at ../data/dataset_0.json.
For the Adaptation, run: python train.py -net sam -mod sam_adpt -exp_name msa-3d-sam-btcv -sam_ckpt ./checkpoint/sam/sam_vit_b_01ec64.pth -image_size 1024 -b 8 -dataset decathlon -thd True -chunk 96 -dataset ../data -num_sample 4
You can modify following parameters to save the memory usage: '-b' the batch size, '-chunk' the 3D depth (channel) for each sample, '-num_sample' number of samples for Monai.RandCropByPosNegLabeld, 'evl_chunk' the 3D channel split step in the evaluation, decrease it if out of memory in the evaluation.

Run on your own dataset

It is simple to run MSA on the other datasets. Just write another dataset class following which in ./dataset.py. You only need to make sure you return a dict with { 'image': A tensor saving images with size [C,H,W] for 2D image, size [C, H, W, D] for 3D data. D is the depth of 3D volume, C is the channel of a scan/frame, which is commonly 1 for CT, MRI, US data. If processing, say like a colorful surgical video, D could the number of time frames, and C will be 3 for a RGB frame. 'label': The target masks. Same size with the images except the resolutions (H and W). 'p_label': The prompt label to decide positive/negative prompt. To simplify, you can always set 1 if don't need the negative prompt function. 'pt': The prompt. Should be the same as that in SAM, e.g., a click prompt should be [x of click, y of click], one click for each scan/frame if using 3d data. 'image_meta_dict': Optional. if you want save/visulize the result, you should put the name of the image in it with the key ['filename_or_obj']. ...(others as you want) } Welcome to open issues if you meet any problem. It would be appreciated if you could contribute your dataset extensions. Unlike natural images, medical images vary a lot depending on different tasks. Expanding the generalization of a method requires everyone's efforts.

TODO LIST

Jupyter tutorials.
Fix bugs in BTCV. Add BTCV example.
Release REFUGE2, BraTs dataloaders and examples
Changable Image Resolution
Fix bugs in Multi-GPU parallel
Sample and Vis in training
Release general data pre-processing and post-processing
Release evaluation
Deploy on HuggingFace
configuration
Release SSL code
Release Medical Adapter Zoo

Cite

 @article{wu2023medical,
   title={Medical sam adapter: Adapting segment anything model for medical image segmentation},
   author={Wu, Junde and Fu, Rao and Fang, Huihui and Liu, Yuanpei and Wang, Zhaowei and Xu, Yanwu and Jin, Yueming and Arbel, Tal},
   journal={arXiv preprint arXiv:2304.12620},
   year={2023}
 }