PRISM-placenta

• Interactive Segmentation Model for Placenta Segmentation from 3D Ultrasound Images (arXiv version)

• PRISM Lite: A lightweight model for interactive 3D placenta segmentation in ultrasound (arXiv version)

• PRISM: A Promptable and Robust Interactive Segmentation Model with Visual Prompts (arXiv version)

News

[07/20/24] PRISM: Lite: A lightweight model version for placenta segmentation is uploaded.

[07/07/24] Repo is created as a version control for the placenta application.

TODO

demo (radio) is on the list

TL;DR

• PRISM is an effective and efficient 3D interactive model for placenta segmentation

• The detailed effective and efficient results are included in our paper, please check if you are interested

• We note this is such an open area, and we keep pushing it with more analyses. Please let me know if you have any thoughts

Workflow

To mimic human behavior, we sample prompts from the FN and FP regions of the current segmentation at each iteration. The initial sampling only has positive prompts.

Effective

The quantitative results can be viewed in our paper, briefly, mean Dice scores are 0.9 for automated and 0.97 for PRISM. Qualitative results with compared methods. We tested the performance of nnUnet, and it is worse than our state-of-the-art model (automated).

Efficient

We consider a Dice score of 0.95 as a bar for success, which is higher than inter-rater variability (0.85-0.90)

Datasets

We used the in-house dataset, the details can be viewed in other papers from our group (you can find them in the Reference section in the paper)

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 ultrasound --data_dir your_data_directory --save_name your_save_name --multiple_outputs --dynamic --use_box --refine

add "--use_scribble" and "--efficient_scribble" if you want to train with scribbles.

Train PRISM Lite

use "--initial_seg" and "--use_penn" to adapt the masks from pretrained model (developed by Penn's group).

python train.py --data ultrasound --data_dir your_data_directory --save_name your_save_name --multiple_outputs --dynamic --use_box --refine --initial_seg --use_penn

Train (Distributed Data Parallel)

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

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

Test

python test.py --data ultrasound --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

This repo is tested to reproduce the results in our paper. However, I haven't cleaned the code, and it has some arguments for past/future efficient or effective analyses.

License

The model is licensed under the Apache 2.0 license

Acknowledgements

Thanks for the code from: SAM, SAM-Med3D, ProMISe, ScribblePrompt, nnU-Net

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

@article{li2024interactive,
  title={Interactive Segmentation Model for Placenta Segmentation from 3D Ultrasound Images},
  author={Li, Hao and Oguz, Baris and Arenas, Gabriel and Yao, Xing and Wang, Jiacheng and Pouch, Alison and Byram, Brett and Schwartz, Nadav and Oguz, Ipek},
  journal={arXiv preprint arXiv:2407.08020},
  year={2024}
}

@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 we are always happy to help! :)