Transformers documentation

Сегментация экземпляров

Transformers

You are viewing main version, which requires installation from source. If you'd like regular pip install, checkout the latest stable version (v5.22.2).

Join the Hugging Face community

and get access to the augmented documentation experience

Collaborate on models, datasets and Spaces

Faster examples with accelerated inference

Switch between documentation themes

to get started

Сегментация экземпляров

Сегментация экземпляров — это задача компьютерного зрения, состоящая в обнаружении объектов на изображении и сегментации каждого из них на уровне пикселей. В отличие от обнаружения объектов (которое выдает ограничивающие рамки), сегментация экземпляров формирует точную маску для каждого обнаруженного объекта, что позволяет различать отдельные экземпляры даже при их пересечении.

В этом руководстве вы узнаете, как:

Загрузите набор данных для сегментации объектов из Hugging Face Hub.
Настройте RF-DETR-Seg, модель сегментации объектов на основе трансформеров, с помощью Trainer из Transformers.
Оцените вашу модель с помощью средней IoU.
Запустите инференс и визуализируйте предсказания.

Мы будем использовать набор данных satellite-building-segmentation, который содержит около 9,6 тыс. спутниковых изображений, размеченных масками инстансов зданий.

Чтобы увидеть все архитектуры и контрольные точки, совместимые с этой задачей, рекомендуем обратиться к странице задачи.

Перед началом установите все необходимые библиотеки:

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

Мы рекомендуем делиться вашей моделью с сообществом. Войдите в свой аккаунт Hugging Face, чтобы загрузить модель на Hub после завершения обучения:

>>> from huggingface_hub import notebook_login

>>> notebook_login()

Загрузите набор данных

Датасет satellite-building-segmentation представлен в нативном формате Hugging Face Datasets, поэтому его можно загрузить напрямую с помощью load_dataset. Каждая строка содержит:

image: спутниковое изображение (PIL)
image_id: уникальный идентификатор изображения
width / height: размеры изображения
objects: словарь с аннотациями для каждого объекта, содержащий id, category, bbox, area, segmentation (координаты полигона) и 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

Рассмотрим один пример. Каждая запись содержит image, image_id и словарь objects с аннотациями для каждого объекта. Для каждого объекта указаны bbox в формате [x, y, width, height] и поле segmentation с координатами полигонов:

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

Визуализация примера с эталонными масками:

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

Загрузите модель и процессор изображений

Используйте AutoImageProcessor и AutoModelForInstanceSegmentation для загрузки модели RF-DETR-Seg. При загрузке модели передайте отображения id2label и label2id, чтобы настроить классификационную головку для единственного класса «building». Поскольку предобученная модель обучалась на COCO (91 класс), установите ignore_mismatched_sizes=True, чтобы инициализировать классификационную головку с правильным количеством выходов.

Обработчик изображений выполняет всю предварительную обработку: изменение размера изображений с сохранением соотношения сторон, нормализацию с использованием статистики ImageNet, выравнивание до единого размера и — что особенно важно для сегментации экземпляров — преобразование полигональных аннотаций в бинарные маски, изменение размера этих масок и нормализацию ограничивающих прямоугольников в формат [cx, cy, w, h] в диапазоне [0, 1], который ожидает модель.

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

Предобработка данных

Чтобы дообучить модель, необходимо предобработать данные в соответствии с форматом, который ожидает модель. RfDetrImageProcessor выполняет всю основную работу, если передать ему изображения и аннотации в формате COCO с параметром return_segmentation_masks=True:

Преобразует сегментации полигонов в бинарные маски
Масштабирует изображения, ограничивающие рамки и маски до размера входных данных модели
Нормализует значения пикселей с использованием среднего и стандартного отклонения ImageNet
Преобразует ограничивающие рамки из формата [x, y, w, h] в нормализованный формат [cx, cy, w, h]
Выравнивает изображения до единого размера и создаёт pixel_mask

Преобразование восстанавливает словари аннотаций в стиле COCO, которые ожидаются процессором изображений из столбца objects набора данных.

Используйте with_transform, чтобы применять предобработку лениво (на лету при загрузке выборок), что позволяет избежать хранения всего обработанного набора данных в памяти.

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

Проверьте предобработанный пример. Он содержит pixel_values (нормализованный тензор изображения), pixel_mask (указывает на реальные пиксели и заполнение) и labels (словарь с class_labels, boxes в нормализованном формате центра и бинарными masks):

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

Коллатор данных

Поскольку в каждом изображении разное количество объектов, метки представляют собой словари переменной длины, которые нельзя объединить в единый тензор. Определите собственную функцию collate, которая обычно складывает pixel_values и pixel_mask, но сохраняет labels в виде списка словарей для каждого изображения:

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

Определение метрики оценки

Для отслеживания качества сегментации в процессе обучения вычислите средний IoU на основе объединения на каждой эпохе оценки. Для каждого изображения:

Объедините все предсказанные маски объектов в единую бинарную карту «здания» (бинаризуя логиты масок каждого запроса с помощью sigmoid > 0.5)
Объедините все маски экземпляров эталонных данных в единую бинарную карту
Вычисление пересечения по объединению (Intersection-over-Union) между двумя картами

Маски запросов в RF-DETR-Seg не фильтруются порогом по классовой оценке, поэтому метрика игнорирует классные оценки и напрямую объединяет логиты масок по каждому запросу. Несоответствующие запросы порождают практически пустые маски, поэтому они не искажают результат объединения.

Это даёт метрику для каждого изображения, показывающую, насколько хорошо модель покрывает здания, которая усредняется по всему набору для валидации.

Бенчмарки сегментации инстансов (например, COCO) обычно сообщают среднюю точность по маскам (mask mAP), которая оценивает каждую предсказанную маску инстанса по сравнению с эталонной по диапазону порогов IoU и, следовательно, поощряет правильное разделение отдельных объектов. Средний IoU на основе объединения, используемый здесь, является более простым и быстрым приближением: он измеряет общее покрытие пикселей, а не качество по каждому инстансу, что делает его удобным для отслеживания прогресса во время обучения. Для стандартной оценки с учётом инстансов вычислите mask mAP, например, с помощью MeanAveragePrecision(iou_type="segm") из torchmetrics.

Передайте это в Trainer в качестве функции compute_metrics, вместо того чтобы создавать подкласс тренера. При eval_do_concat_batches=False (установленном в TrainingArguments ниже) предсказания и метки, полученные в ходе стандартного прохода оценки, передаются в compute_metrics в виде списка результатов по батчам, поэтому метрика переиспользует эти предсказания, и второй проход по валидационной выборке не требуется. В кортеже выходных данных модели по индексу 3 хранятся 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}

Объект Trainer автоматически добавляет префикс eval_ к возвращаемым ключам, поэтому в результате получается метрика eval_mean_iou, которая используется ниже для выбора контрольных точек.

Обучение

Когда данные, модель и метрики готовы, настройте обучение. Несколько важных замечаний по TrainingArguments:

remove_unused_columns=False: необходимо, так как по умолчанию столбцы удаляются до выполнения нашего преобразования.
eval_do_concat_batches=False: Метки для instance segmentation — это словари переменной длины, их нельзя объединять между батчами. Это также позволяет сохранять предсказания сгруппированными по батчам, чтобы compute_metrics мог сопоставить их с соответствующими метками.
metric_for_best_model="eval_mean_iou": выбирает лучшую контрольную точку по качеству сегментации, а не только по значению функции потерь.
fp16=True: Обучение с смешанной точностью значительно ускоряет обучение на современных GPU.

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

Если в аргументах обучения задано push_to_hub=True, контрольные точки обучения будут отправлены на Hugging Face Hub. По завершении обучения отправьте на Hub и финальную модель, вызвав метод push_to_hub().

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

Вывод

Теперь, когда у вас есть дообученная модель, используйте её для инференса на новых спутниковых снимках. Рабочий процесс:

Предобработайте изображение с помощью процессора изображений
Выполнить прямой проход через модель
Используйте post_process_instance_segmentation для постобработки результатов, чтобы получить маски на уровне пикселей

Этап постобработки преобразует необработанные выходные данные запроса (логиты + маски низкого разрешения) в карты сегментации объектов полного разрешения, применяя пороговую фильтрацию по оценкам и бинаризацию масок.

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

Визуализируйте предсказания. Карта сегментации присваивает каждому пикселю идентификатор сегмента (-1 для фона). Наложите каждое обнаруженное здание случайным цветом:

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

←Обнаружение объектов Обнаружение объектов в режиме zero-shot→