点击下方卡片,关注“OpenCV与AI深度学习”
视觉/图像重磅干货,第一时间送达!
SAM 2: Segment Anything Model 2
https://github.com/facebookresearch/segment-anything-2
SAM 2使用步骤与代码演示
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 torch
from sam2.build_sam import build_sam2
from 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 torch
from 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 model
model = SAM("sam2_s.pt")
# Display model information (optional)
model.info()
# Segment image or video
results = model('d.jpg') # 图片推理
#results = model('d.jpg') # 视频推理
# Display results
for result in results:
result.show()
图片推理效果:
from ultralytics import ASSETS, SAM
# Load a model
model = SAM("sam2_s.pt")
# Segment with point prompt
results = model("b.jpg", points=[120, 80], labels=[1], device="cpu")
# Display results
for result in results:
result.show()
图片推理效果:
from ultralytics import ASSETS, SAM
# Load a model
model = SAM("sam2_s.pt")
# Segment with bounding box prompt
results = model("a.jpg", bboxes=[30, 10, 283, 267], labels=[1], device="cpu")
# Display results
for result in results:
result.show()
视频推理效果:
最后附上SAM 2与YOLOv8 seg推理时间对比:
下载1:Pytorch常用函数手册