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

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

demo (gradio)

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.

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 tumor
Iterative correction for multiple tumors
Qualitative results with compared methods

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

• Our preprocessed

  We used four public datasets for 3D tumor segmentation in colon, pancreas, liver and kidney.

• Original

  Here are the links for the datasets: MSD-colon, MSD-pancreas, LiTS2017 and KiTS2021.

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

The model is licensed under the Apache 2.0 license

• SAM
• SAM-Med3D
• ProMISe
• nnU-Net

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 [email protected] for any questions and always happy to help! :)

Original PRISM repository.