UniSeg代码

这是我们MICCAI 2023论文"UniSeg：一个提示驱动的通用分割模型以及强大的表示学习器"的官方PyTorch实现。在这篇论文中，我们提出了一个提示驱动的通用分割模型（UniSeg），用于在具有不同模态和领域的3D医学图像上分割多个器官、肿瘤和椎骨。

新闻

2023.07.17：我们更新了代码以更好地支持新的多任务分割。您只需修改UniSeg_Trainer中的self.task、self.task_class和self.total_task_num。
2023.07.19：我们提供了预测新数据的配置文件。此外，我们更新了新数据预测代码，以限制指定任务的输出类别。
2023.10.13：🎉🎉🎉我们的UniSeg在MICCAI SegRap 2023的两个任务中均获得第二名，仅通过在数据集上进行微调。

要求

CUDA 11.5 Python 3.8 Pytorch 1.11.0 CuDNN 8.3.2.44

使用方法

安装

克隆此仓库。

git clone https://github.com/yeerwen/UniSeg.git
cd UniSeg

数据准备

预处理

步骤1：
- 通过pip install nnunet安装nnunet。
- 设置路径，例如：
  - 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"
步骤2：
- cd Upstream
- 注意，预处理后的数据集的输出路径应在$nnUNet_raw_data_base/nnUNet_raw_data/目录中。
- 运行python prepare_Kidney_Dataset.py以规范化肾脏数据集的体积名称。
- 运行python Convert_MOTS_to_nnUNet_dataset.py以预处理MOTS数据集。
- 运行python Convert_VerSe20_to_nnUNet_dataset.py以预处理VerSe20数据集并生成splits_final.pkl。
- 运行python Convert_Prostate_to_nnUNet_dataset.py以预处理前列腺数据集并生成splits_final.pkl。
- 运行python Convert_BraTS21_to_nnUNet_dataset.py以预处理BraTS21数据集并生成splits_final.pkl。
- 运行python Convert_AutoPET_to_nnUNet_dataset.py以预处理AutoPET2022数据集并生成splits_final.pkl。
步骤3：
- 复制Upstream/nnunet以替换通过pip install nnunet安装的nnunet（地址通常为'anaconda3/envs/你的环境/lib/python3.8/site-packages/nnunet'）。
- 运行nnUNet_plan_and_preprocess -t 91 --verify_dataset_integrity --planner3d MOTSPlanner3D。
- 运行nnUNet_plan_and_preprocess -t 37 --verify_dataset_integrity --planner3d VerSe20Planner3D。
- 运行nnUNet_plan_and_preprocess -t 20 --verify_dataset_integrity --planner3d ProstatePlanner3D。
- 运行nnUNet_plan_and_preprocess -t 21 --verify_dataset_integrity --planner3d BraTS21Planner3D。
- 运行nnUNet_plan_and_preprocess -t 11 --verify_dataset_integrity --planner3d AutoPETPlanner3D。
- 将每个数据集的splits_final.pkl移动到其预处理数据集的地址。例如，'***/nnUNet_preprocessed/Task091_MOTS/splits_final.pkl'。注意，为了遵循DoDNet，我们在Upstream/MOTS_data_split/splits_final.pkl中提供了MOTS数据集的splits_final.pkl。
- 运行python merge_each_sub_dataet.py以形成新的数据集。
- 为确保我们使用相同的数据划分，我们在Upstream/splits_final_11_tasks.pkl中提供了最终的数据划分。

训练和测试

将Upstream/run_ssl.sh和Upstream/UniSeg_Metrics_test.py移动到"***/nnUNet_trained_models/"。
cd ***/nnUNet_trained_models/。
运行sh run_ssl.sh进行训练（GPU内存消耗：约10GB，时间消耗：每个epoch约210秒）。

预训练权重

上游训练模型可在UniSeg_11_Tasks获取。
plans.pkl文件。

下游任务

cd Downstream
下载BTCV数据集。
下载VS数据集。
运行python Convert_BTCV_to_nnUNet_dataset.py预处理BTCV数据集并生成splits_final.pkl。
运行python Convert_VSseg_to_nnUNet_dataset.py预处理VS数据集并生成splits_final.pkl。
更新'Downstream/nnunet/training/network_training/UniSeg_Trainer_DS.py'文件中预训练模型的地址（第97行）。
复制Downstream/nnunet替换通过pip install nnunet安装的nnunet（地址通常为'anaconda3/envs/your envs/lib/python3.8/site-packages/nnunet'）。
运行nnUNet_plan_and_preprocess -t 60 --verify_dataset_integrity。
运行nnUNet_plan_and_preprocess -t 61 --verify_dataset_integrity。
将两个数据集的splits_final.pkl移动到它们预处理后数据集的地址。
为确保我们对下游数据集使用相同的数据划分，我们在Downstream/splits_final_BTCV.pkl和Downstream/splits_final_VS.pkl中提供了最终的数据划分。
训练和测试：
- 对于BTCV数据集：CUDA_VISIBLE_DEVICES=0 nnUNet_n_proc_DA=32 nnUNet_train 3d_fullres UniSeg_Trainer_DS 60 0
- 对于VS数据集：CUDA_VISIBLE_DEVICES=0 nnUNet_n_proc_DA=32 nnUNet_train 3d_fullres UniSeg_Trainer_DS 61 0

对新数据进行预测

下载上游训练模型和配置文件。
将它们移动到./nnUNet_trained_models/UniSeg_Trainer/3d_fullres/Task097_11task/UniSeg_Trainer__DoDNetPlans/fold_0/并分别重命名为model_final_checkpoint.model和model_final_checkpoint.model.pkl。
cd Upstream
复制Upstream/nnunet替换通过pip install nnunet安装的nnunet
运行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：输入图像的路径，输入图像的命名格式：name_0000.nii.gz（name_0001.nii.gz）
- -o：输出掩码的路径
- -task_id 选择的分割任务。
  - -1表示预测特定模态下的所有分割任务。
  - 0："肝脏和肝脏肿瘤分割"
  - 1："肾脏和肾脏肿瘤分割"
  - 2："肝血管和肝脏肿瘤分割"
  - 3："胰腺和胰腺肿瘤分割"
  - 4："结肠肿瘤分割"
  - 5："肺部肿瘤分割"
  - 6："脾脏分割"
  - 7："椎骨分割"
  - 8："前列腺分割"
  - "9"："脑肿瘤：水肿、非增强和增强分割"
  - "10"："全身肿瘤分割"
- -num_image：输入图像的通道数
- -modality："CT"或"MR"（前列腺）或"MR,MR,MR,MR"（脑肿瘤）或"CT,PET"（全身肿瘤）
- -spacing：重采样图像的间距

待办事项

引用

如果此代码对您的研究有帮助，请引用：

@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}
}

致谢

整个框架基于nnUNet v1。

联系方式

叶亦文 (ywye@mail.nwpu.edu.cn)