Awesome
AutoSAM
This repo is pytorch implementation of paper "How to Efficiently Adapt Large Segmentation Model(SAM) to Medical Image Domains" by Xinrong Hu et al.
[Paper
]
Setup
The code requires python>=3.8
, as well as pytorch>=1.7
and torchvision>=0.8
.
clone the repository locally:
git clone git@github.com:xhu248/AutoSAM.git
and install requirements:
cd AutoSAM; pip install -e .
Download the checkpoints from SAM and place them under AutoSAM/
Dataset
The original ACDC data files can be dowonloaded at Automated Cardiac Diagnosis Challenge . The data is provided in nii.gz format. We convert them into PNG files as SAM requires RGB input. The processed data can be downloaded here
How to use
Finetune CNN decoder
python scripts/main_feat_seg.py --src_dir ${ACDC_folder} \
--data_dir ${ACDC_folder}/imgs/ --save_dir ./${output_dir} \
--b 4 --dataset ACDC --gpu ${gpu} \
--fold ${fold} --tr_size ${tr_size} --model_type ${model_type} --num_classes 4
${tr_size} decides how many volumes used in the training; ${model_type} is selected from vit_b (default), vit_l, and vit_h;
Finetune AutoSAM
python scripts/main_autosam_seg.py --src_dir ${ACDC_folder} \
--data_dir ${ACDC_folder}/imgs/ --save_dir ./${output_dir} \
--b 4 --dataset ACDC --gpu ${gpu} \
--fold ${fold} --tr_size ${tr_size} --model_type ${model_type} --num_classes 4
This repo also supports distributed training
python scripts/main_autosam_seg.py --src_dir ${ACDC_folder} --dist-url 'tcp://localhost:10002' \
--data_dir ${ACDC_folder}/imgs/ --save_dir ./${output_dir} \
--multiprocessing-distributed --world-size 1 --rank 0 -b 4 --dataset ACDC \
--fold ${fold} --tr_size ${tr_size} --model_type ${model_type} --num_classes 4
Todo
- Evaluate on more datasets
- Add more baselines
Citation
If you find our codes useful, please cite
@article{hu2023efficiently,
title={How to Efficiently Adapt Large Segmentation Model (SAM) to Medical Images},
author={Hu, Xinrong and Xu, Xiaowei and Shi, Yiyu},
journal={arXiv preprint arXiv:2306.13731},
year={2023}
}