Transformers documentation
Segmentación de imágenes
Segmentación de imágenes
Los modelos de segmentación de imágenes separan las áreas que corresponden a distintas regiones de interés en una imagen. Estos modelos funcionan asignando una etiqueta a cada píxel. Existen varios tipos de segmentación: segmentación semántica, segmentación de instancias y segmentación panóptica.
En esta guía, haremos lo siguiente:
- Echa un vistazo a los diferentes tipos de segmentación.
- Disponer de un ejemplo de ajuste fino de extremo a extremo para la segmentación semántica.
Antes de comenzar, asegúrate de tener instaladas todas las bibliotecas necesarias:
# uncomment to install the necessary libraries
!pip install -q datasets transformers evaluate accelerate trackioTe animamos a iniciar sesión en tu cuenta de Hugging Face para que puedas subir y compartir tu modelo con la comunidad. Cuando se te solicite, introduce tu token para iniciar sesión:
>>> from huggingface_hub import notebook_login
>>> notebook_login()Tipos de segmentación
La segmentación semántica asigna una etiqueta o clase a cada píxel de una imagen. Veamos la salida de un modelo de segmentación semántica. Asignará la misma clase a cada instancia de un objeto que encuentre en la imagen; por ejemplo, todos los gatos se etiquetarían como “cat” en lugar de “cat-1”, “cat-2”. Podemos usar el pipeline de segmentación de imágenes de transformers para inferir rápidamente un modelo de segmentación semántica. Veamos la imagen de ejemplo.
from transformers import pipeline
from PIL import Image
import requests
url = "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/transformers/tasks/segmentation_input.jpg"
image = Image.open(requests.get(url, stream=True).raw)
image
Utilizaremos nvidia/segformer-b1-finetuned-cityscapes-1024-1024.
semantic_segmentation = pipeline("image-segmentation", "nvidia/segformer-b1-finetuned-cityscapes-1024-1024")
results = semantic_segmentation(image)
resultsLa salida del pipeline de segmentación incluye una máscara para cada clase predicha.
[{'score': None,
'label': 'road',
'mask': <PIL.Image.Image image mode=L size=612x415>},
{'score': None,
'label': 'sidewalk',
'mask': <PIL.Image.Image image mode=L size=612x415>},
{'score': None,
'label': 'building',
'mask': <PIL.Image.Image image mode=L size=612x415>},
{'score': None,
'label': 'wall',
'mask': <PIL.Image.Image image mode=L size=612x415>},
{'score': None,
'label': 'pole',
'mask': <PIL.Image.Image image mode=L size=612x415>},
{'score': None,
'label': 'traffic sign',
'mask': <PIL.Image.Image image mode=L size=612x415>},
{'score': None,
'label': 'vegetation',
'mask': <PIL.Image.Image image mode=L size=612x415>},
{'score': None,
'label': 'terrain',
'mask': <PIL.Image.Image image mode=L size=612x415>},
{'score': None,
'label': 'sky',
'mask': <PIL.Image.Image image mode=L size=612x415>},
{'score': None,
'label': 'car',
'mask': <PIL.Image.Image image mode=L size=612x415>}]Al examinar la máscara de la clase de coche, podemos ver que cada coche se clasifica con la misma máscara.
results[-1]["mask"]
En la segmentación de instancias, el objetivo no es clasificar cada píxel, sino predecir una máscara para cada instancia de un objeto en una imagen dada. Funciona de manera muy similar a la detección de objetos, donde en lugar de un recuadro delimitador para cada instancia, se utiliza una máscara de segmentación. Para esto, utilizaremos facebook/mask2former-swin-large-cityscapes-instance.
instance_segmentation = pipeline("image-segmentation", "facebook/mask2former-swin-large-cityscapes-instance")
results = instance_segmentation(image)
resultsComo se puede ver a continuación, hay varios coches clasificados y no hay clasificación para los píxeles que no pertenecen a instancias de coche ni de persona.
[{'score': 0.999944,
'label': 'car',
'mask': <PIL.Image.Image image mode=L size=612x415>},
{'score': 0.999945,
'label': 'car',
'mask': <PIL.Image.Image image mode=L size=612x415>},
{'score': 0.999652,
'label': 'car',
'mask': <PIL.Image.Image image mode=L size=612x415>},
{'score': 0.903529,
'label': 'person',
'mask': <PIL.Image.Image image mode=L size=612x415>}]Revisando una de las máscaras de coches de abajo.
results[2]["mask"]
La segmentación panóptica combina la segmentación semántica y la segmentación de instancias, en la que cada píxel se clasifica en una clase y una instancia de esa clase, y existen múltiples máscaras para cada instancia de una clase. Para esto, podemos usar facebook/mask2former-swin-large-cityscapes-panoptic.
panoptic_segmentation = pipeline("image-segmentation", "facebook/mask2former-swin-large-cityscapes-panoptic")
results = panoptic_segmentation(image)
resultsComo se puede ver a continuación, tenemos más clases. Más adelante ilustraremos que cada píxel se clasifica en una de las clases.
[{'score': 0.999981,
'label': 'car',
'mask': <PIL.Image.Image image mode=L size=612x415>},
{'score': 0.999958,
'label': 'car',
'mask': <PIL.Image.Image image mode=L size=612x415>},
{'score': 0.99997,
'label': 'vegetation',
'mask': <PIL.Image.Image image mode=L size=612x415>},
{'score': 0.999575,
'label': 'pole',
'mask': <PIL.Image.Image image mode=L size=612x415>},
{'score': 0.999958,
'label': 'building',
'mask': <PIL.Image.Image image mode=L size=612x415>},
{'score': 0.999634,
'label': 'road',
'mask': <PIL.Image.Image image mode=L size=612x415>},
{'score': 0.996092,
'label': 'sidewalk',
'mask': <PIL.Image.Image image mode=L size=612x415>},
{'score': 0.999221,
'label': 'car',
'mask': <PIL.Image.Image image mode=L size=612x415>},
{'score': 0.99987,
'label': 'sky',
'mask': <PIL.Image.Image image mode=L size=612x415>}]Hagamos una comparación lado a lado para todos los tipos de segmentación.
Ahora que hemos visto todos los tipos de segmentación, profundicemos en el ajuste fino de un modelo para la segmentación semántica.
Entre las aplicaciones habituales de la segmentación semántica en el mundo real se encuentran el entrenamiento de vehículos autónomos para identificar peatones e información vial importante, la identificación de células y anomalías en imágenes médicas y el monitoreo de cambios ambientales a partir de imágenes satelitales.
Ajuste fino de un modelo para segmentación
Ahora haremos lo siguiente:
- Ajusta SegFormer en el conjunto de datos SceneParse150.
- Usa tu modelo afinado para la inferencia.
Para ver todas las arquitecturas y puntos de control compatibles con esta tarea, recomendamos consultar la página de la tarea
Cargar el conjunto de datos SceneParse150
Comience cargando un subconjunto más pequeño del conjunto de datos SceneParse150 desde la biblioteca 🤗 Datasets. Esto le dará la oportunidad de experimentar y asegurarse de que todo funcione correctamente antes de dedicar más tiempo a entrenar con el conjunto de datos completo.
>>> from datasets import load_dataset
>>> ds = load_dataset("merve/scene_parse_150", split="train[:50]")Divide la partición train del conjunto de datos en un conjunto de entrenamiento y uno de prueba con el método train_test_split:
>>> ds = ds.train_test_split(test_size=0.2)
>>> train_ds = ds["train"]
>>> test_ds = ds["test"]A continuación, echa un vistazo a un ejemplo:
>>> train_ds[0]
{'image': <PIL.JpegImagePlugin.JpegImageFile image mode=RGB size=512x683 at 0x7F9B0C201F90>,
'annotation': <PIL.PngImagePlugin.PngImageFile image mode=L size=512x683 at 0x7F9B0C201DD0>,
'scene_category': 368}
# view the image
>>> train_ds[0]["image"]image: una imagen PIL de la escena.annotation: una imagen PIL del mapa de segmentación, que también es el objetivo del modelo.scene_category: un identificador de categoría que describe la escena de la imagen, como “kitchen” o “office”. En esta guía solo necesitarásimageeannotation, ambos de los cuales son imágenes PIL.
También querrás crear un diccionario que asocie un id de etiqueta con una clase de etiqueta, lo cual será útil al configurar el modelo más adelante. Descarga las correspondencias desde el Hub y crea los diccionarios id2label y label2id:
>>> import json
>>> from pathlib import Path
>>> from huggingface_hub import hf_hub_download
>>> repo_id = "huggingface/label-files"
>>> filename = "ade20k-id2label.json"
>>> id2label = json.loads(Path(hf_hub_download(repo_id, filename, repo_type="dataset")).read_text())
>>> id2label = {int(k): v for k, v in id2label.items()}
>>> label2id = {v: k for k, v in id2label.items()}
>>> num_labels = len(id2label)Preprocesar
El siguiente paso es cargar un procesador de imágenes de SegFormer para preparar las imágenes y las anotaciones para el modelo. Algunos conjuntos de datos, como este, utilizan el índice cero como clase de fondo. Sin embargo, la clase de fondo no está incluida realmente en las 150 clases, por lo que deberá establecer do_reduce_labels=True para restar uno a todas las etiquetas. El índice cero se reemplaza por 255 para que la función de pérdida de SegFormer lo ignore:
>>> from transformers import AutoImageProcessor
>>> checkpoint = "nvidia/mit-b0"
>>> image_processor = AutoImageProcessor.from_pretrained(checkpoint, do_reduce_labels=True)Es habitual aplicar algunas aumentaciones de datos a un conjunto de datos de imágenes para que un modelo sea más robusto frente al sobreajuste. En esta guía, utilizará la función ColorJitter de torchvision para cambiar aleatoriamente las propiedades de color de una imagen, pero también puede usar cualquier biblioteca de imágenes que desee.
>>> from torchvision.transforms import ColorJitter
>>> jitter = ColorJitter(brightness=0.25, contrast=0.25, saturation=0.25, hue=0.1)Ahora crea dos funciones de preprocesamiento para preparar las imágenes y las anotaciones para el modelo. Estas funciones convierten las imágenes en pixel_values y las anotaciones en labels. Para el conjunto de entrenamiento, se aplica jitter antes de pasar las imágenes al procesador de imágenes. Para el conjunto de prueba, el procesador de imágenes recorta y normaliza las images, y solo recorta las labels porque no se aplica aumento de datos durante la prueba.
>>> def train_transforms(example_batch):
... images = [jitter(x) for x in example_batch["image"]]
... labels = [x for x in example_batch["annotation"]]
... inputs = image_processor(images, labels)
... return inputs
>>> def val_transforms(example_batch):
... images = [x for x in example_batch["image"]]
... labels = [x for x in example_batch["annotation"]]
... inputs = image_processor(images, labels)
... return inputsPara aplicar jitter a todo el conjunto de datos, use la función set_transform de 🤗 Datasets. La transformación se aplica sobre la marcha, lo que es más rápido y consume menos espacio en disco:
>>> train_ds.set_transform(train_transforms)
>>> test_ds.set_transform(val_transforms)Evaluar
Incluir una métrica durante el entrenamiento suele ser útil para evaluar el rendimiento de tu modelo. Puedes cargar rápidamente un método de evaluación con la biblioteca 🤗 Evaluate. Para esta tarea, carga la métrica mean Intersection over Union (IoU) (consulta el quick tour de 🤗 Evaluate para obtener más información sobre cómo cargar y calcular una métrica):
>>> import evaluate
>>> metric = evaluate.load("mean_iou")Luego, crea una función para calcular las métricas. Es necesario convertir primero tus predicciones a logits y, a continuación, redimensionarlas para que coincidan con el tamaño de las etiquetas antes de poder llamar a compute:
>>> import numpy as np
>>> import torch
>>> from torch import nn
>>> def compute_metrics(eval_pred):
... with torch.no_grad():
... logits, labels = eval_pred
... logits_tensor = torch.from_numpy(logits)
... logits_tensor = nn.functional.interpolate(
... logits_tensor,
... size=labels.shape[-2:],
... mode="bilinear",
... align_corners=False,
... ).argmax(dim=1)
... pred_labels = logits_tensor.detach().cpu().numpy()
... metrics = metric.compute(
... predictions=pred_labels,
... references=labels,
... num_labels=num_labels,
... ignore_index=255,
... reduce_labels=False,
... )
... for key, value in metrics.items():
... if isinstance(value, np.ndarray):
... metrics[key] = value.tolist()
... return metricsTu función compute_metrics ya está lista y volverás a ella cuando configures el entrenamiento.
Entrenar
Si no está familiarizado con el ajuste fino de un modelo con Trainer, ¡eche un vistazo al tutorial básico aquí!
¡Ya está listo para comenzar a entrenar su modelo! Cargue SegFormer con AutoModelForSemanticSegmentation y pase al modelo el mapeo entre los ids de etiquetas y las clases de etiquetas:
>>> from transformers import AutoModelForSemanticSegmentation, TrainingArguments, Trainer
>>> model = AutoModelForSemanticSegmentation.from_pretrained(checkpoint, id2label=id2label, label2id=label2id)En este punto, solo quedan tres pasos:
- Define los hiperparámetros de entrenamiento en TrainingArguments. Es importante que no elimines las columnas no utilizadas, ya que esto descartaría la columna
image. Sin la columnaimage, no puedes crearpixel_values. Estableceremove_unused_columns=Falsepara evitar este comportamiento. El único otro parámetro obligatorio esoutput_dir, que especifica dónde guardar tu modelo. Subirás este modelo al Hub estableciendopush_to_hub=True(necesitas iniciar sesión en Hugging Face para subir tu modelo). Al final de cada época, Trainer evaluará la métrica IoU y guardará el punto de control de entrenamiento. - Pasa los argumentos de entrenamiento a Trainer junto con el modelo, el conjunto de datos, el tokenizador, el data collator y la función
compute_metrics. - Llama a train() para ajustar tu modelo.
>>> training_args = TrainingArguments(
... output_dir="segformer-b0-scene-parse-150",
... learning_rate=6e-5,
... num_train_epochs=50,
... per_device_train_batch_size=2,
... per_device_eval_batch_size=2,
... save_total_limit=3,
... eval_strategy="steps",
... save_strategy="steps",
... save_steps=20,
... eval_steps=20,
... logging_steps=1,
... report_to="trackio",
... run_name="scene-parse-150",
... eval_accumulation_steps=5,
... remove_unused_columns=False,
... push_to_hub=True,
... )
>>> trainer = Trainer(
... model=model,
... args=training_args,
... train_dataset=train_ds,
... eval_dataset=test_ds,
... compute_metrics=compute_metrics,
... )
>>> trainer.train()Una vez completado el entrenamiento, comparte tu modelo en el Hub con el método push_to_hub() para que todo el mundo pueda usarlo:
>>> trainer.push_to_hub()Inferencia
Genial, ahora que has afinado un modelo, ¡puedes usarlo para inferencia!
Recargue el conjunto de datos y cargue una imagen para la inferencia.
>>> from datasets import load_dataset
>>> ds = load_dataset("merve/scene_parse_150", split="train[:50]")
>>> ds = ds.train_test_split(test_size=0.2)
>>> test_ds = ds["test"]
>>> image = ds["test"][0]["image"]
>>> image
Ahora veremos cómo realizar la inferencia sin una pipeline. Procesa la imagen con un procesador de imágenes y coloca pixel_values en una GPU:
>>> from accelerate import Accelerator
>>> device = Accelerator().device
>>> encoding = image_processor(image, return_tensors="pt")
>>> pixel_values = encoding.pixel_values.to(device)Pasa tu entrada al modelo y devuelve los logits:
>>> outputs = model(pixel_values=pixel_values)
>>> logits = outputs.logits.cpu()A continuación, reescala los logits al tamaño de la imagen original:
>>> upsampled_logits = nn.functional.interpolate(
... logits,
... size=image.size[::-1],
... mode="bilinear",
... align_corners=False,
... )
>>> pred_seg = upsampled_logits.argmax(dim=1)[0]Para visualizar los resultados, cargue la paleta de colores del conjunto de datos como ade_palette(), que asigna a cada clase sus valores RGB.
def ade_palette():
return np.asarray([
[0, 0, 0],
[120, 120, 120],
[180, 120, 120],
[6, 230, 230],
[80, 50, 50],
[4, 200, 3],
[120, 120, 80],
[140, 140, 140],
[204, 5, 255],
[230, 230, 230],
[4, 250, 7],
[224, 5, 255],
[235, 255, 7],
[150, 5, 61],
[120, 120, 70],
[8, 255, 51],
[255, 6, 82],
[143, 255, 140],
[204, 255, 4],
[255, 51, 7],
[204, 70, 3],
[0, 102, 200],
[61, 230, 250],
[255, 6, 51],
[11, 102, 255],
[255, 7, 71],
[255, 9, 224],
[9, 7, 230],
[220, 220, 220],
[255, 9, 92],
[112, 9, 255],
[8, 255, 214],
[7, 255, 224],
[255, 184, 6],
[10, 255, 71],
[255, 41, 10],
[7, 255, 255],
[224, 255, 8],
[102, 8, 255],
[255, 61, 6],
[255, 194, 7],
[255, 122, 8],
[0, 255, 20],
[255, 8, 41],
[255, 5, 153],
[6, 51, 255],
[235, 12, 255],
[160, 150, 20],
[0, 163, 255],
[140, 140, 140],
[250, 10, 15],
[20, 255, 0],
[31, 255, 0],
[255, 31, 0],
[255, 224, 0],
[153, 255, 0],
[0, 0, 255],
[255, 71, 0],
[0, 235, 255],
[0, 173, 255],
[31, 0, 255],
[11, 200, 200],
[255, 82, 0],
[0, 255, 245],
[0, 61, 255],
[0, 255, 112],
[0, 255, 133],
[255, 0, 0],
[255, 163, 0],
[255, 102, 0],
[194, 255, 0],
[0, 143, 255],
[51, 255, 0],
[0, 82, 255],
[0, 255, 41],
[0, 255, 173],
[10, 0, 255],
[173, 255, 0],
[0, 255, 153],
[255, 92, 0],
[255, 0, 255],
[255, 0, 245],
[255, 0, 102],
[255, 173, 0],
[255, 0, 20],
[255, 184, 184],
[0, 31, 255],
[0, 255, 61],
[0, 71, 255],
[255, 0, 204],
[0, 255, 194],
[0, 255, 82],
[0, 10, 255],
[0, 112, 255],
[51, 0, 255],
[0, 194, 255],
[0, 122, 255],
[0, 255, 163],
[255, 153, 0],
[0, 255, 10],
[255, 112, 0],
[143, 255, 0],
[82, 0, 255],
[163, 255, 0],
[255, 235, 0],
[8, 184, 170],
[133, 0, 255],
[0, 255, 92],
[184, 0, 255],
[255, 0, 31],
[0, 184, 255],
[0, 214, 255],
[255, 0, 112],
[92, 255, 0],
[0, 224, 255],
[112, 224, 255],
[70, 184, 160],
[163, 0, 255],
[153, 0, 255],
[71, 255, 0],
[255, 0, 163],
[255, 204, 0],
[255, 0, 143],
[0, 255, 235],
[133, 255, 0],
[255, 0, 235],
[245, 0, 255],
[255, 0, 122],
[255, 245, 0],
[10, 190, 212],
[214, 255, 0],
[0, 204, 255],
[20, 0, 255],
[255, 255, 0],
[0, 153, 255],
[0, 41, 255],
[0, 255, 204],
[41, 0, 255],
[41, 255, 0],
[173, 0, 255],
[0, 245, 255],
[71, 0, 255],
[122, 0, 255],
[0, 255, 184],
[0, 92, 255],
[184, 255, 0],
[0, 133, 255],
[255, 214, 0],
[25, 194, 194],
[102, 255, 0],
[92, 0, 255],
])Luego, puedes combinar y graficar tu imagen y el mapa de segmentación predicho:
>>> import matplotlib.pyplot as plt
>>> import numpy as np
>>> color_seg = np.zeros((pred_seg.shape[0], pred_seg.shape[1], 3), dtype=np.uint8)
>>> palette = np.array(ade_palette())
>>> for label, color in enumerate(palette):
... color_seg[pred_seg == 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.show()