Segmentación de instancias

La segmentación de instancias es la tarea de visión por computadora que consiste en detectar objetos en una imagen y segmentar cada uno a nivel de píxel. A diferencia de la detección de objetos (que produce cajas delimitadoras), la segmentación de instancias genera una máscara precisa para cada objeto detectado, lo que permite distinguir instancias individuales incluso cuando se superponen.

En esta guía aprenderás a:

Cargar un conjunto de datos de segmentación de instancias desde Hugging Face Hub.
Ajusta RF-DETR-Seg, un modelo de segmentación de instancias basado en transformers, utilizando el Trainer de Transformers.
Evalúe su modelo con IoU medio.
Ejecuta la inferencia y visualiza las predicciones.

Utilizaremos el conjunto de datos satellite-building-segmentation, que contiene aproximadamente 9,6 mil imágenes satelitales anotadas con máscaras de instancias de edificios.

Para ver todas las arquitecturas y puntos de control compatibles con esta tarea, recomendamos consultar la página de la tarea.

Antes de comenzar, instale todas las bibliotecas necesarias:

pip install -Uq "transformers>=5.9" datasets torchvision

Te animamos a compartir tu modelo con la comunidad. Inicia sesión en tu cuenta de Hugging Face para subirlo al Hub cuando el entrenamiento esté completo:

>>> from huggingface_hub import notebook_login

>>> notebook_login()

Cargar el conjunto de datos

El conjunto de datos satellite-building-segmentation está en el formato nativo de Hugging Face Datasets, por lo que se puede cargar directamente con load_dataset. Cada fila contiene:

image: la imagen satelital (PIL)
image_id: identificador único de la imagen
width / height: dimensiones de la imagen
objects: un diccionario de anotaciones por instancia con id, category, bbox, area, segmentation (coordenadas del polígono) y iscrowd

>>> from datasets import load_dataset

>>> MODEL_ID = "Roboflow/rf-detr-seg-medium"
>>> DATASET_ID = "merve/satellite-building-segmentation"

>>> ds = load_dataset(DATASET_ID)
>>> train_ds = ds["train"]
>>> valid_ds = ds["validation"]
>>> print(f"Train: {len(train_ds)} images, Valid: {len(valid_ds)} images")
Train: 6764 images, Valid: 1934 images

Inspeccionar un ejemplo individual. Cada registro tiene una image, un image_id y un diccionario objects que contiene anotaciones por instancia. Cada instancia tiene un bbox en formato [x, y, width, height] y un campo segmentation con coordenadas de polígono:

>>> sample = train_ds[0]
>>> print(f"Image ID: {sample['image_id']}")
Image ID: 0
>>> print(f"Image size: {sample['image'].size}")
Image size: (512, 512)
>>> print(f"Number of instances: {len(sample['objects']['id'])}")
Number of instances: 7
>>> print(f"\nObjects keys: {list(sample['objects'].keys())}")
<BLANKLINE>
Objects keys: ['id', 'area', 'bbox', 'segmentation', 'category', 'iscrowd']
>>> print(f"First bbox: {sample['objects']['bbox'][0]}")
First bbox: [80.0, 0.0, 51.0, 51.5]
>>> print(f"First category: {sample['objects']['category'][0]}")
First category: 0

Visualice un ejemplo con sus máscaras de referencia:

>>> import numpy as np
>>> import matplotlib.pyplot as plt
>>> from PIL import Image, ImageDraw

>>> sample = train_ds[0]
>>> image = sample["image"].convert("RGB")

>>> fig, axes = plt.subplots(1, 2, figsize=(14, 6))
>>> _ = axes[0].imshow(image)
>>> _ = axes[0].set_title("Original image")
>>> _ = axes[0].axis("off")

>>> overlay = image.copy()
>>> draw = ImageDraw.Draw(overlay, "RGBA")
>>> objects = sample["objects"]
>>> for seg in objects["segmentation"]:
...     for poly in seg:
...         coords = list(zip(poly[0::2], poly[1::2]))
...         color = tuple(np.random.randint(50, 255, 3)) + (100,)
...         draw.polygon(coords, fill=color, outline="red")

>>> _ = axes[1].imshow(overlay)
>>> _ = axes[1].set_title(f"Ground truth ({len(objects['id'])} buildings)")
>>> _ = axes[1].axis("off")
>>> plt.tight_layout()
>>> plt.show()

Dataset Sample

Cargue el modelo y el procesador de imágenes

Use AutoImageProcessor y AutoModelForInstanceSegmentation para cargar el modelo RF-DETR-Seg. Al cargar el modelo, pase las mapeaciones id2label y label2id para configurar la cabeza de clasificación para la única clase “building”. Dado que el modelo preentrenado se entrenó en COCO (91 clases), establezca ignore_mismatched_sizes=True para reinicializar la cabeza de clasificación con el número correcto de salidas.

El procesador de imágenes se encarga de todo el preprocesamiento: redimensionar las imágenes manteniendo la relación de aspecto, normalizar con las estadísticas de ImageNet, rellenar hasta un tamaño uniforme y, lo cual es crucial para la segmentación de instancias, convertir las anotaciones de polígonos en máscaras binarias, redimensionar esas máscaras y normalizar las cajas delimitadoras al formato [cx, cy, w, h] en el rango [0, 1] que el modelo espera.

>>> from transformers import AutoImageProcessor, AutoModelForInstanceSegmentation

>>> id2label = {0: "building"}
>>> label2id = {"building": 0}

>>> image_processor = AutoImageProcessor.from_pretrained(MODEL_ID)
>>> model = AutoModelForInstanceSegmentation.from_pretrained(
...     MODEL_ID,
...     id2label=id2label,
...     label2id=label2id,
...     ignore_mismatched_sizes=True,
... )

Preprocesar los datos

Para ajustar el modelo, debe preprocesar los datos para que coincidan con el formato que el modelo espera. El RfDetrImageProcessor realiza todo el trabajo pesado cuando se le pasan imágenes y anotaciones en formato COCO con return_segmentation_masks=True:

Rasteriza las segmentaciones de polígonos en máscaras binarias
Redimensiona las imágenes, las cajas delimitadoras y las máscaras al tamaño de entrada del modelo
Normaliza los valores de píxeles con la media/desviación estándar de ImageNet
Convierte las cajas delimitadoras de [x, y, w, h] a [cx, cy, w, h] normalizadas
Rellena las imágenes hasta un tamaño uniforme y crea una pixel_mask

La transformación reconstruye los diccionarios de anotación en formato COCO que el procesador de imágenes espera de la columna objects del conjunto de datos.

Use with_transform para aplicar el preprocesamiento de forma perezosa (al vuelo, cuando se cargan las muestras), lo que evita almacenar el conjunto de datos procesado completo en memoria.

>>> from functools import partial
>>> from typing import Any


>>> def transform_batch(examples: dict[str, Any], image_processor) -> dict[str, Any]:
...     """Convert HF dataset rows into COCO-style dicts and pass to the processor."""
...     images, targets = [], []
...     for image, img_id, objects in zip(
...         examples["image"], examples["image_id"], examples["objects"]
...     ):
...         if not objects["id"]:
...             continue
...         annotations = [
...             {
...                 "id": ann_id,
...                 "image_id": img_id,
...                 "category_id": cat,
...                 "bbox": bbox,
...                 "area": area,
...                 "segmentation": seg,
...                 "iscrowd": crowd,
...             }
...             for ann_id, cat, bbox, area, seg, crowd in zip(
...                 objects["id"],
...                 objects["category"],
...                 objects["bbox"],
...                 objects["area"],
...                 objects["segmentation"],
...                 objects["iscrowd"],
...             )
...         ]
...         images.append(image.convert("RGB"))
...         targets.append({"image_id": img_id, "annotations": annotations})
...
...     if not images:
...         return {}
...     return image_processor(
...         images=images, annotations=targets, return_segmentation_masks=True, return_tensors="pt"
...     )


>>> transform = partial(transform_batch, image_processor=image_processor)
>>> train_ds = train_ds.shuffle(seed=42).with_transform(transform)
>>> valid_ds = valid_ds.with_transform(transform)

Verifique un ejemplo preprocesado. Contiene pixel_values (el tensor de imagen normalizado), pixel_mask (que indica los píxeles reales frente al relleno) y labels (un diccionario con class_labels, boxes en formato de centro normalizado y masks binarias):

>>> example = train_ds[0]
>>> print(f"pixel_values shape: {example['pixel_values'].shape}")
pixel_values shape: torch.Size([3, 432, 432])
>>> print(f"pixel_mask shape: {example['pixel_mask'].shape}")
pixel_mask shape: torch.Size([432, 432])
>>> print(f"labels keys: {list(example['labels'].keys())}")
labels keys: ['size', 'image_id', 'class_labels', 'boxes', 'area', 'iscrowd', 'orig_size', 'masks']
>>> print(f"  class_labels: {example['labels']['class_labels']}")
  class_labels: tensor([0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
        0, 0])
>>> print(f"  boxes shape: {example['labels']['boxes'].shape}")
  boxes shape: torch.Size([26, 4])
>>> print(f"  masks shape: {example['labels']['masks'].shape}")
  masks shape: torch.Size([26, 432, 432])

Colaborador de datos

Dado que cada imagen tiene un número diferente de instancias de objetos, las etiquetas son diccionarios de longitud variable y no se pueden apilar en un único tensor. Defina una función de collate personalizada que apila pixel_values y pixel_mask de forma normal, pero mantiene labels como una lista de diccionarios por imagen:

>>> import torch


>>> def collate_fn(batch: list[dict[str, Any]]) -> dict[str, Any]:
...     out = {
...         "pixel_values": torch.stack([x["pixel_values"] for x in batch]),
...         "labels": [x["labels"] for x in batch],
...     }
...     if "pixel_mask" in batch[0]:
...         out["pixel_mask"] = torch.stack([x["pixel_mask"] for x in batch])
...     return out

Definir la métrica de evaluación

Para hacer seguimiento de la calidad de la segmentación durante el entrenamiento, calcule un IoU medio basado en la unión en cada época de evaluación. Para cada imagen:

Fusiona todas las máscaras de instancias predichas en un único mapa binario de “edificios” (binariza los logits de la máscara de cada consulta con sigmoid > 0.5)
Fusionar todas las máscaras de instancias de referencia en un único mapa binario
Calcular la Intersección sobre la Unión entre los dos mapas

Las máscaras de consulta de RF-DETR-Seg no están controladas por un umbral de puntuación de clase, por lo que la métrica ignora las puntuaciones de clase y une directamente los logits de máscara por consulta. Las consultas no emparejadas producen máscaras casi vacías, por lo que no contaminan la unión.

Esto proporciona una métrica por imagen de “qué tan bien cubre el modelo los edificios”, que se promedia sobre el conjunto de validación completo.

Los benchmarks de segmentación de instancias (como COCO) suelen informar la precisión media promedio de las máscaras (mAP), que puntúa cada máscara de instancia predicha frente a la verdad de referencia en un rango de umbrales de IoU y, por lo tanto, premia la separación correcta de objetos individuales. La IoU media basada en la unión utilizada aquí es un sustituto más simple y rápido: mide la cobertura de píxeles general en lugar de la calidad por instancia, lo que la hace conveniente para seguir el progreso durante el entrenamiento. Para una evaluación estándar y consciente de instancias, calcule la mAP de máscaras en su lugar, por ejemplo con MeanAveragePrecision(iou_type="segm") de torchmetrics.

Pásalo al Trainer como una función compute_metrics en lugar de crear una subclase del trainer. Con eval_do_concat_batches=False (configurado en TrainingArguments a continuación), las predicciones y etiquetas generadas por el paso de evaluación estándar se pasan a compute_metrics como una lista de salidas por lote, de modo que la métrica reutiliza esas predicciones y no se necesita una segunda pasada de adelante sobre el conjunto de validación. En la tupla de salida del modelo, el índice 3 contiene pred_masks:

>>> import torch.nn.functional as F


>>> @torch.no_grad()
... def compute_mean_iou(pred_masks, gt_masks, target_size):
...     """Union-based IoU: merge all instances per image, then compute IoU."""
...     pred_masks = F.interpolate(pred_masks[None], size=target_size, mode="bilinear", align_corners=False)[0]
...     pred_union = (pred_masks.sigmoid() > 0.5).any(dim=0)
...
...     gt_union = gt_masks.any(dim=0).bool()
...
...     intersection = (pred_union & gt_union).sum().float()
...     union = (pred_union | gt_union).sum().float()
...     return (intersection / union.clamp(min=1)).item()


>>> @torch.no_grad()
... def compute_metrics(evaluation_results):
...     predictions, targets = evaluation_results.predictions, evaluation_results.label_ids
...     ious = []
...     for pred_batch, target_batch in zip(predictions, targets):
...         batch_masks = pred_batch[3]
...         for i, gt_label in enumerate(target_batch):
...             gt_masks = torch.as_tensor(gt_label["masks"])
...             if gt_masks.numel() == 0:
...                 continue
...             target_size = gt_masks.shape[-2:]
...             pred_masks = torch.as_tensor(batch_masks[i])
...             ious.append(compute_mean_iou(pred_masks, gt_masks, target_size))
...
...     mean_iou = sum(ious) / len(ious) if ious else 0.0
...     return {"mean_iou": mean_iou}

El Trainer añade automáticamente el prefijo eval_ a las claves devueltas, por lo que esto genera la métrica eval_mean_iou utilizada a continuación para la selección de puntos de control.

Entrenamiento

Con los datos, el modelo y las métricas listos, configura el entrenamiento. Algunas notas importantes sobre TrainingArguments:

remove_unused_columns=False: Es necesario porque el comportamiento predeterminado eliminaría las columnas antes de que se ejecute nuestra transformación.
eval_do_concat_batches=False: Las etiquetas de segmentación de instancias son diccionarios de longitud variable, por lo que no se pueden concatenar entre lotes. Esto también mantiene las predicciones agrupadas por lote para que compute_metrics pueda asociarlas con sus etiquetas.
metric_for_best_model="eval_mean_iou": Selecciona el mejor checkpoint según la calidad de la segmentación, no solo la pérdida.
fp16=True: El entrenamiento con precisión mixta acelera significativamente el entrenamiento en GPUs modernas.

>>> from transformers import Trainer, TrainingArguments

>>> training_args = TrainingArguments(
...     output_dir="rf-detr-seg-satellite-buildings",
...     num_train_epochs=10,
...     per_device_train_batch_size=16,
...     per_device_eval_batch_size=16,
...     learning_rate=1e-4,
...     weight_decay=1e-4,
...     lr_scheduler_type="cosine",
...     warmup_steps=0.1,
...     fp16=True,
...     dataloader_num_workers=4,
...     eval_strategy="epoch",
...     save_strategy="epoch",
...     save_total_limit=2,
...     load_best_model_at_end=True,
...     metric_for_best_model="eval_mean_iou",
...     greater_is_better=True,
...     remove_unused_columns=False,
...     eval_do_concat_batches=False,
...     push_to_hub=False,
... )

>>> trainer = Trainer(
...     model=model,
...     args=training_args,
...     train_dataset=train_ds,
...     eval_dataset=valid_ds,
...     processing_class=image_processor,
...     data_collator=collate_fn,
...     compute_metrics=compute_metrics,
... )

>>> trainer.train()

Si estableces push_to_hub=True en los argumentos de entrenamiento, los puntos de control del entrenamiento se envían al Hugging Face Hub. Al finalizar el entrenamiento, envía también el modelo final al Hub llamando al método push_to_hub().

>>> trainer.push_to_hub(
...     dataset=DATASET_ID,
...     tags=["instance-segmentation", "rf-detr-seg", "vision", "satellite", "building"],
... )

Inferencia

Ahora que tienes un modelo afinado, úsalo para realizar inferencia en nuevas imágenes satelitales. El flujo de trabajo es:

Preprocesa la imagen con el procesador de imágenes
Ejecutar una pasada hacia adelante a través del modelo
Procesa los resultados con post_process_instance_segmentation para obtener máscaras a nivel de píxel

La etapa de posprocesado convierte las salidas crudas de la consulta (logits + máscaras de baja resolución) en mapas de segmentación de instancias a resolución completa, aplicando umbralización de puntuación y binarización de máscaras.

>>> from datasets import load_dataset

>>> test_ds = load_dataset(DATASET_ID, split="test")
>>> sample = test_ds[0]
>>> image = sample["image"].convert("RGB")

>>> device = next(model.parameters()).device
>>> inputs = image_processor(images=image, return_tensors="pt").to(device)

>>> with torch.no_grad():
...     outputs = model(**inputs)

>>> results = image_processor.post_process_instance_segmentation(
...     outputs, threshold=0.5, target_sizes=[(image.height, image.width)]
... )[0]

>>> print(f"Detected {len(results['segments_info'])} buildings")
Detected 0 buildings
>>> for seg_info in results["segments_info"][:5]:
...     print(f"  Building (score: {seg_info['score']:.3f})")

Visualiza las predicciones. El mapa de segmentación asigna a cada píxel un ID de segmento (-1 para el fondo). Superpón cada edificio detectado con un color aleatorio:

>>> fig, axes = plt.subplots(1, 2, figsize=(14, 6))

>>> _ = axes[0].imshow(image)
>>> _ = axes[0].set_title("Input satellite image")
>>> _ = axes[0].axis("off")

>>> seg_map = results["segmentation"].cpu().numpy()
>>> overlay = np.array(image).copy()
>>> for seg_info in results["segments_info"]:
...     mask = seg_map == seg_info["id"]
...     color = np.random.randint(0, 255, 3)
...     overlay[mask] = (overlay[mask] * 0.4 + color * 0.6).astype(np.uint8)

>>> _ = axes[1].imshow(overlay)
>>> _ = axes[1].set_title(f"Predicted masks ({len(results['segments_info'])} buildings)")
>>> _ = axes[1].axis("off")
>>> plt.tight_layout()
>>> plt.show()

Fine-tuning Result

Update on GitHub

Transformers

Segmentación de instancias

Cargar el conjunto de datos

Cargue el modelo y el procesador de imágenes

Preprocesar los datos

Colaborador de datos

Definir la métrica de evaluación

Entrenamiento

Inferencia