MSA【5】：SAM-Med2D

---

前言

由于医学图像和自然图像之间存在较大差异，以及缺少大规模医学图像基准数据集，这是导致AI在医学领域进展缓慢的原因之一。构建大规模基准数据集和可靠的基线模型，能够推动AI在医疗领域的快速发展，加速医疗向更通用的方向转变。上海 AI-Lab 提出了 SAM-Med2D，标志着大模型在医学图像分割的里程碑。

原论文链接：SAM-Med2D
OpenGVLab 官方博客：SAM-Med2D：打破自然图像与医学图像的领域鸿沟，医疗版 SAM 开源了！

---

Abstract

SAM-Med2D

• 构建了一个涵盖各种模式和对象的大规模医学图像分割数据集
• 通过涉及边界框、点和遮罩的更全面的提示，将 SAM 适应于医学图像分割
• 对原始 SAM 的编码器和解码器进行了微调，以获得性能良好的 SAM-Med2D，从而形成了迄今为止最全面的微调策略

---

Introduction

Segment Anything Model (SAM) 代表了自然图像分割领域最先进的研究进展。然而，直接将预训练的 SAM 应用于医学图像分割，效果并不理想。这种限制主要源于自然图像和医学图像之间存在显著的领域差异。

医学图像和自然图像存在显著差异，本质就是成像方式完全不同：自然图像是以自然光为光源，通过手机、相机等设备转换为 RGB 格式的图像，像素值范围也是 0-255 之间；医学图像是通过特定的协议和扫描仪收集的，其展现方式因其特定的临床目的而各异，包括电子、激光、X射线、超声波等，成像都不是 RGB 图，其像素值范围从几千到几十万不等 （见 Fig 1. (a)）。

另外，由于标注成本高且标注质量参差不齐，导致医学图像在不同医院和临床研究中的质量存在很大差异。上述挑战导致医学图像和自然图像在数据量、数据质量上存在显著差异。图1b比较了公开的自然图像数据集和医学图像数据集的数据量规模差异（见 Fig 2. (b)）。

---

Methods

Incorporation of Medical Knowledge into SAM

• 当一个掩码包含多个类别时，本文会生成多个掩码，每个掩码只包含一个类别
• 将具有多个连接组件的掩码拆分为具有单一连接组件的多个掩码
  • 可以将具有 2 个联通区域的脑室拆分为 2 个 mask 分别去预测（在训练的过程中也可以看作 2 张不同的训练数据用于训练）
• 排除了目标区域小于总图像 0.153% （100 256×256）的掩码，这意味着当图像调整为 256×256 时，其目标区域必须超过 100 像素
  • 和通常的 crop 后的处理一样，提高训练效率

Transition from SAM to SAM-Med2D

SAM 由三个主要部分组成：大规模图像编码器、提示编码器和轻量级掩码解码器。该框架允许根据不同的提示为同一图像生成不同的掩码：

• 图像编码器利用预先训练好的视觉变换器（ViT）处理高分辨率输入，并以原始图像 1/16 的比例输出特征图
• 提示编码器包括稀疏提示和密集提示，将点、边界框或文本映射为 256 维向量，并对密集提示进行卷积下采样，在每个级别使用 GELU 激活函数
• 掩码解码器接收来自两个编码器的嵌入信息，并通过交叉关注机制更新图像嵌入和提示嵌入

Adapting Image Encoder

本文从信道和空间两个维度对图像编码器进行调整：

• 在通道维度上
  • 本文首先使用全局平均池化技术将输入特征图的分辨率压缩到 $C\times 1\times 1$
  • 然后，使用一个线性层压缩信道嵌入，再使用另一个线性层还原信道嵌入，压缩比为 0.25
  • 最后，通过一个 sigmoid 函数获得通道维度的权重，并将其与输入特征图相乘，作为下一层的输入
• 在空间维度上
  • 本文使用卷积层将特征图的空间分辨率降低 2 倍，然后使用转置卷积恢复空间分辨率，并保持与输入相同的通道数
• 每个适配器层之后都会添加一个跳转连接

Prompt Encoder and Mask Decoder

Prompt Encoder

• 稀疏提示（点和边界框）
  • 每个点都表示为其位置编码的向量嵌入，以及表示其前景或背景位置的两个学习嵌入之和
  • 每个边界框使用其左上角和右下角的位置编码，以及代表 "左上角 "和 "右下角 "的学习嵌入作为向量嵌入
• 密集提示（掩码）
  • 使用模型第一次迭代后生成的低分辨率特征图作为掩码提示，应用两个卷积嵌入，将输入掩码的比例缩小 4 倍，输出通道分别为原始输入的 1/4 和 1/16
  • 最后，使用 1×1 卷积将通道维度映射为 256

Mask Decoder

没有改变掩码解码器的结构，并在训练过程中不断更新其参数

• 在反向传播过程中，本文选择相对于 GT 具有最高 IoU 得分的预测掩码来计算损失，并相应地传播梯度
• 本文将上一次迭代生成的低分辨率特征图映射到 [0, 1] 范围内，作为当前迭代的密集提示
  • 在实际训练过程中，本文发现即使只使用稀疏提示，模型也能很快收敛，这就削弱了密集提示的影响
  • 因此，本文采用了 SAM 的训练策略，即在最后一次迭代和一次随机的中间迭代中，只提供密集提示，以鼓励模型从所提供的掩码中获益

Fine-tuning Strategy

通过模拟交互式分割来训练 SAM-Med2D，对于每一批数据，本文都会对模型进行 9 次迭代训练：

• 在第一次迭代中，本文以相等的概率随机选择一个前景点或边界框作为稀疏提示
  • 前景点从 GT 中采样
  • 边界框是 GT 的最大包围矩形，每个坐标最多偏移 5 个像素
• 除了第一次迭代同时更新适配器层、提示编码器和掩码解码器的参数外，随后的迭代只更新掩码解码器的参数
  • 第一次更新所有参数是为了让adapter学习到医学图像的知识，后续的 prompt 操作与编码器是无关的，所以只更新解码器
• 从第二次迭代开始，本文从之前的掩码预测和 GT 之间的误差区域随机选择 1、3、5 或 9 个点作为后续稀疏提示
  • follow 了 SAM 的工作，Bbox prompt 非常容易拟合，如果多给几次 bbox 的话很快会过拟合
  • 多次迭代需要保存计算图，需要很大的显存开销，只更新掩码就能减少显存了

Evaluate SAM-Med2D

• Bbox prompt 通过用方框框定目标来引导算法进行分割，直观易懂
  • 它提供了目标的大致位置和形状信息，从而引导算法进行更精确的分割
  • 不过，在某些情况下，用 Bbox 框精确包围目标可能具有挑战性，特别是对于形状复杂的目标、边缘模糊的目标或重叠的目标
• Point prompt 通过标记目标的关键点或区域来引导算法执行分割
  • 这种模式可以提供更精确的分割指导，特别是对于具有复杂形状或局部细节的目标

---

Experiment

Implementation Details

在训练过程中，所有图像都被调整为 256 x 256 的分辨率

对于每幅图像，本文随机选择 5 个相应的掩码。如果可用的掩码较少，本文会随机重复样本

• 每个 GPU 处理 50 幅图像及其对应的 250 个掩码

Quantitative Evaluation

Overall Performance

Performance Evaluation on Different Modalities

---

总结

SAM-Med2D 通过在大规模医学图像数据集上微调来适应各种医学图像场景。与预训练的 SAM 相比取得了令人满意的性能改进和泛化能力。在同等分辨率时，FT-SAM 在 Bbox 提示模式下实现了 11.93％ 的提升，而 SAM-Med2D 实现了 17.67％ 的提升。在单点提示模式下， SAM-Med2D 表现出了压倒性的优势 (18.94% vs. 70.01%)。

作者还对不同模型在解剖学结构，不同成像模式和主要器官维度进行了综合评估，SAM-Med2D 在各方面优于预训练的 SAM 和 FT-SAM。此外，在 9 个 MICCAI 2023 数据集的泛化实验证明了在大规模数据集上预训练的模型具有强大的领域可迁移性。尤其是使用点提示模式时，SAM-Med2D 可以更快地生成所需的 mask，甚至优于其他方法使用 Bbox 提示模式的结果。