ScribblePrompt
ScribblePrompt copied to clipboard
ScribblePrompt: Fast and Flexible Interactive Segmentation for Any Medical Image
ScribblePrompt
Project Page | Paper | Online Demo | Video
Official implementation of ScribblePrompt: Fast and Flexible Interactive Segmentation for any Biomedical Image accepted at ECCV 2024
Hallee E. Wong, Marianne Rakic, John Guttag, Adrian V. Dalca
Updates
- (2024-07-01) ScribblePrompt has been accepted to ECCV 2024!
- (2024-06-17) ScribblePrompt won the Bench-to-Bedside Award at the CVPR 2024 DCAMI Workshop!
- (2024-04-16) Released MedScribble -- a diverse dataset of segmentation tasks with scribble annotations
- (2024-04-15) An updated version of the paper is on arXiv!
- (2024-04-14) Added Google Colab Tutorial
- (2024-01-19) Released scribble simulation code
- (2023-12-15) Released model code and weights
- (2023-12-12) Paper and online demo released
Overview
ScribblePrompt is an interactive segmentation tool that enables users to segment unseen structures in medical images using scribbles, clicks, and bounding boxes.
Try ScribblePrompt
- Interactive online demo on Hugging Face Spaces
- See Installation and Getting Started for how to run the Gradio demo locally
- Jupyter notebook colab tutorial
Models
We provide checkpoints for two versions of ScribblePrompt:
-
ScribblePrompt-UNet with an efficient fully-convolutional architecture
-
ScribblePrompt-SAM based on the Segment Anything Model
Both models have been trained with iterative scribbles, click, and bounding box interactions on a diverse collection of 65 medical imaging datasets with both real and synthetic labels.
MedScribble Dataset
We release MedScribble, a dataset of multi-annotator scribble annotations for diverse biomedical image segmentation tasks, under ./MedScribble
. See the readme for more info and ./MedScribble/tutorial.ipynb
for a preview of the data.
Installation
You can install scribbleprompt
in two ways:
- With pip:
pip install git+https://github.com/halleewong/ScribblePrompt.git
- Manually: cloning it and installing dependencies
git clone https://github.com/halleewong/ScribblePrompt
python -m pip install -r ./ScribblePrompt/requirements.txt
export PYTHONPATH="$PYTHONPATH:$(realpath ./ScribblePrompt)"
The following optional dependencies are necessary for the local demo:
pip install gradio==3.40.1
Getting Started
First, download the model checkpoints to ./checkpoints
.
To run an interactive demo locally:
python demos/app.py
To instantiate ScribblePrompt-UNet and make a prediction:
from scribbleprompt import ScribblePromptUNet
sp_unet = ScribblePromptUNet()
mask = sp_unet.predict(
image, # (B, 1, H, W)
point_coords, # (B, n, 2)
point_labels, # (B, n)
scribbles, # (B, 2, H, W)
box, # (B, n, 4)
mask_input, # (B, 1, H, W)
) # -> (B, 1, H, W)
To instantiate ScribblePrompt-SAM and make a prediction:
from scribbleprompt import ScribblePromptSAM
sp_sam = ScribblePromptSAM()
mask, img_features, low_res_logits = sp_sam.predict(
image, # (B, 1, H, W)
point_coords, # (B, n, 2)
point_labels, # (B, n)
scribbles, # (B, 2, H, W)
box, # (B, n, 4)
mask_input, # (B, 1, 256, 256)
) # -> (B, 1, H, W), (B, 16, 256, 256), (B, 1, 256, 256)
For best results, image
should have spatial dimensions $(H,W) = (128,128)$ and pixel values min-max normalized to the $[0,1]$ range.
For ScribblePrompt-UNet, mask_input
should be the logits from the previous prediction. For ScribblePrompt-SAM, mask_input
should be low_res_logits
from the previous prediction.
To Do
- [x] Release Gradio demo
- [x] Release model code and weights
- [x] Release jupyter notebook tutorial
- [x] Release scribble simulation code
- [x] Release MedScribble dataset
- [ ] Release training code
- [ ] Release segmentation labels collected using ScribblePrompt
Acknowledgements
-
Code for ScribblePrompt-SAM builds on Segment Anything. Thanks to Meta AI for open-sourcing the model.
-
We use functions from voxsynth for applying random deformations during scribble simulation
Citation
If you find our work or any of our materials useful, please cite our paper:
@article{wong2024scribbleprompt,
title={ScribblePrompt: Fast and Flexible Interactive Segmentation for Any Biomedical Image},
author={Hallee E. Wong and Marianne Rakic and John Guttag and Adrian V. Dalca},
journal={European Conference on Computer Vision (ECCV)},
year={2024},
}
License
Code for this project is released under the Apache 2.0 License