Datasets documentation
Семантическая сегментация
Семантическая сегментация
Наборы данных для семантической сегментации используются для обучения модели классифицировать каждый пиксель изображения. Эти наборы данных обеспечивают широкий спектр приложений, таких как удаление фона с изображений, стилизация изображений или понимание сцен для автономного вождения. В этом руководстве показано, как применять преобразования к набору данных для сегментации изображений.
Перед началом убедитесь, что у вас установлены актуальные версии albumentations и cv2:
pip install -U albumentations opencv-python
Albumentations — это библиотека на Python для выполнения аугментации данных в задачах компьютерного зрения. Она поддерживает различные задачи компьютерного зрения, такие как классификация изображений, обнаружение объектов, сегментация и оценка ключевых точек.
В этом руководстве используется набор данных Scene Parsing для сегментации и разбора изображения на различные области, связанные с семантическими категориями, такими как небо, дорога, человек и кровать.
Загрузите разделение 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)
Update on GitHubТеперь, когда вы знаете, как обработать датасет для семантической сегментации, изучите как обучить модель семантической сегментации и использовать её для инференса.