基于SAM的视频标注

在本文中，我们将演示基础模型的应用，例如 Meta 的 Segment Anything 和 YOLOv8，以自动检测、分类和绘制视频中感兴趣对象的蒙版。这是之前指南的后续：使用 Meta 的 Segment Anything 和 YOLOv8 自动分类掩码。在本指南中，我们将自动检测和分割视频中的对象。

视频有很多帧，标记起来很乏味。分割蒙版的标记更加耗时，因为它们在逐帧之间变化很小，每次都需要手动微调。借助基础模型，可以自动执行并显著加快标记过程，从而在更短的时间内标记更多视频数据。这使你可以将宝贵的时间集中在审查上，只需更正 AI 模型的输出即可。

我们将完成一个简单的语义分割任务：在滑板的人周围绘制蒙版。

下面是该过程的高级摘要，我们将在下面逐步完成该过程，并附有代码：

1）加载 YOLOv8、SAM 和 Labelbox Python SDK

2）对于视频的每一帧：

运行对象检测器以生成具有指定类分类的边界框
将边界框作为输入提供给 Meta 的 Segment Anything 模型，该模型将生成分割掩码
以 Labelbox Python SDK 期望的格式准备掩码预测

3）通过预测导入一次性将所有帧上传到 Labelbox

4）打开视频编辑器，像往常一样查看或修改预标签

你可以使用我们的 Colab 笔记本在你的视频上运行上述所有开箱即用的功能。只需加载你的视频，即可在几分钟内获得自动分割的蒙版，并在 Labelbox 中上课！

在本指南中，我们将使用以下视频：

1、加载 YOLOv8

YOLOv8 是一个最先进的对象检测器，它围绕常见对象生成边界框和类。它是 YOLO（You Only Look Once）系列模型的最新版本，拥有一些令人印象深刻的功能。YOLOv8 以其速度和准确性而闻名，使其成为广泛应用的宝贵工具。在这里，我们使用 YOLOv8 来自动检测和定位视频中的滑板人。

import ultralytics
ultralytics.checks()
from ultralytics import YOLO
model = YOLO(f'{HOME}/yolov8n.pt')

# each class id is assigned a different color
colors = np.random.randint(0, 256, size=(len(model.names), 3))
print(model.names)

# Specify which classes you care about. The rest of classes will be filtered out.
chosen_class_ids = [0] # 0 refers to person, as per model.names

2、加载 SAM

Meta 的 SAM 模型是一种最先进的计算机视觉模型，旨在将图像和视频准确地分割成不同的对象。使用先进的深度学习技术，Segment Anything 能够识别和分割图像中的对象，使其成为广泛应用的强大工具。SAM 模型能够根据提示生成分割掩码，包括边界框提示，我们将在下面的代码中使用这些提示。

如需了解 SAM 的编辑体验，请阅读我们的另一篇博文 Auto-Segment 2.0 powered by Meta's Segment Anything Model。

import torch
import matplotlib.pyplot as plt
from segment_anything import sam_model_registry, SamAutomaticMaskGenerator, SamPredictor

DEVICE = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
sam = sam_model_registry["vit_h"](checkpoint=CHECKPOINT_PATH).to(device=DEVICE)
mask_predictor = SamPredictor(sam)

3、加载 Labelbox 的 Python SDK

Labelbox 的 Python SDK 为你提供了创建本体、项目和数据集以及将蒙版上传到视频的简单方法。

import labelbox as lb
import labelbox.types as lb_types

# Create a Labelbox API key for your account by following the instructions here:
# https://docs.labelbox.com/reference/create-api-key
# Then, fill it in here
API_KEY = ""
client = lb.Client(API_KEY)

4、逐帧运行 YOLOv8 和 SAM

在这里，我们在每一帧上运行模型并自动生成蒙版。

cap = cv2.VideoCapture(VIDEO_PATH)

# This will contain the resulting mask predictions for upload to Labelbox
mask_frames = []

frame_num = 1
while cap.isOpened():
  ret, frame = cap.read()
  if not ret:
    break

  # Run frame through YOLOv8 to get detections
  detections = model.predict(frame, conf=0.7)
 
  # Run frame and detections through SAM to get masks
  transformed_boxes = mask_predictor.transform.apply_boxes_torch(detections[0].boxes.xyxy, list(get_video_dimensions(cap)))
  mask_predictor.set_image(frame)
  masks, scores, logits = mask_predictor.predict_torch(
    boxes = transformed_boxes,
    multimask_output=False,
    point_coords=None,
    point_labels=None
  )

  # Combine mask predictions into a single mask, each with a different color
  class_ids = detections[0].boxes.cpu().cls
  merged_with_colors = add_color_to_mask(masks[0][0], colors[int(class_ids[0])]).astype(np.uint8)
  for i in range(1, len(masks)):
    curr_mask_with_colors = add_color_to_mask(masks[i][0], colors[int(class_ids[i])])
    merged_with_colors = np.bitwise_or(merged_with_colors, curr_mask_with_colors)

  # Upload multi-colored combined mask to temp location
  # to get temp instance uri
  instance_uri = get_instance_uri(client, global_key, merged_colored_mask)

  # Create MaskFrame object to be uploaded to Labelbox
  mask_frame = lb_types.MaskFrame(index=frame_num, instance_uri=instance_uri)
  mask_frames.append(mask_frame)

  frame_num += 1

cap.release()

5、将预测的蒙版作为预标签上传到 Labelbox

预测的蒙版可以通过我们的 SDK 轻松无缝地集成到 Labelbox 中。

在视频数据上创建分割掩码可能既繁琐又耗时。利用 Labelbox 中基础模型的强大功能，你可以在几分钟内轻松生成带有分类的蒙版。现在，无需花费数小时标记视频数据，而是可以加快视频标记速度，不仅可以缩短上市时间，还可以降低开发模型的成本。

原文链接：Using Meta's Segment Anything (SAM) model on video with Labelbox's model-assisted labeling

BimAnt翻译整理，转载请标明出处