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 tumor |  |
| Iterative correction for multiple tumors |  |
| Qualitative results with compared methods |  |
Datasets
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.
Models
| colon | pancreas | liver | kidney |
|---|---|---|---|
| Download | Download | Download | Download |
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! :)