开发者指南 / KerasCV / 语义分割与 KerasCV

语义分割与 KerasCV

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

在 Colab 中查看 GitHub 源码


背景

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

KerasCV 提供了 Google 开发的 DeepLabv3+ 模型用于语义分割。本教程演示了如何微调和使用 DeepLabv3+ 模型进行图像语义分割。它的架构结合了空洞卷积、上下文信息聚合和强大的骨干网络,以实现准确和详细的语义分割。DeepLabv3+ 模型在多个图像分割基准上显示出卓越的结果。

参考文献

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


设置和导入

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

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

  • keras-cv
  • keras-core
!pip install -q --upgrade keras-cv
!pip install -q --upgrade keras # 升级到 Keras 3.

安装 keras-corekeras-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 工具函数将输入预处理为一个包含 imagessegmentation_masks 的字典。图像和分割掩码被调整为 512x512。生成的数据集随后被批量分为每组 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,  # oxford iiit pet 数据集的类别数量。VOC 数据集还包括 1 个背景类。
    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


模型配置

请随意修改模型训练的配置,并注意训练结果的变化。这是一个很好的练习,可以更好地理解训练流程。

学习率调度由优化器使用,以计算每个周期的学习率。然后,优化器使用学习率来更新模型的权重。在这种情况下,学习率调度使用余弦衰减函数。余弦衰减函数一开始较高,然后逐渐减少,最终达到零。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,
)

我们实例化一个 DeepLabV3+ 模型,使用在 ImageNet 分类上预训练的 ResNet50 主干: resnet50_v2_imagenet 预训练权重将作为 DeepLabV3Plus 模型的主干特征提取器。num_classes 参数指定模型将训练分割的类别数量。

model = keras_cv.models.DeepLabV3Plus.from_preset(
    "resnet50_v2_imagenet", num_classes=NUM_CLASSES
)
从 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()
模型: "deep_lab_v3_plus_1"
┏━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━┓
┃ 层 (类型)         输出形状       参数 #  连接到         ┃
┡━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━┩
│ input_layer_9       │ (, ,      │       0 │ -                    │
│ (输入层)        │ , 3)          │         │                      │
├─────────────────────┼───────────────────┼─────────┼──────────────────────┤
│ functional_11       │ [(, ,     │ 23,556… │ input_layer_9[0][0]  │
│ (功能性)        │ , 256),       │         │                      │
│                     │ (, ,      │         │                      │
│                     │ , 2048)]      │         │                      │
├─────────────────────┼───────────────────┼─────────┼──────────────────────┤
│ spatial_pyramid_po… │ (, ,      │ 15,538… │ functional_11[0][1]  │
│ (空间金字塔池化, 256)        │         │                      │
├─────────────────────┼───────────────────┼─────────┼──────────────────────┤
│ encoder_output_ups… │ (, ,      │       0 │ spatial_pyramid_poo… │
│ (上采样2D)      │ , 256)        │         │                      │
├─────────────────────┼───────────────────┼─────────┼──────────────────────┤
│ sequential_14       │ (, ,      │  12,480 │ functional_11[0][0]  │
│ (顺序)        │ , 48)         │         │                      │
├─────────────────────┼───────────────────┼─────────┼──────────────────────┤
│ concatenate_1       │ (, ,      │       0 │ encoder_output_upsa… │
│ (连接)       │ , 304)        │         │ sequential_14[0][0]  │
├─────────────────────┼───────────────────┼─────────┼──────────────────────┤
│ sequential_15       │ (, ,      │  84,224 │ concatenate_1[0][0]  │
│ (顺序)        │ , 21)         │         │                      │
└─────────────────────┴───────────────────┴─────────┴──────────────────────┘
 总参数: 39,191,488 (149.50 MB)
 可训练参数: 39,146,464 (149.33 MB)
 非可训练参数: 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/未知  735s 319ms/步 - 类别准确率: 0.7026 - 损失: 1.2143 - 平均IOU: 0.0706

/usr/lib/python3.10/contextlib.py:153: 用户警告: 您的输入数据耗尽;中断训练。确保您的数据集或生成器可以生成至少 `steps_per_epoch * epochs` 批次。您可能需要在构建数据集时使用 `.repeat()` 函数。
  self.gen.throw(typ, value, traceback)

 2124/2124 ━━━━━━━━━━━━━━━━━━━━ 813s 356ms/步 - 类别准确率: 0.7026 - 损失: 1.2143 - 平均IOU: 0.0706 - 验证类别准确率: 0.7768 - 验证损失: 0.8223 - 验证平均IOU: 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 架构,效率更高,准确性优于先前的图像分割模型。