Home

Awesome

PRISM

PRISM: A Promptable and Robust Interactive Segmentation Model with Visual Prompts

The Arxiv PDF file is loading a bit slowly.

News

[03/07/24] The pretrained PRISM models and preprocessed datasets are uploaded.

TODO

demo (gradio)

What is PRISM?

<img src='figs/framework_v1.png' width='600'>

PRISM is a robust model/method for interactive segmentation in medical imaging. We strive for human-level performance, as a human-in-loop interactive segmentation model with prompts should gradually refine its outcomes until they closely match inter-rater variability.

PRISM tumor segmentation examples

Quantitative results is avaliable in paper (provide a link).

Briefly, PRISM produces tumor segmentation with mean Dice values of 93.79 (colon), 94.48 (pancreas), 94.18 (liver), and 96.58 (kidney).

Iterative correction for colon tumoriterative_colon
Iterative correction for multiple tumorsiterative_all
Qualitative results with compared methodsqualitative_results

Datasets

The anatomical differences among individuals and ambiguous boundaries are present in the datasets.

Models

colonpancreasliverkidney
DownloadDownloadDownloadDownload

Get Started

Installation

conda create -n prism python=3.9
conda activate prism
sudo install git
pip install torch==1.12.1+cu113 torchvision==0.13.1+cu113 --extra-index-url https://download.pytorch.org/whl/cu113 # install pytorch
pip install git+https://github.com/facebookresearch/segment-anything.git # install segment anything packages
pip install git+https://github.com/deepmind/surface-distance.git # for normalized surface dice (NSD) evaluation
pip install -r requirements.txt

Train

python train.py --data colon --data_dir your_data_directory --save_name your_save_name --multiple_outputs --dynamic --use_box --refine

Train (Distributed Data Parallel)

the only difference between this and above (train) command is the use of "--ddp".

python train.py --data colon --data_dir your_data_directory --save_name your_save_name -multiple_outputs --dynamic --use_box --refine --ddp

Test

put downloaded pretrained model under the implementation directory

python test.py --data colon --data_dir your_data_directory --split test --checkpoint best --save_name prism_pretrain --num_clicks 1 --iter_nums 11 --multiple_outputs --use_box --use_scribble --efficient_scribble --refine --refine_test

FAQ

if you got the error as AttributeError: module 'cv2' has no attribute 'ximgproc', please check this out

DDP mode has lower Dice and more epoch numbers may solve it

On my end, combining trainer and trainer_basic speeds up

training the model without refine module (as we reported in the paper) has better accuracy than with refine but not using it

License

The model is licensed under the Apache 2.0 license

Acknowledgements

If you find this repository useful, please consider citing this paper:

@article{li2024prism,
  title={PRISM: A Promptable and Robust Interactive Segmentation Model with Visual Prompts},
  author={Li, Hao and Liu, Han and Hu, Dewei and Wang, Jiacheng and Oguz, Ipek},
  journal={arXiv preprint arXiv:2404.15028},
  year={2024}
}

Please send an email to hao.li.1@vanderbilt.edu for any questions and always happy to help! :)