► 開發者指南 / 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。然後將生成的數據集分批為 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

模型配置

請隨時修改模型訓練的配置，並注意訓練結果的變化。這是一個很好的練習，可以更好地理解訓練管道。

學習率計劃由優化器用於計算每個時期的學習率。然後，優化器使用學習率來更新模型的權重。在這種情況下，學習率計劃使用餘弦衰減函數。餘弦衰減函數從高開始，然後隨著時間的推移而減少，最終達到零。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+ 模型執行語義分割
◆ 訓練自訂語義分割模型
◆ 下載數據
◆ 預處理數據
◆ 數據增強
◆ 模型配置
◆ 編譯模型
◆ 使用訓練好的模型進行預測