/UniSeg

MICCAI 2023 Paper (Early Acceptance)

Primary LanguagePythonOtherNOASSERTION

UniSeg-code

This is the official pytorch implementation of our MICCAI 2023 paper "UniSeg: A Prompt-driven Universal Segmentation Model as well as A Strong Representation Learner". In this paper, we propose a Prompt-Driven Universal Segmentation model (UniSeg) to segment multiple organs, tumors, and vertebrae on 3D medical images with diverse modalities and domains.

UniSeg illustration

News

  • 2023.07.17: We have updated the code to better support the new multi-task segmentation. You just need to modify the self.task, self.task_class, and self.total_task_num in the UniSeg_Trainer.
  • 2023.07.19: We have provided the configuration file for predicting new data. In addition, we have updated the new data prediction code to restrict the output categories for specified tasks.
  • 2023.10.13: 🎉🎉🎉Our UniSeg achieved second place on both tasks of MICCAI SegRap 2023 with simply fine-tuning on the dataset.

Requirements

CUDA 11.5
Python 3.8
Pytorch 1.11.0
CuDNN 8.3.2.44

Usage

Installation

  • Clone this repo.
git clone https://github.com/yeerwen/UniSeg.git
cd UniSeg

Data Preparation

Pre-processing

  • Step 1:

    • Install nnunet by pip install nnunet.
    • Set path, for example:
      • export nnUNet_raw_data_base="/data/userdisk0/ywye/nnUNet_raw"
      • export nnUNet_preprocessed="/erwen_SSD/1T/nnUNet_preprocessed"
      • export RESULTS_FOLDER="/data/userdisk0/ywye/nnUNet_trained_models"

  • Step 2:

    • cd Upstream
    • Note that the output paths of the preprocessed datasets should be in the $nnUNet_raw_data_base/nnUNet_raw_data/ directory.
    • Run python prepare_Kidney_Dataset.py to normalize the name of the volumes for the Kidney dataset.
    • Run python Convert_MOTS_to_nnUNet_dataset.py to pre-process the MOTS dataset.
    • Run python Convert_VerSe20_to_nnUNet_dataset.py to pre-process the VerSe20 dataset and generate splits_final.pkl.
    • Run python Convert_Prostate_to_nnUNet_dataset.py to pre-process the Prostate dataset and generate splits_final.pkl.
    • Run python Convert_BraTS21_to_nnUNet_dataset.py to pre-process the BraTS21 dataset and generate splits_final.pkl.
    • Run python Convert_AutoPET_to_nnUNet_dataset.py to pre-process the AutoPET2022 dataset and generate splits_final.pkl.

  • Step 3:

    • Copy Upstream/nnunet to replace nnunet, which is installed by pip install nnunet (the address is usually 'anaconda3/envs/your envs/lib/python3.8/site-packages/nnunet').
    • Run nnUNet_plan_and_preprocess -t 91 --verify_dataset_integrity --planner3d MOTSPlanner3D.
    • Run nnUNet_plan_and_preprocess -t 37 --verify_dataset_integrity --planner3d VerSe20Planner3D.
    • Run nnUNet_plan_and_preprocess -t 20 --verify_dataset_integrity --planner3d ProstatePlanner3D.
    • Run nnUNet_plan_and_preprocess -t 21 --verify_dataset_integrity --planner3d BraTS21Planner3D.
    • Run nnUNet_plan_and_preprocess -t 11 --verify_dataset_integrity --planner3d AutoPETPlanner3D.
    • Move splits_final.pkl of each dataset to the address of its pre-processed dataset. For example, '***/nnUNet_preprocessed/Task091_MOTS/splits_final.pkl'. Note that, to follow DoDNet, we provide splits_final.pkl of the MOTS dataset in Upstream/MOTS_data_split/splits_final.pkl.
    • Run python merge_each_sub_dataet.py to form a new dataset.
    • To make sure that we use the same data split, we provide the final data split in Upstream/splits_final_11_tasks.pkl.

Training and Test

  • Move Upstream/run_ssl.sh and Upstream/UniSeg_Metrics_test.py to "***/nnUNet_trained_models/".
  • cd ***/nnUNet_trained_models/.
  • Run sh run_ssl.sh for training (GPU Memory Cost: ~10GB, Time Cost: ~210s each epoch).

Pretrained weights

Downstream Tasks

  • cd Downstream
  • Download BTCV dataset.
  • Download VS dataset.
  • Run python Convert_BTCV_to_nnUNet_dataset.py to pre-process the BTCV dataset and generate splits_final.pkl.
  • Run python Convert_VSseg_to_nnUNet_dataset.py to pre-process the VS dataset and generate splits_final.pkl.
  • Update the address of the pre-trained model in the 'Downstream/nnunet/training/network_training/UniSeg_Trainer_DS.py' file (line 97)
  • Copy Downstream/nnunet to replace nnunet, which is installed by pip install nnunet (the address is usually 'anaconda3/envs/your envs/lib/python3.8/site-packages/nnunet').
  • Run nnUNet_plan_and_preprocess -t 60 --verify_dataset_integrity.
  • Run nnUNet_plan_and_preprocess -t 61 --verify_dataset_integrity.
  • Move splits_final.pkl of two datasets to the addresses of their pre-processed datasets.
  • To make sure that we use the same data split for the downstream datasets, we provide the final data splits in Downstream/splits_final_BTCV.pkl and Downstream/splits_final_VS.pkl.
  • Training and Test:
    • For the BTCV dataset: CUDA_VISIBLE_DEVICES=0 nnUNet_n_proc_DA=32 nnUNet_train 3d_fullres UniSeg_Trainer_DS 60 0
    • For the VS dataset: CUDA_VISIBLE_DEVICES=0 nnUNet_n_proc_DA=32 nnUNet_train 3d_fullres UniSeg_Trainer_DS 61 0

Prediction on New Data

  • Download the Upstream trained model and configuration file.
  • Move them to ./nnUNet_trained_models/UniSeg_Trainer/3d_fullres/Task097_11task/UniSeg_Trainer__DoDNetPlans/fold_0/ and rename them to model_final_checkpoint.model and model_final_checkpoint.model.pkl, respectively.
  • cd Upstream
  • Copy Upstream/nnunet to replace nnunet, which is installed by pip install nnunet
  • Run CUDA_VISIBLE_DEVICES=1 nnUNet_n_proc_DA=32 nnUNet_predict -i /data/userdisk0/ywye/nnUNet_raw/nnUNet_raw_data/Test/Image/ -o /data/userdisk0/ywye/nnUNet_raw/nnUNet_raw_data/Test/Predict/10/ -t 97 -m 3d_fullres -tr UniSeg_Trainer -f 0 -task_id 7 -exp_name UniSeg_Trainer -num_image 1 -modality CT -spacing 3.0,1.5,1.5
    • -i: Path of the input image(s), name format of the input image: name_0000.nii.gz (name_0001.nii.gz)
    • -o: Path of the output mask(s)
    • -task_id Selected segmentation task.
      • -1 means predicting all segmentation tasks under a specific modality.
      • 0: "liver and liver tumor segmentation"
      • 1: "kidney and kidney tumor segmentation"
      • 2: "hepatic vessel and hepatic tumor segmentation"
      • 3: "pancreas and pancreas tumor segmentation"
      • 4: "colon tumor segmentation"
      • 5: "lung tumor segmentation"
      • 6: "spleen segmentation"
      • 7: "vertebrae segmentation"
      • 8: "prostate segmentation"
      • "9": "brain tumors: edema, non-enhancing, and enhancing segmentation"
      • "10": "whole-body tumors segmentation"
    • -num_image: Channel number of the input image(s)
    • -modality: "CT" or "MR" (prostate) or "MR,MR,MR,MR" (brain tumors) or "CT,PET" (whole-body tumors)
    • -spacing: Spacing of resampled image(s)
UniSeg illustration

To do

  • Dataset Links
  • Pre-processing Code
  • Upstream Code Release
  • Upstream Trained Model
  • Downstream Code Release
  • Inference of Upstream Trained Model on New Data

Citation

If this code is helpful for your study, please cite:

@article{ye2023uniseg,
  title={UniSeg: A Prompt-driven Universal Segmentation Model as well as A Strong Representation Learner},
  author={Yiwen Ye, Yutong Xie, Jianpeng Zhang, Ziyang Chen, and Yong Xia},
  booktitle={International Conference on Medical Image Computing and Computer-Assisted Intervention},
  pages={508--518},
  year={2023},
  organization={Springer}
}

Acknowledgements

The whole framework is based on nnUNet v1.

Contact

Yiwen Ye (ywye@mail.nwpu.edu.cn)