数据集文档

语义分割

Hugging Face's logo
加入 Hugging Face 社区

并获取增强型文档体验

开始

语义分割

语义分割数据集用于训练模型对图像中的每个像素进行分类。这些数据集可以实现多种应用,例如图像的背景去除、图像风格化或自动驾驶的场景理解。本指南将向您展示如何对图像分割数据集应用转换。

在开始之前,请确保您已安装最新的 albumentationscv2 版本

pip install -U albumentations opencv-python

Albumentations 是一个 Python 库,用于对计算机视觉执行数据增强。它支持各种计算机视觉任务,例如图像分类、目标检测、分割和关键点估计。

本指南使用 场景解析 数据集来分割和解析图像,将其分解成与语义类别(例如天空、道路、人、床)相关的不同图像区域。

加载数据集的 train 拆分并查看一个示例

>>> 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:图像的分割掩码。
  • scene_category:图像的标签或场景类别(例如“厨房”或“办公室”)。

接下来,查看带有

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

类似地,您可以查看相应的分割掩码

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

我们还可以添加一个 调色板 到分割掩码上,并将其叠加到原始图像上以可视化数据集

定义调色板后,您就可以开始可视化一些叠加效果了。

>>> 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() 函数对数据集的批次进行实时转换,以减少磁盘空间占用。

>>> 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)

现在您已经了解如何处理语义分割数据集,请学习 如何训练语义分割模型 并将其用于推理。

< > 在 GitHub 上更新