0-Dr-SAM: U-Shape Structure Segment Anything Model for Generalizable Medical Image Segmentation

https://doi.org/10.1007/978-981-97-5600-1_17

注：此文与Link所指向的Dr-SAM不同

背景

医学图像分割在现代医疗诊断中占据关键地位，精准的分割能够助力医师精确判别病情、规划治疗策略以及监测病症演进。然而，医学领域获取大规模标注图像数据困难重重，致使既有方法多依赖特定任务定制模型，普适性欠佳。

自然语言与视觉领域的基础模型成效斐然，但迁移至医学图像分割时难题频现。以 SAM 为例，其虽具强大泛化力，但在医学图像应用中遭遇诸多阻碍，诸如掩码边界粗糙，易遗漏纤细结构分割；于复杂情形（如视网膜血管、无蒂息肉）预测失准、掩码损坏或误差显著，全自动分割性能欠佳，根源在于不良提示的多尺度信息融合及掩码分割难题。

本文提出了一种构建适用于医学分割领域的 DrSAM 模型，在充分利用 SAM 预训练权重的前提下，凭借创新设计的模块，有效攻克医学图像分割难题，实现性能与效率的精妙平衡，达成高精度自动医学图像分割

实验方法

DrSAM 在 SAM 中引入 U 结构和 Med-Output Token，实现更适合医学图像的高质量掩码预测。为了保持 SAM 的零采样能力，Med-Output Token 重用了 SAM 的掩码解码器，通过新的 MLP 层，并使用 U 结构提取的分层特征进行点积。在训练时，固定预训练 SAM 的模型参数。

SAM 基础模型

DrSAM 的编码器沿用 SAM 预训练的 ViT - B 模型，并于训练时锁定图像与提示编码器。如此一来，图像嵌入可提前计算，训练时无需加载图像编码器，GPU 效率得以大幅提升。

U 形结构特征提取

借鉴 U - Net 架构精髓，从图像编码器的全局注意力层采撷多尺度特征以生成图像嵌入。聚焦 ViT - B 编码器，依序提取第 2、5、8、12 块的特征，其尺度由 64×64 经逐次上采样拓展至 512×512、256×256、128×128。借由残差连接，将这些特征与对应解码器特征融合，运用双卷积降维，恰似 U - Net 的上采样操作。特别地，在第 2 阶段解码器连接中融入 SAM 转置卷积上采样特征，力保 SAM 零样本能力，规避遗忘困境5。

Med - Output Token

创新性引入 Med - Output Token（尺寸为 1×256），与输出和提示令牌一并馈入注意力模块。于各注意力层内，Med - Output Token 与其余令牌开展自注意力交互并更新权重。在与 U 形层次特征点积前，借新 MLP 结构由 Med - Output Token 生成动态卷积核，防止过拟合新数据集，捍卫 SAM 零样本分割性能6。

DrSAM 训练

精准定位训练范畴，聚焦 U 形结构、Med - Output Token 及其关联的 MLP 展开训练，提示部分选用混合提示类型，同时稳固 SAM 部分权重。设定初始学习率为 0.001，采用 AdamW 优化器，批次规模设为 6，经 20 个 epoch 的精心训练，于单个 NVIDIA RTX 3090 GPU 上达成高效微调7。

为严密监督 DrSAM 的掩码预测质量，巧妙融合 BCELoss（二元交叉熵损失）与 Dice Loss，二者协同制衡二元分类损失与分割精度，其公式为：
$$
BCELoss(y, \hat{y})=-\frac{1}{N} \sum_{i=1}^{N}\left[y_{i} \cdot log \left(\hat{y}{i}\right)+\left(1-y{i}\right) \cdot log \left(1-\hat{y}{i}\right)\right]
\Dice Loss(y, \hat{y})=1-\frac{2 \cdot \sum{i=1}^{N} y_{i} \cdot \hat{y}{i}}{\sum{i=1}^{N} y_{i}^{2}+\sum_{i=1}^{N} \hat{y}_{i}^{2}}\
\Loss(y, \hat{y})=BCELoss(y, \hat{y})+Dice Loss(y, \hat{y})
$$

DrSAM 推理

DrSAM 推理流程与 SAM 一脉相承，差异在于融入 Med - Output Token 的掩码预测结果。推理之际，先于 512×512 分辨率整合 SAM 与 DrSAM 的预测掩码以校正，而后上采样至原始分辨率（如 1024×1024）输出，确保分割结果精准契合医学图像需求。

对比 SAM

DrSAM 在分割品质上实现跃升，训练环节展现出高效经济特质，仅需单一消费级 GPU 即可精准微调。模型架构轻盈高效，参数、GPU 内存占用及单幅图像推理耗时的增量微乎其微，全方位超越 SAM9。

对比 MedSAM

相较 MedSAM，DrSAM 于垂直领域优势凸显。MedSAM 依赖海量医学图像数据微调 SAM 基础模型，对计算资源索求甚巨，基层医疗与教育机构应用受限。DrSAM 独辟蹊径，仅微调 U 形多尺度特征结构，训练轻快、参数精简，在计算与数据资源局促场景（如基层医疗设施）适应性超群

实验结果

在 Cell、Chase和 Hyper Kvasir - SEG三个医学数据集上微调模型，增强医学图像分割鲁棒性，数据集按 1/2 训练和 1/2 测试划分。

与 SAM 和 MedSAM 对比，用 mIoU、mBIoU 和 mDice 评估。DrSAM 在三个数据集表现最佳，对细小和边界模糊对象分割能力强，且零样本实验显示其具备基础模型能力

消融研究

在 Hyper Kvasir - SEG 数据集消融实验表明，Med - Output Token 和 U 形结构均能提升性能，二者结合时 DrSAM 效果最优

HyperKvasir-SEG 数据集上的 DrSAM 消融实验结果。

总结

DrSAM 是医学图像分割基础模型，用少量数据微调，U 形结构和 Med - Output Token 提升性能，保留 SAM 优势，在医学分割任务通用性强，具自动医学图像分割潜力