精通使用 SAM 2 进行图像和视频分割

Introduction

本指南将带您了解 Segment Anything Model 2（SAM 2）是什么、如何工作，以及如何利用它对图片和视频中的物体进行分割。它在将物体分割到图片中的表现及灵活性方面达到了最先进的水平，成为各种计算机视觉应用中的重要资源。这直接为您提供详细的、逐步的指导，帮助您设置和使用 SAM 2 进行图像分割。通过本指南，您将能够利用框和点提示为图片生成分割掩码。

Learning Objectives

• 描述 Segment Anything Model 2（SAM 2）在图像和视频分割中的关键特性和应用。
• 成功配置支持 CUDA 的环境，安装必要的依赖项，并克隆 Segment Anything Model 2 的代码库以进行图像分割任务。
• 应用 SAM 2 使用框和点提示生成图像的分割掩码，并有效可视化结果。
• 评估 SAM 2 如何通过实现实时分割、自动化复杂任务以及使更广泛的受众能够轻松创建内容来彻底改变照片和视频编辑。

Table of contents

Prerequisites

在您开始之前，请确保您拥有支持 CUDA 的 GPU 以加快处理速度。此外，请验证您的计算机上是否已安装 Python。本指南假设您对 Python 和图像处理的概念有一些基本了解。

What is SAM 2?

Segment Anything Model 2（SAM 2）是由 Facebook AI Research（FAIR）开发的一种先进的图像分割工具。2024年7月29日，Meta AI 发布了 SAM 2，这是一个先进的图像和视频分割基础模型。SAM 2 使用户能够在图像或视频中提供点或框，以为特定物体生成分割掩码。

单击此处访问它。

Key Features of SAM 2

• 先进的掩码生成：SAM 2 根据用户输入（例如点或边界框）生成高质量的分割掩码。
• 灵活性：该模型同时支持图像和视频分割。
• 速度与效率：得益于 CUDA 的支持，SAM 2 能够快速执行分割任务，适用于实时应用。

Core Components of SAM 2

• 图像编码器：对输入图像进行编码以便处理。
• 提示编码器：将用户提供的点或框转换为模型能够使用的格式。
• 掩码解码器：根据编码后的输入生成最终的分割掩码。

Applications of SAM 2

现在让我们看看 SAM 2 的一些应用：

• 照片和视频编辑：SAM 2 提供精确的物体分割，能够在照片和视频中进行详细的编辑和创意效果。
• 自动驾驶车辆：在自动驾驶中，SAM 2 可以用于实时识别和跟踪行人、车辆和交通标志等物体。
• 医学影像：SAM 2 可以帮助分割医学影像中的解剖结构，辅助诊断和治疗规划。

What is Image Segmentation?

图像分割是一种计算机视觉技术，涉及将图像分成多个段或区域，以简化其分析。每个段表示图像内的不同物体或物体的一部分，从而更容易识别和分析特定元素。

图像分割的类型

• 语义分割：将每个像素分类到预定义的类别中。
• 实例分割：区分同一物体类别的不同实例。
• 全景分割：结合语义分割和实例分割。

Setting Up and Utilizing SAM 2 for Image Segmentation

我们将指导您如何在您的环境中设置 Segment Anything Model 2（SAM 2），并利用其强大的功能进行精确的图像分割任务。从确保您的 GPU 准备就绪到配置模型，并将其应用于真实图像，每个步骤将详细介绍，以帮助您充分利用 SAM 2 的全部潜力。

Step 1: Check GPU Availability and Set Up the Environment

首先，让我们确保您的环境正确设置，从检查 GPU 可用性开始，并设置当前工作目录。

# 检查 GPU 可用性和 CUDA 版本
!nvidia-smi
!nvcc --version
# 导入必要模块
import os
# 设置当前工作目录
HOME = os.getcwd()
print("HOME:", HOME)

Explanation

• !nvidia-smi 和 !nvcc –version：这些命令用于检查您的框架是否包含支持 CUDA 的 GPU，并显示 CUDA 版本。
• os.getcwd()：这个函数获取当前工作目录，可以用于管理文件路径。

第 2 步：克隆 SAM 2 仓库并安装依赖项

接下来，我们需要从 GitHub 克隆 SAM 2 仓库并安装所需的依赖项。

# 克隆 SAM 2 仓库
!git clone https://github.com/facebookresearch/segment-anything-2.git
# 切换到仓库目录
%cd segment-anything-2
# 安装 SAM 2 包
!pip install -e .
# 安装其他包
!pip install supervision jupyter_bbox_widget

说明

• !git clone：将 SAM 2 仓库克隆到本地机器。
• %cd：将目录切换到克隆的仓库。
• !pip install -e .：以可编辑模式安装 SAM 2 包。
• !pip install supervision jupyter_bbox_widget：安装用于可视化和边界框小部件支持的额外包。

第 3 步：下载模型检查点

模型检查点至关重要，因为它们包含 SAM 2 的训练参数。我们将下载多个不同模型尺寸的检查点。

# 创建一个用于存储检查点的目录
!mkdir -p checkpoints
# 下载模型检查点
!wget -q https://dl.fbaipublicfiles.com/segment_anything_2/072824/sam2_hiera_tiny.pt -P checkpoints
!wget -q https://dl.fbaipublicfiles.com/segment_anything_2/072824/sam2_hiera_small.pt -P checkpoints
!wget -q https://dl.fbaipublicfiles.com/segment_anything_2/072824/sam2_hiera_base_plus.pt -P checkpoints
!wget -q https://dl.fbaipublicfiles.com/segment_anything_2/072824/sam2_hiera_large.pt -P checkpoints

说明

• !mkdir -p checkpoints：创建一个用于存储模型检查点的目录。
• !wget -q … -P checkpoints：将模型检查点下载到检查点目录中。不同的检查点代表不同尺寸和能力的模型。

第 4 步：下载示例图片

为了演示，我们将使用一些示例图片。您也可以通过相似的步骤使用自己的图片。

# 创建一个用于存储数据的目录
!mkdir -p data
# 下载示例图片
!wget -q https://media.roboflow.com/notebooks/examples/dog.jpeg -P data
!wget -q https://media.roboflow.com/notebooks/examples/dog-2.jpeg -P data
!wget -q https://media.roboflow.com/notebooks/examples/dog-3.jpeg -P data
!wget -q https://media.roboflow.com/notebooks/examples/dog-4.jpeg -P data

说明

• !mkdir -p data：创建一个用于存储示例图片的目录。
• !wget -q … -P data：将示例图片下载到数据目录中。

第 5 步：设置 SAM 2 模型并加载一张图片

现在，我们将设置 SAM 2 模型，加载一张图片，并准备进行分割。

import cv2
import torch
import numpy as np
import supervision as sv
from sam2.build_sam import build_sam2
from sam2.sam2_image_predictor import SAM2ImagePredictor
from sam2.automatic_mask_generator import SAM2AutomaticMaskGenerator
# 如果可用，启用 CUDA
torch.autocast(device_type="cuda", dtype=torch.bfloat16).__enter__()
if torch.cuda.get_device_properties(0).major >= 8:
torch.backends.cuda.matmul.allow_tf32 = True
torch.backends.cudnn.allow_tf32 = True
# 设置设备为 CUDA
DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
# 定义模型检查点和配置
CHECKPOINT = "checkpoints/sam2_hiera_large.pt"
CONFIG = "sam2_hiera_l.yaml"
# 构建 SAM 2 模型
sam2_model = build_sam2(CONFIG, CHECKPOINT, device=DEVICE, apply_postprocessing=False)
# 创建自动掩码生成器
mask_generator = SAM2AutomaticMaskGenerator(sam2_model)
# 加载一张用于分割的图片
IMAGE_PATH = "/content/WhatsApp Image 2024-08-02 at 14.17.11_2b223e01.jpg"
image_bgr = cv2.imread(IMAGE_PATH)
image_rgb = cv2.cvtColor(image_bgr, cv2.COLOR_BGR2RGB)
# 生成分割掩码
sam2_result = mask_generator.generate(image_rgb)

说明

• CUDA 设置：启用 CUDA 以加速处理，并在可用时将设备设置为 GPU。
• 模型设置：使用指定的配置和检查点构建 SAM 2 模型。
• 图片加载：加载并将示例图片转换为 RGB 格式。
• 掩码生成：使用自动掩码生成器为加载的图片生成分割掩码。

第 6 步：可视化分割掩码

我们将现在可视化由 SAM 2 生成的分割掩码。

# 在图片上注释掩码
mask_annotator = sv.MaskAnnotator(color_lookup=sv.ColorLookup.INDEX)
detections = sv.Detections.from_sam(sam_result=sam2_result)
annotated_image = mask_annotator.annotate(scene=image_bgr.copy(), detections=detections)
# 并排显示原始图像和分割图像
sv.plot_images_grid(
images=[image_bgr, annotated_image],
grid_size=(1, 2),
titles=['source image', 'segmented image']
)

# 提取并绘制各个掩码
masks = [
mask['segmentation']
for mask in sorted(sam2_result, key=lambda x: x['area'], reverse=True)
]
sv.plot_images_grid(
images=masks[:16],
grid_size=(4, 4),
size=(12, 12)
)

解释：

• 掩码注释：在原始图像上注释分割掩码。
• 可视化：将原始图像和分割图像并排绘制，并且还绘制各个掩码。

第7步：使用框提示进行分割

框提示允许我们指定图像中的感兴趣区域以进行分割。

# 定义SAM 2图像预测器
predictor = SAM2ImagePredictor(sam2_model)
# 重新加载图像
image_bgr = cv2.imread(IMAGE_PATH)
image_rgb = cv2.cvtColor(image_bgr, cv2.COLOR_BGR2RGB)
# 为边界框输入编码图像
import base64
def encode_image(filepath):
with open(filepath, 'rb') as f:
image_bytes = f.read()
encoded = str(base64.b64encode(image_bytes), 'utf-8')
return "data:image/jpg;base64," + encoded
# 在Colab中启用自定义小部件管理器
IS_COLAB = True
if IS_COLAB:
from google.colab import output
output.enable_custom_widget_manager()
from jupyter_bbox_widget import BBoxWidget
# 创建一个边界框小部件
widget = BBoxWidget()
widget.image = encode_image(IMAGE_PATH)
# 显示小部件
widget

解释

• 图像预测器：定义SAM 2图像预测器。
• 图像编码：对图像进行编码，以便与边界框小部件一起使用。
• 小部件设置：设置边界框小部件以指定感兴趣的区域。

第8步：获取边界框并执行分割

在指定边界框后，我们可以使用它们生成分割掩码。

# 从小部件获取边界框
boxes = widget.bboxes
boxes = np.array([
[
box['x'],
box['y'],
box['x'] + box['width'],
box['y'] + box['height']
] for box in boxes
])
[{'x': 457, 'y': 341, 'width': 0, 'height': 0, 'label': ''},
{'x': 205, 'y': 79, 'width': 0, 'height': 1, 'label': ''}]
# 设置预测器中的图像
predictor.set_image(image_rgb)
# 使用边界框生成掩码
masks, scores, logits = predictor.predict(
box=boxes,
multimask_output=False
)
# 将掩码转换为二进制格式
masks = np.squeeze(masks)
# 注释和可视化掩码
box_annotator = sv.BoxAnnotator(color=sv.Color.white())
mask_annotator = sv.MaskAnnotator(color_lookup=sv.ColorLookup.INDEX)
detections = sv.Detections(
xyxy=boxes,
mask=masks.astype(bool)
)
source_image = box_annotator.annotate(scene=image_bgr.copy(), detections=detections)
segmented_image = mask_annotator.annotate(scene=image_bgr.copy(), detections=detections)
# 绘制注释图像
sv.plot_images_grid(
images=[source_image, segmented_image],
grid_size=(1, 2),
titles=['源图像', '分割图像']
)

解释

• 边界框：检索使用小部件指定的边界框。
• 掩码生成：使用边界框生成分割掩码。
• 可视化：在原始图像上注释和可视化掩码。

第9步：使用点提示进行分割

点提示允许我们指定单个感兴趣点以进行分割。

# 基于边界框创建点提示
input_point = np.array([
[
box['x'] + (box['width'] // 2),
box['y'] + (box['height'] // 2)
] for box in widget.bboxes
])
input_label = np.array([1] * len(input_point))
# 使用点提示生成掩码
masks, scores, logits = predictor.predict(
point_coords=input_point,
point_labels=input_label,
multimask_output=True
)
# 将掩码转换为二进制格式
masks = np.squeeze(masks)
# 注释和可视化掩码
point_annotator = sv.PointAnnotator(color_lookup=sv.ColorLookup.INDEX)
mask_annotator = sv.MaskAnnotator(color_lookup=sv.ColorLookup.INDEX)
detections = sv.Detections(
xyxy=sv.mask_to_xyxy(masks=masks),
mask=masks.astype(bool)
)
source_image = point_annotator.annotate(scene=image_bgr.copy(), detections=detections)
segmented_image = mask_annotator.annotate(scene=image_bgr.copy(), detections=detections)
# 绘制注释图像
sv.plot_images_grid(
images=[source_image, segmented_image],
grid_size=(1, 2),
titles=['源图像', '分割图像']
)

说明

• 点提示：根据边界框创建点提示。
• 掩膜生成：使用点提示生成分割掩膜。
• 可视化：在原始图像上注释和可视化掩膜。

使用 SAM 2 时需要记住的关键点

现在让我们来看一下以下几个重要的关键点：

彻底改变照片和视频编辑

• 有潜力彻底改变照片和视频编辑行业。
• 未来的增强可能包括更高的精确度、更低的计算需求和先进的人工智能集成。

实时分割和编辑

• 演变可能导致实时分割和编辑功能。
• 允许在视频和图像中进行无缝更改，且只需最少努力。

为所有人提供创作可能性

• 为专业人士和业余爱好者开启新的创作可能性。
• 简化视觉内容的处理、创造惊艳效果和制作高质量媒体。

自动化复杂任务

• 自动化复杂的分割任务。
• 显著加速工作流程，使复杂的编辑变得更易于访问和高效。

使内容创作大众化

• 使高水平的编辑工具可供更广泛的受众使用。
• 赋能讲故事的人，并激励各个领域的创新，包括娱乐、广告和教育。

对 VFX 行业的影响

• 通过简化复杂的流程增强视觉效果（VFX）制作。
• 减少创建复杂 VFX 所需的时间和精力，使更有雄心的项目成为可能，同时提高整体质量。

SAM 2 的令人印象深刻的潜力

segment anything model 2 (SAM 2) 准备通过引入精准度和计算效率的重大进步，彻底改变照片和视频编辑领域。通过整合先进的人工智能能力，SAM 2 将使用户互动变得更加直观，并实现实时分割和编辑，让无缝更改轻松实现。这项突破性技术承诺使内容创作大众化，赋予专业人士和业余爱好者操纵视觉内容的能力，创造惊艳效果，并轻松制作高质量媒体。

随着 SAM 2 自动化复杂分割任务，其将加速工作流程，使复杂的编辑对更广泛的受众可及。这一转变将激励各个行业的创新，从娱乐和广告到教育。在视觉效果（VFX）领域，SAM 2 将简化复杂的流程，减少创建复杂 VFX 所需的时间和精力。这将使更有雄心的项目成为可能，提升视觉叙事的质量，并在 VFX 世界中开启新的创作可能性。

结论

通过遵循本指南，您已经了解如何设置和使用 Segment Anything Model 2 (SAM 2) 进行图像分割，使用框和点提示。SAM 2 提供强大而灵活的工具用于图像中的对象分割，使其成为各种计算机视觉任务的宝贵资产。欢迎您在自己的图像上进行实验，进一步探索 SAM 2 的能力。

关键点总结

• SAM 2 是 Meta AI 开发的先进工具，能够使用框和点提示进行精确和灵活的图像和视频分割。
• 该模型可以通过自动化复杂的分割任务显著增强照片和视频编辑，使其更易于访问和高效。
• 设置 SAM 2 需要一个支持 CUDA 的 GPU，以及对 Python 和图像处理概念的基本理解。
• SAM 2 的能力为专业人士和业余爱好者的内容创作开启新的可能性，提供实时分割和创作控制。
• 该模型有潜力通过使高水平编辑工具大众化，彻底改变包括视觉效果、娱乐、广告和教育等各个行业。