Segmentação de instâncias

A segmentação de instâncias é a tarefa de visão computacional que consiste em detectar objetos em uma imagem e segmentar cada um deles no nível de pixel. Diferentemente da detecção de objetos (que produz caixas delimitadoras), a segmentação de instâncias gera uma máscara precisa para cada objeto detectado, permitindo distinguir instâncias individuais mesmo quando elas se sobrepõem.

Neste guia, você aprenderá como:

Carregue um conjunto de dados de segmentação de instâncias do Hugging Face Hub.
Ajuste finamente o RF-DETR-Seg, um modelo de segmentação de instâncias baseado em transformers, usando o Trainer do Transformers.
Avalie seu modelo com IoU médio.
Execute a inferência e visualize as previsões.

Usaremos o conjunto de dados satellite-building-segmentation, que contém cerca de 9,6 mil imagens de satélite anotadas com máscaras de instâncias de edifícios.

Para ver todas as arquiteturas e checkpoints compatíveis com esta tarefa, recomendamos consultar a página da tarefa.

Antes de começar, instale todas as bibliotecas necessárias:

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

Incentivamos você a compartilhar seu modelo com a comunidade. Faça login na sua conta do Hugging Face para enviá-lo ao Hub quando o treinamento estiver concluído:

>>> from huggingface_hub import notebook_login

>>> notebook_login()

Carregue o conjunto de dados

O conjunto de dados satellite-building-segmentation está no formato nativo do Hugging Face Datasets, portanto, carregue-o diretamente com load_dataset. Cada linha contém:

image: a imagem de satélite (PIL)
image_id: identificador único da imagem
width / height: dimensões da imagem
objects: um dicionário de anotações por instância com id, category, bbox, area, segmentation (coordenadas do polígono) e 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

Inspecione um exemplo individual. Cada registro possui uma image, um image_id e um dicionário objects contendo anotações por instância. Cada instância tem um bbox no formato [x, y, width, height] e um campo segmentation com 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

Visualize um exemplo com suas máscaras de referência:

>>> 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

Carregue o modelo e o processador de imagem

Use AutoImageProcessor e AutoModelForInstanceSegmentation para carregar o modelo RF-DETR-Seg. Ao carregar o modelo, passe os mapeamentos id2label e label2id para configurar a cabeça de classificação para a única classe “building”. Como o modelo pré-treinado foi treinado no COCO (91 classes), defina ignore_mismatched_sizes=True para reinitializar a cabeça de classificação com o número correto de saídas.

O processador de imagens trata todo o pré-processamento: redimensiona as imagens mantendo a proporção, normaliza com as estatísticas do ImageNet, aplica padding para um tamanho uniforme e — crucialmente para a segmentação de instâncias — converte anotações de polígonos em máscaras binárias, redimensiona essas máscaras e normaliza as caixas delimitadoras para o formato [cx, cy, w, h] no intervalo [0, 1] que o 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,
... )

Pré-processar os dados

Para ajustar o modelo, você deve pré-processar os dados para corresponder ao formato que o modelo espera. O RfDetrImageProcessor faz todo o trabalho pesado quando você passa imagens e anotações no formato COCO com return_segmentation_masks=True:

Rasteriza segmentações de polígonos em máscaras binárias
Redimensiona imagens, caixas delimitadoras e máscaras para o tamanho de entrada do modelo
Normaliza os valores de pixels com a média/desvio padrão do ImageNet
Converte caixas delimitadoras de [x, y, w, h] para [cx, cy, w, h] normalizadas
Preenche as imagens até um tamanho uniforme e cria uma pixel_mask

O transform reconstrói os dicionários de anotação no estilo COCO que o processador de imagens espera da coluna objects do conjunto de dados.

Use with_transform para aplicar o pré-processamento de forma preguiçosa (no momento em que as amostras são carregadas), o que evita armazenar o conjunto de dados processado inteiro na memória.

>>> 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 um exemplo pré-processado. Ele contém pixel_values (o tensor de imagem normalizado), pixel_mask (indicando pixels reais em relação ao preenchimento) e labels (um dicionário com class_labels, boxes em formato de centro normalizado e masks binárias):

>>> 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 dados

Como cada imagem possui um número diferente de instâncias de objetos, os rótulos são dicionários de comprimento variável e não podem ser empilhados em um único tensor. Defina uma função de colagem personalizada que empilha pixel_values e pixel_mask normalmente, mas mantém labels como uma lista de dicionários por imagem:

>>> 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 métrica de avaliação

Para acompanhar a qualidade da segmentação durante o treinamento, calcule um mean IoU baseado em união em cada época de avaliação. Para cada imagem:

Mescle todas as máscaras de instâncias previstas em um único mapa binário de “edifícios” (binarize os logits de máscara de cada consulta com sigmoid > 0.5)
Mesclar todas as máscaras de instâncias de referência em um único mapa binário
Calcule a Interseção sobre a União entre os dois mapas

As máscaras de consulta do RF-DETR-Seg não são filtradas por um limiar de pontuação de classe, portanto a métrica ignora as pontuações de classe e une diretamente as logits de máscara por consulta. Consultas não correspondidas geram máscaras quase vazias, de modo que não contaminam a união.

Isso fornece uma métrica por imagem de “o quão bem o modelo cobre os prédios”, que é média sobre o conjunto de validação completo.

Os benchmarks de segmentação de instâncias (como o COCO) geralmente relatam a precisão média de máscaras (mAP), que pontua cada máscara de instância prevista em relação ao ground truth em uma faixa de limiares de IoU e, portanto, recompensa a separação correta de objetos individuais. A média de IoU baseada em união usada aqui é um substituto mais simples e rápido: ela mede a cobertura geral de pixels em vez da qualidade por instância, o que a torna conveniente para acompanhar o progresso durante o treinamento. Para uma avaliação padrão e consciente de instâncias, calcule o mAP de máscaras, por exemplo, com MeanAveragePrecision(iou_type="segm") do torchmetrics.

Passe isso para o Trainer como uma função compute_metrics em vez de criar uma subclasse do trainer. Com eval_do_concat_batches=False (definido no TrainingArguments abaixo), as previsões e os rótulos gerados pela passagem de avaliação padrão são passados para compute_metrics como uma lista de saídas por lote, de modo que a métrica reutiliza essas previsões e não é necessária uma segunda passagem de forward no conjunto de validação. Na tupla de saída do modelo, o índice 3 contém 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}

O Trainer adiciona automaticamente o prefixo eval_ às chaves retornadas, de modo que isso gera a métrica eval_mean_iou usada abaixo para a seleção de checkpoints.

Treinamento

Com os dados, o modelo e as métricas prontos, configure o treinamento. Algumas observações importantes sobre o TrainingArguments:

remove_unused_columns=False: Necessário porque o comportamento padrão removeria as colunas antes da execução do nosso transform.
eval_do_concat_batches=False: Os rótulos de segmentação de instância são dicionários de comprimento variável e não podem ser concatenados entre lotes. Isso também mantém as previsões agrupadas por lote, permitindo que compute_metrics as corresponda aos seus rótulos.
metric_for_best_model="eval_mean_iou": Selecione o melhor checkpoint com base na qualidade da segmentação, e não apenas na perda.
fp16=True: O treinamento com precisão mista acelera significativamente o treinamento em 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()

Se você definir push_to_hub=True nos argumentos de treinamento, os checkpoints de treinamento serão enviados para o Hugging Face Hub. Ao concluir o treinamento, envie o modelo final para o Hub também, chamando o método push_to_hub().

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

Inferência

Agora que você tem um modelo ajustado, use-o para inferência em novas imagens de satélite. O fluxo de trabalho é:

Pré-processa a imagem com o processador de imagens
Execute uma passagem direta pelo modelo
Pós-processe as saídas com post_process_instance_segmentation para obter máscaras em nível de pixel

A etapa de pós-processamento converte as saídas brutas da consulta (logits + máscaras de baixa resolução) em mapas de segmentação de instâncias em resolução completa, aplicando o limite de pontuação e a binarização das 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})")

Visualize as previsões. O mapa de segmentação atribui a cada pixel um ID de segmento (-1 para o fundo). Sobreponha cada edifício detectado com uma cor aleatória:

>>> 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

Segmentação de instâncias

Carregue o conjunto de dados

Carregue o modelo e o processador de imagem

Pré-processar os dados

Colaborador de dados

Definir métrica de avaliação

Treinamento

Inferência