SAM模型—探析(PyTorch)

什么是图像分割?图像分割是将数字图像划分为多个区域(或段)的过程,使得属于同一区域(或段)的像素共享一些(语义)特征。应用领域:医学成像(定位肿瘤);物体检测(行人检测、卫星图像中的物体检测);基于内容的图像检索(查找所有包含猫/狗/披萨的图像)。面临的挑战:标记数据集困难且成本高昂(操作员需要创建像素完美的区域);模型通常是特定于应用程序的(例如,仅针对特定类型的医疗应用进行训练,而不能应用于行人检测等其他领域);以前的模型通常不可提示,也就是说,我们无法告诉模型只分割人、汽车或狗。

Vision Transformer(ViT)

Vision Transformer(ViT)是一种将Transformer架构应用于计算机视觉任务的模型。我们将图像分割成固定大小的,线性的嵌入每个块,添加位置嵌入,并将得到的向量序列发送到标准 Transformer编码器。为了执行分类,我们使用标准方法,即在序列中添加一个额外的可学习“分类标记”。

MAE模型

Masked Autoencoder Vision Transformer(MAE)是一种自监督学习方法,用于预训练Vision Transformer模型。MAE架构在预训练期间,会屏蔽掉一大批随机图像块子集(例如75%)。编码器应用于可见图像块的小子集。在编码器之后引入掩码token,全套编码图像块和掩码token由小型解码器处理,解码器会以像素为单位重建原始图像。预训练后,解码器被丢弃,编码器应用于未损坏的图像(全套图像块)以执行识别任务。特点:

  • 非对称编码器-解码器架构:编码器只处理未被遮蔽的图像patches,不包含mask tokens;解码器较轻量级,用于从潜在表示和mask tokens重建原始图像。
  • 高比例遮蔽:随机遮蔽输入图像的大部分区域(约75%patches);这创造了一个非平凡且有意义的自监督任务。
  • 预训练目标:模型需要重建被遮蔽的像素值;这迫使模型学习图像的语义和结构信息。
  • 训练效率:高比例遮蔽和非对称架构使大型模型的训练更加高效;加速训练(3倍或更多)并提高准确性。
  • 迁移学习:预训练后,丢弃解码器,仅保留编码器用于下游任务微调;在下游任务中表现优于有监督预训练,并展现出良好的可扩展性。

SAM模型

Segment Anything Model(SAM)是由Meta AI开发的一种先进的图像分割模型SAM是一个能够对任何图像中的任何对象进行分割的AI模型。它使用多种输入提示(如点击、框选或文本描述)来生成高质量的对象掩码SAM由以下三个主要组件组成:

  • 图像编码器:用于提取图像特征。
  • 提示编码器:用于处理用户提供的提示信息。
  • 掩码解码器:用于生成最终的分割掩码。

特点

  • 通用性SAM能够分割几乎任何类型的对象,而不局限于预定义的类别;通过提示(prompts)机制,SAM可以根据用户提供的提示进行分割。
  • 灵活性:支持多种输入提示,如点、框和文本描述;适用于不同的图像分布和任务,具有零样本迁移能力(zero-shot transfer)。
  • 高效性:通过高效的模型设计和数据循环收集,SAM能够实时生成分割结果;使用了大型数据集进行训练,包括11亿个掩码和1100万张高分辨率图像。

应用场景

  • 图像编辑和处理:如对象移除、背景替换等。
  • 计算机视觉研究:提供一个强大的基础模型,促进各种视觉任务的研究。
  • 医学图像分析:用于自动分割医学图像中的病变区域。
  • 自动驾驶:用于实时分割道路和障碍物。
  • 增强现实:实现实时对象分割和跟踪。

优势

  • 高精度:在各种分割任务中表现优异,零样本性能甚至超过一些完全监督的模型。
  • 快速:能够实时生成分割结果,适用于实时应用。
  • 易用性:提供简单的API和开源代码,易于集成到各种应用中。

可提示分割任务,允许查找给定的掩码:点(例如用户通过鼠标单击选择的点);框(用户定义的矩形)文本提示(例如“找到此图像中的所有狗”)上述内容的组合(例如框和背景点)。快速编码器-解码器模型(在网络浏览器中需要50毫秒才能根据提示生成掩码);歧义感知:例如,给定一个点,它可能对应多个掩码,模型应该返回最有可能的三个(部分、子部分和整体)。Segment Anything引擎收集了11亿个分割掩码;所有掩码均已自动生成(无需人工监督)。总的来说,Segment Anything Model(SAM)是一个强大、灵活且高效的图像分割工具,能够适应广泛的应用场景,为计算机视觉领域带来了显著的进步。

1
2
3
4
5
6
7
8
9
10
11
12
13
import torch
from segment_anything import SamPredictor, sam_model_registry

# 加载预训练模型
sam = sam_model_registry["vit_h"](checkpoint="${path}/to/checkpoint.pth")
predictor = SamPredictor(sam)

# 加载图像并进行分割
image = ... # 读取图像
predictor.set_image(image)

# 使用提示进行分割
masks, _, _ = predictor.predict(point_coords=[[x, y]], point_labels=[1])

稀疏提示被映射到256维向量嵌入,如下所示。一个点表示为该点位置的位置编码 与两个学习到的嵌入的总和,这两个嵌入代表该点在前景还是背景。一个框由一个嵌入对表示:(1)其左上角的位置编码与表示“左上角”的学习到的嵌入相加,(2)表示“右下角”的学习到的嵌入。最后,为了表示文本,我们使用CLIP中的文本编码器(一般来说,任何文本编码器都是可以的)。

SAM提示编码器。这里考虑两组提示:稀疏提示(点、框、文本)和密集提示(掩码)。我们用位置编码表示点和框,并与每个提示类型的学习嵌入和文本相加,使用CLIP的文本编码器。密集提示使用卷积嵌入,并与图像嵌入的每个元素相加。

SAM掩码卷积Dense提示(即掩码)与图​​像具有空间对应性。我们输入的掩码比输入图像的分辨率低了4倍,然后使用两个2×2、步长为2的卷积(输出通道分别为416)将其再缩小4倍。最后的 1×1卷积将通道维度映射到256。每层由GELU激活和层归一化分隔。然后每个元素添加掩码和图像嵌入。如果没有掩码提示,则将“无掩码”的学习嵌入添加到每个图像的嵌入位置。

SAM掩码解码器掩码解码器有效地将图像嵌入、提示嵌入和输出token映射到掩码。这里采用了Transformer解码器块的修改,然后是动态掩码预测头。修改后的解码器块在两个方向使用提示自注意力交叉注意力来更新所有的嵌入。执行这两个块之后,我们对图像嵌入进行了上采样,MLP将输出token映射到动态线性分类器,然后计算每个图像位置掩码的前景概率

SAM:输出token

SAM:输出Transformer。为了确保解码器能够访问关键的几何信息,只要位置编码进入注意力层,就会将其添加到图像嵌入中。此外,只要更新的token进入注意力层,就会将整个原始提示token(包括其位置编码)重新添加到更新的token中。这使得对提示token的几何位置和类型的依赖性很强。

IoU(Intersection over Union)是计算机视觉领域中用于评估目标检测和图像分割算法性能的重要指标。IoU是预测边界框(或分割掩码)与真实边界框(或分割掩码)之间重叠程度的度量。它计算两个区域的相交面积除以它们的合并面积。计算公式:IoU = (预测区域 真实区域) /(预测区域 真实区域),IoU的值在[0,1]之间。1表示完美匹配,0表示没有重叠。应用:目标检测:评估预测边界框与真实边界框的匹配程度;图像分割:评估预测分割掩码与真实分割掩码的匹配程度。通常设置一个IoU阈值(如0.5)来判断检测是否正确。高于阈值视为正确检测。

SAMLoss。我们以focal损失和dice损失的线性组合监督掩码预测,focal损失和dice损失的比例为20:1,我们观察到每个解码器层之后的辅助深度监督是无用的。IoU预测头使用IoU预测和预测掩码的IoU与实际掩码之间的均方误差损失进行训练。它以恒定比例因子1.0添加到掩码损失中。

SAMfocal损失。在物体检测中非常常用,它类似于交叉熵,但针对类别不平衡进行了调整。也就是说,它允许模型将更多的注意力放在困难的例子上,而不是容易的例子上。为什么会出现类别不平衡?因为我们的大多数像素都是非掩码的,只有一小部分是掩码的。

SAMDice损失。Dice分数,也称为Sørensen-Dice系数,是衡量两组数据之间相似度的指标。在图像分割中,Dice分数可用于评估预测分割掩码与实际分割掩码之间的相似度Dice分数的范围从[0,1],0表示无重叠、1表示完全重叠。它等价于F1分值。

SAM:交互式训练。我们在训练期间模拟了一个交互式分割设置。首先,以相等的概率为目标掩码随机选择一个前景点或边界框。从地面实框掩码中均匀采样点。将框作为地面实框掩码的边界框,在每个坐标中添加随机噪声,标准偏差等于框边长的10%,最大为20像素。这种噪声分布是实例分割等应用(在目标对象周围产生紧密的框)和交互式分割(用户可能绘制一个松散的框)之间的合理折衷。从第一个提示进行预测后,从前一个掩码预测和地面实框掩码之间的误差区域中均匀选择后续点。如果误差区域分别为假阴性或假阳性,则每个新点分别为前景或背景。我们还将上一次迭代的掩码预测作为模型的附加提示来提供。为了向下一次迭代提供更多信息,我们提供未阈值化的掩码逻辑,而不是二值化掩码。当返回多个掩码时,传递给下一次迭代并用于对下一个点进行采样的掩码是具有最高预测IoU的掩码。我们发现在8个迭代采样点之后收益递减(我们最多测试了16个)。此外,为了鼓励模型从提供的掩码中受益,我们还使用了另外两次迭代,其中没有采样额外的点。其中一个迭代随机插入8个迭代采样点中,另一个迭代始终位于末尾。这总共提供了11次迭代:一个采样的初始输入提示8个迭代采样点和两次迭代,其中没有向模型提供新的外部信息,因此它可以学习改进自己的掩码预测。我们注意到,使用相对较多的迭代是可行的,因为轻量级掩码解码器需要的图像编码器计算量不到1%,因此每次迭代只会增加很小的开销。这与以前的交互式方法不同,以前的交互式方法每次优化器更新只执行一个或几个交互式步骤。

SAM数据引擎。由于互联网上分割掩码并不是很多,我们构建了一个数据引擎来收集1.1B掩码数据集SA-1B。数据引擎有三个阶段:(1)模型辅助手动注释阶段。(2)半自动阶段,混合了自动预测掩码和模型辅助注释。(3)全自动阶段,其中模型无需对输入注释即可生成掩码。

  • 模型辅助手动注释阶段:在第一阶段,类似于经典的交互式分割,专业注释团队使用由SAM提供的基于浏览器的交互式分割工具,通过单击前景/背景对象点来token蒙版。可以使用像素精确的“画笔”和“橡皮擦”工具来细化蒙版。我们的模型辅助注释直接在浏览器内实时运行(使用预先计算的图像嵌入),从而实现真正的交互式体验。我们没有对token对象施加语义约束,注释者可以自由标记“东西”和“事物”。建议注释者可以标记命名或描述的对象,但没有收集这些名称或描述。注释者被要求按顺序标记对象,并被鼓励在蒙版注释时间超过30秒后继续处理下一张图像。在此阶段开始时,使用常见的公共分割数据集训练SAM。在进行足够的数据注释后,仅使用新注释的蒙版对SAM进行重新训练。随着收集到更多的蒙版,图像编码器从ViT-B扩展到ViT-H;我们总共重新训练了模型 6次。随着模型的改进,每个蒙版的平均注释时间从34秒减少到14秒。我们注意到14秒比COCO的蒙版注释快6.5倍,仅比使用极值点的边界框标记慢2倍。随着SAM的改进,每幅图像的平均蒙版数量从20个增加到44个。总体而言,我们在此阶段从120k张图像中收集了4.3M个蒙版。
  • 半自动阶段:在此阶段,我们旨在增加掩码的多样性,以提高我们的模型分割任何事物的能力。为了让注释者专注于不太突出的物体,我们首先自动检测了可信掩码。然后,我们向注释者展示了预先填充了这些掩码的图像,并要求他们注释任何其他未注释的物体。为了检测可信掩码,我们使用通用的“物体”类别在所有第一阶段的掩码上训练了一个边界框检测器。在此阶段,我们在18万张图像中收集了额外的590万个掩码(总共1020万个掩码)。与第一阶段一样,我们定期在新收集的数据上重新训练我们的模型(5次)。每个掩码的平均注释时间回升至34秒(不包括自动掩码),因为这些物体更难标记。每张图像的平均掩码数量从44个增加到72个(包括自动掩码)。
  • 全自动阶段:在最后阶段,注释是全自动的。因为我们对模型进行了两项重大改进。首先,在这个阶段开始时,我们已经收集了足够的掩码来大大改进了模型,包括来自上一阶段的各种掩码。其次,到这个阶段,已经开发了模糊感知模型,即使在模糊情况下也能预测有效的掩码。具体来说,我们用一个32×32的规则点网格提示模型,并为每个点预测一组可能对应于有效对象的掩码。使用模糊感知模型,如果一个点位于某个部分或子部分上,我们的模型将返回子部分、部分和整个对象。我们模型的IoU预测模块用于选择可信掩码;此外,我们仅识别和选择稳定的掩码(如果在0.5 − δ0.5 + δ处对概率图进行阈值处理导致相似的掩码,则我们认为该掩码是稳定的)。最后,在选择了可信且稳定的蒙版后,我们应用非最大抑制 (NMS) 来过滤重复项。为了进一步提高较小蒙版的质量,我们还处理了多个重叠的放大图像裁剪。我们对数据集中的1100万张图片应用了全自动蒙版生成,共生成了11亿张高质量蒙版。

NMS(Non-Maximum Suppression, 非极大值抑制)是一种在目标检测等计算机视觉任务中广泛使用的后处理技术。NMS用于消除重复的检测结果,只保留最相关的边界框,从而减少假阳性检测。输入:边界框列表及其置信度分数参数:IoU阈值。原理:首先,根据置信度得分对所有预测的边界框进行排序;选择得分最高的边界框,将其添加到最终结果列表中;计算该边界框与所有剩余边界框的IoU(交并比);移除与选中边界框IoU大于某个阈值(通常为0.5)的其他边界框;重复上述步骤,直到处理完所有边界框。输出:在此过程中保留的边界框列表。