Doctor-SAM （ICIC 2024 Oral）

This repo. is the official implementation of 'Dr-SAM: U-shape Structure Segment Anything Model for Generalizable Medical Image Segmentation'.
Authors: Xiangzuo Huo, Shengwei Tian, Bingming Zhou, Long Yu, Aolun Li.

Overview

DrSAM Segmentation Qualitative Results

Run

0. Requirements:

• python3
• pytorch <= 2.0.0

1. Train:

• Run python train_single.py

2. Evaluate:

• Modify parser.add_argument('--eval', default=True) in train_single.py
• Run python train_single.py

3. Visualize:

• Modify parser.add_argument('--visualize', default=True) in train_single.py
• Run python train_single.py

Model Checkpoints

Click the links below to download the checkpoint for the corresponding model type.

• vit_h: ViT-H SAM model.
• vit_l: ViT-L SAM model.
• vit_b: ViT-B SAM model.

Init checkpoint can be downloaded from hugging face link

Expected checkpoint

pretrained_checkpoint
|____sam_vit_b_maskdecoder.pth
|____sam_vit_b_01ec64.pth
|____sam_vit_l_maskdecoder.pth
|____sam_vit_l_0b3195.pth
|____sam_vit_h_maskdecoder.pth
|____sam_vit_h_4b8939.pth

Reference

Some of the codes in this repo are borrowed from:

• SAM
• SAM-HQ

Thank them for their awesome work!

Citation

If you find our paper/code is helpful, please consider citing:

@inproceedings{huo2024dr,
  title={Dr-SAM: U-Shape Structure Segment Anything Model for Generalizable Medical Image Segmentation},
  author={Huo, Xiangzuo and Tian, Shengwei and Zhou, Bingming and Yu, Long and Li, Aolun},
  booktitle={International Conference on Intelligent Computing},
  pages={197--207},
  year={2024},
  organization={Springer}
}