0-Segment Anything

Segment Anything

SAM

Code:https://github.com/facebookresearch/segment-anything

HP:https://segment-anything.com

原理解析ref:【视觉分割大模型SAM(原理解析+代码实践)】link

背景

计算机视觉大模型。构建了迄今为止(迄今为止)最大的分割数据集,在 11M 许可和尊重隐私的图像上拥有超过 10 亿个掩码。该模型经过设计和训练,具有可提示性,因此可以将 Zero-shot 转移到新的图像分布和任务。上发布 1B 掩码和 11M 图像的 Segment Anything Model (SAM) 和相应的数据集 (SA-1B),以促进对计算机视觉基础模型的研究。

大模型起于NLP,自然语言处理领域。在自然语言处理中,有很多称为基础模型的模型,比如谷歌的BERT以及OpenAI的gpt,都是基础模型,这些基础模型输入是一段话,输出也是一段话,将这样的模型称为seq2seq,也就是序列预测,有了这些基础模型之后,就可以很方便地甚至不需要训练就能迁移到翻译或者文本摘要这种常规的自然语言处理的任务上。那这种基础模型是如何来进行训练的呢,训练离不开数据,各位在进行yolo模型训练的时候肯定也被数据标注折磨过,一张张的标注属于你任务的数据之后才能用于模型的训练,如果你标注的数据比较少的话,那训练的效果可能还比较差。对于自然语言处理的这种基础模型而言,有个比较方便的一点是无需人工标注,只需要从网络上搜集大量的文本即可,搜集文本之后,随机扣掉这段话中的部分单词作为输入,你得输出则是尽可能地将这些这些扣掉的词预测出来,这种方式是不是和你在高中英语中做的完型填空类似,由于你是那现成的文本扣掉部分单词来作为输入,你只需要写个简单的脚本随机扣掉这些单词就能制作一份巨大的数据集,完全不需要人工进行大量的标注,其中这种方式也就是大名鼎鼎的bert所使用的方式。

在计算机视觉的领域中,也有使用类似方法来进行训练的模型,就是有何凯明提出的MAE,注意,这里的这个作者同时也是resnet的作者。他的做法是将图像随机扣除部分块,然后将他进行还原,通过这种方式来进行训练也不需要进行标注。

MAE

对于自然语言处理而言,这里的下游任务可以是翻译,可以是文本摘要。对于计算机视觉而言,这里的下游任务可以是分割,可以是检测。其中进行零样本学习的常用方面就是提示词工程,也就是你在chatgpt中经常听到的prompting。

有了NLP任务的参考,研究人员就想能不能把同样的方式应用在计算机视觉的领域上来。对于计算机视觉的任务,希望以图像建立一个类似这样的模型,然后迁移到其他任务上面,其中在计算机视觉的任务中,最难的就是分割了,分割要找到物体的轮廓。视觉和语言不一样,我们要处理的是图像,有了处理的图像,我们还缺少提示,那这样提示可以是什么样子的方式呢。

在NLP的领域中,存在一些被称为基础模型的模型,他们通过预测句子中的下一个词进行训练,称为顺序预测。通过这些基础的模型可以轻松地适应到其他的NLP的任务上,比如翻译或者是文本摘要,这种实现方式也可以称为是零样本迁移学习。其中比较著名的方法就是prompting,通过聊天的形式来进行交互。NLP有效的前提是网络上存在大量的文本,而对于序列的预测,比如说知道一些词然后预测后面的词是什么,这种不需要人工标注的标签就能完成训练。但是问题转化到计算机视觉的任务上,尽管网络上存在数十亿的图像,但是由于缺乏有效标注的mask的信息,所以在计算机视觉的任务上建立这样模型成为了挑战,也就是作者提出的三个问题:

  1. 什么任务将启用零样本泛化?
  2. 对应的模型架构是什么?
  3. 哪些数据可以为这个任务和模型提供支持?

作者的目标是通过引入三个相互关联的组件来构建一个分割的基础模型:一个可提示的分割任务、一个通过数据标注提供动力并能够通过提示工程实现一系列任务零样本迁移的分割模型(SAM),以及一个用于收集我们的数据集SA-1B(包含超过10亿个掩码)的数据引擎。

图1 SAM

实验方法

分割一切的任务

在自然语言处理(NLP)领域,乃至最近的计算机视觉领域,基础模型都是一种很有前景的发展方向,通常通过使用“提示”技术,能够对新的数据集和任务进行零样本和少量样本的学习。受这一工作方向的启发,提出了可提示的分割任务,其目标是在给定任何分割提示的情况下返回一个有效的分割掩码(见图1a)。提示简单地指定了图像中要分割的内容,例如,提示可以包括标识对象的空间或文本信息。有效输出掩码的要求意味着,即使提示含糊不清且可能指代多个对象(例如,衬衫上的一个点可能表示衬衫或穿衬衫的人),输出也应该是其中至少一个对象的合理掩码。我们将可提示的分割任务既用作预训练目标,也用作通过提示工程解决一般的下游分割任务的方法。

参考这种提示词工程的做法,给定一张图片和给定对应的提示信息,得到对应位置的mask区域。这里的提示信息可以是mask、points、box或者是text,如果是mask的话则使用卷积来进行表示,如果是其余的三种提示词则使用位置编码的形式来进行表示,其中text可以通过clip一类的模型获取词嵌入。其中主干网络的部分使用的特征提取能力更强的VIT网络,整体的构思如下:

图4 SAM

Segment Anything Model (SAM) 概述。重量级图像编码器输出图像嵌入,然后可以通过各种输入提示有效地查询该图像嵌入,以分摊的实时速度生成对象掩码。对于对应于多个对象的不明确提示,SAM 可以输出多个有效掩码和关联的置信度分数。

分割一切的模型

可提示的分割任务和实际使用目标对模型架构施加了约束。具体而言,模型必须支持灵活的提示,需要以分摊的实时方式计算掩码以允许交互式使用,并且必须具备处理歧义的能力。发现一个简单的设计就能满足所有这三个约束条件:一个强大的图像编码器计算图像嵌入,一个提示编码器嵌入提示,然后将这两个信息源在一个轻量级的掩码解码器中结合起来,预测分割掩码。我们将此模型称为“任意分割模型”或SAM(见图1b)。通过将SAM分离为图像编码器和快速的提示编码器/掩码解码器,可以使用不同的提示重复使用相同的图像嵌入(并分摊其成本)。给定图像嵌入,提示编码器和掩码解码器可以在网页浏览器中在约50毫秒内从提示预测掩码。专注于点、框和掩码提示,并展示了使用自由形式文本提示的初步结果。为了使SAM具备处理歧义的能力,设计它针对单个提示预测多个掩码,从而使SAM能够自然地处理歧义,例如衬衫与人之间。

  • 图像的编码器:

图像编码器。出于可扩展性和强大的预训练方法的考虑,使用了一个经过最小调整以适应高分辨率输入的MAE预训练视觉Transformer(ViT)。图像编码器每张图像运行一次,并可在提示模型之前应用,这里使用的mae来进行预训练。(PS: 之前有个想法是使用mae先预训练一个模型出来,但是明显这个想法是使用在这个模型里面的,所以多读书是非常必要的)

  • 提示词的编码器:

提示编码器。考虑两组提示:稀疏提示(点、框、文本)和密集提示(掩码)。用位置编码来表示点和框,并将其与每种提示类型的学习嵌入和来自CLIP的现成文本编码器中的自由格式文本相加。密集提示(即掩码)使用卷积进行嵌入,并与图像嵌入进行逐元素相加。

  • 掩码的解码器:

掩码解码器能够高效地将图像嵌入、提示嵌入和输出标记映射到一个掩码。采用了一个经过修改的Transformer解码器块,后面跟着一个动态掩码预测头。修改后的解码器块在两个方向上(从提示到图像嵌入和从图像嵌入到提示)使用提示自注意力和交叉注意力来更新所有嵌入。运行两个块之后,对图像嵌入进行上采样,并且一个多层感知机 (MLP)将输出标记映射到一个动态线性分类器,然后该分类器计算图像每个位置的前景掩码概率。

  • 解决歧义的问题:

解决歧义问题。如果给定一个模糊的提示,模型将平均多个有效的掩码作为一个输出。为了解决这个问题,修改了模型,使其能够针对单个提示预测多个输出掩码(见图3)。我们发现,3个掩码输出足以处理大多数常见情况(嵌套掩码通常最多有三层:整体、部分和子部分)。比如上面的剪刀的图像,其实由三个有效的掩码。

图3 :每列显示 SAM 从单个模糊点提示(绿色圆圈)生成的 3 个有效掩码

Segment Anything 数据引擎

数据引擎分为3个阶段:

第一个阶段

使用公开的数据集进行分割,并让模型和标注专家进行互动,标注专家使用浏览器通过擦除或者绘画画布来纠正模型输出的掩码。在第一阶段,类似于经典的交互式分割,一组专业标注者使用由SAM支持的基于浏览器的交互式分割工具,通过点击前景/背景对象点来标注掩码。可以使用像素精确的“画笔”和“橡皮擦”工具来完善掩码。模型辅助标注在浏览器内实时运行(使用预计算的图像嵌入),从而实现了真正的交互式体验。没有对标注对象施加语义约束,标注者可以自由标注“材质”和“物体”。建议标注者标注他们能够命名或描述的对象,但没有收集这些名称或描述。标注者被要求按照突出程度顺序标注对象,并且如果标注一个掩码需要超过30秒,就会被鼓励继续标注下一张图像。

第二个阶段

第二个阶段在第一个阶段的基础之上更关注模型的细节。在这一阶段,目标是增加掩码的多样性,以提高模型分割任何物体的能力。为了让标注者关注不太突出的物体,首先自动检测置信度高的掩码。然后,向标注者展示预填充了这些掩码的图像,并要求他们标注任何其他未标注的物体。为了检测置信度高的掩码,使用通用的“物体”类别,在所有第一阶段的掩码上训练了一个边界框检测器。在这一阶段,在18万张图像中额外收集了590万个掩码(总共1020万个掩码)。与第一阶段一样,定期在新收集的数据上重新训练模型(5次)。由于这些物体更难标注,每个掩码的平均标注时间回升至34秒(不包括自动掩码)。每张图像的平均掩码数量从44个增加到72个(包括自动掩码)。

第三个阶段

第三个阶段引入32✖️32的常规网格作为输入提示,这个阶段涵盖物体的部分、子部分和整体。在最后阶段,标注是完全自动的。这得益于模型的两大改进。首先,在这一阶段开始时,我们已经收集了足够的掩码来极大地改进模型,包括上一阶段的各种掩码。其次,到这一阶段,已经开发了具有歧义感知能力的模型,这使得我们即使在存在歧义的情况下也能预测出有效的掩码。具体来说,用32×32的常规网格点来提示模型,并为每个点预测一组可能对应有效物体的掩码。借助具有歧义感知能力的模型,如果一个点位于某个部分或子部分上,模型将返回子部分、部分和整个物体。使用模型的交并比(IoU)预测模块来选择置信度高的掩码;此外,我们只识别和选择稳定的掩码(如果在0.5−δ和0.5+δ的概率图上应用阈值处理能得到相似的掩码,则该掩码是稳定的)。最后,在选择置信度高且稳定的掩码后,应用非极大值抑制(NMS)来过滤重复的掩码。为了进一步提高较小掩码的质量,还处理了多个重叠的放大图像裁剪。将全自动掩码生成应用于数据集中的所有1100万张图像,总共生成了11亿个高质量的掩码。最终在1100w张图片上生成了11亿个掩码,称为SA-1B数据集。

实验结果

采用了零样本的学习方式,并且进行预测的图像是没有在文中提到的SA-1B数据集出现过的。展示了一系列实验,这些实验涵盖了低、中、高级别的图像理解,并大致与该领域的历史发展并行。具体来说,提示SAM执行以下任务:(1)执行边缘检测,(2)分割所有内容,即生成对象提议,(3)分割检测到的对象,即实例分割,以及(4)作为概念验证,从自由格式的文本中分割对象。这四个任务与SAM训练时的可提示分割任务显著不同,并通过提示工程实现。我们的实验以消融研究结束。

单点语义分割

按照交互式分割中的标准评估协议,从真实掩码的“中心”(在掩码内部距离变换的最大值处)采样点。由于SAM能够预测多个掩码,默认情况下仅评估模型最具置信度的掩码。所有基线方法都是单掩码方法。我们主要与RITM进行比较,这是一种强大的交互式分割器,与其他强大的基线相比,在本文的基准测试中表现最佳。

单点语义分割

边缘检测

使用了简化版的自动掩码生成流程。具体来说,我们用16×16的常规前景点网格来提示SAM,生成了768个预测的掩码(每个点3个)。通过非极大值抑制(NMS)去除冗余的掩码。然后,使用未阈值化的掩码概率图的Sobel滤波和标准轻量级后处理(包括边缘NMS)来计算边缘图。

简单总结来说就是先通过网格点作为提示来生成掩码,最后通过滤波和后处理完成物体边缘的提取:

边缘检测

提议边界框

将前面的流程稍作修改,变为计算输出掩码的边界框的形式。

提议边界框

实例分割

转向更高级别的视觉任务,使用SAM作为实例分割器的分割模块。实现过程很简单:运行一个目标检测器(之前使用的ViTDet),并用其输出的边界框来提示SAM。这说明了如何在一个更大的系统中组合使用SAM。

实例分割

文本生成图像

首先通过CLIP等模型生成词嵌入,作为提示词输入到SAM模型中,得到对应物体的MASK。考虑了一个更高层次的任务:从自由格式的文本中分割对象。这个实验证明了SAM处理文本提示的能力。虽然在之前的所有实验中都使用了完全相同的SAM,但在这个实验中,对SAM的训练过程进行了修改,使其能够识别文本,但这种方式并不需要新的文本注释。具体来说,对于每个手动收集的面积大于1002的掩码,我们提取CLIP图像嵌入。然后,在训练过程中,将提取的CLIP图像嵌入作为SAM的第一次交互。这里的关键观察是,由于CLIP的图像嵌入是与其文本嵌入对齐的,因此可以使用图像嵌入进行训练,但使用文本嵌入进行推理。也就是说,在推理时,将文本通过CLIP的文本编码器运行,然后将得到的文本嵌入作为提示给SAM。从实验中可以看出,通过文本加点的集合可以帮助模型输出更完美的mask,如果只是通过文本的话分割的情况有限。

文本生成图像

总结

基础模型

自机器学习的早期以来,预训练模型就已经适应了下游任务 。近年来,随着对规模的日益重视,这种范式变得越来越重要,这种模型最近被(重新)命名为“基础模型”:即“在广泛的数据上进行大规模训练并适应广泛的下游任务”的模型 。本文的工作与这个定义密切相关,尽管注意到图像分割的基础模型本质上是一个有限的范围,因为它代表了计算机视觉的一个重要但很小的子集。

组合性

预先训练的模型可以增强新功能,甚至超出训练时想象的功能。一个突出的例子是 CLIP 如何用作大型系统中的组件,例如 DALL·E 。作者的目标是使用 SAM 使这种组合变得简单明了。为了实现这一目标,要求 SAM 为各种分段提示预测有效的掩码。其效果是在 SAM 和其他组件之间创建一个可靠的接口。例如,MCC 可以很容易地使用 SAM 来分割感兴趣的对象,并实现对不可见对象的强烈泛化,以便从单个 RGB-D 图像进行 3D 重建。在另一个示例中,可穿戴设备检测到的注视点可以提示 SAM,从而支持新的应用程序。由于 SAM 能够推广到以自我为中心的图像等新领域,因此此类系统无需额外培训即可工作。

局限性

虽然 SAM 总体上表现良好,但它并不完美。它可能会错过精细的结构,有时会产生小的断开连接的组件产生幻觉,并且不会像 “放大 ”那样清晰地产生边界。一般来说,当提供许多点时,预计专用的交互式分割方法会优于 SAM。与这些方法不同,SAM 旨在实现通用性和广泛性,而不是高 IoU 交互式分割。此外,SAM 可以实时处理提示,但在使用重型图像编码器时,SAM 的整体性能并不是实时的。作者尝试的文本到掩码任务是探索性的,并不完全稳健,尽管相信它可以通过更多的努力来改进。虽然 SAM 可以执行许多任务,但目前尚不清楚如何设计实现语义和全景分割的简单提示。最后,还有一些特定于领域的工具,预计它们在各自的领域中会优于 SAM。

结论

Segment Anything 项目试图将图像分割带入基础模型时代。主要贡献是使这一飞跃成为可能的新任务(可提示分割)、模型 (SAM) 和数据集 (SA-1B)。SAM 是否达到基础模型的地位还有待观察它在社区中的使用方式,但无论我们期待这项工作的前景、超过 1B 掩码的发布以及本文的可提示细分模型将有助于铺平未来的道路。

代码分析

环境:WSL2 Ubuntu22.04

sam自动掩码目标生成:

初始准备:准备代码与环境

1
2
3
4
5
6
7
8
#one by one
gh repo clone facebookresearch/segment-anything
cd segment-anything
conda creat env -n sam python=3.8.5
source activate sam
#寻找合适的cuda版本,我这里cuda version是12.6,选择pytorch-cuda=12.4版本
conda install pytorch torchvision torchaudio pytorch-cuda=12.4 -c pytorch -c nvidia
pip install -v -e .

演示代码:

1
vi test.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
import numpy as np
import torch
import matplotlib.pyplot as plt
import cv2
import sys

sys.path.append("..")
from segment_anything import sam_model_registry, SamAutomaticMaskGenerator, SamPredictor


def show_anns(anns):
if len(anns) == 0:
return
sorted_anns = sorted(anns, key=(lambda x: x['area']), reverse=True)
ax = plt.gca()
ax.set_autoscale_on(False)

img = np.ones((sorted_anns[0]['segmentation'].shape[0], sorted_anns[0]['segmentation'].shape[1], 4))
img[:, :, 3] = 0
for ann in sorted_anns:
m = ann['segmentation']
color_mask = np.concatenate([np.random.random(3), [0.35]])
img[m] = color_mask
ax.imshow(img)

#图片路径
image = cv2.imread('img/01L.jpg')
image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)

# plt.figure(figsize=(20, 20))
# plt.imshow(image)
# plt.axis('off')
# plt.show()

#base检查点路径
# sam_checkpoint = "checkpoint/sam_vit_b_01ec64.pth"
# model_type = "vit_b"
#huge检查点路径
sam_checkpoint = "../sam_checkpoint/sam_vit_h_4b8939.pth"
model_type = "vit_h"

device = "cuda"

sam = sam_model_registry[model_type](checkpoint=sam_checkpoint)
sam.to(device=device)

mask_generator = SamAutomaticMaskGenerator(sam)

masks = mask_generator.generate(image)

print(len(masks))
print(masks[0].keys())

plt.figure(figsize=(20, 20))
plt.imshow(image)
show_anns(masks)
plt.axis('off')
plt.show()

运行结果