论文解读-SAM-Med3D: Towards General-purposeSegmentation Models for Volumetric MedicalImages

论文PDF链接[2310.15161] SAM-Med3D: Towards General-purpose Segmentation Models for Volumetric Medical Imageshttps://arxiv.org/abs/2310.15161论文代码链接GitHub - uni-medical/SAM-Med3D: SAM-Med3D: An Efficient General-purpose Promptable Segmentation Model for 3D Volumetric Medical Image · GitHubhttps://github.com/uni-medical/SAM-Med3D摘要AbstractExisting volumetric medical image segmentation models are typically task-specific: they perform well on designated targets, yet fail to generalize across different anatomical structures or imaging modalities, which greatly limits their wide clinical deployment.In this work, we propose SAM-Med3D, a general-purpose segmentation model tailored for volumetric medical images. Fed with merely a small number of 3D point prompts, SAM-Med3D is capable of precisely segmenting various anatomical tissues and lesions spanning multiple imaging modalities.To support model training, we constructed a large-scale 3D medical imaging dataset named SA-Med3D-140K, aggregated from abundant public datasets and authorized private clinical data. This dataset contains 22,000 volumetric medical scans paired with 143,000 corresponding 3D segmentation masks.SAM-Med3D adopts a fully learnable native 3D architecture with prompt-driven segmentation capacity. We train the model via a two-stage training pipeline on SA-Med3D-140K, and it achieves state-of-the-art segmentation performance on both known and unseen segmentation objects.We conduct comprehensive quantitative and qualitative evaluations across 16 independent datasets covering rich clinical scenarios, including varied anatomical regions, imaging modalities, segmentation targets, as well as zero-shot generalization on novel unseen tasks. Experimental results validate SAM-Med3D’s superior efficiency and segmentation accuracy. Moreover, the pre-trained SAM-Med3D encoder delivers strong transfer learning potential for diverse downstream medical segmentation tasks.Our research proves that large-scale medical annotated data can be leveraged to build universal medical foundation models with broad clinical application prospects.Keywords: General-purpose segmentation; Volumetric medical images; Promptable model; 3D medical foundation model翻译现有的三维体医学图像分割模型大多为任务专用模型这类模型仅能在特定目标上取得优异效果难以在不同解剖结构、不同影像模态之间泛化这一缺陷极大限制了其在临床场景中的大规模落地使用。本文提出面向三维体医学影像的通用分割模型 SAM-Med3D。仅依靠少量三维提示点SAM-Med3D 就能对多模态影像下各类解剖组织与病灶完成高精度分割。为支撑模型训练我们整合海量公开数据集与授权私有临床数据构建了大规模三维医学影像数据集 SA-Med3D-140K数据集包含 22000 份三维医学扫描影像配套 143000 张三维分割掩码。SAM-Med3D 采用全可学习原生 3D 架构具备提示驱动分割能力我们基于 SA-Med3D-140K 数据集采用两阶段训练策略完成模型训练该模型在已见过、未见过的分割目标上均表现出出色分割效果。本文在覆盖多元临床场景的 16 套独立数据集上完成全面评测测试维度包含不同解剖部位、影像模态、分割目标以及模型在全新未知任务上的零样本迁移能力。实验结果证实 SAM-Med3D 兼具更高推理效率与分割精度同时预训练后的 SAM-Med3D 编码器在各类下游医学分割任务中展现出极强的迁移学习潜力。本研究证明依托大规模医学标注数据我们能够构建具备广阔临床应用潜力的通用医疗人工智能基础模型。关键词通用分割三维体医学影像提示式模型3D 医学基础模型SAM-Med3D 模型架构完整解析SAM-Med3D 是首个原生基于三维 Transformer 构建的提示驱动型医学容积图像分割基础模型区别于原始 SAM 先切片 2D 推理、再融合三维结果的方案模型全程在三维空间执行特征提取与交互天然适配 CT、MRI 等三维医学影像[H,W,D]输入。整体架构延续 SAM 经典的三模块范式3D 图像编码器3D Image Encoder、3D 提示编码器3D Prompt Encoder、3D 掩码解码器3D Mask Decoder所有模块均为三维算子设计消除二维切片方式带来的层间特征不一致问题。输入为尺寸(H*W*D)的三维容积医学图像首先送入3D Image Encoder。编码器起始模块为 3D Patch Embedding该模块利用三维卷积将原始三维图像切分为互不重叠的三维 Patch映射为序列形式的三维特征 Token随后接入 3D 绝对位置编码3D Abs PE向特征注入空间位置信息弥补 Transformer 本身不具备位置感知能力的缺陷。特征序列继而堆叠多层3D 多头自注意力模块3D Multi-head Self-Attention Block也就是框架图放大展示的核心单元。在 3D 多头自注意力内部输入特征先经过 Layer Norm 3D 归一化并行生成查询向量Q、键向量K、值向量V。注意力分数计算形式写作与标准多头自注意力不同SAM-Med3D 引入3D 相对位置编码3D Rel PE直接在注意力分数矩阵上融入三维体素之间的相对空间偏移关系相比绝对位置编码能更好适应医学影像多变的图像尺寸与目标解剖尺度。注意力输出经过残差相加后再次归一化送入前馈网络Feed-forward完成通道维度特征变换。经过多层 3D 注意力 Block 迭代编码器最终输出固定维度的图像嵌入Image Embedding作为后续解码器的图像上下文特征。3D Prompt Encoder 负责编码用户提供的交互提示支持三维坐标点提示((x,y,z))以及粗掩码提示两类输入。对于三维点提示模型通过可学习嵌入Learnable Embedding表征提示类型叠加 3D 绝对位置编码编码点的三维空间坐标对于掩码提示则采用堆叠两次的三维卷积、三维层归一化LN 3D与 GELU 激活函数进行特征提取。两种提示分别编码后输出统一维度的提示嵌入Prompt Embeddings传递至掩码解码器实现不同交互方式的统一表征。最右侧的3D Mask Decoder接收图像嵌入与两组提示嵌入依靠双层 Transformer Block 完成图像特征与提示特征的跨模态交互。交互完成后使用转置三维卷积Transposed Conv 3D×2逐步对低分辨率图像嵌入进行上采样恢复三维空间分辨率最后通过多层感知机 MLP 映射输出与输入容积图像相同尺寸的三维分割预测掩膜Prediction。在训练策略层面SAM-Med3D 采用两阶段训练方案依托自建大规模数据集 SA-Med3D-140K22K 三维影像、143K 三维 3D 掩码完成预训练。原生全 3D 架构搭配三维位置感知注意力机制让模型仅依靠少量三维点提示即可分割未见解剖结构、全新影像模态具备强大的零样本迁移能力。作为预训练基础模型编码器提取的三维医学特征能够迁移至各类下游医学分割任务解决传统专用分割模型泛化能力弱、需要大量标注重新训练的痛点。SAM-Med3D 完整前向推理流程整体推理流程分为输入处理、3D 图像编码、3D 提示编码、3D 掩码解码、输出预测五大步骤。模型输入是尺寸为 ([H,W,D]) 的三维容积医学图像CT/MRI 体数据同时接收用户交互提示支持三维空间点 ((x,y,z)) 或者掩码提示两条输入分支。第一步三维图像送入3D Image Encoder3D 图像编码器。先通过 3D Patch Embedding 将连续三维图像划分成三维 Patch 并映射为特征 Token叠加 3D 绝对位置编码3D Abs PE随后特征依次流经多层 3D 多头自注意力 Block依靠带有 3D 相对位置编码的自注意力机制建模三维空间内体素之间的长距离依赖经过多层 Block 特征提取后输出低分辨率的Image Embedding图像嵌入特征保存整幅容积图像的全局上下文信息。第二步交互提示送入3D Prompt Encoder3D 提示编码器进行单独编码。如果输入是三维坐标点提示使用可学习嵌入搭配 3D 绝对位置编码对空间坐标建模如果输入是掩码提示则通过连续 3D 卷积、层归一化与 GELU 激活提取掩码特征。两种提示经过编码器处理后生成统一维度的Prompt Embeddings提示嵌入特征随后一并输送给解码器。第三步3D Mask Decoder3D 掩码解码器同时接收图像嵌入与提示嵌入。内部依靠两层 Transformer Block 实现图像特征与提示特征的深度交互让模型根据提示定位目标解剖结构或病灶交互完成后利用 2 层转置 3D 卷积逐步上采样特征图恢复三维空间分辨率最后经由 MLP 完成特征映射输出与原始输入尺寸一致的三维分割掩膜 Prediction。简单梳理时序链路容积图像 -3D 图像编码器 - 图像嵌入交互提示3D 点 / 掩码- 3D 提示编码器 - 提示嵌入图像嵌入 提示嵌入- 3D 掩码解码器 - 三维分割预测结果补充和原始 2D-SAM 核心流程差异原始 SAM 对三维医学图像需要逐切片做 2D 推理最后融合各层结果容易出现层间分割断裂SAM-Med3D 整套链路全部使用 3D 算子编码、提示交互、解码全程在三维空间运算天然利用 Z 轴层面信息保证容积数据分割的空间连续性。