SAM + YOLOv8 图像分割及对象检测

SAM(Segment Anything Model）是由 Meta 的研究人员团队创建和训练的深度学习模型。该创新成果发表在 2023 年 4 月 5 日发表的一篇研究论文中，它立即引起了公众的广泛兴趣——相关的 Twitter 帖子迄今为止已累积超过 350 万次浏览：

计算机视觉专业人士现在将注意力转向 SAM——但为什么呢？

在这里插入图片描述

推荐：用 NSDT编辑器快速搭建可编程3D场景

1、什么是SAM？

在 Segment everything 研究论文中，SAM 被称为“基础模型”。

基础模型是在大量数据上训练的机器学习模型（通常通过自监督或半监督学习），其目的是在更具体的任务上使用和重新训练。

换句话说，SAM 是一个预训练模型，旨在适应其他任务（特别是通过微调）。

例如，SAM 可以重新训练并用于仅对数据集中的人员进行分割。

人物分割是 SAM 可以执行的一项附件任务，因为它已经在包含此类对象的数据集上进行了训练 - 但不仅如此！

2、SAM 是如何训练的？

SAM 在 SA-1B 数据集上进行了训练，该数据集由 Meta 与 Segment Anything 研究论文并行引入。

Facebook 母公司的数据集包含超过 1100 万张几乎在整个地球上收集的图像——这是开发具有泛化能力的模型的一个重要方面。

在这里插入图片描述

几乎在整个地球上收集的图像 – SA-1B 数据集

这些高质量图像（平均 1500×2250 像素）伴随着与数据集标签相对应的 11 亿个分割掩模。

Meta 使用此数据集的目的是为AI博士创建细分参考。它已获得官方免费许可用于研究目的。

尽管信息量很大，但值得注意的是，掩模与类别无关。换句话说，即使SAM可以生成一个人的掩模，它也无法表明这个掩模代表一个人。

这是需要牢记的重要一点，因为这意味着 SAM 必须与其他算法结合才能真正发挥作用。

让我们仔细看看。

3、如何使用 SAM？

首先，我们需要加载 2 项：

segment-anything GitHub 文件夹，其中包含使用 SAM 的类和函数
使用 Meta 研究人员获得的模型版本的预训练模型权重

!pip install git+https://github.com/facebookresearch/segment-anything.git &> /dev/null
!wget -q https://dl.fbaipublicfiles.com/segment_anything/sam_vit_h_4b8939.pth

接下来，我们创建 3 个全局变量：

MODEL_TYPE：要使用的 SAM 架构
CHECKPOINT_PATH：包含模型权重的文件的路径
DEVICE：使用的处理器，“cpu”或“cuda”（如果 GPU 可用）

MODEL_TYPE = "vit_h"
CHECKPOINT_PATH = "/content/sam_vit_h_4b8939.pth"
DEVICE = "cuda" #cpu,cuda

我们现在可以使用 sam_model_registry 函数加载 SAM 模型，指示模型权重：

from segment_anything import sam_model_registry

sam = sam_model_registry[MODEL_TYPE](checkpoint=CHECKPOINT_PATH).to(device=DEVICE)

模型加载后，Meta 为我们提供了两种使用选项：

生成器选项，允许你从图像中获取模型生成的所有掩模
Predictor选项，它允许我们根据提示从图像中获取一个或多个特定的掩模。

我们将在下面几行中探讨这两个选项。

在此之前，让我们从互联网加载一张图像，我们将在该图像上试验我们的模型：

from urllib.request import urlopen
import cv2
import numpy as np
from google.colab.patches import cv2_imshow

resp = urlopen('https://images.unsplash.com/photo-1615948812700-8828458d368a?ixlib=rb-4.0.3&ixid=M3wxMjA3fDB8MHxwaG90by1wYWdlfHx8fGVufDB8fHx8fA%3D%3D&auto=format&fit=crop&w=2072&q=80')
image = np.asarray(bytearray(resp.read()), dtype='uint8')
image = cv2.imdecode(image, cv2.IMREAD_COLOR)
image = cv2.resize(image, (int(image.shape[1]/2.5), int(image.shape[0]/2.5)))

cv2_imshow(image)

在这里插入图片描述

我们的图像包含几个人、一只狗和一些汽车。

现在我们将使用 SAM 和 Generator 选项来分割该图像。

4、生成器

在本节中，我们将使用 SAM 的生成器版本。这将使我们能够获得由于模型对图像的分析而生成的一组掩模。

让我们初始化 SamAutomaticMaskGenerator 对象：

from segment_anything import SamAutomaticMaskGenerator

mask_generator = SamAutomaticMaskGenerator(sam)

接下来，我们使用 generate()函数启动掩码生成：

masks_generated = mask_generator.generate(image)

该函数为检测到的每个对象生成一个掩码以及其他数据。 SAM 实际上生成一组与其检测到的对象相关的信息（以字典形式）。

5、预测结果

我们可以显示每组信息获得的键：

print(masks_generated[0].keys())

输出：

dict_keys(['segmentation', 'area', 'bbox', 'predicted_iou', 'point_coords', 'stability_score', 'crop_box'])

结果是一组 7 条信息。第一个“分割”表示与检测到的对象的位置相对应的像素：如果像素包含对象，则为 True，否则为 False。

掩码可以显示如下：

cv2_imshow(masks_generated[3]['segmentation'].astype(int)*255)

在这里插入图片描述

该集合中的其他信息对应于以下描述：

area：遮罩区域（以像素为单位）
bbox：XYWH 格式的掩模边界框
Predicted_iou：模型预测的掩模质量得分
point_coords：生成此掩码的采样输入点
stable_score：额外的掩模质量分数
Crop_box：用于生成 XYWH 格式的此蒙版的图像裁剪

大多数从业者不会使用此信息，但对于特定情况，重要的是要知道 SAM 不仅生成掩模，而且还生成诸如此类的附加信息。

以下是为上面显示的掩码获得的其余信息：

print('area :', masks_generated[3]['area'])
print('bbox :',masks_generated[3]['bbox'])
print('predicted_iou :',masks_generated[3]['predicted_iou'])
print('point_coords :',masks_generated[3]['point_coords'])
print('stability_score :',masks_generated[3]['stability_score'])
print('crop_box :',masks_generated[3]['crop_box'])

输出：

area : 5200 bbox : [499, 284, 92, 70]
predicted_iou : 1.005275845527649
point_coords : [[582.1875, 318.546875]]
stability_score : 0.981315553188324
crop_box : [0, 0, 828, 551]

我们还可以显示 SAM 生成的掩码数量：

print(len(masks_generated))

输出：

SAM 从我们的图像中总共生成了 111 个掩模。

6、显示预测

使用这篇文章介绍的 draw_masks_fromDict 函数，我们可以绘制图像上生成的所有蒙版：

segmented_image = draw_masks_fromDict(image, masks_generated)

cv2_imshow(segmented_image)

在这里插入图片描述

开始的图像现在包含 SAM 生成的掩模。

在本节中，我们使用 SAM 的生成器版本。这使我们能够从图像生成 111 个掩模。除了掩模之外，SAM 还生成额外的检测信息。为了可视化模型的预测，我们最终在起始图像上绘制了所有掩模。

因此，SAM 使我们能够执行图像分割。然而，我们可以看到生成的掩码是无序的：没有分类来区分不同的掩码。例如，人们的面具并不与单一颜色相关联。因此，无法对所得分段进行排序。这里获得的唯一信息是对象的位置和界限。

此外，生成的掩模可以重叠。事实上，SAM 可以检测其他物体内部的物体。从积极的一面来看，这证明了 SAM 能够检测图像中几乎所有物体。这意味着我们可以分割狗、汽车、人以及其他物体，例如车轮、窗户或裤子。因此，SAM 的生成器版本能够分割图像中的所有对象，甚至是重叠的对象。

7、超越生成器

但是，此功能也具有不利的一面：它增加了给定区域中的预测数量，这可能会破坏目标的实现。例如，如果要检测图像中的人物，则同时检测与其夹克和裤子对应的面具无关紧要。

此外，由于SAM没有在标记数据上进行训练，因此不可能过滤其预测以保留我们感兴趣的预测。这意味着，即使我们使用 SAM 的生成器版本分割数据集中的所有图像，也不可能轻松提取例如人的掩码。因此，SAM 生成器分割图像中所有对象的能力可能不适合解决某些问题。

因此，对于目标对象检测，不适合使用 SAM 的生成器版本。相反，我们需要使用预测器版本。此版本将使我们能够使用 SAM 并提示指定我们的请求和要检测的目标对象。

8、预测器

在本部分中，我们将使用 SAM 的预测器版本。预测器版本将使我们能够检测目标对象。为此，我们将发送 SAM 提示以指定我们要检测的对象。

目前，可以通过两种方式向 SAM 发送提示：

按兴趣点
通过边界框

SAM 可以将针对代表对象的图像像素的兴趣点（x 和 y 坐标）作为输入。然后，由兴趣点指定的对象将使 SAM 能够生成与此对象关联的掩码。

SAM 还可以将分隔图像中对象轮廓的边界框作为输入。根据这些轮廓，SAM 将生成适当的掩码。

注意：“提示”是一个时髦的术语，在大多数情况下，用于指发送到 ChatGPT 的文本请求。但是，如 SAM 所示，提示不仅限于文本请求。它扩展到从业者可以发送到机器学习模型的一组查询。

需要注意的是，虽然该功能目前尚未公开，但 Meta 已通过其细分任何模型为文本请求理解提供了条件。

也就是说，对于本教程的其余部分，我们需要有一个发送到 SAM 的提示。边界框是计算机视觉标准，因此我们将使用它们。

9、使用边界框提示

如果要继续学习本教程，则必须首先具有与要分割的对象关联的边界框。

如果没有图像的边界框，则可以使用 YOLO 模板在几行代码中轻松生成它们。

你可以了解如何使用此模板快速生成自己的边界框。专门针对最新版本 YOLO 的教程在这里等着你。

一旦在我们的图像上使用 YOLO，我们就会得到这样的结果：

image_bboxes = image.copy()

boxes = np.array(results[0].to('cpu').boxes.data)

plot_bboxes(image_bboxes, boxes, score=False)

在这里插入图片描述

注意：结果变量是模型预测的结果。

使用 YOLO 获得的边界框采用以下形式：

print(boxes)

输出：

[[ 495.96 285.65 589.8 356.48 0.89921 2]
[ 270.63 147.99 403.17 496.82 0.79781 0]
…
[ 235.32 279.23 508.93 399.63 0.3193 2]
[ 612.13 303.94 647.61 333.11 0.2854 2]]

前 4 个值表示边界框坐标，第 5 个值表示预测边界框的置信度分数，第 6 个值表示检测到的类。

现在我们有了提示，让我们初始化 SamPredictor 对象：

from segment_anything import SamPredictor

mask_predictor = SamPredictor(sam)

接下来，我们指定要由SAM分析的图像：

mask_predictor.set_image(image)

从这里开始，本教程分为两部分：

单个对象检测
批量对象检测

让我们从第一个选项开始。

10、检测单个对象

为了预测对象的掩码，我们在 predict()函数中告诉 Predictor 与该对象对应的边界框：

mask, _, _ = mask_predictor.predict(
    box=boxes[1][:-2]
)

我们以布尔数组的形式获得一个掩码，指示检测到的对象的位置（如之前在字典的“分割”键中）：如果像素包含对象，则为 True，否则为 False。

我们可以使用这篇文章中描述的 draw_mask函数在图像上绘制此蒙版：
在这里插入图片描述

我们的凸显现在包含 SAM 检测到的掩码。

多亏了给SAM的提示，我们已经能够获得对象的遮罩并将其显示在我们的图像上。

现在让我们看看如何检测与所有边界框对应的掩码。

11、检测多个对象

为了对一组边界框进行预测，我们需要将它们收集到 PyTorch 张量中。

然后我们使用 transform.apply_boxes_torch()来更新我们的对象。

最后，我们使用 predict_torch来预测相应的掩码。

import torch

input_boxes = torch.tensor(boxes[:, :-2], device=mask_predictor.device)

transformed_boxes = mask_predictor.transform.apply_boxes_torch(input_boxes, image.shape[:2])
masks, _, _ = mask_predictor.predict_torch(
    point_coords=None,
    point_labels=None,
    boxes=transformed_boxes,
    multimask_output=False,
)

结果是在一维 (1, 551, 828) 上编码的一批 13 个掩码。

为了更好地操作这个张量，让我们删除第一个不相关的维度：

print(masks.shape)
masks = torch.squeeze(masks, 1)
print(masks.shape)

输出：

torch.Size([13, 1, 551, 828])
torch.Size([13, 551, 828])

在 SAM 上游使用边界框的优点是，我们可以将每个生成的掩码与边界框对应的标签相关联，从而在显示时使用颜色来区分它们。

让我们定义一个与 YOLO 可以预测的类关联的颜色渐变：

COLORS = [(89, 161, 197),(67, 161, 255),(19, 222, 24),(186, 55, 2),(167, 146, 11),(190, 76, 98),(130, 172, 179),(115, 209, 128),(204, 79, 135),(136, 126, 185),(209, 213, 45),(44, 52, 10),(101, 158, 121),(179, 124, 12),(25, 33, 189),(45, 115, 11),(73, 197, 184),(62, 225, 221),(32, 46, 52),(20, 165, 16),(54, 15, 57),(12, 150, 9),(10, 46, 99),(94, 89, 46),(48, 37, 106),(42, 10, 96),(7, 164, 128),(98, 213, 120),(40, 5, 219),(54, 25, 150),(251, 74, 172),(0, 236, 196),(21, 104, 190),(226, 74, 232),(120, 67, 25),(191, 106, 197),(8, 15, 134),(21, 2, 1),(142, 63, 109),(133, 148, 146),(187, 77, 253),(155, 22, 122),(218, 130, 77),(164, 102, 79),(43, 152, 125),(185, 124, 151),(95, 159, 238),(128, 89, 85),(228, 6, 60),(6, 41, 210),(11, 1, 133),(30, 96, 58),(230, 136, 109),(126, 45, 174),(164, 63, 165),(32, 111, 29),(232, 40, 70),(55, 31, 198),(148, 211, 129),(10, 186, 211),(181, 201, 94),(55, 35, 92),(129, 140, 233),(70, 250, 116),(61, 209, 152),(216, 21, 138),(100, 0, 176),(3, 42, 70),(151, 13, 44),(216, 102, 88),(125, 216, 93),(171, 236, 47),(253, 127, 103),(205, 137, 244),(193, 137, 224),(36, 152, 214),(17, 50, 238),(154, 165, 67),(114, 129, 60),(119, 24, 48),(73, 8, 110)]

最后，我们可以使用本文中开发的 draw_masks_fromList函数来绘制我们所有的蒙版，为每个标签关联一种颜色：

segmented_image = draw_masks_fromList(image, masks.to('cpu'), boxes, COLORS)

cv2_imshow(segmented_image)

在这里插入图片描述

我们使用提供的边界框显示了 YOLO 预测的所有遮罩。此外，每个蒙版都根据边界框指示的类别进行着色。这使得区分各种分割对象变得容易。

原文连接：SAM+YOLOv8简明教程 — BimAnt