Keras 3 API 文档 / KerasCV / 模型 / 任务 / 分割一切

分割一切

[源代码]

SegmentAnythingModel

keras_cv.models.SegmentAnythingModel(
    backbone, prompt_encoder, mask_decoder, **kwargs
)

分割一切 (SAM) 模型。

参数

  • backbone (keras_cv.models.Backbone): 输入图像的特征提取器。
  • prompt_encoder (keras_cv.models.SAMPromptEncoder): 用于计算点、框和掩码提示嵌入的 Keras 层。
  • mask_decoder (keras_cv.models.SAMMaskDecoder): 一个 Keras 层,用于在给定由骨干网络和提示编码器生成的嵌入的情况下生成分割掩码。

参考文献

示例

>>> import numpy as np
>>> from keras_cv.src.models import ViTDetBBackbone
>>> from keras_cv.src.models import SAMPromptEncoder
>>> from keras_cv.src.models import SAMMaskDecoder

创建 SAM 模型的所有组件

>>> backbone = ViTDetBBackbone()
>>> prompt_encoder = SAMPromptEncoder()
>>> mask_decoder = SAMMaskDecoder()

实例化模型

>>> sam = SegmentAnythingModel(
...     backbone=backbone,
...     prompt_encoder=prompt_encoder,
...     mask_decoder=mask_decoder
... )

定义骨干网络的输入。对于我们正在使用的 ViT 骨干网络,它必须是形状为 (1024, 1024, 3) 的图像批次。

>>> image = np.ones((1, 1024, 1024, 3))

SAM 通过提示输入图像来工作。有三种提示方式

(1) 带标签的点:前景点(标签为 1 的点)被编码,使得掩码解码器生成的输出掩码包含它们,而背景点(标签为 0 的点)被编码,使得生成的掩码不包含它们。(2) 框:框告诉模型分割图像的哪个部分/裁剪。 (3) 掩码:输入掩码可用于细化掩码解码器的输出。

这些提示可以混合和匹配,但必须至少存在一个提示。要关闭特定提示,只需将其从模型的输入中排除。

TODO(ianstenbit):删除对 1 轴的需要,并修复框的形状

.

(1) 对于点提示,预期形状为 (batch, num_points, 2)。标签必须具有相应的形状 (batch, num_points)。(2) 对于框提示,预期形状为 (batch, 1, 2, 2)。(3) 同样,掩码提示的形状为 (batch, 1, H, W, 1)

例如,要传递所有提示,请执行以下操作

>>> points = np.array([[[512., 512.], [100., 100.]]])
>>> # For labels: 1 means foreground point, 0 means background
>>> labels = np.array([[1., 0.]])
>>> box = np.array([[[[384., 384.], [640., 640.]]]])
>>> input_mask = np.ones((1, 1, 256, 256, 1))

准备输入字典

>>> inputs = {
...     "images": image,
...     "points": points,
...     "labels": labels,
...     "boxes": box,
...     "masks": input_mask
... }
...
>>> outputs = sam.predict(inputs)
>>> masks, iou_pred = outputs["masks"], outputs["iou_pred"]

输出 masks 中的第一个掩码(即 masks[:, 0, ...])是模型根据提示预测的最佳掩码。其他 masks(即 masks[:, 1:, ...])是备选预测,如果需要,可以优先于第一个掩码使用。

现在,在只有点和框提示的情况下,只需排除掩码即可

>>> inputs = {
...     "images": image,
...     "points": points,
...     "labels": labels,
...     "boxes": box,
... }
...
>>> outputs = sam.predict(inputs)
>>> masks, iou_pred = outputs["masks"], outputs["iou_pred"]

TODO(ianstenbit):删除对此填充的需要

.

另一个示例是只有点提示存在。请注意,如果存在点提示但不存在框提示,则必须使用零点和 -1 标签对点进行填充。

>>> padded_points = np.concatenate(
...     [points, np.zeros((1, 1, 2))], axis=1
... )
...
>>> padded_labels = np.concatenate(
...     [labels, -np.ones((1, 1))], axis=1
... )
>>> inputs = {
...     "images": image,
...     "points": padded_points,
...     "labels": padded_labels,
... }
...
>>> outputs = sam.predict(inputs)
>>> masks, iou_pred = outputs["masks"], outputs["iou_pred"]

请注意,分割一切模型仅支持推理,尚不支持训练。因此,目前调用 fit 方法将会失败。


[源代码]

from_preset 方法

SegmentAnythingModel.from_preset()

从预设配置和权重实例化 SegmentAnythingModel 模型。

参数

  • preset: 字符串。必须是以下之一:"resnet18"、"resnet34"、"resnet50"、"resnet101"、"resnet152"、"resnet18_v2"、"resnet34_v2"、"resnet50_v2"、"resnet101_v2"、"resnet152_v2"、"mobilenet_v3_small"、"mobilenet_v3_large"、"csp_darknet_tiny"、"csp_darknet_s"、"csp_darknet_m"、"csp_darknet_l"、"csp_darknet_xl"、"efficientnetv1_b0"、"efficientnetv1_b1"、"efficientnetv1_b2"、"efficientnetv1_b3"、"efficientnetv1_b4"、"efficientnetv1_b5"、"efficientnetv1_b6"、"efficientnetv1_b7"、"efficientnetv2_s"、"efficientnetv2_m"、"efficientnetv2_l"、"efficientnetv2_b0"、"efficientnetv2_b1"、"efficientnetv2_b2"、"efficientnetv2_b3"、"densenet121"、"densenet169"、"densenet201"、"efficientnetlite_b0"、"efficientnetlite_b1"、"efficientnetlite_b2"、"efficientnetlite_b3"、"efficientnetlite_b4"、"yolo_v8_xs_backbone"、"yolo_v8_s_backbone"、"yolo_v8_m_backbone"、"yolo_v8_l_backbone"、"yolo_v8_xl_backbone"、"vitdet_base"、"vitdet_large"、"vitdet_huge"、"videoswin_tiny"、"videoswin_small"、"videoswin_base"、"resnet50_imagenet"、"resnet50_v2_imagenet"、"mobilenet_v3_large_imagenet"、"mobilenet_v3_small_imagenet"、"csp_darknet_tiny_imagenet"、"csp_darknet_l_imagenet"、"efficientnetv2_s_imagenet"、"efficientnetv2_b0_imagenet"、"efficientnetv2_b1_imagenet"、"efficientnetv2_b2_imagenet"、"densenet121_imagenet"、"densenet169_imagenet"、"densenet201_imagenet"、"yolo_v8_xs_backbone_coco"、"yolo_v8_s_backbone_coco"、"yolo_v8_m_backbone_coco"、"yolo_v8_l_backbone_coco"、"yolo_v8_xl_backbone_coco"、"vitdet_base_sa1b"、"vitdet_large_sa1b"、"vitdet_huge_sa1b"、"videoswin_tiny_kinetics400"、"videoswin_small_kinetics400"、"videoswin_base_kinetics400"、"videoswin_base_kinetics400_imagenet22k"、"videoswin_base_kinetics600_imagenet22k"、"videoswin_base_something_something_v2"、"sam_base_sa1b"、"sam_large_sa1b"、"sam_huge_sa1b"。如果正在寻找具有预训练权重的预设,请选择以下之一:"resnet50_imagenet"、"resnet50_v2_imagenet"、"mobilenet_v3_large_imagenet"、"mobilenet_v3_small_imagenet"、"csp_darknet_tiny_imagenet"、"csp_darknet_l_imagenet"、"efficientnetv2_s_imagenet"、"efficientnetv2_b0_imagenet"、"efficientnetv2_b1_imagenet"、"efficientnetv2_b2_imagenet"、"densenet121_imagenet"、"densenet169_imagenet"、"densenet201_imagenet"、"yolo_v8_xs_backbone_coco"、"yolo_v8_s_backbone_coco"、"yolo_v8_m_backbone_coco"、"yolo_v8_l_backbone_coco"、"yolo_v8_xl_backbone_coco"、"vitdet_base_sa1b"、"vitdet_large_sa1b"、"vitdet_huge_sa1b"、"videoswin_tiny_kinetics400"、"videoswin_small_kinetics400"、"videoswin_base_kinetics400"、"videoswin_base_kinetics400_imagenet22k"、"videoswin_base_kinetics600_imagenet22k"、"videoswin_base_something_something_v2"、"sam_base_sa1b"、"sam_large_sa1b"、"sam_huge_sa1b"。
  • load_weights: 是否将预训练权重加载到模型中。默认为 None,遵循预设是否有可用的预训练权重。
  • input_shape : 将传递给骨干网络初始化的输入形状,默认为 None。如果为 None,则使用预设值。

示例

# Load architecture and weights from preset
model = keras_cv.models.SegmentAnythingModel.from_preset(
    "resnet50_imagenet",
)

# Load randomly initialized model from preset architecture with weights
model = keras_cv.models.SegmentAnythingModel.from_preset(
    "resnet50_imagenet",
    load_weights=False,
预设名称 参数 描述
sam_base_sa1b 93.74M 在 SA1B 数据集上训练的基本 SAM 模型。
sam_large_sa1b 312.34M 在 SA1B 数据集上训练的大型 SAM 模型。
sam_huge_sa1b 641.09M 在 SA1B 数据集上训练的超大型 SAM 模型。

[源代码]

SAMMaskDecoder

keras_cv.models.SAMMaskDecoder(
    transformer_dim=256,
    transformer=None,
    num_multimask_outputs=3,
    iou_head_depth=3,
    iou_head_hidden_dim=256,
    activation="gelu",
    **kwargs
)

分割一切模型 (SAM) 的掩码解码器。

这个轻量级模块有效地将图像嵌入和一组提示嵌入映射到输出掩码。在应用 Transformer 解码器之前,该层首先将学习的输出标记嵌入插入到提示嵌入集中,该嵌入将在解码器的输出处使用。为简单起见,这些嵌入(不包括图像嵌入)统称为“标记”。

图像嵌入、位置图像嵌入和标记通过 Transformer 解码器传递。解码器运行后,该层使用两个转置卷积层将更新后的图像嵌入上采样 4 倍(现在它相对于输入图像缩小了 4 倍)。然后,标记再次关注图像嵌入,并更新输出标记嵌入传递到一个小的 3 层 MLP,该 MLP 输出一个与上采样图像嵌入的通道维度匹配的向量。最后,通过上采样图像嵌入和 MLP 输出之间的空间逐点乘积来预测掩码。

参数

  • transformer_dim (int, 可选): Transformer 解码器的输入特征数量。默认为 256
  • transformer (keras.layers.Layer, 可选): Transformer 解码器。默认为 None。当为 None 时,使用 keras_cv.models.TwoWayTransformer 层。
  • num_multimask_outputs (int, 可选): 多掩码输出的数量。模型将生成这么多额外的掩码。模型生成的总掩码数为 1 + num_multimask_outputs。默认为 3
  • iou_head_depth (int, 可选): 用于预测 IoU 置信度的密集网络的深度。默认为 3
  • iou_head_hidden_dim (int, 可选): 用于预测 IoU 置信度的密集网络中隐藏层的单元数量。默认为 256
  • activation (str, 可选): 在掩码上采样网络中使用的激活函数。默认为 "gelu"

参考文献


[源代码]

SAMPromptEncoder

keras_cv.models.SAMPromptEncoder(
    embed_dim=256,
    image_embedding_size=(64, 64),
    input_image_size=(1024, 1024),
    mask_in_chans=16,
    activation="gelu",
    **kwargs
)

分割一切模型 (SAM) 的提示编码器。

提示编码器为三种类型的提示生成编码

  • 点提示:图像上的点以及指示该点是前景(掩码的一部分)还是背景(不是掩码的一部分)的标签。
  • 框提示:一批格式为 [(x1, y1), (x2, y2)] 的边界框,用于确定图像中掩码的位置。
  • 掩码:可以传递输入掩码以细化输出掩码的位置嵌入。

首先,将点提示和框提示连接起来,并使用随机空间频率生成位置编码。一个点表示为点位置的位置编码与指示该点是前景还是背景的两个学习嵌入之一的总和。一个框由一对嵌入表示

(1) 其左上角的位置编码与表示“左上角”的学习嵌入的总和,以及 (2) 相同的结构,但使用表示“右下角”的学习嵌入。

框和点编码被称为“稀疏编码”。

如果传递了掩码提示,则使用卷积神经网络将其缩小以生成“密集编码”。如果没有传递掩码提示,则改为使用嵌入层生成“无掩码”嵌入。

参数

  • embed_dim (int, 可选): 输出嵌入中的特征数量。默认为 256
  • image_embedding_size (int, 可选): 图像编码器生成的图像嵌入中的特征数量。默认为 (64, 64)
  • input_image_size (tuple[int], 可选): 被提示的图像的高度和宽度的元组。默认为 (1024, 1024)
  • mask_in_chans (int, 可选): 掩码提示的通道数。默认为 16
  • activation (str, 可选): 在掩码降尺度神经网络中使用的激活函数。默认为 "gelu"

参考文献


[源代码]

TwoWayTransformer

keras_cv.models.TwoWayTransformer(
    depth=2,
    embed_dim=256,
    num_heads=8,
    mlp_dim=2048,
    activation="relu",
    attention_downsample_rate=2,
    **kwargs
)

双向交叉注意力 Transformer 解码器。

一个 Transformer 解码器,它使用查询来关注输入图像,这些查询的位置嵌入由提供。

变压器解码器的设计如图 [1]_ 所示。每个解码器层执行 4 个步骤:(1) 对标记进行自注意力,(2) 从标记(作为查询)到图像嵌入进行交叉注意力,(3) 点式 MLP 更新每个标记,以及 (4) 从图像嵌入(作为查询)到标记进行交叉注意力。最后一步使用提示信息更新图像嵌入。每个自注意力/交叉注意力和 MLP 都有一个残差连接和层归一化。

为了确保解码器能够访问关键的几何信息,每当图像嵌入参与注意力层时,都会向其添加位置编码。此外,每当原始提示标记(包括其位置编码)参与注意力层时,都会将其重新添加到更新后的标记中。这使得对提示标记的几何位置和类型都有强烈的依赖性。

参数

  • depth(int,可选):注意力块的深度(要使用的注意力块的数量)。默认为 2
  • embed_dim(int,可选):输入图像和点嵌入的特征数量。默认为 256
  • num_heads(int,可选):注意力层中使用的头部数量。默认为 8
  • mlp_dim(int,可选):注意力层中使用的 MLP 块的隐藏层中的单元数量。默认为 2048
  • activation(str,可选):注意力层中使用的 MLP 块的输出层的激活函数。默认为 "relu"
  • attention_downsample_rate(int,可选):注意力层的下采样率。默认为 2

参考文献


[源代码]

MultiHeadAttentionWithDownsampling

keras_cv.layers.MultiHeadAttentionWithDownsampling(
    num_heads, key_dim, downsample_rate=1, **kwargs
)

具有下采样的多头注意力。

一个允许在投影到查询、键和值后缩放下嵌入大小的注意力层。

此层首先使用密集层缩放下输入查询、键和值的特征。然后执行多头注意力,并将注意力图投影回(上采样)到输入特征的数量。

参数

  • num_heads(int):注意力头的数量。
  • key_dim(int):查询、键和值每个注意力头的尺寸。
  • downsample_rate(int,可选):缩放下输入特征的因子,即大小为 key_dim 的输入特征被投影到 key_dim // downsample_rate

参考文献


[源代码]

TwoWayMultiHeadAttention

keras_cv.layers.TwoWayMultiHeadAttention(
    num_heads,
    key_dim,
    mlp_dim,
    skip_first_layer_pe,
    attention_downsample_rate=2,
    activation="relu",
    **kwargs
)

双向多头注意力层。

参数

  • num_heads(int):注意力头的数量。
  • key_dim(int):查询、键和值每个注意力头的尺寸。
  • mlp_dim(int):在 mlp 块中使用的隐藏维度数。
  • skip_first_layer_pe(bool):一个布尔值,指示是否跳过第一层位置嵌入。
  • attention_downsample_rate(int,可选):在注意力层中使用的下采样率。默认为 2。
  • activation(str,可选):mlp 块输出层的激活函数。默认为“relu”。

参考文献


[源代码]

RandomFrequencyPositionalEmbeddings

keras_cv.layers.RandomFrequencyPositionalEmbeddings(
    num_positional_features, scale, **kwargs
)

使用随机空间频率的位置编码。

此层使用随机空间频率将 2D 空间中的坐标/点映射到位置编码。

参数

  • num_positional_features(int):输出中的位置特征数量。
  • scale(float):随机频率的标准差。

参考文献