► 开发者指南 / KerasCV / KerasCV中的任意分割！

KerasCV中的任意分割！

作者： Tirth Patel, Ian Stenbit
创建日期： 2023/12/04
最后修改日期： 2023/12/19
描述： 在KerasCV中使用文本、框和点提示进行任意分割。

概述

Segment Anything Model (SAM) 从输入提示（如点或框）生成高质量的对象掩膜，并可以用于为图像中的所有对象生成掩膜。它已在一个包含1100万张图片和11亿个掩膜的数据集上进行训练，并在各种分割任务中表现出强大的零-shot性能。

在本指南中，我们将展示如何使用KerasCV实现的 Segment Anything Model 并展示TensorFlow和JAX性能提升的强大之处。

首先，让我们获取所有依赖项和演示用的图像。

!pip install -Uq keras-cv
!pip install -Uq keras

!wget -q https://raw.githubusercontent.com/facebookresearch/segment-anything/main/notebooks/images/truck.jpg

选择您的后端

使用Keras 3，您可以选择使用您喜欢的后端！

import os

os.environ["KERAS_BACKEND"] = "jax"

import timeit
import numpy as np
import matplotlib.pyplot as plt
import keras
from keras import ops
import keras_cv

辅助函数

让我们定义一些辅助函数以可视化图像、提示和分割结果。

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)


def show_points(coords, labels, ax, marker_size=375):
    pos_points = coords[labels == 1]
    neg_points = coords[labels == 0]
    ax.scatter(
        pos_points[:, 0],
        pos_points[:, 1],
        color="green",
        marker="*",
        s=marker_size,
        edgecolor="white",
        linewidth=1.25,
    )
    ax.scatter(
        neg_points[:, 0],
        neg_points[:, 1],
        color="red",
        marker="*",
        s=marker_size,
        edgecolor="white",
        linewidth=1.25,
    )


def show_box(box, ax):
    box = box.reshape(-1)
    x0, y0 = box[0], box[1]
    w, h = box[2] - box[0], box[3] - box[1]
    ax.add_patch(
        plt.Rectangle((x0, y0), w, h, edgecolor="green", facecolor=(0, 0, 0, 0), lw=2)
    )


def inference_resizing(image, pad=True):
    # 计算预处理形状
    image = ops.cast(image, dtype="float32")
    old_h, old_w = image.shape[0], image.shape[1]
    scale = 1024 * 1.0 / max(old_h, old_w)
    new_h = old_h * scale
    new_w = old_w * scale
    preprocess_shape = int(new_h + 0.5), int(new_w + 0.5)

    # 调整图像大小
    image = ops.image.resize(image[None, ...], preprocess_shape)[0]

    # 填充较短的一侧
    if pad:
        pixel_mean = ops.array([123.675, 116.28, 103.53])
        pixel_std = ops.array([58.395, 57.12, 57.375])
        image = (image - pixel_mean) / pixel_std
        h, w = image.shape[0], image.shape[1]
        pad_h = 1024 - h
        pad_w = 1024 - w
        image = ops.pad(image, [(0, pad_h), (0, pad_w), (0, 0)])
        # KerasCV现在重新缩放图像并对其进行规范化。
        # 只需反规范化，以便当KerasCV重新规范化时
        # 填充值映射到0。
        image = image * pixel_std + pixel_mean
    return image

获取预训练的SAM模型

我们可以使用KerasCV的from_preset工厂方法初始化一个训练好的SAM模型。在这里，我们使用在SA-1B数据集上训练的巨型ViT主干（sam_huge_sa1b）来获取高质量的分割掩膜。您还可以使用sam_large_sa1b或sam_base_sa1b之一以获得更好的性能（代价是降低分割掩膜的质量）。

model = keras_cv.models.SegmentAnythingModel.from_preset("sam_huge_sa1b")

理解提示

Segment Anything允许使用点、框和掩膜对图像进行提示：

点提示是最基本的提示：模型尝试根据图像中的一个点来猜测对象。该点可以是前景点（即所需分割掩膜中包含该点）或背景点（即该点位于所需掩膜之外）。
另一个提示模型的方法是使用边界框。给定一个边界框，模型尝试分割其中包含的对象。
最后，模型也可以通过掩码本身进行提示。这在需要细化之前预测或已知的分割掩码边界时非常有用。

使该模型极其强大的原因是能够结合上述提示。点、框和掩码提示可以以多种不同方式组合，以达到最佳效果。

让我们看看如何将这些提示传递给 KerasCV 中的 Segment Anything 模型。SAM 模型的输入是一个字典，包含以下键：

"images"：要分割的一组图像。必须具有形状 (B, 1024, 1024, 3)。
"points"：一组点提示。每个点是一个 (x, y) 坐标，起源于图像的左上角。换句话说，每个点的形式是 (r, c) 其中 r 和 c 是图像中像素的行和列。必须具有形状 (B, N, 2)。
"labels"：给定点的一组标签。1 代表前景点，0 代表背景点。必须具有形状 (B, N)。
"boxes"：一组框。注意模型每个批次只接受一个框。因此，期望的形状是 (B, 1, 2, 2)。每个框由两个点组成：框的左上角和右下角。这里的点遵循与点提示相同的语义。这里第二维中的 1 表示框提示的存在。如果框提示缺失，必须传递形状为 (B, 0, 2, 2) 的占位符输入。
"masks"：一组掩码。与框提示一样，每个图像只能有一个掩码提示。如果存在，输入掩码的形状必须是 (B, 1, 256, 256, 1)，如果掩码提示缺失，则形状为 (B, 0, 256, 256, 1)。

占位符提示仅在直接调用模型时（即 model(...)）是必需的。当调用 predict 方法时，可以从输入字典中省略缺失的提示。

点提示

首先，让我们使用点提示对图像进行分割。我们加载图像并将其调整为形状 (1024, 1024)，这是预训练 SAM 模型所期望的图像大小。

# 加载我们的图像
image = np.array(keras.utils.load_img("truck.jpg"))
image = inference_resizing(image)

plt.figure(figsize=(10, 10))
plt.imshow(ops.convert_to_numpy(image) / 255.0)
plt.axis("on")
plt.show()

png

接下来，我们将定义我们想要分割的对象上的点。让我们尝试在坐标 (284, 213) 处分割卡车的窗户玻璃。

# 定义输入点提示
input_point = np.array([[284, 213.5]])
input_label = np.array([1])

plt.figure(figsize=(10, 10))
plt.imshow(ops.convert_to_numpy(image) / 255.0)
show_points(input_point, input_label, plt.gca())
plt.axis("on")
plt.show()

png

现在让我们调用模型的 predict 方法来获取分割掩码。

注意：我们不会直接调用模型（model(...)），因为这样做需要占位符提示。缺失的提示会被 predict 方法自动处理，因此我们调用它。而且，当没有框提示时，点和标签需要用零点提示和 -1 标签提示进行填充。下面的单元演示了这如何工作。

outputs = model.predict(
    {
        "images": image[np.newaxis, ...],
        "points": np.concatenate(
            [input_point[np.newaxis, ...], np.zeros((1, 1, 2))], axis=1
        ),
        "labels": np.concatenate(
            [input_label[np.newaxis, ...], np.full((1, 1), fill_value=-1)], axis=1
        ),
    }
)

 1/1 ━━━━━━━━━━━━━━━━━━━━ 48s 48s/step

SegmentAnythingModel.predict 返回两个输出。第一个是形状为 (1, 4, 256, 256) 的 logits（分割掩码），另一个是每个预测掩码的 IoU 置信度分数（形状为 (1, 4)）。预训练的 SAM 模型预测了四个掩码：第一个是模型针对给定提示所能给出的最佳掩码，其他三个是备用掩码，可以在最佳预测不包含所需对象时使用。用户可以选择他们偏好的掩码。

让我们可视化模型返回的掩码！

# 将掩码调整为我们的图像形状，即 (1024, 1024)
# 将 logits 转换为 numpy 数组
# 并将 logits 转换为布尔掩码

png

正如预期，模型返回了卡车窗户的分割掩码。然而，我们的点提示也可能意味着其他许多东西。例如，另一种可能的掩码包含我们的点，仅仅是窗户右侧或整个卡车。

让我们也可视化模型预测的其他掩码。

fig, ax = plt.subplots(1, 3, figsize=(20, 60))
masks, scores = outputs["masks"][0][1:], outputs["iou_pred"][0][1:]
for i, (mask, score) in enumerate(zip(masks, scores)):
    mask = inference_resizing(mask[..., None], pad=False)[..., 0]
    mask, score = map(ops.convert_to_numpy, (mask, score))
    mask = 1 * (mask > 0.0)
    ax[i].imshow(ops.convert_to_numpy(image) / 255.0)
    show_mask(mask, ax[i])
    show_points(input_point, input_label, ax[i])
    ax[i].set_title(f"Mask {i+1}, Score: {score:.3f}", fontsize=12)
    ax[i].axis("off")
plt.show()

png

很好！SAM 能够捕捉到我们点提示的模糊性，并且返回了其他可能的分割掩码。

盒子提示

现在，让我们看看如何使用盒子来提示模型。盒子由两个点指定，左上角和右下角的边界框，格式为 xyxy。让我们使用围绕卡车左前轮的边界框提示模型。

# 让我们指定盒子
input_box = np.array([[240, 340], [400, 500]])

outputs = model.predict(
    {"images": image[np.newaxis, ...], "boxes": input_box[np.newaxis, np.newaxis, ...]}
)
mask = inference_resizing(outputs["masks"][0][0][..., None], pad=False)[..., 0]
mask = ops.convert_to_numpy(mask) > 0.0

plt.figure(figsize=(10, 10))
plt.imshow(ops.convert_to_numpy(image) / 255.0)
show_mask(mask, plt.gca())
show_box(input_box, plt.gca())
plt.axis("off")
plt.show()

 1/1 ━━━━━━━━━━━━━━━━━━━━ 13s 13s/step

png

太好了！模型完美地分割出了我们的边界框中的左前轮。

组合提示

为了充分发挥模型的真正潜力，让我们结合盒子和点提示，看看模型会怎么做。

# 让我们指定盒子
input_box = np.array([[240, 340], [400, 500]])
# 让我们指定点并将其标记为背景
input_point = np.array([[325, 425]])
input_label = np.array([0])

outputs = model.predict(
    {
        "images": image[np.newaxis, ...],
        "points": input_point[np.newaxis, ...],
        "labels": input_label[np.newaxis, ...],
        "boxes": input_box[np.newaxis, np.newaxis, ...],
    }
)
mask = inference_resizing(outputs["masks"][0][0][..., None], pad=False)[..., 0]
mask = ops.convert_to_numpy(mask) > 0.0

plt.figure(figsize=(10, 10))
plt.imshow(ops.convert_to_numpy(image) / 255.0)
show_mask(mask, plt.gca())
show_box(input_box, plt.gca())
show_points(input_point, input_label, plt.gca())
plt.axis("off")
plt.show()

 1/1 ━━━━━━━━━━━━━━━━━━━━ 16s 16s/step

png

好极了！模型理解我们想要从掩码中排除的对象是轮胎的轮圈。

文本提示

最后，让我们看看文本提示如何与 KerasCV 的 SegmentAnythingModel 一起使用。

在这个演示中，我们将使用官方的 Grounding DINO 模型。 Grounding DINO 是一种模型，输入为 (image, text) 对，生成一个围绕 text 描述的对象的边界框。有关模型实现的更多细节，您可以参考论文。

在演示的这一部分，我们需要从源代码安装 groundingdino 包：

pip install -U git+https://github.com/IDEA-Research/GroundingDINO.git

然后，我们可以安装预训练模型的权重和配置：

!wget -q https://github.com/IDEA-Research/GroundingDINO/releases/download/v0.1.0-alpha/groundingdino_swint_ogc.pth
!wget -q https://raw.githubusercontent.com/IDEA-Research/GroundingDINO/v0.1.0-alpha2/groundingdino/config/GroundingDINO_SwinT_OGC.py

from groundingdino.util.inference import Model as GroundingDINO

CONFIG_PATH = "GroundingDINO_SwinT_OGC.py"
WEIGHTS_PATH = "groundingdino_swint_ogc.pth"

grounding_dino = GroundingDINO(CONFIG_PATH, WEIGHTS_PATH)

/home/tirthp/oss/virtualenvs/keras-io-dev/lib/python3.10/site-packages/torch/functional.py:504: 用户警告: torch.meshgrid: 在即将发布的版本中，必须传递索引参数。 (在 ../aten/src/ATen/native/TensorShape.cpp:3526 内部触发。)
  return _VF.meshgrid(tensors, **kwargs)  # type: ignore[attr-defined]

最终文本编码器类型: bert-base-uncased

让我们为这一部分加载一张狗的图片！

filepath = keras.utils.get_file(
    origin="https://storage.googleapis.com/keras-cv/test-images/mountain-dog.jpeg"
)
image = np.array(keras.utils.load_img(filepath))
image = ops.convert_to_numpy(inference_resizing(image))

plt.figure(figsize=(10, 10))
plt.imshow(image / 255.0)
plt.axis("on")
plt.show()

将输入数据裁剪到有效范围以用于带RGB数据的imshow（浮点数范围为[0..1]或整数范围为[0..255]）。

png

我们首先使用Grounding DINO模型预测我们要分割的对象的边界框。然后，我们使用边界框提示SAM模型以获取分割掩码。

让我们尝试分割狗的安全带。请更改下面的图像和文本，以使用图像中的文本分割你想要的任何内容！

# 让我们预测狗的安全带的边界框
boxes = grounding_dino.predict_with_caption(image.astype(np.uint8), "harness")
boxes = np.array(boxes[0].xyxy)

outputs = model.predict(
    {
        "images": np.repeat(image[np.newaxis, ...], boxes.shape[0], axis=0),
        "boxes": boxes.reshape(-1, 1, 2, 2),
    },
    batch_size=1,
)

/home/tirthp/oss/virtualenvs/keras-io-dev/lib/python3.10/site-packages/transformers/modeling_utils.py:942: FutureWarning: `device` 参数已弃用，将在Transformers版本5中移除。
  warnings.warn(
/home/tirthp/oss/virtualenvs/keras-io-dev/lib/python3.10/site-packages/torch/utils/checkpoint.py:429: UserWarning: torch.utils.checkpoint：请明确传入use_reentrant=True或use_reentrant=False。use_reentrant的默认值将来将更新为False。为了保持当前行为，请传入use_reentrant=True。建议您使用use_reentrant=False。有关两个变种之间差异的更多详细信息，请参阅文档。
  warnings.warn(
/home/tirthp/oss/virtualenvs/keras-io-dev/lib/python3.10/site-packages/torch/utils/checkpoint.py:61: UserWarning: 没有输入需要计算梯度。梯度将为None
  warnings.warn(

 1/1 ━━━━━━━━━━━━━━━━━━━━ 13s 13s/step

就这样！我们用Gounding DINO + SAM的组合为我们的文本提示得到了一个分割掩码！这是一种非常强大的技术，可以结合不同的模型来扩展应用！

让我们可视化结果。

plt.figure(figsize=(10, 10))
plt.imshow(image / 255.0)

for mask in outputs["masks"]:
    mask = inference_resizing(mask[0][..., None], pad=False)[..., 0]
    mask = ops.convert_to_numpy(mask) > 0.0
    show_mask(mask, plt.gca())
    show_box(boxes, plt.gca())

plt.axis("off")
plt.show()

将输入数据裁剪到有效范围以用于带RGB数据的imshow（浮点数范围为[0..1]或整数范围为[0..255]）。

png

优化SAM

您可以使用mixed_float16或bfloat16数据类型策略，以在相对较低的精度损失下获得巨大的速度提升和内存优化。

# 加载我们的图像
image = np.array(keras.utils.load_img("truck.jpg"))
image = inference_resizing(image)

# 指定提示
input_box = np.array([[240, 340], [400, 500]])

# 让我们先看看模型在float32数据类型下的速度
time_taken = timeit.repeat(
    'model.predict({"images": image[np.newaxis, ...], "boxes": input_box[np.newaxis, np.newaxis, ...]}, verbose=False)',
    repeat=3,
    number=3,
    globals=globals(),
)
print(f"使用float32数据类型所需时间: {min(time_taken) / 3:.10f}s")

# 在Keras中设置数据类型策略
keras.mixed_precision.set_global_policy("mixed_float16")

model = keras_cv.models.SegmentAnythingModel.from_preset("sam_huge_sa1b")

time_taken = timeit.repeat(
    'model.predict({"images": image[np.newaxis, ...], "boxes": input_box[np.newaxis, np.newaxis, ...]}, verbose=False)',
    repeat=3,
    number=3,
    globals=globals(),
)
print(f"使用float16数据类型所需时间: {min(time_taken) / 3:.10f}s")

使用float32数据类型所需时间: 0.5304666963s
使用float16数据类型所需时间: 0.1586400040s

这是KerasCV实现与原始PyTorch实现的比较！

benchmark

用于生成基准测试的脚本在这里。

结论

KerasCV的SegmentAnythingModel支持多种应用，并且在Keras 3的帮助下，可以在TensorFlow、JAX和PyTorch上运行该模型！借助XLA 在JAX和TensorFlow中，模型的运行速度比原始实现快好几倍。此外，使用Keras的混合精度支持只需一行代码即可优化内存使用和计算时间！

有关更高级的用法，请查看自动掩码生成器演示。

KerasCV中的任意分割！

◆ 概述

◆ 选择您的后端

◆ 辅助函数

◆ 获取预训练的SAM模型

◆ 理解提示

◆ 点提示

◆ 盒子提示

◆ 组合提示

◆ 文本提示

◆ 优化SAM

◆ 结论