MedSAM 论文详解

论文名称:Segment Anything in Medical Images

论文地址:https://arxiv.org/pdf/2304.12306.pdf

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

立个flag,要快快之前把这个玩意搞好,用在我自己的数据集上!

1. SAM 拆解分析

      基础模型有很强的泛化能力,这种能力通过提示工程(prompt engineering)实现,想要实现提示分割任务,需要解决三个问题:任务目标、模型结构和数据。

      分割提示指的是完成图像分割需要的 prompt 或者提示语,最常见的 prompt 包括点、bbox(框)、掩模图和文本描述。任务要求收到任意提示符时均需要输出至少一个有效 mask,即使提示符存在歧义。比如衣服上的点可能是想分割衣服,也可能是想区分人体。这两者至少需要输出一个。如下图所示,SAM 结构简单,包括一个图像编码器,一个 prompt 编码器和一个轻量级的解码器。

      此外,一个大规模的数据集是必不可少的,现有的数据集无法胜任这个任务。故设计了一个数据引擎去制造大量的高质量数据来解决这个问题,通过数据引擎可以得到数据集 SA-1B,总共包含了 1100 万张高分辨率图片和 11 亿个 mask。

      关于 Image encoder,输入图像输入前被预处理为 1024*1024,Image encoder 采用MAE VIT-H/16,是经典的视觉 Transformer 结构,最后输出(256,64,64)的图像 embedding。

       关于 Prompt encoder,根据不同的 prompt 会有不同的编码方式,点和框的 embedding 通过位置编码获得,Mask 的 embedding 通过卷积操作获得,文本的 embedding 则是通过 Clip 的 encoder 获得。

       关于 Mask decoder,首先做 prompt 的 self-attention, prompt 到图像 embedding 的 Cross-attention。Cross-attention 是指在有两个相同维度序列的情况下,而其中一个序列用作查询 Q 输入,而另一个序列用作键 K 和值 V 输入。将两个相同维度的嵌入序列不对称地组合在一起。然后,右侧 MLP 均为三层,输出维度与图像 embedding channel 相同的向量,左侧 MLP 为 2048 个神经元,主要作用为聚合全局特征。使用 MLP 更新 token,再做图像 embedding 到 prompt 的 Cross-attention。经过两轮 decode layer 之后,token 再次与图像 embedding 进行 Cross-attention,output token 作为可训练参数在 decoder 前加入到 prompt 中,分别通过两个 MLP,得到 mask 和 mask 的 IOU。

2. 方法

2.1 A teardown analysis of SAM

        SAM利用了基于transformer的架构,该架构已被证明在自然语言处理和图像识别任务中非常有效。具体来说,SAM使用基于transformer的图像编码器来提取图像特征,并使用提示编码器来合并用户交互,然后使用掩码解码器来基于图像嵌入、提示嵌入和输出令牌生成分割结果和置信度分数。

       图像编码器中的视觉转换器采用掩码自动编码器建模进行预训练,该可以处理高分辨率图像(即1024×1024)。获得的图像嵌入缩小比例为16×(64×64)。提示编码器是为不同的用户输入而定制的。SAM支持四种不同的提示:点、框、文本和掩码。每个点都由傅里叶位置编码和两个可学习的标记进行编码,分别用于指定前景和背景。边界框由其左上角和右下角的点编码进行编码。自由形式的文本由CLIP中的预先训练过的文本编码器进行编码。掩码提示符与输入图像具有相同的空间分辨率,并由卷积特征图进行编码。最后,掩模解码器采用轻量级设计,由两层动态掩模预测头和一个交叉联合(IoU)得分回归头组成。掩模预测头可以生成3个4×的缩小掩模,分别对应于对象的整个对象、部分和子部分。

       SAM支持三种主要的分割模式:以全自动的方式分割一切,边界框模式,和点模式。然而,由于两个主要原因,这种分割结果的使用范围有限。一方面,分割结果没有语义标签。另一方面,临床医生主要关注临床场景中有意义的roi,如肝脏、肾脏、脾脏和病变。基于边界框的分割模式只需给出左上角和右下角的点,就可以对右肾产生良好的结果(图2c)。对于基于点的分割模式(图2d),我们首先给出右肾中心的一个前景点,但分割结果包括整个腹部组织。然后,我们在过分割区域上添加一个背景点。分割Mask缩小到肝脏和右肾。在肝脏上添加另一个背景点后,我们最终得到了预期的肾脏分割。

        综上所述,在应用SAM进行医学图像分割时,分段-全模式容易产生无用的区域分区,而基于点的模式则不明确,需要多次预测校正迭代。相比之下,基于边界盒的模式可以明确指定ROI,不需要多次试验和错误,获得合理的分割结果。此外,目前常用的注释方法之一是标记放射学中最长的直径,如实体肿瘤(RECIST)的反应评价标准。我们可以很容易地根据RECIST注释获得目标的边界框提示符。因此,我们认为在医学图像分割任务中使用SAM时,基于边界盒的分割模式比基于分段一切分割模式和基于点的分割模式具有更广泛的实用价值。

2.2 MedSAM:专用的医学图像分割基础模型

       为了使SAM适应于医学图像的分割,有必要选择一个合适的用户提示符和网络的组成部分来进行微调。基于以上分析,边界框提示符是指定分割目标的正确选择。SAM的网络架构包含三个主要组成部分:图像编码器、提示编码器和掩码解码器。人们可以选择调整它们中的任何一种组合。该图像编码器是基于一个视觉变压器,它在SAM中具有最大的计算开销。为了最小化计算成本,我们保持图像编码器冻结。提示符编码器会对边界框的位置信息进行编码,并且可以是从SAM中预先训练的边界框编码器中重用,所以我们也冻结了这个组件。需要微调的其余部分是掩码解码器,如图所示。

(MedSAM:微调SAM的医学图像分割。我们冻结了图像编码器和提示编码器,并只微调掩码解码器。)

       由于图像编码器可以在提示模型之前应用,我们可以对所有训练图像进行图像嵌入,以避免每个提示对图像嵌入的重复计算,从而显著提高训练效率。掩码解码器只需要生成一个掩码,而不是生成三个掩码,因为在大多数情况下,边界框提示符可以清楚地指定预期的分割目标。

       然后就是实验巴拉巴拉,感觉很简单,就是整不出来。
 

猜你喜欢

转载自blog.csdn.net/Zosse/article/details/131631978