Skip to content

UyJZ/MedSAM-BUAA

 
 

Folders and files

NameName
Last commit message
Last commit date

Latest commit

 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 

Repository files navigation

MedSAM

This is the official repository for MedSAM: Segment Anything in Medical Images.

Welcome to join our mailing list to get updates.

News

Installation

  1. Create a virtual environment conda create -n medsam python=3.10 -y and activate it conda activate medsam
  2. Install Pytorch 2.0
  3. git clone https://github.com/bowang-lab/MedSAM
  4. Enter the MedSAM folder cd MedSAM and run pip install -e .

Get Started

Download the model checkpoint and place it at e.g., work_dir/MedSAM/medsam_vit_b

We provide three ways to quickly test the model on your images

  1. Command line
python MedSAM_Inference.py # segment the demo image

Segment other images with the following flags

-i input_img
-o output path
--box bounding box of the segmentation target
  1. Jupyter-notebook

We provide a step-by-step tutorial on CoLab

You can also run it locally with tutorial_quickstart.ipynb.

  1. GUI

Install PyQt5 with pip: pip install PyQt5 or conda: conda install -c anaconda pyqt

python gui.py

Load the image to the GUI and specify segmentation targets by drawing bounding boxes.

MedSAM-Demo.mp4

Model Training

Data preprocessing

Download SAM checkpoint and place it at work_dir/SAM/sam_vit_b_01ec64.pth .

Download the demo dataset and unzip it to data/FLARE22Train/.

This dataset contains 50 abdomen CT scans and each scan contains an annotation mask with 13 organs. The names of the organ label are available at MICCAI FLARE2022.

Run pre-processing

Install cc3d: pip install connected-components-3d

python pre_CT_MR.py
  • split dataset: 80% for training and 20% for testing
  • adjust CT scans to soft tissue window level (40) and width (400)
  • max-min normalization
  • resample image size to 1024x1024
  • save the pre-processed images and labels as npy files

Training on multiple GPUs (Recommend)

The model was trained on five A100 nodes and each node has four GPUs (80G) (20 A100 GPUs in total). Please use the slurm script to start the training process.

sbatch train_multi_gpus.sh

When the training process is done, please convert the checkpoint to SAM's format for convenient inference.

python utils/ckpt_convert.py # Please set the corresponding checkpoint path first

Training on one GPU

python train_one_gpu.py

Combined parameter-efficient variants

Set the -model_type argument on any training script to pick one of the new hybrids:

  • vit_lora_prompt: LoRA-tuned attention with deep prompt residuals.
  • vit_adapter_lora: Adapter backbone with frozen ViT blocks plus LoRA heads.
  • vit_adapter_prompt: Adapter backbone augmented with prompt tuning.
  • vit_adapter_lora_prompt: Stacks adapter prompts and LoRA on the same model.

Example:

python train_one_gpu_sam_adapter.py -dataset_name busi -model_type vit_adapter_lora_prompt

开发原则

dataloader

dataloader统一放在dataloader.py中

模型

模型统一放到segment_anything/modeling目录下,类的命名和image_encoder等保持一致,确保ckpt可以正确加载,如果需要改loss,新加一个train_one_gpu.py,模型需要注册到segment_anything/build_sam.py中

Acknowledgements

  • We highly appreciate all the challenge organizers and dataset owners for providing the public dataset to the community.
  • We thank Meta AI for making the source code of segment anything publicly available.
  • We also thank Alexandre Bonnet for sharing this great blog

Reference

@article{MedSAM,
  title={Segment Anything in Medical Images},
  author={Ma, Jun and He, Yuting and Li, Feifei and Han, Lin and You, Chenyu and Wang, Bo},
  journal={Nature Communications},
  volume={15},
  pages={654},
  year={2024}
}

About

Segment Anything in Medical Images in BUAA

Resources

License

Stars

Watchers

Forks

Releases

No releases published

Packages

No packages published

Languages

  • Jupyter Notebook 80.7%
  • Python 19.1%
  • Shell 0.2%