SegmentAnythingModel
类keras_cv.models.SegmentAnythingModel(
backbone, prompt_encoder, mask_decoder, **kwargs
)
分割一切 (SAM) 模型。
参数
参考文献
示例
>>> 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 模型。
参数
None
,遵循预设是否有可用的预训练权重。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 输出之间的空间逐点乘积来预测掩码。
参数
256
。None
。当为 None
时,使用 keras_cv.models.TwoWayTransformer
层。1 + num_multimask_outputs
。默认为 3
。3
。256
。"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) 的提示编码器。
提示编码器为三种类型的提示生成编码
首先,将点提示和框提示连接起来,并使用随机空间频率生成位置编码。一个点表示为点位置的位置编码与指示该点是前景还是背景的两个学习嵌入之一的总和。一个框由一对嵌入表示
(1) 其左上角的位置编码与表示“左上角”的学习嵌入的总和,以及 (2) 相同的结构,但使用表示“右下角”的学习嵌入。
框和点编码被称为“稀疏编码”。
如果传递了掩码提示,则使用卷积神经网络将其缩小以生成“密集编码”。如果没有传递掩码提示,则改为使用嵌入层生成“无掩码”嵌入。
参数
256
。(64, 64)
。(1024, 1024)
。16
。"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 都有一个残差连接和层归一化。
为了确保解码器能够访问关键的几何信息,每当图像嵌入参与注意力层时,都会向其添加位置编码。此外,每当原始提示标记(包括其位置编码)参与注意力层时,都会将其重新添加到更新后的标记中。这使得对提示标记的几何位置和类型都有强烈的依赖性。
参数
2
。256
。8
。2048
。"relu"
。2
。参考文献
MultiHeadAttentionWithDownsampling
类keras_cv.layers.MultiHeadAttentionWithDownsampling(
num_heads, key_dim, downsample_rate=1, **kwargs
)
具有下采样的多头注意力。
一个允许在投影到查询、键和值后缩放下嵌入大小的注意力层。
此层首先使用密集层缩放下输入查询、键和值的特征。然后执行多头注意力,并将注意力图投影回(上采样)到输入特征的数量。
参数
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
)
双向多头注意力层。
参数
参考文献
RandomFrequencyPositionalEmbeddings
类keras_cv.layers.RandomFrequencyPositionalEmbeddings(
num_positional_features, scale, **kwargs
)
使用随机空间频率的位置编码。
此层使用随机空间频率将 2D 空间中的坐标/点映射到位置编码。
参数
参考文献