語意分割

語意分割資料集用於訓練模型，以對影像中的每個像素進行分類。這些資料集可應用於廣泛的領域，例如影像背景移除、影像風格化，或自動駕駛中的場景理解。本指南將向您展示如何對影像分割資料集進行轉換 (transformations)。

開始之前，請確保您已安裝最新版本的 albumentations 和 cv2

pip install -U albumentations opencv-python

Albumentations 是一個用於執行電腦視覺資料增強 (data augmentation) 的 Python 函式庫。它支援多種電腦視覺任務，如影像分類、物件偵測、分割以及關鍵點估計。

本指南使用場景解析 (Scene Parsing) 資料集，將影像分割並解析為與語意類別相關聯的不同影像區域，例如天空、道路、人員和床鋪。

載入資料集的 train 切分 (split) 並查看其中一個範例

>>> from datasets import load_dataset

>>> dataset = load_dataset("scene_parse_150", split="train")
>>> index = 10
>>> dataset[index]
{'image': <PIL.JpegImagePlugin.JpegImageFile image mode=RGB size=683x512 at 0x7FB37B0EC810>,
 'annotation': <PIL.PngImagePlugin.PngImageFile image mode=L size=683x512 at 0x7FB37B0EC9D0>,
 'scene_category': 927}

該資料集包含三個欄位

image：PIL 影像物件。
annotation：影像的分割遮罩 (segmentation mask)。
scene_category：影像的標籤或場景類別（例如“廚房”或“辦公室”）。

接下來，查看一張影像：

>>> dataset[index]["image"]

同樣地，您可以查看對應的分割遮罩：

>>> dataset[index]["annotation"]

我們還可以在分割遮罩上添加顏色調色盤 (color palette)，並將其疊加在原始影像上，以便視覺化資料集。

定義好顏色調色盤後，您就可以開始視覺化一些疊加效果了。

>>> import matplotlib.pyplot as plt

>>> def visualize_seg_mask(image: np.ndarray, mask: np.ndarray):
...    color_seg = np.zeros((mask.shape[0], mask.shape[1], 3), dtype=np.uint8)
...    palette = np.array(create_ade20k_label_colormap())
...    for label, color in enumerate(palette):
...        color_seg[mask == label, :] = color
...    color_seg = color_seg[..., ::-1]  # convert to BGR

...    img = np.array(image) * 0.5 + color_seg * 0.5  # plot the image with the segmentation map
...    img = img.astype(np.uint8)

...    plt.figure(figsize=(15, 10))
...    plt.imshow(img)
...    plt.axis("off")
...    plt.show()


>>> visualize_seg_mask(
...     np.array(dataset[index]["image"]),
...     np.array(dataset[index]["annotation"])
... )

現在使用 albumentations 套用一些增強功能。您將首先調整影像大小並調整其亮度。

>>> import albumentations

>>> transform = albumentations.Compose(
...     [
...         albumentations.Resize(256, 256),
...         albumentations.RandomBrightnessContrast(brightness_limit=0.3, contrast_limit=0.3, p=0.5),
...     ]
... )

建立一個函數來對影像套用轉換

>>> def transforms(examples):
...     transformed_images, transformed_masks = [], []
...
...     for image, seg_mask in zip(examples["image"], examples["annotation"]):
...         image, seg_mask = np.array(image), np.array(seg_mask)
...         transformed = transform(image=image, mask=seg_mask)
...         transformed_images.append(transformed["image"])
...         transformed_masks.append(transformed["mask"])
...
...     examples["pixel_values"] = transformed_images
...     examples["label"] = transformed_masks
...     return examples

使用 set_transform() 函數，在資料集批次處理時即時 (on-the-fly) 套用轉換，以節省硬碟空間。

>>> dataset.set_transform(transforms)

您可以透過索引範例中的 pixel_values 和 label 來驗證轉換是否成功。

>>> image = np.array(dataset[index]["pixel_values"])
>>> mask = np.array(dataset[index]["label"])

>>> visualize_seg_mask(image, mask)

在本指南中，您使用了 albumentations 來增強資料集。您也可以使用 torchvision 來套用類似的轉換。

>>> from torchvision.transforms import Resize, ColorJitter, Compose

>>> transformation_chain = Compose([
...     Resize((256, 256)),
...     ColorJitter(brightness=0.25, contrast=0.25, saturation=0.25, hue=0.1)
... ])
>>> resize = Resize((256, 256))

>>> def train_transforms(example_batch):
...     example_batch["pixel_values"] = [transformation_chain(x) for x in example_batch["image"]]
...     example_batch["label"] = [resize(x) for x in example_batch["annotation"]]
...     return example_batch

>>> dataset.set_transform(train_transforms)

>>> image = np.array(dataset[index]["pixel_values"])
>>> mask = np.array(dataset[index]["label"])

>>> visualize_seg_mask(image, mask)

既然您已經知道如何處理語意分割的資料集，接下來請學習如何訓練語意分割模型並將其用於推論 (inference)。

在 GitHub 上更新

資料集 (Datasets)

語意分割