支持图片和视频分割，SAM2最新分割一切大模型分享

置顶杰说新技术

于 2024-09-06 06:00:00 发布

阅读量393

点赞数 1

分类专栏： AIGC 文章标签： AIGC 人工智能

本文链接：https://blog.csdn.net/m0_71062934/article/details/141790639

版权

AIGC 专栏收录该内容

37 篇文章 1 订阅

订阅专栏

Segment Anything Model 2（简称SAM 2）是由Meta（Facebook AI）开发的最新一代图像和视频分割模型。

SAM2能够实现对静态图像和动态视频中的对象进行实时、可提示的分割，将图像与视频分割功能整合到了同一个系统中。

SAM2的一个核心特点是其交互式分割过程，用户可以通过点击选择和细化目标对象，模型会根据这些提示自动将分割传播到视频的后续帧。

此外，SAM2引入了流式记忆模块，这使得模型能够利用先前帧的信息来辅助当前帧的分割任务。

与第一代模型相比，SAM2在多个方面实现了显著的改进，包括支持视频分割、实时处理任意长视频、Zero-shot泛化、提高分割和追踪的准确性，以及解决遮挡问题等。

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

一、环境安装

1、python环境

建议安装python版本在3.10以上。

2、pip库安装

cd segment-anything-2

pip install -e . -i https://pypi.tuna.tsinghua.edu.cn/simple

pip install -U ultralytics -i https://pypi.tuna.tsinghua.edu.cn/simple

3、模型下载：

cd checkpoints && ./download_ckpts.sh

二、功能测试

1、运行测试：

（1）通过sam2接口进行图片分割测试

import numpy as np
import torch
from PIL import Image
import cv2
from sam2.build_sam import build_sam2
from sam2.automatic_mask_generator import SAM2AutomaticMaskGenerator
import time

# Enable CUDA optimizations
if torch.cuda.is_available():
    torch.autocast(device_type="cuda", dtype=torch.float16).__enter__()

    if torch.cuda.get_device_properties(0).major >= 8:
        torch.backends.cuda.matmul.allow_tf32 = True
        torch.backends.cudnn.allow_tf32 = True

def apply_color_mask(image, mask, color, color_dark=0.5):
    """Apply colored mask to the image."""
    for c in range(3):
        image[:, :, c] = np.where(mask == 1, image[:, :, c] * (1 - color_dark) + color_dark * color[c], image[:, :, c])
    return image

def main():
    sam2_checkpoint = "checkpoints/sam2_hiera_large.pt"
    model_cfg = "sam2_hiera_l.yaml"

    # Load image
    try:
        image = Image.open('image.jpg')
    except FileNotFoundError:
        print("Image file not found.")
        return

    image = np.array(image.convert("RGB"))

    # Load SAM2 Model
    sam2 = build_sam2(model_cfg, sam2_checkpoint, device='cuda', apply_postprocessing=False)

    # Initialize mask generator
    mask_generator = SAM2AutomaticMaskGenerator(sam2)

    # Generate masks
    start = time.time()
    masks = mask_generator.generate(image)
    print(f"sam2 infer: {time.time() - start:.3f}s")

    # Apply masks and save result
    image_select = image.copy()
    for mask_info in masks:
        color = tuple(np.random.randint(0, 256, 3).tolist())
        selected_mask = mask_info['segmentation']
        image_select = apply_color_mask(image_select, selected_mask, color)

    cv2.imwrite("result.jpg", image_select)
    print("Result saved to result.jpg")

if __name__ == "__main__":
    main()

未完......

更多详细的欢迎关注：杰哥新技术