► 开发者指南 / KerasCV / 使用 KerasCV 进行语义分割

使用 KerasCV 进行语义分割

作者： Divyashree Sreepathihalli，Ian Stenbit
创建日期 2023/08/22
上次修改日期 2023/08/24
描述：使用 KerasCV 训练和使用 DeepLabv3+ 分割模型。

在 Colab 中查看 • GitHub 源代码

背景

语义分割是一种计算机视觉任务，它涉及为图像的每个像素分配一个类别标签，例如人、自行车或背景，有效地将图像划分为对应于不同对象类别或类别的区域。

KerasCV 提供了 Google 开发的用于语义分割的 DeepLabv3+ 模型。本指南演示了如何使用 KerasCV 对 DeepLabv3+ 模型进行微调并将其用于图像语义分割。其架构结合了空洞卷积、上下文信息聚合和强大的骨干网络，以实现准确而详细的语义分割。DeepLabv3+ 模型已被证明可以在各种图像分割基准测试中取得最先进的结果。

参考文献

用于语义图像分割的具有空洞可分离卷积的编码器-解码器
重新思考用于语义图像分割的空洞卷积

设置和导入

让我们安装依赖项并导入必要的模块。

要运行本教程，您需要安装以下软件包

  • keras-cv
  • keras-core
!pip install -q --upgrade keras-cv
!pip install -q --upgrade keras # Upgrade to Keras 3.

安装 keras-core 和 keras-cv 后，设置 keras-core 的后端。本指南可以使用任何后端（Tensorflow、JAX、PyTorch）运行。

import os

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

import keras
from keras import ops

import keras_cv
import numpy as np

from keras_cv.datasets.pascal_voc.segmentation import load as load_voc

使用预训练的 DeepLabv3+ 模型执行语义分割

KerasCV 语义分割 API 中的最高级 API 是 keras_cv.models API。此 API 包括完全预训练的语义分割模型，例如 keras_cv.models.DeepLabV3Plus。

让我们从构建一个在 pascalvoc 数据集上预训练的 DeepLabv3+ 模型开始。

model = keras_cv.models.DeepLabV3Plus.from_preset(
    "deeplab_v3_plus_resnet50_pascalvoc",
    num_classes=21,
    input_shape=[512, 512, 3],
)

让我们可视化此预训练模型的结果

filepath = keras.utils.get_file(origin="https://i.imgur.com/gCNcJJI.jpg")
image = keras.utils.load_img(filepath)

resize = keras_cv.layers.Resizing(height=512, width=512)
image = resize(image)
image = keras.ops.expand_dims(np.array(image), axis=0)
preds = ops.expand_dims(ops.argmax(model(image), axis=-1), axis=-1)
keras_cv.visualization.plot_segmentation_mask_gallery(
    image,
    value_range=(0, 255),
    num_classes=1,
    y_true=None,
    y_pred=preds,
    scale=3,
    rows=1,
    cols=1,
)

png

训练自定义语义分割模型

在本指南中，我们将为 KerasCV DeepLabV3 语义分割模型组装一个完整的训练管道。这包括数据加载、增强、训练、指标评估和推理！

下载数据

我们使用 KerasCV 数据集下载 Pascal VOC 数据集，并将其拆分为训练数据集 train_ds 和 eval_ds。

train_ds = load_voc(split="sbd_train")
eval_ds = load_voc(split="sbd_eval")

预处理数据

preprocess_tfds_inputs 实用程序函数将输入预处理为包含 images 和 segmentation_masks 的字典。图像和分割掩码被调整为 512x512 的大小。然后，生成的 dataset 被批处理为 4 个图像和分割掩码对的组。

可以使用 keras_cv.visualization.plot_segmentation_mask_gallery 函数可视化此预处理输入训练数据的批次。此函数以图像和分割掩码的批次作为输入，并将其显示在网格中。

def preprocess_tfds_inputs(inputs):
    def unpackage_tfds_inputs(tfds_inputs):
        return {
            "images": tfds_inputs["image"],
            "segmentation_masks": tfds_inputs["class_segmentation"],
        }

    outputs = inputs.map(unpackage_tfds_inputs)
    outputs = outputs.map(keras_cv.layers.Resizing(height=512, width=512))
    outputs = outputs.batch(4, drop_remainder=True)
    return outputs


train_ds = preprocess_tfds_inputs(train_ds)
batch = train_ds.take(1).get_single_element()
keras_cv.visualization.plot_segmentation_mask_gallery(
    batch["images"],
    value_range=(0, 255),
    num_classes=21,  # The number of classes for the oxford iiit pet dataset. The VOC dataset also includes 1 class for the background.
    y_true=batch["segmentation_masks"],
    scale=3,
    rows=2,
    cols=2,
)

png

预处理应用于评估数据集 eval_ds。

eval_ds = preprocess_tfds_inputs(eval_ds)

数据增强

KerasCV 提供了各种图像增强选项。在本例中，我们将使用 RandomFlip 增强来增强训练数据集。RandomFlip 增强随机水平或垂直翻转训练数据集中的图像。这有助于提高模型对图像中对象方向变化的鲁棒性。

train_ds = train_ds.map(keras_cv.layers.RandomFlip())
batch = train_ds.take(1).get_single_element()

keras_cv.visualization.plot_segmentation_mask_gallery(
    batch["images"],
    value_range=(0, 255),
    num_classes=21,
    y_true=batch["segmentation_masks"],
    scale=3,
    rows=2,
    cols=2,
)

png

模型配置

请随意修改模型训练的配置，并注意训练结果的变化。这是一个更好地了解训练管道的绝佳练习。

学习率调度程序用于优化器计算每个 epoch 的学习率。然后，优化器使用学习率更新模型的权重。在本例中，学习率调度程序使用余弦衰减函数。余弦衰减函数从高开始，然后随着时间的推移逐渐减小，最终达到零。VOC 数据集的基数为 2124，批次大小为 4。数据集基数对于学习率衰减很重要，因为它决定了模型将训练多少步。初始学习率与 0.007 成正比，衰减步数为 2124。这意味着学习率将从 INITIAL_LR 开始，然后在 2124 步内衰减至零。 png

BATCH_SIZE = 4
INITIAL_LR = 0.007 * BATCH_SIZE / 16
EPOCHS = 1
NUM_CLASSES = 21
learning_rate = keras.optimizers.schedules.CosineDecay(
    INITIAL_LR,
    decay_steps=EPOCHS * 2124,
)

我们使用在 ImageNet 分类上预训练的 ResNet50 骨干网络实例化 DeepLabV3+ 模型：将 resnet50_v2_imagenet 预训练权重用作 DeepLabV3Plus 模型的骨干特征提取器。num_classes 参数指定模型将训练分割的类别数。

model = keras_cv.models.DeepLabV3Plus.from_preset(
    "resnet50_v2_imagenet", num_classes=NUM_CLASSES
)

Downloading data from https://storage.googleapis.com/keras-cv/models/resnet50v2/imagenet/classification-v2-notop.h5
 94687928/94687928 ━━━━━━━━━━━━━━━━━━━━ 1s 0us/step

编译模型

model.compile() 函数设置模型的训练过程。它定义了 - 优化算法 - 随机梯度下降 (SGD) - 损失函数 - 分类交叉熵 - 评估指标 - 平均 IoU 和分类准确率

语义分割评估指标

平均交并比 (MeanIoU)：MeanIoU 衡量语义分割模型在多大程度上准确地识别和描绘图像中不同的对象或区域。它计算预测对象边界和实际对象边界之间的重叠，提供 0 到 1 之间的分数，其中 1 表示完美匹配。

分类准确率：分类准确率衡量图像中正确分类像素的比例。它给出一个简单的百分比，指示模型在多大程度上准确地预测整个图像中像素的类别。

从本质上讲，MeanIoU 强调识别特定对象边界的准确性，而分类准确率则提供对整体像素级正确性的总体概述。

model.compile(
    optimizer=keras.optimizers.SGD(
        learning_rate=learning_rate, weight_decay=0.0001, momentum=0.9, clipnorm=10.0
    ),
    loss=keras.losses.CategoricalCrossentropy(from_logits=False),
    metrics=[
        keras.metrics.MeanIoU(
            num_classes=NUM_CLASSES, sparse_y_true=False, sparse_y_pred=False
        ),
        keras.metrics.CategoricalAccuracy(),
    ],
)

model.summary()

Model: "deep_lab_v3_plus_1"

┏━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━┓
┃ Layer (type)        ┃ Output Shape      ┃ Param # ┃ Connected to         ┃
┡━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━┩
│ input_layer_9       │ (None, None,      │       0 │ -                    │
│ (InputLayer)        │ None, 3)          │         │                      │
├─────────────────────┼───────────────────┼─────────┼──────────────────────┤
│ functional_11       │ [(None, None,     │ 23,556… │ input_layer_9[0][0]  │
│ (Functional)        │ None, 256),       │         │                      │
│                     │ (None, None,      │         │                      │
│                     │ None, 2048)]      │         │                      │
├─────────────────────┼───────────────────┼─────────┼──────────────────────┤
│ spatial_pyramid_po… │ (None, None,      │ 15,538… │ functional_11[0][1]  │
│ (SpatialPyramidPoo… │ None, 256)        │         │                      │
├─────────────────────┼───────────────────┼─────────┼──────────────────────┤
│ encoder_output_ups… │ (None, None,      │       0 │ spatial_pyramid_poo… │
│ (UpSampling2D)      │ None, 256)        │         │                      │
├─────────────────────┼───────────────────┼─────────┼──────────────────────┤
│ sequential_14       │ (None, None,      │  12,480 │ functional_11[0][0]  │
│ (Sequential)        │ None, 48)         │         │                      │
├─────────────────────┼───────────────────┼─────────┼──────────────────────┤
│ concatenate_1       │ (None, None,      │       0 │ encoder_output_upsa… │
│ (Concatenate)       │ None, 304)        │         │ sequential_14[0][0]  │
├─────────────────────┼───────────────────┼─────────┼──────────────────────┤
│ sequential_15       │ (None, None,      │  84,224 │ concatenate_1[0][0]  │
│ (Sequential)        │ None, 21)         │         │                      │
└─────────────────────┴───────────────────┴─────────┴──────────────────────┘

 Total params: 39,191,488 (149.50 MB)

 Trainable params: 39,146,464 (149.33 MB)

 Non-trainable params: 45,024 (175.88 KB)

实用程序函数 dict_to_tuple 有效地将训练和验证数据集的字典转换为图像和独热编码分割掩码的元组，这在 DeepLabv3+ 模型的训练和评估期间使用。

def dict_to_tuple(x):
    import tensorflow as tf

    return x["images"], tf.one_hot(
        tf.cast(tf.squeeze(x["segmentation_masks"], axis=-1), "int32"), 21
    )


train_ds = train_ds.map(dict_to_tuple)
eval_ds = eval_ds.map(dict_to_tuple)

model.fit(train_ds, validation_data=eval_ds, epochs=EPOCHS)

   2124/Unknown  735s 319ms/step - categorical_accuracy: 0.7026 - loss: 1.2143 - mean_io_u: 0.0706

/usr/lib/python3.10/contextlib.py:153: UserWarning: Your input ran out of data; interrupting training. Make sure that your dataset or generator can generate at least `steps_per_epoch * epochs` batches. You may need to use the `.repeat()` function when building your dataset.
  self.gen.throw(typ, value, traceback)

 2124/2124 ━━━━━━━━━━━━━━━━━━━━ 813s 356ms/step - categorical_accuracy: 0.7026 - loss: 1.2143 - mean_io_u: 0.0706 - val_categorical_accuracy: 0.7768 - val_loss: 0.8223 - val_mean_io_u: 0.1593

<keras.src.callbacks.history.History at 0x7f261a534640>

使用训练好的模型进行预测

现在 DeepLabv3+ 的模型训练已经完成，让我们通过对一些示例图像进行预测来测试它。

test_ds = load_voc(split="sbd_eval")
test_ds = preprocess_tfds_inputs(test_ds)

images, masks = next(iter(train_ds.take(1)))
images = ops.convert_to_tensor(images)
masks = ops.convert_to_tensor(masks)
preds = ops.expand_dims(ops.argmax(model(images), axis=-1), axis=-1)
masks = ops.expand_dims(ops.argmax(masks, axis=-1), axis=-1)

keras_cv.visualization.plot_segmentation_mask_gallery(
    images,
    value_range=(0, 255),
    num_classes=21,
    y_true=masks,
    y_pred=preds,
    scale=3,
    rows=1,
    cols=4,
)

png

以下是一些使用 KerasCV DeepLabv3+ 模型的其他技巧

  • 该模型可以在各种数据集上进行训练，包括 COCO 数据集、PASCAL VOC 数据集和 Cityscapes 数据集。
  • 该模型可以在自定义数据集上进行微调，以提高其在特定任务上的性能。
  • 该模型可用于对图像执行实时推理。
  • 此外，还可以尝试 KerasCV 的 SegFormer 模型 keras_cv.models.segmentation.SegFormer。SegFormer 模型是一个较新的模型，已被证明可以在各种图像分割基准测试中取得最先进的结果。它基于 Swin Transformer 架构，并且比以前的图像分割模型更有效和准确。
使用 KerasCV 进行语义分割
◆ 背景
参考文献
◆ 设置和导入
◆ 使用预训练的 DeepLabv3+ 模型执行语义分割
◆ 训练自定义语义分割模型
◆ 下载数据
◆ 预处理数据
◆ 数据增强
◆ 模型配置
◆ 编译模型
◆ 使用训练好的模型进行预测