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/adrenal/ct_WORD",
  "data/train/liver/ct_WORD",
  # ...
]

（3）启动训练


# 单卡
bash train.sh

# 多卡 DDP（按需修改 GPU 配置）
bash train_ddp.sh

常用参数：--lr（学习率）、--batch_size、--num_epochs 等可在脚本或命令行指定。

5）评估流程

单样本：val_for_single_data.sh数据集：val_on_dataset.sh 批量评估，并用 utils/compute_dataset_metrics.py 统计 Dice、NSD 等指标。

6）项目特点（速记）

原生支持 3D 医学图像分割，覆盖 CT/MRI 等体数据。支持交互式提示（点/框）与自动分割，以少量提示获得高精度结果。提供大规模数据集与预训练权重，可快速迁移到自定义任务。

二、SAM-Med3D 代码仓库结构总览

1) 根目录核心文件

配置与说明文件

.gitignore：指定Git版本控制中需忽略的文件/文件夹（如数据、缓存、日志等），避免提交冗余内容。LICENSE：采用Apache License 2.0，明确代码的使用、复制、修改、分发等权限及法律条款。readme.md：项目核心说明文档，包含：
项目简介（SAM-Med3D的3D医学图像通用分割模型）、新闻动态（数据/模型更新、竞赛信息）。亮点特性（143K 3D掩码数据集、高效提示分割、多数据集验证）。模型 checkpoint 下载链接（Hugging Face/Google Drive/百度网盘）。使用指南（环境搭建、单数据测试、训练/微调步骤）。

脚本文件

lint.sh：代码格式化与规范检查脚本，通过isort工具排序和分组Python导入语句，保证代码风格一致。medim_val_single.py：单数据验证脚本，用于测试单张3D医学图像的分割效果。用户可指定输入图像、Ground Truth（GT）和输出路径，快速验证模型性能。medim_val_dataset.py：数据集批量验证脚本（推测），用于在整个测试集上评估模型，生成批量结果。train.py：模型训练主脚本，核心功能包括：
命令行参数解析（批处理大小、学习率、GPU配置等）。模型构建（基于segment_anything.build_sam3D注册3D SAM模型）。数据加载（结合utils.data_loader处理数据集，支持分布式采样）。训练循环（定义损失函数DiceCELoss、优化器AdamW、学习率调度器，支持checkpoint保存与加载）。支持单GPU和多GPU分布式训练（DDP）。 train.sh：单GPU训练启动脚本，预设训练参数（如批处理大小、工作进程数、预训练权重路径），直接调用train.py执行训练。train_ddp.sh：多GPU分布式训练启动脚本，指定GPU IDs和总批处理大小（单卡批处理×GPU数量），适用于大规模数据和模型训练。val_for_single_data.sh：单数据验证快捷脚本，直接调用medim_val_single.py执行单数据测试。val_on_dataset.sh：数据集批量验证快捷脚本（推测），用于启动medim_val_dataset.py进行全量测试。

2) 文件夹功能

assets/：项目资源文件夹，存放可视化与说明相关文件，如：

示意图（architecture.png模型架构、motivation.png项目动机、comparison.png实验对比）。

分割结果可视化图（vis_anat.png、vis_modal.png）。

paper/：论文相关文件存放处，包含：

主论文SAM_Med3D_ECCV.pdf（ECCV 2024 BIC会议录用论文）。补充材料SAM_Med3D_ECCV_Supplementary.pdf，提供实验细节、附加结果等。

test_data/：测试数据文件夹，包含：

示例测试数据（如amos_val_toy_data/，AMOS数据集的小规模验证数据）。其他测试案例（如ACDC数据集的测试图像与GT），用于模型调试和效果验证。

utils/：工具函数集合，支撑数据处理、模型训练、评估等核心流程：

click_method.py：交互式分割的提示点生成工具，实现多种策略（如基于GT的错误区域选点、无GT的图像阈值选点、基于距离变换的核心错误区域选点），生成用于模型优化的正负点击点。compute_dataset_metrics.py：批量评估指标计算工具，对比预测结果与GT，计算DSC（Dice系数）、NSD（归一化表面距离）等指标，输出单文件、单类别及整体统计结果，支持JSON保存。data_loader.py：数据加载工具，定义Dataset_Union_ALL（多数据集联合加载）和Union_Dataloader（支持分布式采样的数据加载器），包含数据变换（如裁剪、翻转）。data_paths.py：数据路径配置工具，定义img_datas列表指定训练数据目录，供数据加载器读取。infer_utils.py：推理辅助工具（推测），包含模型推理、结果后处理（如掩码阈值化、格式转换）等功能。metric_utils.py：基础指标计算工具（被compute_dataset_metrics.py调用），实现DSC、NSD等分割评估指标的底层计算。prepare_data_from_nnUNet.py：数据格式转换工具，将nnU-Net风格的数据集（多模态图像、原始掩码）转换为项目所需的单模态、二值化掩码格式，包含重采样等预处理。

.github/：GitHub仓库配置文件夹（内容未详细提供），可能包含CI/CD工作流配置、issue模板、贡献指南等，用于项目协作与自动化管理。

segment_anything/：模型核心代码文件夹，定义3D SAM模型结构：

modeling/common.py：包含模型通用组件（如归一化层的前向传播实现）。build_sam3D.py：3D SAM模型注册与构建逻辑（通过sam_model_registry3D注册模型类型，加载预训练权重），是模型实例化的核心入口。

3) 总结

SAM-Med3D仓库围绕3D医学图像分割任务，提供了从数据准备、模型训练（单/多GPU）、交互式提示生成到性能评估的完整工具链。核心逻辑集中在train.py（训练）、utils/（工具函数）和segment_anything/（模型结构），通过readme.md和脚本文件（train.sh、medim_val_single.py等）降低用户使用门槛，支持快速测试与定制化训练。

三、segment_anything文件夹脚本及组成部分作用总结

定位：segment_anything/ 为 SAM-Med3D 的核心实现目录，统一封装 2D/3D 模型构建、自动掩码生成、交互式推理 与 部署/工具，对外以注册表与预测器提供稳定接口。

1）目录结构与关键文件

顶层核心（目录根）

__init__.py：模块导出入口；统一暴露 SamAutomaticMaskGenerator、SamPredictor、sam_model_registry（2D）、sam_model_registry3D（3D）、build_sam/build_sam3D 等。automatic_mask_generator.py：全图自动掩码生成；网格点自动提示 → 批量候选掩码 → 依据 IoU/稳定性/NMS 过滤与后处理；支持二值/RLE/COCO-RLE 输出。build_sam.py（2D）：按 ViT-H/L/B 构建 2D 模型；_build_sam(ImageEncoderViT + PromptEncoder + MaskDecoder)；权重尺寸不匹配时对 pos_embed/rel_pos 做插值自适配；sam_model_registry 管理版本映射。build_sam3D.py（3D）：按 ViT-H/L/B 与 vit_b_ori（常见 128×128×128）构建 3D 模型；_build_sam3D(_ori)(ImageEncoderViT3D + PromptEncoder3D + MaskDecoder3D)；sam_model_registry3D 管理版本映射。predictor.py：交互式分割；set_image/set_torch_image 预处理与嵌入缓存，predict/predict_torch 接收 point_coords/point_labels/box/mask_input 生成掩码并还原到原图尺度；提供 get_image_embedding、reset_image。

核心模型组件（modeling/）

文件	作用	关键要点
`common.py`	通用层	`MLPBlock`、`LayerNorm2d` 等基础模块
`image_encoder.py`	2D 图像编码器	ViT-based；可配置 `global_attn_indexes`；输出 2D 特征嵌入
`image_encoder3D.py`	3D 图像编码器	体数据 3D patch + 3D Transformer；输出 3D 特征嵌入
`prompt_encoder.py`	2D 提示编码	点/框/掩码 → 稀疏/密集嵌入；前景/背景点标签
`prompt_encoder3D.py`	3D 提示编码	3D 点/框/体素提示对齐 3D 特征空间
`mask_decoder.py`	2D 掩码解码	融合图像/提示嵌入；输出低分辨率掩码与 IoU；支持 `multimask_output`
`mask_decoder3D.py`	3D 掩码解码	三维掩码与 IoU 质量预测；含多头注意力 `_separate_heads`
`sam.py`	2D 主类	端到端：预处理 → 编码 → 解码 → 后处理（回到原图尺寸）
`sam3D.py`	3D 主类	端到端体分割：3D 预/后处理（填充/缩放）与尺寸对齐
`sam_model.py`	2D 备选实现	更扁平的 `forward`（字典式批输入），便于特定调用场景

工具与部署（utils/）

文件	作用	关键要点
`__init__.py`	工具导出	聚合并导出常用 3D 变换函数
`amg.py`	掩码数据管理	`MaskData.cat`：高效合并掩码/分数/框（Tensor/ndarray/list）
`onnx.py`	ONNX 部署	ONNX 友好前向：图像嵌入 → 提示编码 → 掩码生成/后处理/稳定性分数
`transforms.py`*	2D 预处理	（若存在） `ResizeLongestSide`、坐标/归一化等标准化流程
`transforms3D.py`	3D 预处理/增强	体素级缩放、裁剪、旋转、归一化、重采样，保持三维一致性

2）工作机制（从构建到推理）

模型构建

2D：build_sam_vit_* → _build_sam 组装组件 → 必要时对 pos_embed/rel_pos 做尺寸适配 → 产出可用实例。3D：build_sam3D_vit_* → _build_sam3D(_ori) 组装组件 → 加载 3D 预训练权重（常用固定体素尺寸，减少位置嵌入调整）。

推理模式

自动分割：SamAutomaticMaskGenerator 进行分块/多尺度 + 网格提示 → 批量候选 → IoU/稳定性/NMS 过滤 → 小连通域/孔洞后处理 → 导出多格式掩码。交互分割：SamPredictor 预缓存图像嵌入 → 接收点/框/掩码提示 → 生成掩码并按原始空间复原；支持多轮迭代修正。

数据与部署配套

预处理/增强：transforms(3D)；批结果管理：amg.py；工程化上线：onnx.py 提供导出与兼容前向。

3）接口与注册表速查

自动推理：SamAutomaticMaskGenerator交互推理：SamPredictor注册表：sam_model_registry（2D）、sam_model_registry3D（3D）组件族：ImageEncoderViT(3D) / PromptEncoder(3D) / MaskDecoder(3D) / sam(3D)常见 3D 输入尺寸：vit_b_ori 常用 128×128×128；其余配置多见 256×256×256（以仓库实际实现为准）

4）要点小结

清晰分层：2D 与 3D 组件并行设计，统一由 build_* 与注册表管理。两种使用范式：自动（全图）与交互（点/框）推理均可闭环运行。工程完备：提供 ONNX 友好前向与掩码批管理，便于规模化部署与评估。

四、utils 文件夹脚本总结

定位：utils/ 提供数据路径管理、加载与预处理、指标计算、推理辅助、命令行封装与数据格式转换等通用工具，支撑训练/评估/部署的全流程。

1）脚本速查表

脚本	作用	典型函数/要点
`__init__.py`	包初始化与导出入口	可能通过 `__all__` 聚合常用 API，或保持空文件仅作包标识
`click_method.py`	命令行接口（CLI）封装	使用 `click` 定义命令/选项，将数据处理、推理等能力一键化
`compute_dataset_metrics.py`	数据集统计与基础指标	计算样本数、维度/模态分布、像素强度统计；可扩展医学影像特有统计
`data_loader.py`	数据读取与增广	`Dataset/DataLoader` 实现；归一化、重采样、裁剪/翻转等（训练启用增强）
`data_paths.py`	路径集中管理	定义数据根目录与任务子路径；`img_datas` 等清单；目录存在性检查/创建
`infer_utils.py`	推理阶段工具	加载权重、`model.eval()`、输入预处理与后处理、批量推理封装
`metric_utils.py`	训练/评估指标计算	Dice、IoU、Accuracy/Precision/Recall/F1、NSD；支持聚合与日志输出
`prepare_data_from_nnUNet.py`	数据格式转换适配	将 nnU-Net 结构转换为项目所需；重采样、二值化、`dataset.json`/分割列表生成

2）关键脚本要点

data_paths.py
统一维护训练/验证/测试目录与任务列表（如 img_datas = [...]）；提供路径拼接、目录创建与健壮性检查，避免硬编码。

data_loader.py

读取：NIfTI/DICOM/图像等到张量；可按任务装载多模态。预处理/增强：强度归一化、重采样（spacing 对齐）、裁剪/旋转/翻转；区分训练与评估管线。性能：num_workers/pin_memory、按需缓存/内存映射。

compute_dataset_metrics.py
数据概览（样本数、体素尺寸/spacing 分布、强度直方图）、类别像素占比；结果输出为 CSV/JSON 便于复现与可视化。

metric_utils.py

分割：Dice、IoU、NSD（归一化表面距离）、HD95（若实现）。分类/检出：Accuracy/Precision/Recall/F1、AUC（视任务而定）。聚合：批次/病例级统计与日志记录，便于曲线对比与早停。

infer_utils.py

加载：权重映射到设备；多卡/AMP 兼容。前后处理：尺寸/强度规范化、补丁滑窗（大体素）、阈值/CRF/连通域过滤；掩码回写原始空间。批推理：目录级推理与结果落盘（nii.gz/png/json）。

click_method.py
将常用流程（数据准备 / 评估 / 推理）暴露为 CLI：
python -m utils.click_method prepare --src ... --dst ...
python -m utils.click_method infer --ckpt ... --in_dir ... --out_dir ...

prepare_data_from_nnUNet.py

结构：imagesTr/labelsTr 等目录规范转换；生成或适配 dataset.json。标签：类别映射/二值化、1/2 对调修正；确保训练时标签一致性。重采样：统一 spacing，记录元数据以便推理回写。

3）实践建议

路径/配置集中化：仅在 data_paths.py 修改根路径与任务列表，避免重复改动。记录数据统计：先用 compute_dataset_metrics.py 了解数据强度与尺寸分布，指导重采样与归一化策略。区分训练/评估管线：增强仅用于训练；评估与推理保持确定性（固定随机种子）。体素一致性：重采样与回写时保留原始 spacing/affine，保证指标可比与临床可用性。

五、小节

SAM-Med3D 仓库面向三维医学体数据（CT/MRI 等）的通用可提示（promptable）分割任务，是论文 “SAM-Med3D: Towards General-purpose Segmentation Models for Volumetric Medical Images” 的官方实现与配套资源。项目以将 2D SAM 的思路扩展到“完全可学习的 3D 结构”为核心，围绕公开与授权数据汇聚而成的大规模 3D 医学分割数据集展开训练与评测：论文与 README 均强调 SA-Med3D-140K 的规模与覆盖面，并给出仓库、模型与数据的入口链接；项目同时公布了更强的微调版本 “SAM-Med3D-turbo”，并在 2024 年 ECCV BIC 获口头报告（Oral）。整体上，仓库既是方法实现，也是一套涵盖数据、权重、使用说明与脚本的工程化落地方案。(GitHub)

从目录结构看，根目录提供训练、验证与评估的主要脚本（如 train.py、train.sh、train_ddp.sh、medim_val_single.py、medim_val_dataset.py 及两类评估入口脚本），并配有 assets/（示意图与可视化）、paper/（论文与补充材料 PDF）、test_data/（示例数据）、utils/（工具与流程脚本）和 segment_anything/（核心模型代码）等子目录。README 的文件清单与导航足以定位到这些资源：如 utils/prepare_data_from_nnUNet.py 用于将 nnU-Net 风格数据转换为项目所需结构，utils/data_paths.py 负责训练数据目录清单（img_datas），而 segment_anything/ 则承载 3D 模型构建与推理的主实现。仓库采用 Apache-2.0 许可证，便于在研究与产业环境中复用与再发布。(GitHub)

上手路径被设计得相当直接。README 给出了完整的环境搭建与“单病例快速验证”示例：创建 Conda 环境后使用 uv 安装指定版本的 PyTorch 及医学影像常用依赖，然后修改 medim_val_single.py 中的图像与标签路径即可运行测试。需要特别注意的是，快速测试默认依赖 ground-truth 掩码来生成提示点；若没有 GT，可用“伪标签”临时标注目标区域以驱动推理。此外，项目已集成到 MedIM 中，用户可以一行代码从 Hugging Face 拉取权重并构建 PyTorch 模型，降低了尝鲜门槛。(GitHub)

训练与微调流程同样围绕 README 给出的“数据 → 配置 → 启动”三步走：首先推荐准备好预训练权重（尤其是 SAM-Med3D-turbo），然后按示例把自有数据组织成 imagesTr/labelsTr 的目录结构；如果原始数据是 nnU-Net 风格，则用提供的脚本完成重采样与掩码二值化的转换。接着在 utils/data_paths.py 中填写或自动收集 img_datas 后，直接参照 train.sh 或 train_ddp.sh 启动训练；命令行参数支持从 checkpoint 继续，或在未提供权重时从头训练。该流程既兼顾了医用体数据的工程细节（如 spacing 统一与掩码处理），也照顾到分布式训练的可扩展性。(GitHub)

评估与验证方面，仓库提供了“单病例脚本”和“数据集级评估脚本”，并配套 utils/compute_dataset_metrics.py 汇总 Dice、NSD 等指标，支持按病例、按类别与整体的多维度统计；test_data/ 下还给出了小规模验证数据，便于快速复现实验。更宏观的性能与特性在 README 的 Highlights 中也有归纳：作者在 16 个常用 3D 医学分割数据集上进行了系统评估，并强调 SAM-Med3D 的交互效率优势——达到满意效果所需的提示点数量比传统做法少 10–100 倍。(GitHub)

方法与实现层面，仓库遵循论文思路，在核心代码中以 3D 图像编码器、3D 提示编码器与 3D 掩码解码器作为主干，面向三维空间直接建模，实现了“少量 3D 点/框提示 → 体素级掩码”的可提示分割范式。论文进一步说明了以“两阶段训练”与“完全可学习 3D 结构”达成跨解剖结构与跨模态的泛化目标；配合仓库中的自动掩码生成与交互式预测接口，既能做无提示的全图候选掩码生成，也能在人机交互场景下快速收敛到目标器官或病灶。作为一个工程化成品，它同时覆盖了数据准备、模型构建、推理导出与指标评估的完整闭环。(arXiv)

数据与权重资源是这一闭环的基石。作者在 Hugging Face 公布了 SA-Med3D-140K 的汇总页，其中明确列出了数据规模（约 2.17 万例 3D 影像与 14.35 万个 3D 掩码）及来源构成；仓库侧则在 Checkpoint 区块集中给出 SAM-Med3D-turbo 的下载地址与镜像，并在“News”中说明 turbo 版本基于 44 个数据集做了更广泛的微调，以提升泛化表现。结合开放许可证，这些资源让用户可以平滑地从“示例验证”过渡到“特定科室与任务”的落地适配。(Hugging Face)
SAM-Med3D：面向三维医疗体数据的通用分割模型（代码仓库笔记）

总的来说，SAM-Med3D 仓库把“方法、数据、权重、脚本与文档”拧成一根线：segment_anything/ 提供模型与推理的 3D 实现，utils/ 管理数据转换与评估，根目录脚本统一了训练与验证入口，README 则在环境、步骤与注意事项上给出可直接执行的配方。这种“研究-工程一体化”的结构，使其既是复现论文结果的权威参考，也是将 3D 可提示分割引入临床/科研工作流的实用起点。(GitHub)

代码仓库链接: link