🚀 分割一切模型(SAM) - ViT基础版(ViT - B)模型卡片
分割一切模型(SAM)能够根据输入的提示(如点或框)生成高质量的对象掩码,还可以为图像中的所有对象生成掩码。它在大规模数据集上进行训练,在各种分割任务中展现出强大的零样本性能。
🚀 快速开始
模型简介
分割一切模型(SAM)由3个模块组成:
VisionEncoder
:基于VIT的图像编码器,通过对图像块进行注意力计算来生成图像嵌入,使用了相对位置嵌入。
PromptEncoder
:为点和边界框生成嵌入。
MaskDecoder
:双向变压器,在图像嵌入和点嵌入之间进行交叉注意力计算,输出结果会被进一步处理。
Neck
:根据MaskDecoder
生成的上下文掩码预测输出掩码。
使用示例
基础用法
from PIL import Image
import requests
from transformers import SamModel, SamProcessor
model = SamModel.from_pretrained("facebook/sam-vit-base")
processor = SamProcessor.from_pretrained("facebook/sam-vit-base")
img_url = "https://huggingface.co/ybelkada/segment-anything/resolve/main/assets/car.png"
raw_image = Image.open(requests.get(img_url, stream=True).raw).convert("RGB")
input_points = [[[450, 600]]]
inputs = processor(raw_image, input_points=input_points, return_tensors="pt").to("cuda")
outputs = model(**inputs)
masks = processor.image_processor.post_process_masks(outputs.pred_masks.cpu(), inputs["original_sizes"].cpu(), inputs["reshaped_input_sizes"].cpu())
scores = outputs.iou_scores
在生成掩码时,除了其他参数,你还可以传入感兴趣对象的大致二维位置、围绕感兴趣对象的边界框(格式应为边界框右上角和左下角的x、y坐标)或分割掩码。根据官方仓库,目前官方模型不支持将文本作为输入。
高级用法
模型可以以“零样本”的方式为输入图像生成分割掩码。模型会自动使用一个包含1024
个点的网格进行提示,并将这些点输入到模型中。
以下代码展示了如何轻松运行自动掩码生成(可在任何设备上运行,只需提供合适的points_per_batch
参数):
from transformers import pipeline
generator = pipeline("mask-generation", device = 0, points_per_batch = 256)
image_url = "https://huggingface.co/ybelkada/segment-anything/resolve/main/assets/car.png"
outputs = generator(image_url, points_per_batch = 256)
以下代码用于显示图像:
import matplotlib.pyplot as plt
from PIL import Image
import numpy as np
def show_mask(mask, ax, random_color=False):
if random_color:
color = np.concatenate([np.random.random(3), np.array([0.6])], axis=0)
else:
color = np.array([30 / 255, 144 / 255, 255 / 255, 0.6])
h, w = mask.shape[-2:]
mask_image = mask.reshape(h, w, 1) * color.reshape(1, 1, -1)
ax.imshow(mask_image)
plt.imshow(np.array(raw_image))
ax = plt.gca()
for mask in outputs["masks"]:
show_mask(mask, ax=ax, random_color=True)
plt.axis("off")
plt.show()
📚 详细文档
模型架构图
分割一切模型(SAM)的详细架构。
快速浏览
原始仓库链接
分割一切模型(SAM) 可以根据输入的提示(如点或框)生成高质量的对象掩码,并且可以为图像中的所有对象生成掩码。它在包含1100万张图像和11亿个掩码的数据集上进行了训练,在各种分割任务中具有强大的零样本性能。
论文摘要指出:
我们推出了分割一切(SA)项目:一个用于图像分割的新任务、模型和数据集。通过在数据收集循环中使用高效的模型,我们构建了迄今为止最大的分割数据集,在1100万张经过授权且尊重隐私的图像上拥有超过11亿个掩码。该模型的设计和训练使其具有可提示性,因此可以零样本迁移到新的图像分布和任务中。我们在众多任务上评估了其能力,发现其零样本性能令人印象深刻 —— 通常与之前的全监督结果相媲美,甚至更优。我们在 https://segment-anything.com 上发布了分割一切模型(SAM)以及包含11亿个掩码和1100万张图像的相应数据集(SA - 1B),以促进计算机视觉基础模型的研究。
免责声明:此模型卡片的内容由Hugging Face团队编写,部分内容从原始的 SAM模型卡片 复制粘贴而来。
📄 许可证
本模型采用Apache - 2.0许可证。
📚 引用
如果您使用此模型,请使用以下BibTeX条目进行引用。
@article{kirillov2023segany,
title={Segment Anything},
author={Kirillov, Alexander and Mintun, Eric and Ravi, Nikhila and Mao, Hanzi and Rolland, Chloe and Gustafson, Laura and Xiao, Tete and Whitehead, Spencer and Berg, Alexander C. and Lo, Wan - Yen and Doll{\'a}r, Piotr and Girshick, Ross},
journal={arXiv:2304.02643},
year={2023}
}