Segment Anything Model (SAM) - ViT Huge (ViT-H)版本的模型卡片

Segment Anything模型 (SAM) 的详细架构。

目录

TL;DR

概述

Segment Anything模型 (SAM) 可以根据输入的提示（如点或框）生成高质量的目标掩码，并用于生成图像中所有对象的掩码。它经过训练，使用了1100万张图像和10亿个掩码，对各种分割任务具有强大的零样本性能。该论文的摘要陈述如下：

我们介绍了Segment Anything (SA)项目：一个新的图像分割任务、模型和数据集。通过将我们高效的模型应用于数据收集循环中，我们建立了迄今为止最大的分割数据集，包括1100万张经过许可并尊重隐私的图像上的10亿个掩码。该模型的设计和训练支持提示，因此它可以零样本地传输到新的图像分布和任务上。我们对其能力进行了评估，并发现其零样本性能令人印象深刻，往往与或甚至优于先前的全面监督结果相竞争。我们正在发布1B掩码和1100万图像的Segment Anything模型 (SAM)和对应的数据集 (SA-1B)，以促进计算机视觉领域基础模型的研究。

免责声明：这个模型卡片的内容是由Hugging Face团队编写的，并且其中的部分内容是从原始文档复制粘贴的。

模型细节

SAM模型由3个模块组成：

• 视觉编码器：基于VIT的图像编码器。它使用对图像的补丁进行注意力计算图像嵌入。使用相对位置嵌入。
• 提示编码器：为点和边界框生成嵌入。
• 掩码解码器：双向转换器，对图像嵌入和点嵌入进行交叉注意力（->）和点嵌入和图像嵌入。输出被馈送给Neck进行基于上下文的掩码预测。
• Neck：根据MaskDecoder生成的上下文掩码进行输出掩码的预测。

使用方法

提示掩码生成

from PIL import Image
import requests
from transformers import SamModel, SamProcessor

model = SamModel.from_pretrained("facebook/sam-vit-huge")
processor = SamProcessor.from_pretrained("facebook/sam-vit-huge")

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]]] # 2D localization of a window

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

除了生成掩码的其他参数，您可以传递对象感兴趣的位置的2D坐标，包围对象感兴趣的边界框（格式应为边界框右上角和左下角点的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()

这应该给出以下结果

引用

如果您使用该模型，请使用以下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}
}