Transformers documentation

EdgeTAM

Hugging Face's logo
Join the Hugging Face community

and get access to the augmented documentation experience

to get started

该模型于 2025-01-13 在 HF papers 上发布,并于 2025-09-29 贡献至 Hugging Face Transformers。

SDPA FlashAttention

EdgeTAM

概述

EdgeTAM 模型在 EdgeTAM: On-Device Track Anything Model 中提出,作者为 Chong Zhou, Chenchen Zhu, Yunyang Xiong, Saksham Suri, Fanyi Xiao, Lemeng Wu, Raghuraman Krishnamoorthi, Bo Dai, Chen Change Loy, Vikas Chandra, Bilge Soran。

EdgeTAM 是 SAM 2 的一种高效适配方案,它引入了 2D Spatial Perceiver 架构,以优化内存注意力机制,从而在移动设备上实现实时视频分割。

该论文的摘要如下:

在 Segment Anything Model (SAM) 的基础上,SAM 2 通过记忆库机制将其能力从图像输入扩展到视频输入,与以往方法相比取得了显著的性能提升,使其成为视频分割任务的基础模型。本文旨在大幅提升 SAM 2 的效率,使其在保持可比性能的同时,甚至能在移动设备上运行。尽管已有若干工作致力于优化 SAM 以提升效率,但我们发现这些方法对于 SAM 2 而言并不充分,因为它们都聚焦于压缩图像编码器,而我们的基准测试表明,新引入的记忆注意力模块同样是延迟瓶颈。基于这一观察,我们提出了 EdgeTAM,它利用一种新颖的 2D Spatial Perceiver 来降低计算成本。具体而言,所提出的 2D Spatial Perceiver 使用一个包含固定数量可学习查询的轻量级 Transformer,对密集存储的帧级记忆进行编码。鉴于视频分割是一项密集预测任务,我们发现保留记忆的空间结构至关重要,因此将查询划分为全局级和块级两组。我们还提出了一种蒸馏流程,在不增加推理开销的情况下进一步提升性能。最终,EdgeTAM 在 DAVIS 2017、MOSE、SA-V val 和 SA-V test 上分别取得了 87.7、70.0、72.3 和 71.7 的 J&F 分数,同时在 iPhone 15 Pro Max 上以 16 FPS 的速度运行。

该模型由 yonigozlan 贡献。 原始代码可在此处找到。

使用示例

使用 Pipeline 自动生成掩码

EdgeTAM 可用于自动掩码生成,以分割图像中的所有对象,使用 mask-generation 流水线:

from transformers import pipeline


generator = pipeline("mask-generation", model="yonigozlan/edgetam-1", device=0)
image_url = "https://huggingface.co/datasets/hf-internal-testing/sam2-fixtures/resolve/main/truck.jpg"
outputs = generator(image_url, points_per_batch=64)

len(outputs["masks"])  # Number of masks generated
39

基础图像分割

单击

你可以通过在想要分割的对象上单击一个点来分割对象:

from transformers import Sam2Processor, EdgeTamModel
import torch
from PIL import Image
import requests


model = EdgeTamModel.from_pretrained("yonigozlan/edgetam-1", device_map="auto")
processor = Sam2Processor.from_pretrained("yonigozlan/edgetam-1")

image_url = "https://huggingface.co/datasets/hf-internal-testing/sam2-fixtures/resolve/main/truck.jpg"
raw_image = Image.open(requests.get(image_url, stream=True).raw).convert("RGB")

input_points = [[[[500, 375]]]]  # Single point click, 4 dimensions (image_dim, object_dim, point_per_object_dim, coordinates)
input_labels = [[[1]]]  # 1 for positive click, 0 for negative click, 3 dimensions (image_dim, object_dim, point_label)

inputs = processor(images=raw_image, input_points=input_points, input_labels=input_labels, return_tensors="pt").to(model.device)

with torch.no_grad():
    outputs = model(**inputs)

masks = processor.post_process_masks(outputs.pred_masks.cpu(), inputs["original_sizes"])[0]

# The model outputs multiple mask predictions ranked by quality score
print(f"Generated {masks.shape[1]} masks with shape {masks.shape}")
Generated 3 masks with shape torch.Size([1, 3, 1200, 1800])
print(f"IoU scores: {outputs.iou_scores.squeeze()}")
IoU scores: tensor([0.0463, 0.4859, 0.7616], device='cuda:0')

多个用于优化的点

您可以提供多个点来细化分割:

# Add both positive and negative points to refine the mask
input_points = [[[[500, 375], [1125, 625]]]]  # Multiple points for refinement
input_labels = [[[1, 1]]]  # Both positive clicks

inputs = processor(images=raw_image, input_points=input_points, input_labels=input_labels, return_tensors="pt").to(model.device)

with torch.no_grad():
    outputs = model(**inputs)

masks = processor.post_process_masks(outputs.pred_masks.cpu(), inputs["original_sizes"])[0]
print(f"IoU scores: {outputs.iou_scores.squeeze()}")
IoU scores: tensor([0.8362, 0.6900, 0.2120], device='cuda:0')

边界框输入

EdgeTAM 还支持用于分割的边界框输入:

# Define bounding box as [x_min, y_min, x_max, y_max]
input_boxes = [[[75, 275, 1725, 850]]]

inputs = processor(images=raw_image, input_boxes=input_boxes, return_tensors="pt").to(model.device)

with torch.no_grad():
    outputs = model(**inputs)

masks = processor.post_process_masks(outputs.pred_masks.cpu(), inputs["original_sizes"])[0]
print(f"IoU scores: {outputs.iou_scores.squeeze()}")
IoU scores: tensor([0.9301, 0.9348, 0.6605], device='cuda:0')

多目标分割

您可以同时分割多个对象:

# Define points for two different objects
input_points = [[[[500, 375]], [[650, 750]]]]  # Points for two objects in same image
input_labels = [[[1], [1]]]  # Positive clicks for both objects

inputs = processor(images=raw_image, input_points=input_points, input_labels=input_labels, return_tensors="pt").to(model.device)

with torch.no_grad():
    outputs = model(**inputs, multimask_output=False)

# Each object gets its own mask
masks = processor.post_process_masks(outputs.pred_masks.cpu(), inputs["original_sizes"])[0]
print(f"Generated masks for {masks.shape[0]} objects")
Generated masks for 2 objects
print(f"IoU scores: {outputs.iou_scores.squeeze()}")
IoU scores: tensor([0.7616, 0.9465], device='cuda:0')

批量推理

批量图像

同时处理多张图像以提高效率:

from transformers import Sam2Processor, EdgeTamModel
import torch
from PIL import Image
import requests


model = EdgeTamModel.from_pretrained("yonigozlan/edgetam-1", device_map="auto")
processor = Sam2Processor.from_pretrained("yonigozlan/edgetam-1")

# Load multiple images
image_urls = [
    "https://huggingface.co/datasets/hf-internal-testing/sam2-fixtures/resolve/main/truck.jpg",
    "https://huggingface.co/datasets/hf-internal-testing/transformers-synthetic-assets/resolve/main/images/dog_sam.png"
]
raw_images = [Image.open(requests.get(url, stream=True).raw).convert("RGB") for url in image_urls]

# Single point per image
input_points = [[[[500, 375]]], [[[770, 200]]]]  # One point for each image
input_labels = [[[1]], [[1]]]  # Positive clicks for both images

inputs = processor(images=raw_images, input_points=input_points, input_labels=input_labels, return_tensors="pt").to(model.device)

with torch.no_grad():
    outputs = model(**inputs, multimask_output=False)

# Post-process masks for each image
all_masks = processor.post_process_masks(outputs.pred_masks.cpu(), inputs["original_sizes"])
print(f"Processed {len(all_masks)} images, each with {all_masks[0].shape[0]} objects")
Processed 2 images, each with 1 objects
print(f"IoU scores: {outputs.iou_scores.squeeze()}")
IoU scores: tensor([0.7618, 0.7999], device='cuda:0')

每张图像的批量对象

使用批量推理对每张图片中的多个对象进行分割:

# Multiple objects per image - different numbers of objects per image
input_points = [
    [[[500, 375]], [[650, 750]]],  # Truck image: 2 objects
    [[[770, 200]]]  # Dog image: 1 object
]
input_labels = [
    [[1], [1]],  # Truck image: positive clicks for both objects
    [[1]]  # Dog image: positive click for the object
]

inputs = processor(images=raw_images, input_points=input_points, input_labels=input_labels, return_tensors="pt").to(model.device)

with torch.no_grad():
    outputs = model(**inputs, multimask_output=False)

all_masks = processor.post_process_masks(outputs.pred_masks.cpu(), inputs["original_sizes"])

批量图像与批量对象及多个点

处理每个对象包含多个点的复杂批次场景:

# Add groceries image for more complex example
groceries_url = "https://huggingface.co/datasets/hf-internal-testing/sam2-fixtures/resolve/main/groceries.jpg"
groceries_image = Image.open(requests.get(groceries_url, stream=True).raw).convert("RGB")
raw_images = [raw_images[0], groceries_image]  # Use truck and groceries images

# Complex batching: multiple images, multiple objects, multiple points per object
input_points = [
    [[[500, 375]], [[650, 750]]],  # Truck image: 2 objects with 1 point each
    [[[400, 300]], [[630, 300], [550, 300]]]  # Groceries image: obj1 has 1 point, obj2 has 2 points
]
input_labels = [
    [[1], [1]],  # Truck image: positive clicks
    [[1], [1, 1]]  # Groceries image: positive clicks for refinement
]

inputs = processor(images=raw_images, input_points=input_points, input_labels=input_labels, return_tensors="pt").to(model.device)

with torch.no_grad():
    outputs = model(**inputs, multimask_output=False)

all_masks = processor.post_process_masks(outputs.pred_masks.cpu(), inputs["original_sizes"])

批量边界框

使用边界框输入处理多张图像:

# Multiple bounding boxes per image (using truck and groceries images)
input_boxes = [
    [[75, 275, 1725, 850], [425, 600, 700, 875], [1375, 550, 1650, 800], [1240, 675, 1400, 750]],  # Truck image: 4 boxes
    [[450, 170, 520, 350], [350, 190, 450, 350], [500, 170, 580, 350], [580, 170, 640, 350]]  # Groceries image: 4 boxes
]

# Update images for this example
raw_images = [raw_images[0], groceries_image]  # truck and groceries

inputs = processor(images=raw_images, input_boxes=input_boxes, return_tensors="pt").to(model.device)

with torch.no_grad():
    outputs = model(**inputs, multimask_output=False)

all_masks = processor.post_process_masks(outputs.pred_masks.cpu(), inputs["original_sizes"])
print(f"Processed {len(input_boxes)} images with {len(input_boxes[0])} and {len(input_boxes[1])} boxes respectively")
Processed 2 images with 4 and 4 boxes respectively
print(f"IoU scores: {outputs.iou_scores.squeeze()}")
IoU scores: tensor([0.9301, 0.9348, 0.6605, 0.9465], device='cuda:0')

使用之前的掩码作为输入

EdgeTAM 可以使用来自先前预测的掩码作为输入,以优化分割:

# Get initial segmentation
input_points = [[[[500, 375]]]]
input_labels = [[[1]]]
inputs = processor(images=raw_image, input_points=input_points, input_labels=input_labels, return_tensors="pt").to(model.device)

with torch.no_grad():
    outputs = model(**inputs)

# Use the best mask as input for refinement
mask_input = outputs.pred_masks[:, :, torch.argmax(outputs.iou_scores.squeeze())]

# Add additional points with the mask input
new_input_points = [[[[500, 375], [450, 300]]]]
new_input_labels = [[[1, 1]]]
inputs = processor(
    input_points=new_input_points,
    input_labels=new_input_labels,
    original_sizes=inputs["original_sizes"],
    return_tensors="pt",
).to(model.device)

with torch.no_grad():
    refined_outputs = model(
        **inputs,
        input_masks=mask_input,
        image_embeddings=outputs.image_embeddings,
        multimask_output=False,
    )

EdgeTamConfig

class transformers.EdgeTamConfig

< >

( transformers_version: str | None = Nonearchitectures: list[str] | None = Noneoutput_hidden_states: bool | None = Falsereturn_dict: bool | None = Truedtype: str | torch.dtype | None = Nonechunk_size_feed_forward: int = 0is_encoder_decoder: bool = Falseid2label: dict[int, str] | dict[str, str] | None = Nonelabel2id: dict[str, int] | dict[str, str] | None = Noneproblem_type: Literal['regression', 'single_label_classification', 'multi_label_classification'] | None = Nonevision_config: dict | transformers.configuration_utils.PreTrainedConfig | None = Noneprompt_encoder_config: dict | transformers.configuration_utils.PreTrainedConfig | None = Nonemask_decoder_config: dict | transformers.configuration_utils.PreTrainedConfig | None = Noneinitializer_range: float = 0.02 )

参数

  • vision_config (Union[dict, ~configuration_utils.PreTrainedConfig],可选) — 视觉骨干网络的配置对象或字典。
  • prompt_encoder_config (Union[dict, EdgeTamPromptEncoderConfig], 可选) — 用于初始化 EdgeTamPromptEncoderConfig 的配置选项字典。
  • mask_decoder_config (Union[dict, EdgeTamMaskDecoderConfig],可选) — 用于初始化 EdgeTamMaskDecoderConfig 的配置选项字典。
  • initializer_range (float,可选,默认为 0.02) — 用于初始化所有权重矩阵的 truncated_normal_initializer 的标准差。

这是用于存储 EdgeTamModel 配置的类。它根据指定的参数实例化 Edgetam 模型,并定义模型架构。使用默认值实例化配置,将得到与 yonigozlan/EdgeTAM-hf 类似的配置。

配置对象继承自 PreTrainedConfig,可用于控制模型输出。有关更多信息,请参阅 PreTrainedConfig 的文档。

示例:

>>> from transformers import (
...     EdgeTamVisionConfig,
...     EdgeTamPromptEncoderConfig,
...     EdgeTamMaskDecoderConfig,
...     EdgeTamModel,
... )

>>> # Initializing a EdgeTamConfig with `"facebook/edgetam.1_hiera_tiny"` style configuration
>>> configuration = EdgeTamConfig()

>>> # Initializing a EdgeTamModel (with random weights) from the `"facebook/edgetam.1_hiera_tiny"` style configuration
>>> model = EdgeTamModel(configuration)

>>> # Accessing the model configuration
>>> configuration = model.config

>>> # We can also initialize a EdgeTamConfig from a EdgeTamVisionConfig, EdgeTamPromptEncoderConfig, and EdgeTamMaskDecoderConfig
>>> # Initializing EDGETAM vision encoder, memory attention, and memory encoder configurations
>>> vision_config = EdgeTamVisionConfig()
>>> prompt_encoder_config = EdgeTamPromptEncoderConfig()
>>> mask_decoder_config = EdgeTamMaskDecoderConfig()

>>> config = EdgeTamConfig(vision_config, prompt_encoder_config, mask_decoder_config)

EdgeTamVisionConfig

class transformers.EdgeTamVisionConfig

< >

( transformers_version: str | None = Nonearchitectures: list[str] | None = Noneoutput_hidden_states: bool | None = Falsereturn_dict: bool | None = Truedtype: str | torch.dtype | None = Nonechunk_size_feed_forward: int = 0is_encoder_decoder: bool = Falseid2label: dict[int, str] | dict[str, str] | None = Nonelabel2id: dict[str, int] | dict[str, str] | None = Noneproblem_type: Literal['regression', 'single_label_classification', 'multi_label_classification'] | None = Nonebackbone_config: dict | transformers.configuration_utils.PreTrainedConfig | None = Nonebackbone_channel_list: list[int] | None = Nonebackbone_feature_sizes: list | None = Nonefpn_hidden_size: int = 256fpn_kernel_size: int = 1fpn_stride: int = 1fpn_padding: int = 0fpn_top_down_levels: list[int] | None = Nonenum_feature_levels: int = 3hidden_act: str = 'gelu'layer_norm_eps: float = 1e-06initializer_range: float = 0.02 )

参数

  • backbone_config (Union[dict, ~configuration_utils.PreTrainedConfig],可选) — 主干模型的配置。
  • backbone_channel_list (List[int],可选,默认为 [384, 192, 96, 48]) — 主干网络的通道维度列表。
  • backbone_feature_sizes (List[List[int]],可选,默认为 [[256, 256], [128, 128], [64, 64]]) — 主干网络输出的特征图的空间尺寸。
  • fpn_hidden_size (int,可选,默认为 256) — FPN 的隐藏维度。
  • fpn_kernel_size (int,可选,默认为 1) — neck 中卷积的核大小。
  • fpn_stride (int,可选,默认为 1) — neck 中卷积的步长。
  • fpn_padding (int,可选,默认为 0) — neck 中卷积的填充值。
  • fpn_top_down_levels (List[int],可选,默认为 [2, 3]) — 用于自顶向下 FPN 连接的层级。
  • num_feature_levels (int,可选,默认为 3) — 要使用的 FPN 特征层数。
  • hidden_act (str,可选,默认为 gelu) — 解码器中的非线性激活函数(函数或字符串)。例如,"gelu"、 "relu"、"silu" 等。
  • layer_norm_eps (float,可选,默认为 1e-06) — 层归一化层所使用的 epsilon 值。
  • initializer_range (float,可选,默认为 0.02) — 用于初始化所有权重矩阵的 truncated_normal_initializer 的标准差。

这是用于存储 EdgeTamModel 配置的类。它根据指定的参数实例化 Edgetam 模型,并定义模型架构。使用默认值实例化配置,将得到与 yonigozlan/EdgeTAM-hf 类似的配置。

配置对象继承自 PreTrainedConfig,可用于控制模型输出。有关更多信息,请参阅 PreTrainedConfig 的文档。

EdgeTamMaskDecoderConfig

class transformers.EdgeTamMaskDecoderConfig

< >

( transformers_version: str | None = Nonearchitectures: list[str] | None = Noneoutput_hidden_states: bool | None = Falsereturn_dict: bool | None = Truedtype: str | torch.dtype | None = Nonechunk_size_feed_forward: int = 0is_encoder_decoder: bool = Falseid2label: dict[int, str] | dict[str, str] | None = Nonelabel2id: dict[str, int] | dict[str, str] | None = Noneproblem_type: Literal['regression', 'single_label_classification', 'multi_label_classification'] | None = Nonehidden_size: int = 256hidden_act: str = 'gelu'mlp_dim: int = 2048num_hidden_layers: int = 2num_attention_heads: int = 8attention_downsample_rate: int = 2num_multimask_outputs: int = 3iou_head_depth: int = 3iou_head_hidden_dim: int = 256dynamic_multimask_via_stability: bool = Truedynamic_multimask_stability_delta: float = 0.05dynamic_multimask_stability_thresh: float = 0.98 )

参数

  • hidden_size (int,可选,默认为 256) — 隐藏表示的维度。
  • hidden_act (str,可选,默认为 gelu) — 解码器中的非线性激活函数(函数或字符串)。例如,"gelu"、 "relu"、"silu" 等。
  • mlp_dim (int,可选,默认为 2048) — 双向 transformer 中 MLP 的维度。
  • num_hidden_layers (int,可选,默认为 2) — Transformer 解码器中的隐藏层数。
  • num_attention_heads (int,可选,默认为 8) — Transformer 解码器中每个注意力层的注意力头数量。
  • attention_downsample_rate (int,可选,默认为 2) — 注意力层的下采样率。
  • num_multimask_outputs (int,可选,默认为 3) — 多掩码输出的数量。
  • iou_head_depth (int, 可选,默认为 3) — IoU 头的深度。
  • iou_head_hidden_dim (int, 可选,默认为 256) — IoU 头的隐藏维度。
  • dynamic_multimask_via_stability (bool,可选,默认为 True) — 是否通过稳定性使用动态多掩码。
  • dynamic_multimask_stability_delta (float,可选,默认值为 0.05) — 动态多掩码的稳定性增量。
  • dynamic_multimask_stability_thresh (float,可选,默认值为 0.98) — 动态多掩码的稳定性阈值。

这是用于存储 EdgeTamModel 配置的类。它根据指定的参数实例化 Edgetam 模型,并定义模型架构。使用默认值实例化配置,将得到与 yonigozlan/EdgeTAM-hf 类似的配置。

配置对象继承自 PreTrainedConfig,可用于控制模型输出。有关更多信息,请参阅 PreTrainedConfig 的文档。

EdgeTamPromptEncoderConfig

class transformers.EdgeTamPromptEncoderConfig

< >

( transformers_version: str | None = Nonearchitectures: list[str] | None = Noneoutput_hidden_states: bool | None = Falsereturn_dict: bool | None = Truedtype: str | torch.dtype | None = Nonechunk_size_feed_forward: int = 0is_encoder_decoder: bool = Falseid2label: dict[int, str] | dict[str, str] | None = Nonelabel2id: dict[str, int] | dict[str, str] | None = Noneproblem_type: Literal['regression', 'single_label_classification', 'multi_label_classification'] | None = Nonehidden_size: int = 256image_size: int | list[int] | tuple[int, int] = 1024patch_size: int | list[int] | tuple[int, int] = 16mask_input_channels: int = 16num_point_embeddings: int = 4hidden_act: str = 'gelu'layer_norm_eps: float = 1e-06scale: int = 1 )

参数

  • hidden_size (int,可选,默认为 256) — 隐藏表示的维度。
  • image_size (Union[int, list[int], tuple[int, int]],可选,默认为 1024) — 每张图片的尺寸(分辨率)。
  • patch_size (Union[int, list[int], tuple[int, int]],可选,默认为 16) — 每个 patch 的大小(分辨率)。
  • mask_input_channels (int,可选,默认为 16) — 输入到 MaskDecoder 模块的通道数。
  • num_point_embeddings (int,可选,默认为 4) — 要使用的点嵌入数量。
  • hidden_act (str,可选,默认为 gelu) — 解码器中的非线性激活函数(函数或字符串)。例如,"gelu"、 "relu"、"silu" 等。
  • layer_norm_eps (float,可选,默认为 1e-06) — 层归一化层所使用的 epsilon 值。
  • scale (float,可选,默认为 1) — 提示编码器的缩放因子。

这是用于存储 EdgeTamModel 配置的类。它根据指定的参数实例化 Edgetam 模型,并定义模型架构。使用默认值实例化配置,将得到与 yonigozlan/EdgeTAM-hf 类似的配置。

配置对象继承自 PreTrainedConfig,可用于控制模型输出。有关更多信息,请参阅 PreTrainedConfig 的文档。

EdgeTamVisionModel

class transformers.EdgeTamVisionModel

< >

( config: EdgeTamVisionConfig )

参数

  • config (EdgeTamVisionConfig) — 模型的配置类,包含模型的所有参数。使用配置文件进行初始化时,不会加载与模型关联的权重,仅加载配置。请查看 from_pretrained() 方法以加载模型权重。

EdgeTAM 的视觉模型,顶部未添加任何头或投影层。

该模型继承自 PreTrainedModel。请查阅父类文档,了解该库为所有模型实现的通用方法(例如下载或保存、调整输入嵌入大小、剪枝头等)。

该模型也是 PyTorch torch.nn.Module 的子类。 请将其作为常规的 PyTorch Module 使用,有关通用用法和行为的所有事项,请参阅 PyTorch 文档。

forward

< >

( pixel_values: typing.Optional[torch.Tensor] = None**kwargs: *<class 'transformers.utils.generic.TransformersKwargs'> )

EdgeTamModel

class transformers.EdgeTamModel

< >

( config: EdgeTamConfig )

参数

  • config (EdgeTamConfig) — 模型的配置类,包含模型的所有参数。使用配置文件进行初始化时,不会加载与模型关联的权重,仅加载配置。请查看 from_pretrained() 方法以加载模型权重。

Segment Anything Model 2 (SAM 2) 用于根据输入图像以及输入的点及标签、边界框或掩码来生成分割掩码。

该模型继承自 PreTrainedModel。请查阅父类文档,了解该库为所有模型实现的通用方法(例如下载或保存、调整输入嵌入大小、剪枝头等)。

该模型也是 PyTorch torch.nn.Module 的子类。 请将其作为常规的 PyTorch Module 使用,有关通用用法和行为的所有事项,请参阅 PyTorch 文档。

forward

< >

( pixel_values: typing.Optional[torch.Tensor] = Noneinput_points: typing.Optional[torch.Tensor] = Noneinput_labels: typing.Optional[torch.LongTensor] = Noneinput_boxes: typing.Optional[torch.Tensor] = Noneinput_masks: typing.Optional[torch.LongTensor] = Noneimage_embeddings: typing.Optional[torch.Tensor] = Nonemultimask_output: bool = Trueattention_similarity: typing.Optional[torch.Tensor] = Nonetarget_embedding: typing.Optional[torch.Tensor] = None**kwargs: *<class 'transformers.utils.generic.TransformersKwargs'> ) → EdgeTamImageSegmentationOutput or tuple(torch.Tensor)

参数

  • pixel_values (torch.Tensor,形状为 (batch_size, num_channels, image_size, image_size),可选) — 与输入图像对应的张量。可以使用 Sam2ImageProcessor 获取像素值。有关详细信息,请参阅 Sam2ImageProcessor.call()(Sam2Processor 使用 Sam2ImageProcessor 来处理图像)。
  • input_points (torch.Tensor,形状为 (batch_size, num_points, 2)) — 输入的二维空间点,提示编码器使用它来编码提示。通常能带来好得多的结果。可以通过向处理器传入列表的列表的列表来获取这些点,处理器将创建维度为 4 的相应 torch 张量。第一个维度是图像批次大小,第二个维度是点批次大小(即对于每个输入点,我们希望模型预测多少个分割掩码),第三个维度是每个分割掩码对应的点数(可以为单个掩码传入多个点),最后一个维度是点的 x(垂直)和 y(水平)坐标。如果为每张图像或每个掩码传入的点数不同,处理器将创建“PAD”点,这些点对应 (0, 0) 坐标,并且使用标签跳过对这些点的嵌入计算。
  • input_labels (torch.LongTensor 形状为 (batch_size, point_batch_size, num_points)) — 点的输入标签,用于提示编码器对提示进行编码。根据官方实现,共有 3 种标签类型
    • 1:该点是一个包含目标对象的点
    • 0:该点不包含目标对象
    • -1:该点对应背景

    我们添加了标签:

    • -10:该点为填充点,因此应被 prompt encoder 忽略

    填充标签应由处理器自动完成。

  • input_boxes (torch.Tensor 形状为 (batch_size, num_boxes, 4)) — 用于点的输入框,提示编码器会用它来编码提示。通常能生成质量更好的掩码。可以通过向处理器传入列表的列表来获取这些框,这将生成一个 torch 张量,其各维度分别对应图像批次大小、每张图像的框数以及框的左上角和右下角点的坐标。顺序为 (x1, y1, x2, y2):
    • x1:输入框左上角点的 x 坐标
    • y1:输入框左上角点的 y 坐标
    • x2:输入框右下角点的 x 坐标
    • y2:输入框右下角点的 y 坐标
  • input_masks(形状为 (batch_size, image_size, image_size) 的 torch.Tensor) — SAM 模型也接受分割掩码作为输入。掩码将由提示编码器嵌入,以生成相应的嵌入,随后输入到掩码解码器中。这些掩码需要由用户手动提供,且其形状必须为 (batch_size, image_size, image_size)。
  • image_embeddings (torch.Tensor,形状为 (batch_size, output_channels, window_size, window_size)) — 图像嵌入,供掩码解码器生成掩码和 iou 分数。为了更节省内存,用户可以先使用 get_image_embeddings 方法获取图像嵌入,然后将它们传入 forward 方法,而不是传入 pixel_values。
  • multimask_output (bool, 可选) — 在原始实现和论文中,模型始终为每张图像输出 3 个掩码(如果相关,则为每个点或每个边界框输出 3 个掩码)。不过,可以通过指定 multimask_output=False 来仅输出单个掩码,该掩码对应于“最佳”掩码。
  • attention_similarity (torch.Tensor, 可选) — 注意力相似度张量,当模型用于如 PerSAM 中所述的个性化时,需将其提供给掩码解码器以实现目标引导的注意力。
  • target_embedding (torch.Tensor, 可选) — 目标概念的嵌入,当模型用于 PerSAM 中介绍的个人化时,需将其提供给掩码解码器以进行目标语义提示。

返回值

EdgeTamImageSegmentationOutput 或 tuple(torch.Tensor)

EdgeTamImageSegmentationOutput 或 torch.Tensor 元组(当传入 return_dict=False 或 config.return_dict=False 时),其包含的元素取决于配置(EdgeTamConfig)和输入。

EdgeTamModel 的 forward 方法重写了 __call__ 特殊方法。

尽管前向传播的实现需要在此函数中定义,但之后应调用 Module 实例,而不是直接调用此函数,因为前者会负责执行预处理和后处理步骤,而后者则会静默地忽略这些步骤。

  • iou_scores(形状为 (batch_size, point_batch_size, num_masks) 的 torch.Tensor)— 预测掩码的交并比(IoU)分数。
  • pred_masks(形状为 (batch_size, point_batch_size, num_masks, height, width) 的 torch.Tensor)— 预测的低分辨率掩码。这是 low_res_masks 的别名。这些掩码需要经过处理器进行后处理,才能调整到原始图像尺寸。
  • object_score_logits(形状为 (batch_size, point_batch_size, 1) 的 torch.Tensor)— 对象得分的 logits,指示是否存在对象。
  • image_embeddings (tuple(torch.Tensor)) — 来自 FPN 的特征,供 mask decoder 使用。这是一个 torch.Tensor 元组,其中每个张量的形状为 (batch_size, channels, height, width)。
  • vision_hidden_states (tuple(torch.Tensor),可选,当 output_hidden_states=True 时返回) — 一个 torch.Tensor 元组(每个阶段一个输出),形状为 (batch_size, height, width, hidden_size)。 视觉模型在每个阶段输出处的隐藏状态。
  • vision_attentions (tuple(torch.Tensor),可选,当 output_attentions=True 时返回) — 一个 torch.Tensor 元组(每层一个),形状为 (batch_size, num_heads, sequence_length, sequence_length)。 视觉模型的注意力权重。
  • mask_decoder_attentions (tuple(torch.Tensor),可选,当 output_attentions=True 时返回) — 一个 torch.Tensor 元组(每层一个),形状为 (batch_size, num_heads, sequence_length, sequence_length)。 掩码解码器的注意力权重。

示例:

>>> from PIL import Image
>>> from huggingface_hub.utils import httpx
>>> from io import BytesIO
>>> from transformers import AutoModel, AutoProcessor

>>> model = AutoModel.from_pretrained("danelcsb/edgetam.1_hiera_tiny")
>>> processor = AutoProcessor.from_pretrained("danelcsb/edgetam.1_hiera_tiny")

>>> url = "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/transformers/model_doc/sam-car.png"
>>> with httpx.stream("GET", url) as response:
...     raw_image = Image.open(BytesIO(response.read())).convert("RGB")
>>> input_points = [[[400, 650]]]  # 2D location of a window on the car
>>> inputs = processor(images=raw_image, input_points=input_points, return_tensors="pt")

>>> # Get segmentation mask
>>> outputs = model(**inputs)

>>> # Postprocess masks
>>> masks = processor.post_process_masks(
...     outputs.pred_masks, inputs["original_sizes"]
... )

get_image_features

< >

( pixel_values: Tensor**kwargs: *<class 'transformers.utils.generic.TransformersKwargs'> ) → EdgeTamVisionEncoderOutput or tuple(torch.Tensor)

参数

  • pixel_values (torch.Tensor) — 形状为 (batch_size, num_channels, height, width) 的输入像素值。

返回值

EdgeTamVisionEncoderOutput 或 tuple(torch.Tensor)

EdgeTamVisionEncoderOutput 或 torch.Tensor 元组(当传入 return_dict=False 或 config.return_dict=False 时),其包含的元素取决于配置(EdgeTamConfig)和输入。

  • last_hidden_state(形状为 (batch_size, height, width, hidden_size) 的 torch.Tensor)— 模型最后一层输出的隐藏状态序列。

  • pooler_output (torch.Tensor,形状为 (batch_size, hidden_size)) — 序列中第一个 token(分类 token)的最后一层隐藏状态,经过用于辅助预训练任务的层进一步处理后得到。例如,对于 BERT 系列模型,这返回的是分类 token 经过线性层和 tanh 激活函数处理后的结果。线性层的权重在预训练期间通过下一句预测(分类)目标进行训练。

  • hidden_states (tuple(torch.Tensor),可选,当传入 output_hidden_states=True 或 config.output_hidden_states=True 时返回) — 一个 torch.Tensor 元组(如果模型具有嵌入层,则包含嵌入层的输出,以及每一层的输出),形状为 (batch_size, sequence_length, hidden_size)。

    模型在每一层输出处的隐藏状态,以及可选的初始嵌入输出。

  • attentions (tuple(torch.Tensor),可选,当传入 output_attentions=True 或 config.output_attentions=True 时返回) — 一个 torch.Tensor 元组(每层一个),形状为 (batch_size, num_heads, sequence_length, sequence_length)。

    注意力权重,即经过注意力 softmax 之后的权重,用于在自注意力头中计算加权平均值。

  • fpn_hidden_states (tuple(torch.Tensor)) — 一个 torch.Tensor 元组(每个特征层一个,从高分辨率到低分辨率),形状为 (batch_size, hidden_size, height, width)。来自特征金字塔网络(Feature Pyramid Network)颈部的特征图。

  • fpn_position_encoding (tuple(torch.Tensor)) — torch.Tensor 的元组(每个特征层一个,从高分辨率到低分辨率),形状为 (batch_size, hidden_size, height, width)。与 fpn_hidden_states 对应的位置编码。

示例:

Update on GitHub