SAMed

本仓库包含以下论文的实现:

定制化的Segment Anything Model用于医学图像分割
张凯东和刘东
技术报告
[论文]

Colab在线演示:

:star: 新闻

• 感谢我的导师的大力投入,我得以微调SAM的vit_h版本,以实现更精确的医学图像分割。现在,我们发布了SAMed的vit_h版本(我们将此版本称为SAMed_h),下表展示了SAMed和SAMed_h的对比。

无需任何技巧,SAMed_h的性能就大幅超越了SAMed。尽管vit_h版本的模型大小比vit_b版本大得多(超过2G,而vit_b约350M),但SAMed_h的LoRA检查点并没有增加很多(从18M增加到21M)。因此,SAMed_h的部署和存储成本与SAMed几乎相当。鉴于业界倾向于部署更大、性能更好的模型,我们相信SAMed_h在实际的计算机辅助诊断和术前规划中更有前景。有关SAMed_h的更多细节,请访问SAMed_h目录。

概述

待办事项

• 制作演示。
• 在更多数据集上微调
• 基于SAM的vit_l或vit_h模式制作SAMed

先决条件

• Linux (我们在Ubuntu 18.04上测试了我们的代码)
• Anaconda
• Python 3.7.11
• Pytorch 1.9.1

首先,请克隆仓库

git clone https://github.com/hitachinsk/SAMed.git

然后,请运行以下命令:

conda create -n SAMed python=3.7.11
conda activate SAMed
pip install -r requirements.txt

如果您有原始的Synapse数据集,我们提供了预处理脚本来处理和规范化用于训练的数据。请参考该文件夹了解更多详情。

快速开始

我们强烈建议您尝试我们的在线演示 。

目前,我们提供了SAMed和SAMed_s模型以便快速复现我们的结果。LoRA检查点及其对应的配置如下表所示。

以下是使用说明:

1. 切换目录到此仓库的根目录。
2. 请下载预训练的SAM模型（由SAM的原始仓库提供）和SAMed的LoRA检查点。将它们放在./checkpoints文件夹中。
3. 请下载测试集并将其放在./testset文件夹中。然后解压并删除该文件。
4. 运行以下命令来测试SAMed的性能。

python test.py --is_savenii --output_dir <您的输出目录>

如果一切正常，您会发现平均DSC为0.8188（81.88），HD为20.64，这与论文中的表1相对应。并在<您的输出目录>中查看测试结果。

此外，我们还提供了SAMed_s模型，它使用LoRA对图像编码器和掩码解码器中的transformer块进行微调。与SAMed相比，SAMed_s的模型尺寸更小，但性能也略有下降。如果您想使用这个模型，请下载并将其放在./checkpoints_s文件夹中，然后运行以下命令来测试其性能。

python test.py --is_savenii --output_dir <您的输出目录> --lora_ckpt checkpoints_s/epoch_159.pth --module sam_lora_image_encoder_mask_decoder

SAMed_s的平均DSC为0.7778（77.78），HD为31.72，这与论文中的表3相对应。

训练

我们使用2个RTX 3090 GPU进行训练。

1. 请下载经过处理的训练集，其分辨率为224x224，并将其放在<您的文件夹>中。然后解压并删除该文件。我们还准备了分辨率为512x512的训练集供参考，224x224版本的训练集是从512x512版本下采样得到的。
2. 运行以下命令来训练SAMed。

python train.py --root_path <您的文件夹> --output <您的输出路径> --warmup --AdamW 

在<您的输出路径>中查看结果。

许可证

本工作采用MIT许可证。详情请参见LICENSE。

引用

如果我们的工作启发了您的研究，或者代码的某些部分对您的工作有用，请引用我们的论文：

@article{samed,
  title={Customized Segment Anything Model for Medical Image Segmentation},
  author={Kaidong Zhang, and Dong Liu},
  journal={arXiv preprint arXiv:2304.13785},
  year={2023}
}

联系方式

如果您有任何问题，请通过以下方式联系我们：

• richu@mail.ustc.edu.cn

致谢

我们感谢Segment Anything Model的开发者和Synapse多器官分割数据集的提供者。SAMed的代码基于TransUnet和SAM LoRA构建，我们对这些优秀的项目表示感谢。