SAM图像分割模型

[source]

SAMImageSegmenter

keras_hub.models.SAMImageSegmenter(backbone, preprocessor=None, **kwargs)

分割一切 (SAM) 图像分割模型。

SAM 通过提示输入图像来工作。有三种提示方式:(1)标记点:前景点(标签为 1 的点)被编码,使得掩码解码器生成的输出掩码包含它们,背景点(标签为 0 的点)被编码,使得生成的掩码不包含它们。(2)框:框告诉模型要分割的图像的哪个部分/裁剪区域。(3)掩码:输入掩码可用于细化掩码解码器的输出。这些提示可以混合和匹配,但必须至少存在一个提示。要关闭特定提示,只需将其从模型的输入中排除即可。(1) 对于点提示,期望的形状是 (batch, num_points, 2)。标签必须具有相应的形状 (batch, num_points)。(2) 对于框提示,期望的形状是 (batch, 1, 2, 2)。(3) 类似地,掩码提示具有形状 (batch, 1, H, W, 1)

参数

示例

使用 from_preset 加载预训练模型。

image_size=128
batch_size=2
input_data = {
    "images": np.ones(
        (batch_size, image_size, image_size, 3),
        dtype="float32",
    ),
    "points": np.ones((batch_size, 1, 2), dtype="float32"),
    "labels": np.ones((batch_size, 1), dtype="float32"),
    "boxes": np.ones((batch_size, 1, 2, 2), dtype="float32"),
    "masks": np.zeros(
        (batch_size, 0, image_size, image_size, 1)
    ),
}
sam = keras_hub.models.SAMImageSegmenter.from_preset('sam_base_sa1b')
outputs = sam.predict(input_data)
masks, iou_pred = outputs["masks"], outputs["iou_pred"]

加载带有自定义骨干网络的分割一切图像分割器

image_size = 128
batch_size = 2
images = np.ones(
    (batch_size, image_size, image_size, 3),
    dtype="float32",
)
image_encoder = keras_hub.models.ViTDetBackbone(
    hidden_size=16,
    num_layers=16,
    intermediate_dim=16 * 4,
    num_heads=16,
    global_attention_layer_indices=[2, 5, 8, 11],
    patch_size=16,
    num_output_channels=8,
    window_size=2,
    image_shape=(image_size, image_size, 3),
)
prompt_encoder = keras_hub.layers.SAMPromptEncoder(
    hidden_size=8,
    image_embedding_size=(8, 8),
    input_image_size=(
        image_size,
        image_size,
    ),
    mask_in_channels=16,
)
mask_decoder = keras_hub.layers.SAMMaskDecoder(
    num_layers=2,
    hidden_size=8,
    intermediate_dim=32,
    num_heads=8,
    embedding_dim=8,
    num_multimask_outputs=3,
    iou_head_depth=3,
    iou_head_hidden_dim=8,
)
backbone = keras_hub.models.SAMBackbone(
    image_encoder=image_encoder,
    prompt_encoder=prompt_encoder,
    mask_decoder=mask_decoder,
)
sam = keras_hub.models.SAMImageSegmenter(
    backbone=backbone
)

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

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))
# Prepare an input dictionary:
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"]

另一个示例是仅存在点提示。请注意,如果存在点提示但不存在框提示,则必须使用零点和 -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"]

[source]

from_preset 方法

SAMImageSegmenter.from_preset(preset, load_weights=True, **kwargs)

从模型预设实例化 keras_hub.models.Task

预设是用于保存和加载预训练模型的配置、权重和其他文件资产的目录。preset 可以作为以下之一传递

  1. 内置预设标识符,例如 'bert_base_en'
  2. Kaggle Models 句柄,例如 'kaggle://user/bert/keras/bert_base_en'
  3. Hugging Face 句柄,例如 'hf://user/bert_base_en'
  4. 本地预设目录的路径,例如 './bert_base_en'

对于任何 Task 子类,您可以运行 cls.presets.keys() 以列出该类上所有可用的内置预设。

此构造函数可以通过两种方式之一调用。可以从特定于任务的基类(如 keras_hub.models.CausalLM.from_preset())调用,也可以从模型类(如 keras_hub.models.BertTextClassifier.from_preset())调用。如果从基类调用,则返回对象的子类将从预设目录中的配置推断出来。

参数

  • preset: 字符串。内置预设标识符、Kaggle Models 句柄、Hugging Face 句柄或本地目录的路径。
  • load_weights: 布尔值。如果为 True,则保存的权重将加载到模型架构中。如果为 False,则所有权重将随机初始化。

示例

# Load a Gemma generative task.
causal_lm = keras_hub.models.CausalLM.from_preset(
    "gemma_2b_en",
)

# Load a Bert classification task.
model = keras_hub.models.TextClassifier.from_preset(
    "bert_base_en",
    num_classes=2,
)
预设 参数 描述
sam_base_sa1b 93.74M 在 SA1B 数据集上训练的基础 SAM 模型。
sam_huge_sa1b 312.34M 在 SA1B 数据集上训练的巨型 SAM 模型。
sam_large_sa1b 641.09M 在 SA1B 数据集上训练的大型 SAM 模型。

backbone 属性

keras_hub.models.SAMImageSegmenter.backbone

具有核心架构的 keras_hub.models.Backbone 模型。


preprocessor 属性

keras_hub.models.SAMImageSegmenter.preprocessor

用于预处理输入的 keras_hub.models.Preprocessor 层。