/PromptNucSeg

[ECCV 2024] Code for "Unleashing the Power of Prompt-driven Nucleus Instance Segmentation"

Primary LanguagePython

Requirements

torch 2.0.1, mmcv, mmdet

Dataset

  1. download Kumar, CPM-17 and PanNuke datasets.

  2. run extract_data.py for data pre-processing.

    datasets
    ├── cpm17
    │   ├── extract_data.py
    │   ├── test
    │   └── train
    ├── cpm17_test_files.npy
    ├── cpm17_train_files.npy
    ├── kumar
    │   ├── extract_data.py
    │   ├── images
    │   └── labels
    ├── kumar_test_files.npy
    ├── kumar_train_files.npy
    ├── pannuke
    │   ├── extract_data.py
    │   ├── Fold 1
    │   ├── Fold 2
    │   ├── Fold 3
    │   ├── Images
    │   └── Masks
    ├── pannuke123_test_files.npy
    ├── pannuke123_train_files.npy
    ├── pannuke123_val_files.npy
    ├── pannuke213_test_files.npy
    ├── pannuke213_train_files.npy
    ├── pannuke213_val_files.npy
    ├── pannuke321_test_files.npy
    ├── pannuke321_train_files.npy
    └── pannuke321_val_files.npy

Training

  1. Train the prompter

    cd prompter
    python main.py --config dpa_pannuke123.py --output_dir dpa_pannuke123
    # python main.py --config dpa_pannuke213.py --output_dir dpa_pannuke123
    # python main.py --config dpa_pannuke321.py --output_dir dpa_pannuke123
  2. Use the trained prompter to generate nuclei prompts for the validation and test sets.

    python predict_prompts.py --config dpa_pannuke123.py --resume checkpoint/dpa_pannuke123/best.pth
    # python predict_prompts.py --config dpa_pannuke213.py --resume checkpoint/dpa_pannuke213/best.pth
    # python predict_prompts.py --config dpa_pannuke321.py --resume checkpoint/dpa_pannuke321/best.pth
  3. Download SAM's pre-trained weights into segmentor/pretrained and train the segmentor.

    cd segmentor
    torchrun --nproc_per_node=4 main.py --config pannuke123_b.py --output_dir pannuke123_b
    # torchrun --nproc_per_node=4 main.py --config pannuke213_b.py --output_dir pannuke213_b
    # torchrun --nproc_per_node=4 main.py --config pannuke321_b.py --output_dir pannuke321_b

Evaluation

see test.sh

Checkpoints

Kumar CPM-17 PanNuke123 PanNuke213 PanNuke321
Prompter OneDrive OneDrive OneDrive OneDrive OneDrive
Segmentor-B OneDrive OneDrive OneDrive OneDrive OneDrive
Segmentor-L OneDrive OneDrive OneDrive OneDrive OneDrive
Segmentor-H OneDrive OneDrive OneDrive OneDrive OneDrive

Contact

If you have any questions or concerns, feel free to report issues or directly contact us (Zhongyi Shui shuizhongyi@westlake.edu.cn).

Citation

If you find this code useful for your research, please cite us using the following BibTeX entry.

@article{shui2023unleashing,
  title={Unleashing the Power of Prompt-driven Nucleus Instance Segmentation},
  author={Shui, Zhongyi and Zhang, Yunlong and Yao, Kai and Zhu, Chenglu and Sun, Yuxuan and Yang, Lin},
  journal={arXiv preprint arXiv:2311.15939},
  year={2023}
}