Segment Anything Model 2(SAM 2)作为Meta公司发布的Segment Anything Model(SAM)的升级版本,在图像和视频分割领域展现出了显著的优点和特性。
论文连接:https://arxiv.org/pdf/2408.00714
Demo: https://sam2.metademolab.com Code: https://github.com/facebookresearch/segment-anything-2 Website: https://ai.meta.com/sam2
分割精度提升:
处理速度加快:
支持视频分割:
实时处理能力:
Zero-Shot泛化能力:
广泛的训练数据:
可提示的模型架构:
编码器-解码器架构:
高效网络设计:
适应性强的分割头:
多模态输入支持:
记忆注意模块:
PVS任务支持:
开源与社区支持:
SAM 2以其高精度、高速度、广泛的支持能力以及强大的zero-shot泛化能力,在图像和视频分割领域展现出了显著的优势和广泛的应用前景。
实现SAM推理有两种方法,一种是直接使用官方的SAM2模型,另一种使用Ultralytics 。
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 .
然后,下载模型: sam2_hiera_tiny.pt:
https://dl.fbaipublicfiles.com/segment_anything_2/072824/sam2_hiera_tiny.pt
sam2_hiera_small.pt:
https://dl.fbaipublicfiles.com/segment_anything_2/072824/sam2_hiera_small.pt
sam2_hiera_base_plus.pt:
https://dl.fbaipublicfiles.com/segment_anything_2/072824/sam2_hiera_base_plus.pt
sam2_hiera_large.pt:
https://dl.fbaipublicfiles.com/segment_anything_2/072824/sam2_hiera_large.pt
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_or_box(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并确保其版本>=8.2.70,torch版本也需>=2.0或直接使用最新版本
pip install -U ultralytics
下载分割模型。链接如下:
SAM 2 tiny链接:
https://github.com/ultralytics/assets/releases/download/v8.2.0/sam2_t.pt
SAM 2 small:
https://github.com/ultralytics/assets/releases/download/v8.2.0/sam2_s.pt
SAM 2 base
https://github.com/ultralytics/assets/releases/download/v8.2.0/sam2_b.pt
SAM 2 large
https://github.com/ultralytics/assets/releases/download/v8.2.0/sam2_l.pt
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('car.jpg') # 图片推理
# Display results
for result in results:
result.show()
from ultralytics import SAM
# Load a model
model = SAM("sam2_b.pt")
# Display model information (optional)
model.info()
# Segment with bounding box prompt
results = model("path/to/image.jpg", bboxes=[100, 100, 200, 200])
# Segment with point prompt
results = model("path/to/image.jpg", points=[150, 150], labels=[1])
for result in results:
result.show()
from ultralytics import SAM
# Load a model
model = SAM("sam2_b.pt")
# Display model information (optional)
model.info()
# Run inference
model("path/to/video.mp4")
在这里,我们将 Meta 的最小 SAM 2 模型 SAM2-t 与 Ultralytics 的最小分割模型 YOLOv8n-seg 进行比较:
这种比较显示了模型大小和速度之间的数量级差异。虽然 SAM 具有自动分割的独特能力,但它并不是 YOLOv8 分割模型的直接竞争对手,因为 YOLOv8 模型的体积更小、速度更快、效率更高。
自动标注是SAM 2的一个强大功能,它使用户能够利用预训练模型快速且准确地生成分割数据集。这一功能在无需大量手动工作的情况下创建大型、高质量数据集时尤其有用。
使用SAM 2自动标注您的数据集,代码如下:
from ultralytics.data.annotator import auto_annotate
auto_annotate(data="path/to/images", det_model="yolov8x.pt", sam_model="sam2_b.pt")
参数 | 类型 | 描述 | 默认值 |
---|---|---|---|
data | str | 包含待标注图像的文件夹的路径。 | |
det_model | str, optional | 预训练的YOLO检测模型。默认为'yolov8x.pt'。 | 'yolov8x.pt' |
sam_model | str, optional | 预训练的SAM 2分割模型。默认为'sam2_b.pt'。 | 'sam2_b.pt' |
device | str, optional | 运行模型的设备。默认为空字符串(如果有的话,则为CPU或GPU)。 | |
output_dir | str, None, optional | 保存标注结果的目录。默认为与'data'同目录下的'labels'文件夹。 | None |