UniSeg代码

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

<div align="center"> <img width="100%" alt="UniSeg示意图" src="https://yellow-cdn.veclightyear.com/ab5030c0/5c5f5a83-13ae-420d-8c3b-efdf5eedd685.png"> </div>

新闻

• 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

数据准备

• 下载MOTS数据集。
• 下载VerSe20数据集。
• 下载前列腺数据集。
• 下载BraTS21数据集。
• 下载AutoPET2022数据集。

预处理

• 步骤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：重采样图像的间距

<div align="center"> <img width="100%" alt="UniSeg插图" src="https://yellow-cdn.veclightyear.com/ab5030c0/435e2172-ba27-45a2-a76b-763f6bcba3be.png"> </div>

待办事项

• 数据集链接
• 预处理代码
• 上游代码发布
• 上游训练模型
• 下游代码发布
• 在新数据上使用上游训练模型进行推理

引用

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

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