【保姆级教程】SAM2（Segemnt Anything2）实现视频单目标与多目标的一键分割与跟踪：无需模型训练

阿_旭

已于 2024-08-12 15:26:48 修改

阅读量1.2k

点赞数 7

分类专栏：前沿视觉模型指南文章标签： SAM2 视频分割目标分割与跟踪深度学习人工智能

于 2024-08-12 15:13:42 首次发布

本文链接：https://blog.csdn.net/qq_42589613/article/details/141130514

版权

前沿视觉模型指南专栏收录该内容

14 篇文章 14 订阅

订阅专栏

《------往期经典推荐------》

一、【100个深度学习实战项目】【链接】，持续更新~~

二、机器学习实战专栏【链接】，已更新31期，欢迎关注，持续更新中~~
三、深度学习【Pytorch】专栏【链接】
四、【Stable Diffusion绘画系列】专栏【链接】
五、YOLOv8改进专栏【链接】，持续更新中~~
六、YOLO性能对比专栏【链接】，持续更新中~

《------正文------》

引言

本文主要是介绍如何使用最新的SAM2分割一切模型，实现对视频的单目标与多目标的一键分割与跟踪。包含详细的步骤与代码说明。
视频单目标分割效果如下：
在这里插入图片描述

视频多目标分割效果如下：
在这里插入图片描述

环境配置

SAM2项目地址:https://github.com/facebookresearch/segment-anything-2

下载源码后，通过pip install --no-build-isolation -e .命令安装项目环境。在安装和运行时会出现各种坑，我已经整理好，并写了一篇文章供小伙伴们参考，帮助大家避坑。文章将《SAM2环境配置问题汇总：CUDA_HOME environment variable is not set；RuntimeError: No available kernel. Aborting》

SAM2一键分割与跟踪视频步骤

下面介绍如何使用 SAM 2 在视频中进行交互式分割。它将涵盖以下内容：

1.通过在一帧上添加提示点创建masklets【时空掩码】

2.在整个视频中传播点击以获得掩码

3.单目标跟踪分割、多目标跟踪与分割

我们使用术语 segment 或 mask 来指代单帧上对象的模型预测，并使用 masklet【时空掩码】来指代整个视频中的时空掩码。

导入相关库

import os
import torch
import numpy as np
import matplotlib.pyplot as plt
from PIL import Image

# 张量计算时的精度设置
torch.autocast(device_type="cuda", dtype=torch.bfloat16).__enter__()

if torch.cuda.get_device_properties(0).major >= 8:
    # turn on tfloat32 for Ampere GPUs (https://pytorch.org/docs/stable/notes/cuda.html#tensorfloat-32-tf32-on-ampere-devices)
    torch.backends.cuda.matmul.allow_tf32 = True
    torch.backends.cudnn.allow_tf32 = True

加载SAM2模型

官方提供了4种大小的SAM2模型，下载地址如下，我这里直接使用最小一个sam2_hiera_tiny.pt进行后续的分割演示。

模型下载地址：
https://dl.fbaipublicfiles.com/segment_anything_2/072824/sam2_hiera_tiny.pt
https://dl.fbaipublicfiles.com/segment_anything_2/072824/sam2_hiera_small.pt
https://dl.fbaipublicfiles.com/segment_anything_2/072824/sam2_hiera_base_plus.pt
https://dl.fbaipublicfiles.com/segment_anything_2/072824/sam2_hiera_large.pt

from sam2.build_sam import build_sam2_video_predictor

# 模型路径与配置文件
sam2_checkpoint = "checkpoints/sam2_hiera_tiny.pt"
model_cfg = "sam2_hiera_t.yaml"

predictor = build_sam2_video_predictor(model_cfg, sam2_checkpoint)

定义相关显示函数

# 定义相关显示函数
def show_mask(mask, ax, obj_id=None, random_color=False):
    if random_color:
        color = np.concatenate([np.random.random(3), np.array([0.6])], axis=0)
    else:
        cmap = plt.get_cmap("tab10")
        cmap_idx = 0 if obj_id is None else obj_id
        color = np.array([*cmap(cmap_idx)[:3], 0.6])
    h, w = mask.shape[-2:]
    mask_image = mask.reshape(h, w, 1) * color.reshape(1, 1, -1)
    ax.imshow(mask_image)


def show_points(coords, labels, ax, marker_size=200):
    pos_points = coords[labels==1]
    neg_points = coords[labels==0]
    ax.scatter(pos_points[:, 0], pos_points[:, 1], color='green', marker='*', s=marker_size, edgecolor='white', linewidth=1.25)
    ax.scatter(neg_points[:, 0], neg_points[:, 1], color='red', marker='*', s=marker_size, edgecolor='white', linewidth=1.25)

选择一个视频，将视频保存为图片帧

我们首先需要将分割的视频保存为图片帧。即将视频存储为 JPEG 帧列表，文件名类似于 <frame_index>.jpg。
使用的是ffmpeg（https://ffmpeg.org/）进行视频帧保存的。命令如下：

ffmpeg -i test1.mp4 -q:v 2 -start_number 0 output/%05d.jpg

其中 -q：v 生成高质量的 JPEG 帧，-start_number 0 要求 ffmpeg 从 00000.jpg 开始 JPEG 文件。

ffmpeg使用方法参考我的另一篇博客： https://a-xu-ai.blog.csdn.net/article/details/141004104

保存结果如下：
在这里插入图片描述

读取所有视频帧

#  读取包含所有视频帧的图片路径
video_dir = "./videos/output"
if not os.path.exists(video_dir):
    os.makedirs(video_dir)

# scan all the JPEG frame names in this directory
frame_names = [
    p for p in os.listdir(video_dir)
    if os.path.splitext(p)[-1] in [".jpg", ".jpeg", ".JPG", ".JPEG"]
]
frame_names.sort(key=lambda p: int(os.path.splitext(p)[0]))

# take a look the first video frame
frame_idx = 0
plt.figure(figsize=(12, 8))
plt.title(f"frame {frame_idx}")
plt.imshow(Image.open(os.path.join(video_dir, frame_names[frame_idx])))

展示视频第一帧图片：
在这里插入图片描述

初始化推理状态

SAM 2 需要有状态推理才能进行交互式视频分割，因此我们需要在此视频上初始化推理状态。

在初始化过程中，它会以 video_path 格式加载所有 JPEG 帧，并将其像素存储在 inference_state中（如下面的进度条所示）。

inference_state = predictor.init_state(video_path=video_dir)

frame loading (JPEG): 100%|██████████████████████████████████████████████████████████| 300/300 [00:07<00:00, 38.12it/s]
D:\1MyShare\3.SAM2\segment-anything-2-main\sam2\modeling\backbones\hieradet.py:72: UserWarning: 1Torch was not compiled with flash attention. (Triggered internally at ..\aten\src\ATen\native\transformers\cuda\sdp_utils.cpp:455.)
  x = F.scaled_dot_product_attention(

示例1：分割与追踪单个目标

第 1 步：在模型中添加分割提示点

首先，分割最中间的赛车。我们需要在模型中添加用于目标分割的提示点。

在这里，我们通过将它们的坐标和标签发送到 add_new_points API 中，在标签 1 处进行正向点击。

注意：标签 1 表示正面点击（添加区域），而标签 0 表示负面点击（删除区域）。

# 视频帧序号
ann_frame_idx = 0  
# 为每一个分割目标建立一个唯一id, 任意整数
ann_obj_id = 1  

# 添加一个前景点
points = np.array([[639, 445]], dtype=np.float32)
# 1表示前景点，0表示背景点
labels = np.array([1], np.int32)
_, out_obj_ids, out_mask_logits = predictor.add_new_points(
    inference_state=inference_state,
    frame_idx=ann_frame_idx,
    obj_id=ann_obj_id,
    points=points,
    labels=labels,
)

# 显示分割结果
plt.figure(figsize=(12, 8))
plt.title(f"frame {ann_frame_idx}")
plt.imshow(Image.open(os.path.join(video_dir, frame_names[ann_frame_idx])))
show_points(points, labels, plt.gca())
show_mask((out_mask_logits[0] > 0.0).cpu().numpy(), plt.gca(), obj_id=out_obj_ids[0])

在这里插入图片描述

第 2 步：在视频中传播获取 masklet(时空掩码) 的提示

为了在整个视频中获取 masklet【时空掩码】，我们使用 propagate_in_video API 传播提示。

# 在整个视频中运行传播，并在字典中存储结果
video_segments = {}  # Video_segments包含每帧分割结果
for out_frame_idx, out_obj_ids, out_mask_logits in predictor.propagate_in_video(inference_state):
    video_segments[out_frame_idx] = {
        out_obj_id: (out_mask_logits[i] > 0.0).cpu().numpy()
        for i, out_obj_id in enumerate(out_obj_ids)
    }

# 每隔几帧显示分割结果
vis_frame_stride = 15
plt.close("all")
for out_frame_idx in range(0, len(frame_names), vis_frame_stride):
    plt.figure(figsize=(6, 4))
    plt.title(f"frame {out_frame_idx}")
    plt.imshow(Image.open(os.path.join(video_dir, frame_names[out_frame_idx])))
    for out_obj_id, out_mask in video_segments[out_frame_idx].items():
        show_mask(out_mask, plt.gca(), obj_id=out_obj_id)

propagate in video: 100%|████████████████████████████████████████████████████████████| 300/300 [00:26<00:00, 11.20it/s]

部分帧图片分割结果展示：
在这里插入图片描述

第 3 步：保存每一帧分割结果

对上述每一帧的分割结果进行保存。

import cv2

colors = [np.array([0, 128, 255], dtype=np.uint8),
         np.array([255, 128, 0], dtype=np.uint8)]

def add_mask2(image, mask, color_id):
    # 单通道
    int_mask = mask.astype(np.uint8)
    int_mask_3d = np.dstack((int_mask, int_mask, int_mask))
    # 创建橙色掩码图像
    # 橙色 (B, G, R)
    mask_color = colors[color_id]
    mask = np.full_like(image, mask_color)

    # 将掩码应用于橙色掩码图像
    mask[int_mask == 0] = 0

    # 使用 cv2.addWeighted 叠加原始图像和橙色掩码图像
    alpha = 0.6  # 原始图像权重
    beta = 1 - 0.6  # 橙色掩码权重
    gamma = 0  # 偏移量

    # 使用掩码矩阵来控制叠加
    res = cv2.addWeighted(image, alpha, mask, beta, gamma, dtype=cv2.CV_8U)

    # 将mask中为黑色部分保留原图,0的区域为True, 非零区域为False
    # 获取黑色区域
    black_areas = int_mask_3d == 0
    res[black_areas] = image[black_areas]
    return res


for out_frame_idx in range(0, len(frame_names)):
    color_num = 0
    now_img = cv2.imread(os.path.join(video_dir, frame_names[out_frame_idx]))
    for out_obj_id, out_mask in video_segments[out_frame_idx].items():
        now_img = add_mask2(now_img, np.squeeze(out_mask),color_num)
        color_num += 1
    mask_out_name = 'masks2/{}.jpg'.format(str(out_frame_idx).zfill(5))
    cv2.imwrite(mask_out_name, now_img)

保存结果如下：
在这里插入图片描述

第4步：合成视频

将上述保存的每一帧分割结果，合成为最终的分割视频。合成命令如下：

ffmpeg -framerate 30 -i frame_%05d.jpg -c:v libx264 -pix_fmt yuv420p output.mp4

视频分割结果如下：
在这里插入图片描述

示例 2：同时分割多个对象

注意：如果您使用此inference_state运行过任何以前的跟踪，请先通过reset_state重置它。

predictor.reset_state(inference_state)

第一步：添加目标

SAM 2 还可以同时分割和跟踪两个或多个对象。当然，一种方法是逐个完成它们。但是，将它们批处理在一起会更有效（例如，这样我们就可以在对象之间共享图像特征以降低计算成本）。

这一次，让我们在这个视频中专注于物体部分并分割两个赛车。在这里，我们为这两个对象添加提示，并为每个对象分配一个唯一的对象 ID。

prompts = {}  # 保存所有提示点信息

在第 0 帧的（x， y） = （200， 300）处单击第一个对象（左孩子的衬衫）。

我们将其分配给对象 ID 2（它可以是任意整数，并且只需要对要跟踪的每个对象是唯一的），该 ID 被传递给 add_new_points API 以区分我们正在单击的对象。

# 添加第一个目标
ann_frame_idx = 0  # the frame index we interact with
ann_obj_id = 2  # give a unique id to each object we interact with (it can be any integers)

# Let's add a positive click at (x, y) = (200, 300) to get started on the first object
points = np.array([[639, 445]], dtype=np.float32)
# for labels, `1` means positive click and `0` means negative click
labels = np.array([1], np.int32)
prompts[ann_obj_id] = points, labels
_, out_obj_ids, out_mask_logits = predictor.add_new_points(
    inference_state=inference_state,
    frame_idx=ann_frame_idx,
    obj_id=ann_obj_id,
    points=points,
    labels=labels,
)

# 添加第二个目标
ann_frame_idx = 0  # the frame index we interact with
ann_obj_id = 3  # give a unique id to each object we interact with (it can be any integers)

points = np.array([[356, 402]], dtype=np.float32)
# for labels, `1` means positive click and `0` means negative click
labels = np.array([1], np.int32)
prompts[ann_obj_id] = points, labels

# `add_new_points` returns masks for all objects added so far on this interacted frame
_, out_obj_ids, out_mask_logits = predictor.add_new_points(
    inference_state=inference_state,
    frame_idx=ann_frame_idx,
    obj_id=ann_obj_id,
    points=points,
    labels=labels,
)

# show the results on the current (interacted) frame on all objects
plt.figure(figsize=(12, 8))
plt.title(f"frame {ann_frame_idx}")
plt.imshow(Image.open(os.path.join(video_dir, frame_names[ann_frame_idx])))
show_points(points, labels, plt.gca())
for i, out_obj_id in enumerate(out_obj_ids):
    show_points(*prompts[out_obj_id], plt.gca())
    show_mask((out_mask_logits[i] > 0.0).cpu().numpy(), plt.gca(), obj_id=out_obj_id)

在这里插入图片描述
发现通过单个点提示进行分割，效果不佳，此时，我们可以通过多个提示点来对每个分割目标进行提示分割。

优化分割结果

每个目标添加多个提示点。

# 添加第一个目标
ann_frame_idx = 0  # the frame index we interact with
ann_obj_id = 2  # give a unique id to each object we interact with (it can be any integers)

# Let's add a positive click at (x, y) = (200, 300) to get started on the first object
points = np.array([[639, 445],[569,451],[441,561],[829,595]], dtype=np.float32)
# for labels, `1` means positive click and `0` means negative click
labels = np.array([1,1,1,1], np.int32)
prompts[ann_obj_id] = points, labels
_, out_obj_ids, out_mask_logits = predictor.add_new_points(
    inference_state=inference_state,
    frame_idx=ann_frame_idx,
    obj_id=ann_obj_id,
    points=points,
    labels=labels,
)

# 添加第二个目标
ann_frame_idx = 0  # the frame index we interact with
ann_obj_id = 3  # give a unique id to each object we interact with (it can be any integers)

points = np.array([[377, 415],[475,403]], dtype=np.float32)
# for labels, `1` means positive click and `0` means negative click
labels = np.array([1,1], np.int32)
prompts[ann_obj_id] = points, labels

# `add_new_points` returns masks for all objects added so far on this interacted frame
_, out_obj_ids, out_mask_logits = predictor.add_new_points(
    inference_state=inference_state,
    frame_idx=ann_frame_idx,
    obj_id=ann_obj_id,
    points=points,
    labels=labels,
)

# show the results on the current (interacted) frame on all objects
plt.figure(figsize=(12, 8))
plt.title(f"frame {ann_frame_idx}")
plt.imshow(Image.open(os.path.join(video_dir, frame_names[ann_frame_idx])))
show_points(points, labels, plt.gca())
for i, out_obj_id in enumerate(out_obj_ids):
    show_points(*prompts[out_obj_id], plt.gca())
    show_mask((out_mask_logits[i] > 0.0).cpu().numpy(), plt.gca(), obj_id=out_obj_id)

在这里插入图片描述
可以发现，每个目标添加多个提示点，可以达到很好的目标分割效果，接下来我们继续后面的步骤。

第 2 步：在视频中传播获取 masklets 的提示

现在，我们传播两个对象的提示，以便在整个视频中获取它们的掩码。

注意：当有多个对象时，propagate_in_video API 将返回每个对象的掩码列表。

# run propagation throughout the video and collect the results in a dict
video_segments = {}  # video_segments contains the per-frame segmentation results
for out_frame_idx, out_obj_ids, out_mask_logits in predictor.propagate_in_video(inference_state):
    video_segments[out_frame_idx] = {
        out_obj_id: (out_mask_logits[i] > 0.0).cpu().numpy()
        for i, out_obj_id in enumerate(out_obj_ids)
    }

# render the segmentation results every few frames
vis_frame_stride = 15
plt.close("all")
for out_frame_idx in range(0, len(frame_names), vis_frame_stride):
    plt.figure(figsize=(6, 4))
    plt.title(f"frame {out_frame_idx}")
    plt.imshow(Image.open(os.path.join(video_dir, frame_names[out_frame_idx])))
    for out_obj_id, out_mask in video_segments[out_frame_idx].items():
        show_mask(out_mask, plt.gca(), obj_id=out_obj_id)

propagate in video: 100%|████████████████████████████████████████████████████████████| 300/300 [00:40<00:00,  7.44it/s]

部分帧的分割结果显示如下：
在这里插入图片描述

可以发现，模型可以对视频的每一帧的指定的两个赛车目标进行很好的分割。

第 3 步：保存每帧分割图片

for out_frame_idx in range(0, len(frame_names)):
    color_num = 0
    now_img = cv2.imread(os.path.join(video_dir, frame_names[out_frame_idx]))
    for out_obj_id, out_mask in video_segments[out_frame_idx].items():
        now_img = add_mask2(now_img, np.squeeze(out_mask),color_num)
        color_num += 1
    mask_out_name = 'masks3/{}.jpg'.format(str(out_frame_idx).zfill(5))
    cv2.imwrite(mask_out_name, now_img)

保存结果如下：
在这里插入图片描述

第4步：将图片视频合成

将上述每一帧的多个目标分割结果，保存为视频，命令如下：

ffmpeg -framerate 30 -i %05d.jpg -c:v libx264 -pix_fmt yuv420p output.mp4

视频分割结果如下：
在这里插入图片描述

视频教程见：
https://www.bilibili.com/video/BV1dDeFeBEcV/?vd_source=b80770ed7412867134a95927c288f50b

好了，这篇文章就介绍到这里，如果对你有帮助，感谢点赞关注！

资料获取

本文的相关资料都已打包好上传，需要的小伙伴可以自行获取学习。
在这里插入图片描述

好了，这篇文章就介绍到这里，喜欢的小伙伴感谢给点个赞和关注，更多精彩内容持续更新~~
关于本篇文章大家有任何建议或意见，欢迎在评论区留言交流！

阿_旭

关注

7
点赞
踩
22

收藏

觉得还不错? 一键收藏
打赏
0
评论
复制链接

分享到 QQ

分享到新浪微博

扫一扫

专栏目录