Datasets documentation
Segmentación semántica
Segmentación semántica
Los conjuntos de datos de segmentación semántica se utilizan para entrenar un modelo que clasifique cada píxel de una imagen. Existen una amplia variedad de aplicaciones habilitadas por estos conjuntos de datos, como la eliminación del fondo de imágenes, la estilización de imágenes o la comprensión de escenas para la conducción autónoma. Esta guía le mostrará cómo aplicar transformaciones a un conjunto de datos de segmentación de imágenes.
Antes de comenzar, asegúrate de tener instaladas las versiones más recientes de albumentations y cv2:
pip install -U albumentations opencv-python
Albumentations es una biblioteca de Python para realizar aumentación de datos para visión por computadora. Admite diversas tareas de visión por computadora, como clasificación de imágenes, detección de objetos, segmentación y estimación de puntos clave.
Esta guía utiliza el conjunto de datos Scene Parsing para segmentar y analizar una imagen en diferentes regiones asociadas con categorías semánticas, como cielo, carretera, persona y cama.
Cargue la división train del conjunto de datos y examine un ejemplo:
>>> 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}El conjunto de datos tiene tres campos:
image: un objeto de imagen PIL.annotation: máscara de segmentación de la imagen.scene_category: la etiqueta o categoría de escena de la imagen (como “kitchen” o “office”).
A continuación, comprueba una imagen con:
>>> dataset[index]["image"]
De manera similar, puedes consultar la máscara de segmentación correspondiente:
>>> dataset[index]["annotation"]
También podemos añadir una paleta de colores a la máscara de segmentación y superponerla sobre la imagen original para visualizar el conjunto de datos:
Después de definir la paleta de colores, estará listo para visualizar algunas superposiciones.
>>> 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"])
... )
Ahora aplica algunas aumentaciones con albumentations. Primero, redimensiona la imagen y ajusta su brillo.
>>> import albumentations
>>> transform = albumentations.Compose(
... [
... albumentations.Resize(256, 256),
... albumentations.RandomBrightnessContrast(brightness_limit=0.3, contrast_limit=0.3, p=0.5),
... ]
... )Cree una función para aplicar la transformación a las imágenes:
>>> 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 examplesUse la función set_transform() para aplicar la transformación al vuelo a los lotes del conjunto de datos y consumir menos espacio en disco:
>>> dataset.set_transform(transforms)Puedes verificar que la transformación funcionó indexando en pixel_values y label de un ejemplo:
>>> image = np.array(dataset[index]["pixel_values"])
>>> mask = np.array(dataset[index]["label"])
>>> visualize_seg_mask(image, mask)
En esta guía, has utilizado albumentations para aumentar el conjunto de datos. También es posible usar torchvision para aplicar algunas transformaciones similares.
>>> 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 GitHubAhora que sabe cómo procesar un conjunto de datos para segmentación semántica, aprenda cómo entrenar un modelo de segmentación semántica y usarlo para inferencia.