Segmentation d'instances

La segmentation d'instances est la tâche de vision par ordinateur qui consiste à détecter les objets dans une image et à segmenter chacun d'entre eux au niveau du pixel. Contrairement à la détection d'objets (qui produit des boîtes englobantes), la segmentation d'instances génère un masque précis pour chaque objet détecté, ce qui permet de distinguer les instances individuelles même lorsqu'elles se chevauchent.

Dans ce guide, vous apprendrez à :

Charger un jeu de données de segmentation d'instances depuis le Hugging Face Hub.
Effectuez l'ajustement fin de RF-DETR-Seg, un modèle de segmentation d'instances basé sur des transformers, à l'aide du Trainer de Transformers.
Évaluez votre modèle avec l'IoU moyen.
Lancer l'inférence et visualiser les prédictions.

Nous utiliserons le jeu de données satellite-building-segmentation, qui contient environ 9,6 k d’images satellites annotées avec des masques d’instances de bâtiments.

Pour consulter toutes les architectures et points de contrôle compatibles avec cette tâche, nous vous recommandons de consulter la page de la tâche.

Avant de commencer, installez toutes les bibliothèques nécessaires :

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

Nous vous encourageons à partager votre modèle avec la communauté. Connectez-vous à votre compte Hugging Face pour le téléverser sur le Hub une fois l’entraînement terminé :

>>> from huggingface_hub import notebook_login

>>> notebook_login()

Charger le jeu de données

Le jeu de données satellite-building-segmentation est au format natif Hugging Face Datasets, il peut donc être chargé directement avec load_dataset. Chaque ligne contient :

image : l’image satellite (PIL)
image_id : identifiant unique de l'image
width / height : dimensions de l'image
objects : un dictionnaire d’annotations par instance avec id, category, bbox, area, segmentation (coordonnées du polygone) et 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

Inspecter un exemple unique. Chaque enregistrement comporte une image, un image_id et un dictionnaire objects contenant des annotations par instance. Chaque instance possède un bbox au format [x, y, width, height] et un champ segmentation avec des coordonnées de polygone :

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

Visualiser un exemple avec ses masques de référence :

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

Charger le modèle et le processeur d’images

Utilisez AutoImageProcessor et AutoModelForInstanceSegmentation pour charger le modèle RF-DETR-Seg. Lors du chargement du modèle, passez les correspondances id2label et label2id pour configurer la tête de classification pour la classe unique « building ». Étant donné que le modèle pré-entraîné a été entraîné sur COCO (91 classes), définissez ignore_mismatched_sizes=True pour réinitialiser la tête de classification avec le bon nombre de sorties.

Le processeur d’images gère tout le prétraitement : redimensionnement des images en conservant le ratio d’aspect, normalisation avec les statistiques d’ImageNet, bourrage à une taille uniforme et — point crucial pour la segmentation d’instances — conversion des annotations polygonales en masques binaires, redimensionnement de ces masques et normalisation des boîtes englobantes au format [cx, cy, w, h] dans l’intervalle [0, 1] attendu par le modèle.

>>> 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étraiter les données

Pour affiner le modèle, vous devez prétraiter les données afin qu’elles correspondent au format attendu par le modèle. Le RfDetrImageProcessor effectue tout le travail lourd lorsque vous lui transmettez des images et des annotations au format COCO avec return_segmentation_masks=True :

Rasterise les segmentations de polygones en masques binaires
Redimensionne les images, les boîtes englobantes et les masques à la taille d'entrée du modèle
Normalise les valeurs de pixels avec la moyenne/écart-type d'ImageNet
Convertit les boîtes englobantes de [x, y, w, h] en [cx, cy, w, h] normalisées
Remplit les images pour leur donner une taille uniforme et crée un pixel_mask

La transformation reconstruit les dictionnaires d’annotation au format COCO attendus par le processeur d’images à partir de la colonne objects du jeu de données.

Utilisez with_transform pour appliquer le prétraitement de manière paresseuse (à la volée lors du chargement des échantillons), ce qui évite de stocker l’ensemble du jeu de données traité en mémoire.

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

Vérifier un exemple prétraité. Il contient pixel_values (le tenseur d’image normalisé), pixel_mask (indiquant les pixels réels par rapport au bourrage) et labels (un dictionnaire avec class_labels, boxes au format centre normalisé et des masks binaires) :

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

Collateur de données

Chaque image ayant un nombre différent d'instances d'objets, les étiquettes sont des dictionnaires de longueur variable et ne peuvent pas être empilées en un seul tenseur. Définissez une fonction de collation personnalisée qui empile pixel_values et pixel_mask normalement, mais qui conserve labels sous forme de liste de dictionnaires par image :

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

Définir la métrique d'évaluation

Pour suivre la qualité de la segmentation pendant l’entraînement, calculez un IoU moyen basé sur l’union à chaque époque d’évaluation. Pour chaque image :

Fusionner tous les masques d’instances prédits en une seule carte binaire « buildings » (binariser les logits de masque de chaque requête avec sigmoid > 0.5)
Fusionner tous les masques d'instances de référence en une seule carte binaire
Calculer l'Intersection-over-Union entre les deux cartes

Les masques de requêtes de RF-DETR-Seg ne sont pas filtrés par un seuil de score de classe, la métrique ignore donc les scores de classe et fusionne directement les logits de masque par requête. Les requêtes non appariées produisent des masques quasi vides, de sorte qu’elles ne polluent pas l’union.

Cela fournit une métrique par image de « dans quelle mesure le modèle couvre les bâtiments », qui est moyennée sur l’ensemble de validation complet.

Les benchmarks de segmentation d’instances (tels que COCO) rapportent généralement le mean average precision des masques (mAP), qui évalue chaque masque d’instance prédit par rapport à la vérité terrain sur une plage de seuils IoU et récompense donc la séparation correcte des objets individuels. Le mean IoU basé sur l’union utilisé ici est un substitut plus simple et plus rapide : il mesure la couverture globale des pixels plutôt que la qualité par instance, ce qui le rend pratique pour suivre les progrès pendant l’entraînement. Pour une évaluation standard, tenant compte des instances, calculez plutôt le mAP des masques, par exemple avec MeanAveragePrecision(iou_type="segm") de torchmetrics.

Transmettez ceci au Trainer en tant que fonction compute_metrics plutôt que de créer une sous-classe du trainer. Avec eval_do_concat_batches=False (défini dans les TrainingArguments ci-dessous), les prédictions et les étiquettes produites par le passage d'évaluation standard sont transmises à compute_metrics sous forme d'une liste de sorties par lot, de sorte que la métrie réutilise ces prédictions et qu'aucun second passage avant sur l'ensemble de validation n'est nécessaire. Dans le tuple de sortie du modèle, l'index 3 contient 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}

Le Trainer préfixe automatiquement les clés retournées avec eval_, ce qui génère la métrique eval_mean_iou utilisée ci-dessous pour la sélection des points de contrôle.

Entraînement

Une fois les données, le modèle et les métriques prêts, configurez l’entraînement. Quelques notes importantes sur TrainingArguments :

remove_unused_columns=False : requis, car le comportement par défaut supprime les colonnes avant l’exécution de notre transformation.
eval_do_concat_batches=False : Les étiquettes de segmentation d’instances sont des dictionnaires de longueur variable, ils ne peuvent pas être concaténés entre les lots. Cela permet également de regrouper les prédictions par lot afin que compute_metrics puisse les associer à leurs étiquettes.
metric_for_best_model="eval_mean_iou" : sélectionne le meilleur point de contrôle en fonction de la qualité de la segmentation, et non uniquement de la perte.
fp16=True : l’entraînement en précision mixte accélère considérablement l’entraînement sur les GPU modernes.

>>> 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 vous définissez push_to_hub=True dans les arguments d’entraînement, les points de contrôle d’entraînement sont poussés vers le Hugging Face Hub. Une fois l’entraînement terminé, poussez également le modèle final vers le Hub en appelant la méthode push_to_hub().

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

Inférence

Maintenant que vous disposez d'un modèle affiné, utilisez-le pour l'inférence sur de nouvelles images satellites. Le flux de travail est le suivant :

Prétraitez l'image avec le processeur d'images
Effectuer un passage avant à travers le modèle
Post-traitez les sorties avec post_process_instance_segmentation pour obtenir des masques au niveau des pixels

L'étape de post-traitement convertit les sorties brutes de la requête (logits + masques basse résolution) en cartes de segmentation d'instances pleine résolution, en appliquant un seuillage des scores et une binarisation des masques.

>>> 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})")

Visualisez les prédictions. La carte de segmentation attribue à chaque pixel un identifiant de segment (-1 pour l'arrière-plan). Superposez chaque bâtiment détecté avec une couleur aléatoire :

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

Segmentation d'instances

Charger le jeu de données

Charger le modèle et le processeur d’images

Prétraiter les données

Collateur de données

Définir la métrique d'évaluation

Entraînement

Inférence