Ajustar un Modelo de Segmentación Semántica con un Conjunto de Datos Personalizado

'Ajustar modelo de segmentación semántica con conjunto de datos personalizado'

Esta guía muestra cómo se puede ajustar finamente Segformer, un modelo de segmentación semántica de última generación. Nuestro objetivo es construir un modelo para un robot de entrega de pizza, para que pueda ver dónde conducir y reconocer obstáculos 🍕🤖. Primero etiquetaremos un conjunto de imágenes de aceras en Segments.ai . Luego ajustaremos finamente un modelo pre-entrenado de SegFormer utilizando 🤗 transformers , una biblioteca de código abierto que ofrece implementaciones fáciles de usar de modelos de última generación. En el proceso, aprenderás cómo trabajar con el Hugging Face Hub, el catálogo de modelos y conjuntos de datos de código abierto más grande.

La segmentación semántica es la tarea de clasificar cada píxel en una imagen. Puedes verlo como una forma más precisa de clasificar una imagen. Tiene una amplia gama de casos de uso en campos como la medicina y la conducción autónoma. Por ejemplo, para nuestro robot de entrega de pizza, es importante saber exactamente dónde está la acera en una imagen, no solo si hay una acera o no.

Debido a que la segmentación semántica es un tipo de clasificación, las arquitecturas de red utilizadas para la clasificación de imágenes y la segmentación semántica son muy similares. En 2014, un documento seminal de Long et al. utilizó redes neuronales convolucionales para la segmentación semántica. Más recientemente, los Transformers se han utilizado para la clasificación de imágenes (por ejemplo, ViT ), y ahora también se utilizan para la segmentación semántica, llevando el estado de la técnica aún más lejos.

• Anunciando el Programa de Residencia de Investigación en IA de 🤗
• Habana Labs y Hugging Face se asocian para acelerar el entrenamiento de modelos Transformer
• Emisiones de CO2 y el 🤗 Hub Liderando la carga

SegFormer es un modelo de segmentación semántica presentado por Xie et al. en 2021. Tiene un codificador jerárquico de Transformers que no utiliza codificaciones posicionales (a diferencia de ViT) y un decodificador simple de perceptrón multicapa. SegFormer logra un rendimiento de última generación en múltiples conjuntos de datos comunes. Veamos cómo se desempeña nuestro robot de entrega de pizza en imágenes de aceras.

Comencemos instalando las dependencias necesarias. Debido a que vamos a subir nuestro conjunto de datos y modelo al Hugging Face Hub, necesitamos instalar Git LFS e iniciar sesión en Hugging Face.

La instalación de git-lfs puede ser diferente en tu sistema. Ten en cuenta que Google Colab tiene Git LFS preinstalado.

pip install -q transformers datasets evaluate segments-ai
apt-get install git-lfs
git lfs install
huggingface-cli login

El primer paso en cualquier proyecto de aprendizaje automático es reunir un buen conjunto de datos. Para entrenar un modelo de segmentación semántica, necesitamos un conjunto de datos con etiquetas de segmentación semántica. Podemos usar un conjunto de datos existente del Hugging Face Hub, como ADE20k , o crear nuestro propio conjunto de datos.

Para nuestro robot de entrega de pizza, podríamos usar un conjunto de datos existente de conducción autónoma como CityScapes o BDD100K . Sin embargo, estos conjuntos de datos fueron capturados por automóviles que conducen en la carretera. Dado que nuestro robot de entrega conducirá en la acera, habrá una discrepancia entre las imágenes en estos conjuntos de datos y los datos que nuestro robot verá en el mundo real.

No queremos que nuestro robot de entrega se confunda, por lo que crearemos nuestro propio conjunto de datos de segmentación semántica utilizando imágenes capturadas en aceras. Mostraremos cómo puedes etiquetar las imágenes que capturamos en los siguientes pasos. Si solo deseas utilizar nuestro conjunto de datos etiquetado y terminado, puedes omitir la sección “Crear tu propio conjunto de datos” y continuar desde “Usar un conjunto de datos del Hub” .

Crea tu propio conjunto de datos

Para crear tu conjunto de datos de segmentación semántica, necesitarás dos cosas:

1. imágenes que cubran las situaciones que tu modelo encontrará en el mundo real
2. etiquetas de segmentación, es decir, imágenes donde cada píxel representa una clase/categoría.

Ya hemos capturado mil imágenes de aceras en Bélgica. Recopilar y etiquetar dicho conjunto de datos puede llevar mucho tiempo, por lo que puedes comenzar con un conjunto de datos más pequeño y ampliarlo si el modelo no funciona lo suficientemente bien.

Algunos ejemplos de las imágenes sin procesar en el conjunto de datos de aceras.

Para obtener las etiquetas de segmentación, debemos indicar las clases de todas las regiones/objetos en estas imágenes. Esto puede ser una tarea que consume mucho tiempo, pero usar las herramientas adecuadas puede acelerar significativamente la tarea. Para la etiquetación, utilizaremos Segments.ai , ya que tiene herramientas inteligentes de etiquetación para la segmentación de imágenes y un SDK de Python fácil de usar.

Configura la tarea de etiquetación en Segments.ai

Primero, crea una cuenta en https://segments.ai/join . A continuación, crea un nuevo conjunto de datos y carga tus imágenes. Puedes hacer esto desde la interfaz web o a través del SDK de Python (consulta el cuaderno ).

Etiquetar las imágenes

Ahora que los datos sin procesar están cargados, ve a segments.ai/home y abre el conjunto de datos recién creado. Haz clic en “Comenzar a etiquetar” y crea máscaras de segmentación. Puedes usar las herramientas de superpíxeles y autosegmentación con inteligencia artificial para etiquetar más rápido.

Consejo: cuando uses la herramienta de superpíxeles, desplázate para cambiar el tamaño de los superpíxeles, y haz clic y arrastra para seleccionar segmentos.

Sube el resultado al Hugging Face Hub

Cuando hayas terminado de etiquetar, crea una nueva versión del conjunto de datos que contenga los datos etiquetados. Puedes hacerlo en la pestaña “Versiones” en Segments.ai, o programáticamente a través del SDK como se muestra en el notebook.

Ten en cuenta que la creación de la versión puede tardar unos segundos. Puedes verificar la pestaña “Versiones” en Segments.ai para comprobar si tu versión aún está siendo creada.

Ahora, convertiremos la versión en un conjunto de datos de Hugging Face utilizando el SDK de Python de Segments.ai. Si aún no has configurado el cliente de Python de Segments, sigue las instrucciones en la sección “Configurar la tarea de etiquetado en Segments.ai” del notebook.

Ten en cuenta que la conversión puede llevar un tiempo, dependiendo del tamaño de tu conjunto de datos.

from segments.huggingface import release2dataset

release = segments_client.get_release(dataset_identifier, release_name)
hf_dataset = release2dataset(release)

Si inspeccionamos las características del nuevo conjunto de datos, podemos ver la columna de imagen y la etiqueta correspondiente. La etiqueta consta de dos partes: una lista de anotaciones y un mapa de segmentación. La anotación corresponde a los diferentes objetos en la imagen. Para cada objeto, la anotación contiene un id y un category_id. El mapa de segmentación es una imagen donde cada píxel contiene el id del objeto en ese píxel. Se puede encontrar más información en la documentación relevante.

Para la segmentación semántica, necesitamos un mapa semántico que contenga un category_id para cada píxel. Usaremos la función get_semantic_bitmap del SDK de Segments.ai para convertir los mapas en mapas semánticos. Para aplicar esta función a todas las filas de nuestro conjunto de datos, usaremos dataset.map.

from segments.utils import get_semantic_bitmap

def convert_segmentation_bitmap(example):
    return {
        "label.segmentation_bitmap":
            get_semantic_bitmap(
                example["label.segmentation_bitmap"],
                example["label.annotations"],
                id_increment=0,
            )
    }


semantic_dataset = hf_dataset.map(
    convert_segmentation_bitmap,
)

También puedes reescribir la función convert_segmentation_bitmap para usar lotes y pasar batched=True a dataset.map. Esto acelerará significativamente el mapeo, pero es posible que debas ajustar el batch_size para asegurarte de que el proceso no se quede sin memoria.

El modelo SegFormer que vamos a ajustar más adelante espera nombres específicos para las características. Para mayor comodidad, igualaremos este formato ahora. Por lo tanto, renombraremos la característica image a pixel_values y la label.segmentation_bitmap a label y descartaremos las otras características.

semantic_dataset = semantic_dataset.rename_column('image', 'pixel_values')
semantic_dataset = semantic_dataset.rename_column('label.segmentation_bitmap', 'label')
semantic_dataset = semantic_dataset.remove_columns(['name', 'uuid', 'status', 'label.annotations'])

Ahora podemos subir el conjunto de datos transformado al Hugging Face Hub. De esta manera, tu equipo y la comunidad de Hugging Face podrán aprovecharlo. En la siguiente sección, veremos cómo puedes cargar el conjunto de datos desde el Hub.

hf_dataset_identifier = f"{hf_username}/{dataset_name}"

semantic_dataset.push_to_hub(hf_dataset_identifier)

Usar un conjunto de datos del Hub

Si no deseas crear tu propio conjunto de datos, pero encontraste un conjunto de datos adecuado para tu caso de uso en el Hugging Face Hub, puedes definir el identificador aquí.

Por ejemplo, puedes usar el conjunto de datos completo de aceras etiquetadas. Ten en cuenta que puedes ver los ejemplos directamente en tu navegador.

hf_dataset_identifier = "segments/sidewalk-semantic"

Ahora que hemos creado un nuevo conjunto de datos y lo hemos enviado al Hugging Face Hub, podemos cargar el conjunto de datos en una sola línea.

from datasets import load_dataset

ds = load_dataset(hf_dataset_identifier)

Comencemos mezclando el conjunto de datos y dividiéndolo en un conjunto de entrenamiento y un conjunto de prueba.

ds = ds.shuffle(seed=1)
ds = ds["train"].train_test_split(test_size=0.2)
train_ds = ds["train"]
test_ds = ds["test"]

Extraeremos el número de etiquetas y los identificadores legibles por humanos, para poder configurar correctamente el modelo de segmentación más adelante.

import json
from huggingface_hub import hf_hub_download

repo_id = f"datasets/{hf_dataset_identifier}"
filename = "id2label.json"
id2label = json.load(open(hf_hub_download(repo_id=hf_dataset_identifier, filename=filename, repo_type="dataset"), "r"))
id2label = {int(k): v for k, v in id2label.items()}
label2id = {v: k for k, v in id2label.items()}

num_labels = len(id2label)

Extractor de características y aumento de datos

Un modelo SegFormer espera que la entrada tenga una forma determinada. Para transformar nuestros datos de entrenamiento y que coincidan con la forma esperada, podemos usar SegFormerFeatureExtractor. Podríamos usar la función ds.map para aplicar el extractor de características a todo el conjunto de datos de entrenamiento de antemano, pero esto puede ocupar mucho espacio en disco. En su lugar, utilizaremos una transformación, que solo preparará un lote de datos cuando esos datos se utilicen realmente (sobre la marcha). De esta manera, podemos comenzar el entrenamiento sin esperar más procesamiento de datos.

En nuestra transformación, también definiremos algunas aumentaciones de datos para que nuestro modelo sea más resistente a diferentes condiciones de iluminación. Utilizaremos la función ColorJitter de torchvision para cambiar aleatoriamente el brillo, el contraste, la saturación y el tono de las imágenes en el lote.

from torchvision.transforms import ColorJitter
from transformers import SegformerFeatureExtractor

feature_extractor = SegformerFeatureExtractor()
jitter = ColorJitter(brightness=0.25, contrast=0.25, saturation=0.25, hue=0.1) 

def train_transforms(example_batch):
    images = [jitter(x) for x in example_batch['pixel_values']]
    labels = [x for x in example_batch['label']]
    inputs = feature_extractor(images, labels)
    return inputs


def val_transforms(example_batch):
    images = [x for x in example_batch['pixel_values']]
    labels = [x for x in example_batch['label']]
    inputs = feature_extractor(images, labels)
    return inputs


# Establecer transformaciones
train_ds.set_transform(train_transforms)
test_ds.set_transform(val_transforms)

Cargar el modelo para ajustar

Los autores de SegFormer definen 5 modelos con tamaños crecientes: B0 a B5. El siguiente gráfico (tomado del artículo original) muestra el rendimiento de estos diferentes modelos en el conjunto de datos ADE20K, en comparación con otros modelos.

Fuente

Aquí cargaremos el modelo SegFormer más pequeño (B0), pre-entrenado en ImageNet-1k. ¡Solo ocupa alrededor de 14 MB! Utilizar un modelo pequeño garantizará que nuestro modelo pueda funcionar sin problemas en nuestro robot de entrega de pizzas.

from transformers import SegformerForSemanticSegmentation

pretrained_model_name = "nvidia/mit-b0" 
model = SegformerForSemanticSegmentation.from_pretrained(
    pretrained_model_name,
    id2label=id2label,
    label2id=label2id
)

Configurar el entrenador

Para ajustar el modelo en nuestros datos, utilizaremos la API del Entrenador de Hugging Face. Necesitamos configurar la configuración de entrenamiento y una métrica de evaluación para usar un Entrenador.

Primero, configuraremos los TrainingArguments. Esto define todos los hiperparámetros de entrenamiento, como la tasa de aprendizaje y el número de épocas, la frecuencia para guardar el modelo, etc. También especificamos enviar el modelo al hub después del entrenamiento (push_to_hub=True) y especificamos un nombre de modelo (hub_model_id).

from transformers import TrainingArguments

epochs = 50
lr = 0.00006
batch_size = 2

hub_model_id = "segformer-b0-finetuned-segments-sidewalk-2"

training_args = TrainingArguments(
    "segformer-b0-finetuned-segments-sidewalk-outputs",
    learning_rate=lr,
    num_train_epochs=epochs,
    per_device_train_batch_size=batch_size,
    per_device_eval_batch_size=batch_size,
    save_total_limit=3,
    evaluation_strategy="steps",
    save_strategy="steps",
    save_steps=20,
    eval_steps=20,
    logging_steps=1,
    eval_accumulation_steps=5,
    load_best_model_at_end=True,
    push_to_hub=True,
    hub_model_id=hub_model_id,
    hub_strategy="end",
)

A continuación, definiremos una función que calcule la métrica de evaluación con la que queremos trabajar. Debido a que estamos realizando segmentación semántica, utilizaremos la Intersección sobre Unión media (mIoU), directamente accesible en la biblioteca evaluate. IoU representa la superposición de las máscaras de segmentación. La mIoU media es el promedio de la IoU de todas las clases semánticas. Echa un vistazo a esta publicación de blog para obtener una visión general de las métricas de evaluación para la segmentación de imágenes.

Debido a que nuestro modelo produce logits con dimensiones altura/4 y ancho/4, debemos ampliarlos antes de poder calcular la mIoU.

import torch
from torch import nn
import evaluate

metric = evaluate.load("mean_iou")

def compute_metrics(eval_pred):
  with torch.no_grad():
    logits, labels = eval_pred
    logits_tensor = torch.from_numpy(logits)
    # escalar los logits al tamaño de la etiqueta
    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()
    # actualmente usando _compute en lugar de compute
    # consulta este problema para obtener más información: https://github.com/huggingface/evaluate/pull/328#issuecomment-1286866576
    metrics = metric._compute(
            predictions=pred_labels,
            references=labels,
            num_labels=len(id2label),
            ignore_index=0,
            reduce_labels=feature_extractor.do_reduce_labels,
        )
    
    # agregar métricas por categoría como pares clave-valor individuales
    per_category_accuracy = metrics.pop("per_category_accuracy").tolist()
    per_category_iou = metrics.pop("per_category_iou").tolist()

    metrics.update({f"accuracy_{id2label[i]}": v for i, v in enumerate(per_category_accuracy)})
    metrics.update({f"iou_{id2label[i]}": v for i, v in enumerate(per_category_iou)})
    
    return metrics

Finalmente, podemos instanciar un objeto Trainer.

from transformers import Trainer

trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=train_ds,
    eval_dataset=test_ds,
    compute_metrics=compute_metrics,
)

Ahora que nuestro entrenador está configurado, el entrenamiento es tan simple como llamar a la función train. No necesitamos preocuparnos por administrar nuestra(s) GPU(s), el entrenador se encargará de eso.

trainer.train()

Cuando hayamos terminado con el entrenamiento, podemos subir nuestro modelo ajustado y el extractor de características a Hub.

Esto también creará automáticamente una tarjeta de modelo con nuestros resultados. Proporcionaremos información adicional en kwargs para que la tarjeta de modelo sea más completa.

kwargs = {
    "tags": ["vision", "segmentación-de-imágenes"],
    "ajustado_desde": pretrained_model_name,
    "conjunto_de_datos": hf_dataset_identifier,
}

feature_extractor.push_to_hub(hub_model_id)
trainer.push_to_hub(**kwargs)

Ahora viene la parte emocionante, ¡usar nuestro modelo ajustado! En esta sección, mostraremos cómo puedes cargar tu modelo desde Hub y usarlo para la inferencia.

Sin embargo, también puedes probar tu modelo directamente en Hugging Face Hub, gracias a los widgets geniales impulsados por la API de inferencia alojada. Si subiste tu modelo a Hub en el paso anterior, deberías ver un widget de inferencia en la página de tu modelo. Puedes agregar ejemplos predeterminados al widget definiendo las URL de las imágenes de ejemplo en la tarjeta de modelo. Consulta esta tarjeta de modelo como ejemplo.

Usar el modelo desde Hub

Primero cargaremos el modelo desde el Hub utilizando SegformerForSemanticSegmentation.from_pretrained() .

from transformers import SegformerFeatureExtractor, SegformerForSemanticSegmentation

feature_extractor = SegformerFeatureExtractor.from_pretrained("nvidia/segformer-b0-finetuned-ade-512-512")
model = SegformerForSemanticSegmentation.from_pretrained(f"{hf_username}/{hub_model_id}")

A continuación, cargaremos una imagen de nuestro conjunto de datos de prueba.

image = test_ds[0]['pixel_values']
gt_seg = test_ds[0]['label']
image

Para segmentar esta imagen de prueba, primero necesitamos preparar la imagen utilizando el extractor de características. Luego la enviamos a través del modelo.

También debemos recordar aumentar el tamaño de los logits de salida al tamaño original de la imagen. Para obtener las predicciones de categoría reales, solo tenemos que aplicar un argmax en los logits.

from torch import nn

inputs = feature_extractor(images=image, return_tensors="pt")
outputs = model(**inputs)
logits = outputs.logits  # forma (tamaño_lote, num_etiquetas, altura/4, ancho/4)

# Primero, redimensionamos los logits al tamaño original de la imagen
upsampled_logits = nn.functional.interpolate(
    logits,
    size=image.size[::-1], # (altura, ancho)
    mode='bilinear',
    align_corners=False
)

# Segundo, aplicamos argmax en la dimensión de clase
pred_seg = upsampled_logits.argmax(dim=1)[0]

Ahora es el momento de mostrar el resultado. Mostraremos el resultado junto a la máscara de la verdad absoluta.

¿Qué piensas? ¿Enviarías nuestro robot de entrega de pizza a la carretera con esta información de segmentación?

El resultado puede que aún no sea perfecto, pero siempre podemos ampliar nuestro conjunto de datos para hacer que el modelo sea más robusto. Ahora también podemos entrenar un modelo SegFormer más grande y ver cómo se compara.

¡Eso es todo! Ahora sabes cómo crear tu propio conjunto de datos de segmentación de imágenes y cómo usarlo para ajustar finamente un modelo de segmentación semántica.

Te presentamos algunas herramientas útiles en el camino, como:

• Segments.ai para etiquetar tus datos
• 🤗 datasets para crear y compartir un conjunto de datos
• 🤗 transformers para ajustar fácilmente un modelo de segmentación de vanguardia
• Hugging Face Hub para compartir nuestro conjunto de datos y modelo, y para crear un widget de inferencia para nuestro modelo

Esperamos que hayas disfrutado de esta publicación y hayas aprendido algo nuevo. No dudes en compartir tu propio modelo con nosotros en Twitter (@TobiasCornille, @NielsRogge y @huggingface).

We will continue to update Zepes; if you have any questions or suggestions, please contact us!