文章目录
- 一、Segment Anything Task
- 二、SAM模型的架构
- (一)图像编码器image encoder
- (二)prompt encoder
- (三)mask decoder
- 三、SAM 模型的思想
- (一)Transformer的自注意力机制
- (二)用户提示的融入
- 四、Segment Anything Data Engine
Segment Anything(SA)项目:一个用于图像分割的新任务、新模型和新数据集
通过FM(基础模型)+prompt解决了CV中难度较大的分割任务,给计算机视觉实现基础模型+提示学习+指令学习提供了一种思路
关键:加大模型容量(构造海量的训练数据,或者构造合适的自监督任务来预训练)
一、Segment Anything Task
- 目标:创建一个类似NLP领域的foundation model(例如GPT),它是一种使用prompt engineering机制的预训练模型并可以灵活地解决下游任务;
- SAM 模型是一个提示型模型,可以在不需要任何标注的情况下,对任意图像进行分割:
- 输入包括一个图像和一个提示,提示可以是点、框、文本或者掩码,用来指示要分割的目标;
- 输出是一个分割掩码,表示图像中每个像素属于前景或背景的概率;
- 与以往的分割模型相比,SAM 的本质区别在于它不需要为每个任务或数据集进行专门的训练或微调,而是可以直接使用预训练的模型来完成分割任务,从而节省了大量的时间和资源。
二、SAM模型的架构
SAM(segment anything model)模型总体上分为3大块:
- 绿色的image encoder,将图像编码为向量;
- 紫色prompt encoder,将支持的prompt 内容编码为向量;
- 橙色的mask解码器,将两个特征融合并解码出最终的分割掩码(使用产生的text embedding和image embedding 在mask decoder模块进行对齐)。
(一)图像编码器image encoder
SAM中的图像编码器使用 MAE预训练的Vision Transformer (ViT) 作为图像编码器,可以适应高分辨率输入。具体的图像编码器支持vit-h,vit-l,vit-b,vit经过了MAE方式的预训练,MAE属于图像中自监督学习的一种,如下图:
- 在MAE中原始图像如vit切割成不重叠的patch
- 保留部分patch进入vit架构的encoder进行学习patch的表示,学习到的patch表示和mask(灰色)的表示
- 按照原始的patch顺序输入到vit架构的decoder,得到复原图像。
训练完成后我们只使用encoder来提取图像特征。从原始模型可以看出,图像的表征embedding是不变的,也就可以在已经编码好的图像embedding多次进行不同的prompt输入得到期望的结果,这个对交互式分割的场景是非常有用的。
(二)prompt encoder
基于分割的任务需求,SAM 支持的prompt可以分为稀疏类型和稠密类型两类。
-
稀疏类型(sparse prompt):包含 point,bbox,free text。
对于point和bbox两类标记:采用sincos的位置编码+学习的类别编码
对于free text标记:采用clip预训练好的text encoder 作为文本编码器 -
稠密类型(dense prompt)
掩码标记:卷积+使用图像编码进行逐像素求和
(在SAM中使用粗略的分割结果作为稠密的prompt输入,这里通过卷积进行特征提取)
(三)mask decoder
- mask decoder的核心:
用transformer 学习和prompt对齐后的image embedding 以及额外4个token的embedding。
4个token embedding:
iou token embedding和3个分割结果 token的embedding,经过transformer 学习得到token 的embedding 会用于最终的任务头,得到目标结果。
- transformer的三个输入:
- token embedding(将文本中的单词或token映射到连续向量空间的过程);
- src(image embedding 和dense prompt embedding的求和);
- pos_src(图像的位置编码)。
- 在transformer中每层做以下四件事情:
- token embedding 做self attention 计算;
- token embedding 和src 之间做cross attention 计算;
- src 和token embedding 之间做cross attention 计算;
- 第2和3之间有前馈mlp网络;cross attention的结果通过残差方式相加并norm。
- transform最终输出前,token embedding 还需要和src 做一次cross attention。transform返回的3个mask token的embedding经过3层mlp后,与对齐后的图像embedding点击得到3个最终的分割结果;iou token 经过mlp得到3个分割结果置信度得分。
三、SAM 模型的思想
(一)Transformer的自注意力机制
SAM 模型是一个基于 Transformer 的视觉模型,它可以将图像和用户的提示作为输入,输出一个分割掩码,表示每个像素的标签。SAM 模型的关键是如何利用 Transformer 的自注意力机制,来学习图像和提示之间的关系,以及图像内部的结构和语义信息。
Transformer的自注意力机制:
Transformer 的特点是使用自注意力机制,来计算序列中每个元素与其他元素的相关性,从而捕捉序列的全局依赖关系。自注意力机制的公式如下:
A t t e n t i o n ( Q , K , V ) = s o f t m a x ( Q K T d k ) V Attention(Q,K,V)=softmax(\frac{QK^T}{\sqrt {d_k}})V Attention(Q,K,V)=softmax(dkQKT)V
其中, Q Q Q, K K K, V V V分别表示查询(query),键(key),值(value)矩阵, d k d_k dk是键的维度,softmax是归一化函数。
-
为了将 Transformer 应用到图像上,SAM 模型首先将图像划分为多个小块,每个小块对应一个像素块,然后将每个像素块的颜色值转化成一个向量,作为 Transformer 的输入序列。这样,每个像素块就相当于 Transformer 中的一个词,而图像就相当于 Transformer 中的一个句子。
-
为了将用户的提示也作为 Transformer 的输入,SAM 模型使用了一种称为提示编码器(prompt encoder)的模块,它可以将不同类型的提示转换为统一的向量表示,然后与图像的向量表示拼接在一起,作为 Transformer 的输入序列。这样,每个提示就相当于 Transformer 中的一个词,而图像和提示的组合就相当于 Transformer 中的一个句子。
通过 Transformer 的自注意力机制,SAM 模型可以计算图像中每个像素块与其他像素块以及用户的提示之间的相关性,从而学习图像的结构和语义信息,以及用户的意图信息。然后,SAM 模型使用一个线性层,将 Transformer 的输出序列映射为一个分割掩码,表示每个像素块的标签。
(二)用户提示的融入
SAM 模型的重点在于它能够理解图像内容、用户意图以及它们之间的关联。具体来说,SAM 模型是通过一个掩码解码器 (Mask decoder) 来融合图像编码器和提示编码器的输出的。掩码解码器的关键思想在于Transformer 的自注意力机制和交叉注意力机制。
四、Segment Anything Data Engine
数据产生的三个阶段:
- 辅助人工标注:模型辅助的手动注释阶段
- 刚开始的SAM是采用公开的分割数据训练;
- 标注时人工采用点击前景点、背景点作为SAM的prompt输入,对分割的结果进行标注和修正;
- 随着标注数据的增多,会采用新标注的数据来重训SAM模型,这个阶段模型反复重训了6次。
- 人工半监督标注:混合了自动预测掩码和模型辅助注释的半自动阶段
- 利用第一阶段的数据训练一个faster-rcnn的目标检测模型;
- 用检测框作为SAM的prompt输入(目标检测难度比分割小的多),输出的分割结果中,人工只需要关注置信度得分低的分割图进行修正,并补充SAM遗漏的结果;
- 同样在这个阶段,随着标注数据的增多,SAM模型会持续的重训,一共进行5次训练。
- 全自动标注:没有注释器输入的情况下生产掩码的全自动阶段
- 第三个阶段类似我们产生伪标签训练的过程
- 用前边数据训练好的SAM在海量数据上产生分割的结果,然后再通过规则过滤掉部分可能错误的结果
参考:
五分钟看明白分割一切模型 SAM
Segment Anything论文和源码解读
【模型解读】【代码复现】Segment Anything Model(SAM)