论文阅读 MedSAM (Segment Anything in Medical Images)

论文阅读 MedSAM (Segment Anything in Medical Images)

发表于《Nature Communications》

论文地址：https://arxiv.org/abs/2304.12306

代码地址：https://github.com/bowang-lab/MedSAM

一、提出背景

医学图像分割是临床实践中的关键组成部分，有助于准确诊断，治疗计划和疾病监测。然而，现有的方法，往往是针对特定的方式或疾病类型，缺乏通用的医学图像分割大模型。

各种医学图像中的SAM进行了全面评估，这些SAM主要在具有明显边界的目标上实现了令人满意的分割结果。然而，模型对于弱边界或低对比度的典型医学目标时表现出很大的局限性。

二、本文创新

MedSAM通过在一个前所未有的大规模的医学图像数据集上微调SAM：

（1）这篇论文通过微调SAM提出了医学图像领域的通用分割大模型MedSAM , MedSAM模型在大规模的医学图像数据集（CT、MR、X射线、超声等）上进行训练，包含1,570,263个图像-掩码对，涵盖了10种成像模态和30多种癌症类型。

（2）考虑到存在3D和2D不同维度、不同格式的图像，作者更实际的方法是开发一个快速的2D分割模型。该模型可根据用户提供的提示轻松适应特定任务，从而提供更高的灵活性和适应性。它还能够通过将3D图像处理为一系列2D切片来处理2D和3D图像。

（3）CT、MR、内窥镜、超声超声图像强度差异非常大，进行了强度归一化，确保了模型对所有图像的一致性和兼容性。

（下图表示 MedSAM在能够处理各种分割任务的大规模数据集上训练。该数据集涵盖了各种解剖结构、病理状况、以及医学成像模式。洋红色轮廓和掩模分别表示专家注释和MedSAM分割结果。）

作者研究不同的分割提示发现：边界框为目标区域指明了更明确的空间背景，使得算法能够更精确地辨别目标。这与基于点的提示形成对比，基于点的提示会引入歧义，尤其是当邻近结构彼此相似时。此外，绘制边界框是高效的，尤其是在涉及多对象分割的场景中。

MedSAM包含一个基于ViT的图像编码器、一个用于整合用户交互（边界框）的prompt编码器以及一个使用图像嵌入、prompt嵌入和输出令牌生成分割结果和置信度得分的掩膜解码器。图像编码器将输入图像映射到高维图像嵌入空间；提示编码器通过位置编码将用户绘制的边界框转换为特征表示；最后，掩模解码器使用交叉注意机制来融合图像嵌入和提示特征。

具体来说：MedSAM 模型的架构及其各个组件的工作原理：

（1）基于视觉Transformer（ViT）的图像编码器

 * **功能** ：图像编码器负责从输入图像中提取高层次的特征表示。
 * **结构** ：使用基础ViT模型作为图像编码器。
 * **原因** ：基础ViT模型在分割性能和计算效率之间取得了平衡。较大的ViT模型（如ViT-Large和ViT-Huge）虽然在准确性上有轻微提升，但显著增加了计算需求。
 * **具体配置** ：基础ViT模型由12个Transformer层组成，每个层包含一个多头自注意（Multi-Head Self-Attention, MHSA）块和一个包含层规范化的多层感知器（MLP）块。

（2）提示编码器

 * **功能** ：提示编码器用于集成用户的交互信息，例如边界框。
 * **结构** ：将边界框的角点映射到256维的向量嵌入。
 * **具体配置** ：每个边界框由左上角点和右下角点的嵌入对表示。

（3）掩码解码器

 * **功能** ：掩码解码器使用图像嵌入、提示嵌入和输出令牌生成最终的分割结果和置信度分数。
 * **结构** ：采用轻量级架构，以便支持实时用户交互。
 * **具体配置** ： 
   * **输入** ：输入图像的尺寸为1024 × 1024 × 3。
   * **图像嵌入** ：输入图像被重塑为16 × 16 × 3的平坦二维块序列，通过图像编码器后，特征尺寸缩减为64 × 64。
   * **提示嵌入** ：提示编码器将边界框提示的角点映射到256维的向量嵌入。
   * **融合** ：掩码解码器由两个Transformer层组成，用于融合图像嵌入和提示嵌入。
   * **上采样** ：通过两个转置卷积层将嵌入分辨率提高到256 × 256。
   * **激活和插值** ：对嵌入进行S形（Sigmoid）激活，然后进行双线性插值以匹配输入图像的大小。

具体步骤

（1）图像预处理 ：

 * 输入图像尺寸为1024 × 1024 × 3。
 * 图像被重塑为16 × 16 × 3的平坦二维块序列。

（2）图像编码 ：

 * 使用基础ViT模型（12个Transformer层）对图像进行编码。
 * 生成的图像嵌入特征尺寸为64 × 64，相比输入图像尺寸缩减了16倍。

（3）提示编码 ：

 * 用户提供的边界框提示被编码为256维的向量嵌入。
 * 每个边界框由左上角点和右下角点的嵌入对表示。

（4）掩码解码 ：

 * 使用两个Transformer层融合图像嵌入和提示嵌入。
 * 通过两个转置卷积层将嵌入分辨率从64 × 64提高到256 × 256。
 * 对嵌入进行S形激活，得到0到1之间的置信度分数。
 * 进行双线性插值，将结果调整回输入图像的大小（1024 × 1024）。

特点：

 * **实时交互** ：通过采用轻量级的掩码解码器架构，确保模型能够支持实时用户交互。
 * **高精度** ：在保持计算效率的同时，确保模型在分割任务中具有高精度。

作者还进行了一项人工注释研究，以评估两种方法的时间成本。对于第一种方法，两个人类专家以逐切片的方式手动注释3D肾上腺肿瘤。对于第二种方法，专家们首先每隔3-10个切片用直线标记（初始标记）画出肿瘤长短轴，这是肿瘤缓解评价中的常用做法。然后，使用MedSAM基于这些稀疏的线性标注来分割肿瘤。最后，专家手动修正分割结果，直到满意为止。作者定量比较了两种方法之间的注释时间成本。实验结果表明，在MedSAM的辅助下，两个专家的标注时间分别减少了82.37%和82.95%。

作者通过86项内部验证任务和60项外部验证任务的综合实验，全面评估了MedSAM，涵盖了各种解剖结构、病理条件和医学成像模式，实验结果表明，MedSAM在多个数据集上的结果优于最先进的分割基础模型。

三、模型初始化与训练

• 初始化 ：模型使用预训练的SAM模型（基于ViT-Base模型）进行初始化。
• 固定部分 ：提示编码器被固定，因为它已经能够编码边界框提示。
• 可训练参数 ：图像编码器和掩码解码器的可训练参数分别为89,670,912和4,058,340。
• 边界框提示 ：专家注释的边界框提示通过0-20像素的随机扰动进行模拟。
• 损失函数 ：使用Dice损失和交叉熵损失的无权重和，这两种损失函数在各种分割任务中表现出色。
• 优化器 ：使用AdamW优化器（β1 = 0.9, β2 = 0.999），初始学习率为1e-4，权重衰减为0.01。
• 批处理和硬件 ：全局批处理大小为160，未使用数据增强。模型在20个A100（80G）GPU上训练150个epoch，最后一个检查点被选为最终模型。

在验证集上的定量和定性实验结果图

总结

MedSAM 模型通过结合基于ViT的图像编码器、提示编码器和轻量级的掩码解码器，实现了在多种成像模态和疾病类型中的高效和准确的分割。这种设计不仅提高了模型的泛化能力，还支持了实时用户交互，使其在实际应用中更具实用价值。