SAM:分割任意模型.
- paper:Segment Anything
- project:Segment Anything Model (SAM): a new AI model from Meta AI that can “cut out” any object, in any image, with a single click
SAM (Segment Anything Model)是一个图像分割的基础模型,模型的设计和训练是通过提示工程实现的,预训练完成后可以对不同的图像分割数据实现零样本的泛化。文章主要讨论了三个问题:
- 什么任务可以实现零样本泛化?
- 什么模型结构能够实现该任务?
- 什么数据可以用于模型预训练?
1. 可提示的分割任务 promptable segmentation task
可提示的分割任务的目标是在给定任何分割提示的情况下返回有效的分割掩码。提示应该指定在图像中分割什么,例如提示可以包括标识对象的空间坐标或文本信息。有效输出掩码意味着即使提示不明确,并且可能涉及多个对象,输出也应该是这些对象中至少一个的合理掩码。
对于分割任务,提示可以设置为一组前景/背景点、一个粗略的边界框或mask、一段文本等能够提供指示在图像中分割什么内容的信息。可提示的分割任务提出了一种自然的预训练算法,该算法模拟每个训练样本的提示序列(例如点、框、掩码),并将模型的掩码预测与真实标签进行比较。预训练任务赋予了模型在推理时对任何提示做出适当响应的能力,因此下游任务可以通过设计适当的提示来解决。
2. 分割任意模型 segment anything model
可提示的分割模型必须支持灵活的提示输入形式,能够实时计算掩码以允许交互式使用,并且能够处理分割语义的模糊性。SAM采用一种简单的设计:一个图像编码器生成图像嵌入,一个提示编码器生成提示嵌入,然后将这两种嵌入组合后通过一个轻量级掩码解码器预测分割掩码。
- 图像编码器(image encoder):采用MAE预训练的ViT模型
- 提示编码器(prompt encoder):稀疏提示(点、框)通过位置编码嵌入后与对应提示类型的学习嵌入相加,文本提示通过CLIP的文本编码器嵌入,密集提示(掩码)使用卷积嵌入,并与图像嵌入相加。
- 掩码解码器(mask decoder):采用Transformer解码器的修改版。解码器块在两个方向上(提示到图像嵌入,图像到提示嵌入)使用提示自注意力和交叉注意力来更新所有嵌入。在两个解码器块之后对图像嵌入进行上采样,MLP把输出映射到动态线性分类器,然后动态线性分类器计算每个图像位置的掩码前景概率。
对于一个输出,如果给出不明确的提示,模型将平均多个有效掩码。为了解决分割结果的模糊性,对单个提示预测多个输出掩码。具体地,为每个提示设置$3$个输出掩码(整体、部分、子部分),模型同时预测每个掩码的置信度得分(估计的IoU)。训练时只计算预测掩码的最小损失,采用Focal损失和Dice损失的线性组合作为损失函数。
3. 数据引擎 data engine
作者收集了一个SA-1B数据集,并设计如下数据引擎完成收集过程:
- 辅助标注阶段:一组专业标注人员通过交互式分割工具点击前景/背景对象点来进行标注,标注人员可以自由地标记目标。在这个阶段从120k张图像中收集了430万个标注。
- 半自动阶段:首先使用一个通用的“对象”类别在所有第一阶段图像上训练了一个Faster R-CNN,使用该边界框检测器构造置信度较高的掩码,并要求标注人员标注图像中任何其他未标注的目标对象。此阶段在180k张图像中额外收集了590万个标注。
- 全自动阶段:完全使用模型产生新的标注。使用$32×32$的规则网格提示模型,并为每个点预测一组可能对应于有效对象的掩码,并从中选择置信度高、稳定的掩码,应用非极大值抑制过滤重复掩码。为了进一步提高较小掩码的质量,模型处理了多个重叠的放大图像裁剪。该方法应用于数据集中的所有1100万张图像,总共生成了11亿个高质量标注。
SA-1B数据集包含1100万张图像,共计11亿个高质量掩码标注,这些图像的平均分辨率为$3300 \times 4950$。该数据集是迄今为止最大的图像分割数据集。