SAM-Med3D:面向三维医疗体数据的通用分割模型 (代码仓库笔记)
SAM-Med3D 代码仓库组成及运行流程总结
目录
- SAM-Med3D 代码仓库组成及运行流程总结
-
- 一、代码仓库简介
-
- 1)目录结构与关键文件
- 2)数据与预训练权重
- 3)快速上手(模型测试)
- 4)训练 / 微调流程
- 5)评估流程
- 6)项目特点(速记)
- 二、SAM-Med3D 代码仓库结构总览
-
- 1) 根目录核心文件
- 2) 文件夹功能
- 3) 总结
- 三、segment\_anything文件夹脚本及组成部分作用总结
-
- 1)目录结构与关键文件
- 2)工作机制(从构建到推理)
- 3)接口与注册表速查
- 4)要点小结
- 四、utils 文件夹脚本总结
-
- 1)脚本速查表
- 2)关键脚本要点
- 3)实践建议
- 五、小节
一、代码仓库简介
项目定位:SAM-Med3D 面向体数据(CT/MRI)的通用分割,基于 SAM 思想扩展到 3D。仓库内容可分为「核心代码、数据与权重、运行与评估脚本」三大块,便于快速上手测试、训练/微调与评估。
1)目录结构与关键文件
-
核心模型(
segment_anything/
)image_encoder3D.py
:3D 图像编码器prompt_encoder3D.py
:3D 提示(点/框等)编码器mask_decoder3D.py
:3D 掩码解码器build_sam3D.py
:模型构建与加载工具
-
训练与微调
train.py
:训练主程序(支持交互式提示模拟与损失计算)train.sh
:单卡训练启动脚本train_ddp.sh
:多卡分布式训练脚本
-
评估与验证
medim_val_single.py
:单样本快速验证medim_val_dataset.py
:数据集批量评估val_for_single_data.sh
/val_on_dataset.sh
:评估脚本入口utils/compute_dataset_metrics.py
:Dice、NSD 等指标计算
-
数据处理与配置(
utils/
)prepare_data_from_nnUNet.py
:将 nnU-Net 格式转换为本项目所需格式(重采样、掩码二值化)data_loader.py
:数据加载与基础增强data_paths.py
:训练/测试数据路径集中配置(img_datas
列表)
2)数据与预训练权重
- 数据集:提供 SA-Med3D-140K(约 143K 个 3D 掩码,245 类),可从 Hugging Face 获取。
- 预训练模型:提供
SAM-Med3D-turbo
等权重(Hugging Face / Google Drive / 百度网盘),建议微调以提升下游性能。
3)快速上手(模型测试)
(1)环境准备
# 创建并激活虚拟环境
conda create --name sammed3d python=3.10
conda activate sammed3d# 安装依赖管理工具 uv
pip install uv# 安装 PyTorch 及相关库(按需替换为你的 CUDA 对应版本)
uv pip install torch==2.6.0 torchvision==0.21.0 torchaudio==2.6.0# 安装医学影像处理依赖
uv pip install torchio opencv-python-headless matplotlib prefetch_generator monai edt surface-distance medim
(2)单样本验证
在 medim_val_single.py
中设置路径:
img_path = "./test_data/Seg_Exps/ACDC/ACDC_test_cases/patient101_frame01_0000.nii.gz"
gt_path = "./test_data/Seg_Exps/ACDC/ACDC_test_gts/patient101_frame01.nii.gz"
out_path = "./test_data/Seg_Exps/ACDC/ACDC_test_SAM_Med3d/patient101_frame01.nii.gz"
运行:
python medim_val_single.py
提示:需要 ground-truth 掩码用于生成提示点;若无真实标签,可临时手工标注“伪标签”。
4)训练 / 微调流程
(1)准备权重(推荐)
ckpt_path = "https://huggingface.co/blueyo0/SAM-Med3D/blob/main/sam_med3d_turbo.pth"
# 示例:使用 MedIM 快速创建(如已安装)
model = medim.create_model("SAM-Med3D", pretrained=True, checkpoint_path=ckpt_path)
将下载的 .pth
放入 ckpt/
并在训练脚本中通过 --checkpoint
指定。
(2)准备数据
- 建议结构(每个任务一套图像与二值掩码):
data/train├── adrenal/ct_WORD│ ├── imagesTr/word_0025.nii.gz│ └── labelsTr/word_0025.nii.gz├── liver/ct_WORD│ └── ...
- nnU-Net 数据可用
utils/prepare_data_from_nnUNet.py
转换。 - 在
utils/data_paths.py
设置训练目录:
img_datas = ["data/train/adr