SAM2(Segment Anything Model 2)新一代分割一切大模型介绍与使用(步骤 + 代码)

Color Space OpenCV与AI深度学习 2024年08月05日 07:45

点击下方卡片,关注“OpenCV与AI深度学习

视觉/图像重干货,第一时间送达!

图片

SAM 2: Segment Anything Model 2

    Segment Anything Model 2(SAM 2)是由Meta公司发布的一个先进的图像和视频分割模型。它是Segment Anything Model(SAM)的升级版本,SAM是Meta的FAIR实验室发布的一款用于图像分割的基础模型,能够在给定提示的情况下生成高质量的对象掩模。
    项目地址:
https://github.com/facebookresearch/segment-anything-2
图片
    特点
    • 准确性提升:SAM 2相比原始的SAM模型在分割精度上有所提高。
    • 速度加快:SAM 2的处理速度提高了大约六倍,这意味着它可以更快地生成分割掩模。
    • 支持视频分割:除了图像分割之外,SAM 2还支持视频中的对象分割。
    • 实时处理:SAM 2可以实现实时处理,这使得它非常适合于需要快速响应的应用场景,如增强现实(AR)和虚拟现实(VR)应用。
    • Zero-Shot泛化:SAM 2具有良好的zero-shot迁移能力,即可以在未见过的数据上工作而不需要额外的训练。
    • 可提示的模型架构:SAM 2继承了SAM的特性,可以根据不同的提示(如点、框、甚至是文本)来生成分割结果。
    结构
    • 编码器-解码器架构:SAM 2很可能会继续使用编码器-解码器架构,其中编码器负责提取特征,解码器则用于生成分割掩模。
    • 高效网络设计:为了达到更高的处理速度,SAM 2可能采用了优化过的网络结构或计算效率更高的组件。
    • 适应性强的分割头:模型可能包含了一个高度灵活的分割头,能够根据不同的提示生成相应的掩模。
    • 多模态输入支持:除了传统的图像输入外,SAM 2还可能支持视频帧序列作为输入,以实现视频分割。
    • 训练数据集:SAM 2的训练数据集可能包含了大量多样化的图像和视频样本,以确保模型的泛化能力和鲁棒性。
    应用场景
    • 增强现实(AR)和虚拟现实(VR):SAM 2可以用于实时分割用户周围的环境,从而增强用户体验。
    • 自动驾驶:在自动驾驶系统中,SAM 2可以帮助车辆识别和理解道路场景中的不同元素。
    • 医学影像分析:SAM 2可以用来自动分割医学影像中的器官或病变区域。

SAM 2使用步骤与代码演示

    方法一:使用github项目
https://github.com/facebookresearch/segment-anything-2

    使用前需要先安装 SAM 2。代码需要python>=3.10,以及torch>=2.3.1和。请按照此处的torchvision>=0.18.1说明安装 PyTorch 和 TorchVision 依赖项。您可以使用以下方式在 GPU 机器上安装 SAM 2:

git clone https://github.com/facebookresearch/segment-anything-2.git
cd segment-anything-2; pip install -e .
    下载模型:
图片

    图像预测:

import torchfrom sam2.build_sam import build_sam2from sam2.sam2_image_predictor import SAM2ImagePredictor
checkpoint = "./checkpoints/sam2_hiera_large.pt"model_cfg = "sam2_hiera_l.yaml"predictor = SAM2ImagePredictor(build_sam2(model_cfg, checkpoint))
with torch.inference_mode(), torch.autocast("cuda", dtype=torch.bfloat16): predictor.set_image(<your_image>) masks, _, _ = predictor.predict(<input_prompts>)

    视频预测:

import torchfrom sam2.build_sam import build_sam2_video_predictor
checkpoint = "./checkpoints/sam2_hiera_large.pt"model_cfg = "sam2_hiera_l.yaml"predictor = build_sam2_video_predictor(model_cfg, checkpoint)
with torch.inference_mode(), torch.autocast("cuda", dtype=torch.bfloat16): state = predictor.init_state(<your_video>)
# add new prompts and instantly get the output on the same frame frame_idx, object_ids, masks = predictor.add_new_points(state, <your_prompts>):
# propagate the prompts to get masklets throughout the video for frame_idx, object_ids, masks in predictor.propagate_in_video(state): ...

    方法二:基于ultralytics包的封装来调用

    【1】安装必要的包。安装ultralytics并确保其版本>=8.2.70,torch版本也需>=2.0或使用最新版本

pip install -U ultralytics

    【2】下载分割模型。下面网址中提供了4个模型,大家可以根据自己需要下载:

https://docs.ultralytics.com/models/sam-2/#how-can-i-use-sam-2-for-real-time-video-segmentation
图片

    【3】全局目标分割。

from ultralytics import ASSETS, SAM
# Load a modelmodel = SAM("sam2_s.pt")
# Display model information (optional)model.info()
# Segment image or videoresults = model('d.jpg'# 图片推理#results = model('d.jpg') # 视频推理
# Display resultsfor result in results: result.show()

    图片推理效果:

图片
图片
图片
图片
    【4】引导分割制定目标(指定点或者矩形区域)。
from ultralytics import ASSETS, SAM
# Load a modelmodel = SAM("sam2_s.pt")
# Segment with point promptresults = model("b.jpg", points=[120, 80], labels=[1], device="cpu")
# Display resultsfor result in results: result.show()

    图片推理效果:

图片

图片

from ultralytics import ASSETS, SAM
# Load a modelmodel = SAM("sam2_s.pt")
# Segment with bounding box promptresults = model("a.jpg", bboxes=[30, 10, 283, 267], labels=[1], device="cpu")
# Display resultsfor result in results: result.show()
图片
图片

图片

图片

图片

图片

    视频推理效果:

图片

图片

    最后附上SAM 2与YOLOv8 seg推理时间对比:

图片

—THE END—

下载1:Pytorch常用函数手册

在「OpenCV与AI深度学习公众号后台回复:Pytorch函数手册即可下载学习全网第一份Pytorch函数常用手册,包括Tensors介绍、基础函数介绍、数据处理函数、优化函数、CUDA编程、多处理等十四章内容。

下载2:145个OpenCV实例应用代码
在「OpenCV与AI深度学习」公众号后台回复:OpenCV145即可下载学习145个OpenCV实例应用代码(Python和C++双语言实现)。