Segment Anything Model(SAM)
能分割一切的模型 2023
SAM是一个提示型模型,其在1100万张图像上训练了超过10亿个掩码,实现了强大的零样本泛化。许多研究人员认为「这是 CV 的 GPT-3 时刻,因为 SAM 已经学会了物体是什么的一般概念,甚至是未知的物体、不熟悉的场景(如水下、细胞显微镜)和模糊的情况」,并展示了作为 CV 基本模型的巨大潜力。
2023年4月6号,Meta AI公开了Segment Anything Model(SAM),使用了有史以来最大的分割数据集Segment Anything 1-Billion mask dataset(SA-1B),其内包含了1100万张图像,总计超过10亿张掩码图,模型在训练时被设计为交互性的可提示模型,因此可以通过零样本学习转移到新的图像分布和任务中。在其中他们提出一个用于图像分割的基础模型,名为SAM。该模型被发现在NLP和CV领域中表现出较强的性能,研究人员试图建立一个类似的模型来统一整个图像分割任务。SAM在设计上可以同时输入原图和特定提示(点、框、阴影、文本),然后根据不同的提示输出不同的分割结果图,并且SAM支持不同提示的交互式分割。
项目关键的三部分包括组件:任务、模型、数据。
Meta AI提出一个大规模多样化的图像分割数据集:SA-1B(包含1100万张图片以及10亿个Mask图)。在这项工作中,SAM的目标是建立一个图像分割的基础模型(Foundation Models)。
目标: 是在给定任何分割提示下返回一个有效的分割掩码,并在一个大规模且支持强大泛化能力的数据集上对其进行预训练,然后用提示工程解决一系列新的数据分布上的下游分割问题。
任务: 在NLP和CV中,基础模型是一个很有前途的发展,受到启发,研究者提出了提示分割任务,其目标是在给定任何分割提示下返回一个有效的分割掩码。
网络结构
SAM 架构主要包含三个部分:图像编码器;提示编码器;以及掩码解码器。
Segment Anything Model(SAM),包含一个强大的图像编码器(计算图像嵌入),一个提示编码器(计算提示嵌入),一个轻量级掩码解码器(实时预测掩码)。在使用时,只需要对图像提取一次图像嵌入,可以在不同的提示下重复使用。给定一个图像嵌入,提示编码器和掩码解码器可以在浏览器中在~50毫秒内根据提示预测掩码。
-
图像编码器:基于可扩展和强大的预训练方法,研究者使用MAE预训练的ViT,最小限度地适用于处理高分辨率输入。图像编码器对每张图像运行一次,在提示模型之前进行应用。
-
提示编码器:考虑两组prompt:稀疏(点、框、文本)和密集(掩码)。对于点和框在提取embedding的时候加入了位置编码信息的embedding,然后将两者相加得到最终的embedding。对于文本信息的编码采用了clip模型。
-
掩码解码器:分割mask解码器模块采用Transformer的解码器部分实现,并在后面接入动态的头部预测模块。同时在训练过程中,从图片到提示和从提示到图片都引入了自监督(self-attention)和跨监督(cross-attention)操作。 最后上采样图片的embedding,通过一个MLP模块,经线性分类器,得到最终的概率图。
数据引擎:为了对庞大数据的图像进行掩码标注,作者开发了数据引擎。如图所示,它是一个模型、数据的闭环系统。
模型标注数据:标注好的数据用来优化模型。以此循环,迭代优化模型以及数据质量。
在像自然图像这样的常见场景中,它具有优秀的泛化能力,而在低对比度的场景中,它的效果较差,而且在复杂场景中需要强大的先验知识。
Learning with Explicit Shape Priors for Medical Image Segmentation
基于显式形状先验的医学图像分割学习 2023 IEEE
医学图像分割是医学图像分析和手术计划的基础任务。近年来,基于unet的网络在医学图像分割领域占据主导地位。然而,卷积神经网络(cnn)的感受野有限,无法模拟器官或肿瘤的长期依赖。此外,这些模型在很大程度上依赖于最终分割头的训练。而现有的方法并不能很好地同时解决这两个问题。因此,在我们的工作中,我们提出了一种新的形状先验模块(SPM),它可以显式地引入形状先验,以提高基于unet的模型的分割性能。显式形状先验包括全局形状先验和局部形状先验。 前者具有粗糙形状表示,为网络提供了对全局上下文建模的能力。后者具有更精细的形状信息作为附加指导,提高了分割性能,减轻了分割头对可学习原型的严重依赖。为了评估SPM的有效性,我们在三个具有挑战性的公共数据集上进行了实验。我们提出的模型实现了最先进的性能。此外,SPM在经典cnn和最近基于transformer的骨干网上表现出出色的泛化能力,可以作为一种即插即播的结构,用于不同数据集的分割任务。
本文集中探讨了三种带有形状先验的分割模型:
- 基于图谱的模型(atlas-based models)
- 基于统计的模型(statistical-based models)
- 基于U-Net的模型(UNet-based models)
前两种方法的泛化能力较差,而 UNet-based 模型由于相比于前两者泛化性能要好,但由于它是倾向于使用隐式形状先验,这在不同形状的器官上缺乏良好的可解释性和泛化能力。综上所述,本文提出了一种新的形状先验模块(Shape Prior Module, SPM),它可以显示地引入形状先验,以促进 UNet-based 模型的分割性能。
关于隐式形状先验和显示形状先验?
- 隐式形状先验通常是通过在模型中加入先验信息,例如特定的损失函数或正则化项来实现的。这些隐式的形状先验通常难以解释,因为它们是通过一些特殊的方式集成到模型中的,而不是直接考虑目标物体的形状信息。例如,在基于 UNet 的模型中,可以通过使用 Dice 损失函数来强制模型更加注重目标物体的轮廓信息,从而隐式地考虑了形状先验信息。
- 显式形状先验则直接将形状先验信息作为输入提供给模型。例如,在本文中,作者提出了一个新的形状先验模块,它明确地将形状先验信息作为输入,并利用这些信息来引导模型更好地分割目标物体。这种显式的形状先验可以更好地解释和调整,因为它们直接考虑了目标物体的形状和结构信息。
SPM
本文的主要贡献是在U-Net网络基础上设计了一个形状先验模块(SPM),通过引入可学习形状先验,为每个类别施加解剖形状约束来增强网络的表示能力。与DETR类似,作者设定了 n 个(分割类别数目)可学习原型,类比于DETR中Transformer解码器中的对象查询。
SPM的输入包括原始跳跃特征 和原始形状先验 ,经过“特征提纯”后会生成对应的增强跳跃特征和增强形状先验 。最终,通过这些增强后的特征和先验,模型会生成更加精准的分割掩膜。与DETR不同,SPM会与多尺度特征进行交互,而不仅仅是来自编码器最深层的特征。因此,在跳跃连接之前的分层编码特征在经过SPM处理后将获得更多的形状信息。增强形状先验由两个部分组成:
- 全局形状先验
- 局部形状先验
它们分别由自更新块(self-update)和交叉更新块(cross-update block) 生成。
Self-update block (SUB)
为了建模类别之间的关系,本文提出了自更新块(参照上图第一行)来生成具有 N 个通道之间交互的形状先验。从流程图上可以看出,输入特征是来自编码层对应层级输出的长跳跃连接,作者首先整了两个Resblock做下特征提取并进行维度压缩。其次,便是借鉴ViT中自注意力机制(Self-Attention) 的启发,通过引入来自深层的特征,构建对应的QKV。如此一来,网络便可以有效的建模远距离的上下文依赖,从而获得与全局区域相关的丰富纹理信息。
Cross-update block (CUB)
SUB带来了全局上下文信息,但最关键的形状和轮廓信息缺忽略了,这是由于缺乏相应的归纳偏置来建模局部视觉结构和定位具有不同尺度的对象。为了解决这个限制,CUB 应运而生。受到卷积核固有的局部性和尺度不变性的归纳偏置的启发,基于卷积的 CUB 为 SPM 注入归纳偏置,以获得更细致的局部形状信息。此外,基于编码器中卷积特征具有定位区分性区域的显著潜力的事实,作者尝试在原始跳跃特征和形状先验之间进行交互。
总结:
形状先验模块通过引入全局和局部特征对原始跳跃特征和形状先验进行了增强。 全局形状先验能够通过自注意力块模拟类间关系,包含粗糙的形状先验和丰富的纹理信息以建立全局上下文。局部形状先验通过卷积引入归纳偏差,从而能够捕捉具有不同尺度的局部视觉结构和定位目标。此外,SPM还通过与全局形状先验的交互来增强原始跳跃特征,从而促进生成具有较好形状表示和全局上下文的特征,最终获得更准确的预测分割掩模。
End
以上仅作个人学习记录使用