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))

输出:

111

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

猜你喜欢

转载自blog.csdn.net/shebao3333/article/details/132787767
今日推荐