Segment-Anything-Model-for-Medical-Images

Implementation of "Segment Anything Model for Medical Images?" in pytorch --for finetuning the SAM with box prompts.

Arxiv link: https://arxiv.org/pdf/2304.14660.pdf

MIA version link: https://www.sciencedirect.com/science/article/pii/S1361841523003213

Our work has been accepted by Medical Image Analysis (MedIA) 2023!

Usage

1. Prepare data and segmentation information json files corresponding to each test set.

Json file content.

"Info" refers to the segmentation target in this dataset, while "color" is the ground truth pixel value corresponding to the target.

{
   "info": {
       "1": "LeftVentricle",
       "2": "LeftVentricularMyocardium",
       "3": "RightVentricle"
   },
   "color": {
       "1": 85,
       "2": 170,
       "3": 255
   }
}

Dataset distribution.

 train: ../data/train_data/images/
 val: ../data/train_data/images/
 test: ../data/test_data/dataset_name/images/
 
 ├── train_data          
 │   ├── images        
 │   │   ├── 000001.png
 │   │   ├── 000002.png
 │   │   └── 000003.png
 │   └── labels         
 │       ├── 00001.png
 │       ├── 00002.png
 │       └── 00003.png
 └── val_data           
 |   ├── images        
 |   │   ├── 000001.png
 |   │   ├── 000002.png
 |   │   └── 000003.png
 |   └── labels         
 |       ├── 00001.png
 |       ├── 00002.png
 |       └── 00003.png
 └── test_data          
     ├── dataset1        
     │   ├── images
     |   |     ├── 000001.png
     |   │     ├── 000002.png
     |   |     └── 000003.png
     │   └── labels
     |         ├── 000001.png
     |         ├── 000002.png
     |         └── 000003.png
     └── dataset2         
         ├── images
         |     ├── 000001.png
         │     ├── 000002.png
         |     └── 000003.png
         └── labels
               ├── 000001.png
               ├── 000002.png
               └── 000003.png

2. Generate embedding for each single image.

$ python pre_grey_rgb2D.py         --img_path  data/train_data/images    --gt_path data/train_data/images                        --checkpoint sam_vit_b_01ec64.pth              #for preparing training data (embeddings) with ViT-B  
$ python pre_grey_rgb2D_Huge.py    --img_path  data/test_data            --gt_path data/test_data     --task_name 22_Heart       --checkpoint sam_vit_b_01ec64.pth              #for preparing testing data (embeddings) with ViT-B

3. Finetune SAM with your own data.

$ python train_only_box.py    --tr_npz_path data/precompute_vit_b/train  --val_npz_path data/precompute_vit_b/valid --model_type vit_b # finetune ViT-B
                              --tr_npz_path data/precompute_vit_h/train  --val_npz_path data/precompute_vit_h/valid --model_type vit_h # finetune ViT-H

4. Test on finetuned models and output the Dice results.

$ python test_only_box.py

5. Calculation of all the indicators (Dice, IOU, HD, etc.).

$ python cal_matric.py

Our pretrained weights

Checkpoints download path: https://drive.google.com/drive/folders/1jry-07RxGYQnT9cQE8weuCZurDCu58pj?usp=sharing

COSMOS 1050K Dataset

We collected and sorted 53 public datasets to build the large COSMOS 1050K medical image segmentation dataset. Following are the links to the datasets used in our paper.

Ownership and license of the datasets belong to their corresponding original papers, authors, or competition organizers. If you use the datasets, please cite the corresponding paper or links.

Links

Acknowledgments

Our code is based on Segment Anything and MedSAM. We appreciate the authors for their great works. We also sincerely appreciate all the challenge organizers and owners for providing the public medical image segmentation datasets.

Citation

If you find the code useful for your research, please cite our paper.

@article{huang2024segment,
  title={Segment anything model for medical images?},
  author={Huang, Yuhao and Yang, Xin and Liu, Lian and Zhou, Han and Chang, Ao and Zhou, Xinrui and Chen, Rusi and Yu, Junxuan and Chen, Jiongquan and Chen, Chaoyu and others},
  journal={Medical Image Analysis},
  volume={92},
  pages={103061},
  year={2024},
  publisher={Elsevier}
}

ksoh97/Segment-Anything-Model-for-Medical-Images

Segment-Anything-Model-for-Medical-Images

Our work has been accepted by Medical Image Analysis (MedIA) 2023!

Usage

1. Prepare data and segmentation information json files corresponding to each test set.

2. Generate embedding for each single image.

3. Finetune SAM with your own data.

4. Test on finetuned models and output the Dice results.

5. Calculation of all the indicators (Dice, IOU, HD, etc.).

Our pretrained weights

COSMOS 1050K Dataset

AbdomenCT-1K

ACDC

AMOS 2022

AutoLaparo

BrainPTM 2021

BraTS20

CAMUS

CellSeg Challenge-NeurIPS 2022

CHAOS

CHASE-DB1

Chest CT Segmentation

CRAG

crossMoDA

CVC-ClinicDB

DRIVE

EndoTect 2020

ETIS-Larib Polyp DB

FeTA

HaN-Seg

I2CVB

iChallenge-AMD

iChallenge-PALM

IDRiD 2018

iSeg 2019

ISIC 2018

IXI

KiPA22

KiTS19

KiTS21

Kvasir-Instrumen

Kvasir-SEG

LiVScar

LUNA16

M&Ms

Montgomery County CXR Set

MRSpineSeg

MSD

Multi-Atlas Labeling Beyond the Cranial Vault（Abdomen）：MALBCV-Abdomen

NCI-ISBI 2013

PROMISE12

QUBIQ 2021

SIIM-ACR

SKI10

SLIVER07

STARE

TN-SCUI 2020

VerSe19&VerSe20

Warwick-QU

WORD

EPFL_EM

ssTEM

TotalSegmentator

4C2021 C04 TLS01

Acknowledgments

Citation