Сегментация экземпляров — Mask R-CNN

Добавьте к детектору Faster R-CNN небольшую ветвь масок — и получите сегментацию экземпляров. Самая сложная часть — RoIAlign, и она сложнее, чем кажется.

Тип: Сборка + изучение Языки: Python Предварительные требования: Фаза 4, урок 06 (YOLO), фаза 4, урок 07 (U-Net) Время: ~75 минут

Цели обучения

• Проследить архитектуру Mask R-CNN от начала до конца: backbone, FPN, RPN, RoIAlign, голову рамок и голову масок
• Реализовать RoIAlign с нуля и объяснить, почему RoIPool больше не используется
• Использовать предобученную модель torchvision maskrcnn_resnet50_fpn_v2 для получения экземплярных масок продакшен-качества и правильно читать формат её выхода
• Дообучить Mask R-CNN на небольшом пользовательском датасете, заменив головы рамок и масок и сохранив backbone замороженным

Проблема

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

Mask R-CNN (He et al., 2017) решила эту задачу, переформулировав сегментацию экземпляров как обнаружение плюс маска. Конструкция оказалась настолько удачной, что в следующие пять лет почти каждая статья по сегментации экземпляров была вариантом Mask R-CNN, а реализация torchvision до сих пор остаётся продакшен-стандартом для малых и средних датасетов.

Сложная инженерная проблема здесь — выборка: как вырезать область признаков фиксированного размера из предлагаемой рамки, углы которой не совпадают с границами пикселей? Ошибка в этом стоит десятых долей пункта mAP повсюду. Ответ — RoIAlign.

Концепция

Архитектура

Нужно понять пять компонентов:

1. Backbone — ResNet-50 или ResNet-101, обученная на ImageNet. Создаёт иерархию карт признаков с шагами 4, 8, 16, 32.
2. FPN (Feature Pyramid Network) — нисходящие и боковые соединения, которые дают каждому уровню C каналов семантически насыщенных признаков. Для обнаружения выбирают уровень FPN, соответствующий размеру объекта.
3. RPN (Region Proposal Network) — небольшая свёрточная голова, которая в каждой позиции якоря предсказывает «есть ли здесь объект?» и «как уточнить рамку?». Создаёт около 1000 предложений на изображение.
4. RoIAlign — извлекает участок признаков фиксированного размера (например, 7x7) из любой рамки на любом уровне FPN. Билинейная выборка без квантования.
5. Головы — двухслойная голова рамок уточняет рамку и выбирает класс, а небольшая свёрточная голова выдаёт двоичную маску 28x28 для каждого предложения.

Почему RoIAlign, а не RoIPool

Исходный Fast R-CNN использовал RoIPool: он делит предлагаемую рамку на сетку, берёт максимальный признак в каждой ячейке и округляет все координаты до целых. Из-за такого округления карта признаков смещается относительно координат пикселей входа вплоть до целого пикселя карты признаков — немного для изображения 224x224, но катастрофично, когда шаг карты признаков равен 32.

RoIPool:
  рамка (34.7, 51.3, 98.2, 142.9)
  округлить -> (34, 51, 98, 142)
  разделить сетку -> округлить границу каждой ячейки
  смещение накапливается на каждом шаге

RoIAlign:
  рамка (34.7, 51.3, 98.2, 142.9)
  выборка в точных вещественных координатах с билинейной интерполяцией
  нигде нет округления

RoIAlign бесплатно повышает mask AP на COCO на 3–4 пункта. Теперь его использует каждый детектор, для которого важна локализация: YOLOv7 seg, RT-DETR и Mask2Former.

RPN в одном абзаце

В каждой позиции карты признаков разместите K якорных рамок разных размеров и соотношений сторон. Для каждого якоря предскажите оценку objectness и смещение регрессии, превращающее якорь в более подходящую рамку. Оставьте примерно 1 000 рамок с наилучшими оценками, примените NMS при IoU 0.7 и передайте выжившие головы. RPN обучается собственной мини-потерей — с той же структурой, что и потеря YOLO из урока 6, только с двумя классами (объект / нет объекта).

Голова масок

Для каждого предложения (после RoIAlign) голова масок — это небольшая FCN: четыре свёртки 3x3, одна deconv 2x и финальная свёртка 1x1, создающая num_classes выходных каналов с разрешением 28x28. Сохраняется только канал, соответствующий предсказанному классу; остальные игнорируются. Это отделяет предсказание маски от классификации.

Увеличьте маску 28x28 до исходного размера предложения в пикселях, чтобы получить итоговую двоичную маску.

Потери

В Mask R-CNN суммируются четыре потери:

L = L_rpn_cls + L_rpn_box + L_box_cls + L_box_reg + L_mask

• L_rpn_cls, L_rpn_box — objectness и регрессия рамок для предложений RPN.
• L_box_cls — кросс-энтропия по классам (C+1), включая фон, в классификаторе головы.
• L_box_reg — smooth L1 для уточнения рамки головой.
• L_mask — попиксельная двоичная кросс-энтропия на выходе маски 28x28.

У каждой потери есть собственный вес по умолчанию; реализация torchvision предоставляет их как аргументы конструктора.

Формат выхода

torchvision.models.detection.maskrcnn_resnet50_fpn_v2 возвращает список словарей, по одному на изображение:

{
    "boxes":  (N, 4) в пиксельных координатах (x1, y1, x2, y2),
    "labels": (N,) идентификаторы классов, 0 = фон, поэтому индексы начинаются с 1,
    "scores": (N,) оценки уверенности,
    "masks":  (N, 1, H, W) вещественные маски в [0, 1] — для двоичной маски порог 0.5,
}

Маска уже имеет полное разрешение изображения. Выход головы 28x28 увеличивается внутри модели.

Соберите сами

Шаг 1: RoIAlign с нуля

Это единственный компонент Mask R-CNN, который проще понять по коду, чем по тексту.

import torch
import torch.nn.functional as F

def roi_align_single(feature, box, output_size=7, spatial_scale=1 / 16.0):
    """
    feature: (C, H, W) single-image feature map
    box: (x1, y1, x2, y2) in original image pixel coordinates
    output_size: side of the output grid (7 for box head, 14 for mask head)
    spatial_scale: reciprocal of the feature map stride
    """
    C, H, W = feature.shape
    x1, y1, x2, y2 = [c * spatial_scale - 0.5 for c in box]
    bin_w = (x2 - x1) / output_size
    bin_h = (y2 - y1) / output_size

    grid_y = torch.linspace(y1 + bin_h / 2, y2 - bin_h / 2, output_size)
    grid_x = torch.linspace(x1 + bin_w / 2, x2 - bin_w / 2, output_size)
    yy, xx = torch.meshgrid(grid_y, grid_x, indexing="ij")

    gx = 2 * (xx + 0.5) / W - 1
    gy = 2 * (yy + 0.5) / H - 1
    grid = torch.stack([gx, gy], dim=-1).unsqueeze(0)
    sampled = F.grid_sample(feature.unsqueeze(0), grid, mode="bilinear",
                            align_corners=False)
    return sampled.squeeze(0)

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

Шаг 2: Сравните с RoIAlign из torchvision

from torchvision.ops import roi_align

feature = torch.randn(1, 16, 50, 50)
boxes = torch.tensor(0, 10, 20, 100, 90, dtype=torch.float32)  # (batch_idx, x1, y1, x2, y2)

ours = roi_align_single(feature[0], boxes[0, 1:].tolist(), output_size=7, spatial_scale=1/4)
theirs = roi_align(feature, boxes, output_size=(7, 7), spatial_scale=1/4, sampling_ratio=1, aligned=True)[0]

print(f"shape ours:   {tuple(ours.shape)}")
print(f"shape theirs: {tuple(theirs.shape)}")
print(f"max|diff|:    {(ours - theirs).abs().max().item():.3e}")

При sampling_ratio=1 и aligned=True результаты совпадают с точностью до 1e-5.

Шаг 3: Загрузите предобученную Mask R-CNN

import torch
from torchvision.models.detection import maskrcnn_resnet50_fpn_v2, MaskRCNN_ResNet50_FPN_V2_Weights

model = maskrcnn_resnet50_fpn_v2(weights=MaskRCNN_ResNet50_FPN_V2_Weights.DEFAULT)
model.eval()
print(f"params: {sum(p.numel() for p in model.parameters()):,}")
print(f"classes (including background): {len(model.roi_heads.box_predictor.cls_score.out_features * [0])}")

46M параметров, 91 класс (COCO). Первый класс (id 0) — фон; всё, что модель действительно обнаруживает, начинается с id 1.

Шаг 4: Запустите инференс

with torch.no_grad():
    x = torch.randn(3, 400, 600)
    predictions = model([x])
p = predictions[0]
print(f"boxes:  {tuple(p['boxes'].shape)}")
print(f"labels: {tuple(p['labels'].shape)}")
print(f"scores: {tuple(p['scores'].shape)}")
print(f"masks:  {tuple(p['masks'].shape)}")

Тензор масок имеет форму (N, 1, H, W). Примените порог 0.5, чтобы получить двоичную маску на объект:

binary_masks = (p['masks'] > 0.5).squeeze(1)  # (N, H, W) boolean

Шаг 5: Замените головы для своего числа классов

Обычный рецепт тонкой настройки: переиспользовать backbone, FPN и RPN; заменить две классификационные головы.

from torchvision.models.detection.faster_rcnn import FastRCNNPredictor
from torchvision.models.detection.mask_rcnn import MaskRCNNPredictor

def build_custom_maskrcnn(num_classes):
    model = maskrcnn_resnet50_fpn_v2(weights=MaskRCNN_ResNet50_FPN_V2_Weights.DEFAULT)
    in_features = model.roi_heads.box_predictor.cls_score.in_features
    model.roi_heads.box_predictor = FastRCNNPredictor(in_features, num_classes)
    in_features_mask = model.roi_heads.mask_predictor.conv5_mask.in_channels
    hidden_layer = 256
    model.roi_heads.mask_predictor = MaskRCNNPredictor(in_features_mask, hidden_layer, num_classes)
    return model

custom = build_custom_maskrcnn(num_classes=5)
print(f"custom cls_score.out_features: {custom.roi_heads.box_predictor.cls_score.out_features}")

num_classes обязан включать класс фона, поэтому для датасета с 4 классами объектов нужно использовать num_classes=5.

Шаг 6: Заморозьте то, чему не нужно обучаться

На небольших датасетах заморозьте backbone и FPN. Обучаются только objectness и регрессия RPN, а также две головы.

def freeze_backbone_and_fpn(model):
    # torchvision Mask R-CNN packs the FPN inside `model.backbone` (as
    # `model.backbone.fpn`), so iterating `model.backbone.parameters()` covers
    # both the ResNet feature layers and the FPN lateral/output convs.
    for p in model.backbone.parameters():
        p.requires_grad = False
    return model

custom = freeze_backbone_and_fpn(custom)
trainable = sum(p.numel() for p in custom.parameters() if p.requires_grad)
print(f"trainable after freeze: {trainable:,}")

Для датасетов из 500 изображений это определяет, будет ли модель сходиться или переобучаться.

Используйте

Полный цикл обучения Mask R-CNN в torchvision занимает 40 строк и практически не меняется от задачи к задаче — достаточно заменить датасет.

def train_step(model, images, targets, optimizer):
    model.train()
    loss_dict = model(images, targets)
    losses = sum(loss for loss in loss_dict.values())
    optimizer.zero_grad()
    losses.backward()
    optimizer.step()
    return {k: v.item() for k, v in loss_dict.items()}

Список targets должен содержать по одному словарю на изображение с boxes, labels и masks (двоичные тензоры формы (num_instances, H, W)). При обучении модель возвращает словарь из четырёх потерь, а при оценке — список предсказаний; выбор определяется model.training.

Оценщик pycocotools вычисляет mAP@IoU=0.5:0.95 и для рамок, и для масок; нужны оба числа, чтобы понять, что является узким местом: голова рамок или голова масок.

Подготовьте к поставке

Этот урок создаёт:

• outputs/prompt-instance-vs-semantic-router.md — промпт, который задаёт три вопроса и выбирает сегментацию экземпляров, семантическую или паноптическую сегментацию, а также точную начальную модель.
• outputs/skill-mask-rcnn-head-swapper.md — навык, который по новому num_classes генерирует 10 строк кода для замены голов в любой модели обнаружения torchvision.

Упражнения

1. (Легко) Проверьте свой RoIAlign относительно torchvision.ops.roi_align на 100 случайных рамках. Сообщите максимальную абсолютную разность. Также запустите RoIPool (поведение до 2017 года) и покажите, что для рамок около границы он расходится примерно на 1–2 пикселя карты признаков.
2. (Средне) Дообучите maskrcnn_resnet50_fpn_v2 на пользовательском датасете из 50 изображений (любые два класса: воздушные шары, рыбы, выбоины, логотипы). Заморозьте backbone, обучайте 20 эпох и сообщите mask AP@0.5.
3. (Сложно) Замените голову масок Mask R-CNN на голову, которая предсказывает 56x56 вместо 28x28. Измерьте mAP@IoU=0.75 до и после. Объясните, почему прирост (или его отсутствие) соответствует ожидаемому компромиссу между точностью границ и памятью.

Ключевые термины

Термин	Как говорят	Что это на самом деле означает
Mask R-CNN	«Обнаружение плюс маски»	Faster R-CNN + небольшая голова FCN, предсказывающая маску 28x28 для каждого предложения и класса
FPN	«Пирамида признаков»	Нисходящие и боковые соединения, дающие каждому уровню шага C каналов семантически насыщенных признаков
RPN	«Генератор предложений регионов»	Небольшая свёрточная голова, создающая около 1000 предложений объект / нет объекта на изображение
RoIAlign	«Вырезка без округления»	Билинейно выбирает сетку признаков фиксированного размера из любой рамки с вещественными координатами
RoIPool	«Вырезка до 2017 года»	Имеет ту же цель, что RoIAlign, но округляет координаты рамки; устарел
Mask AP	«Экземплярный mAP»	Average precision, вычисленная с IoU масок вместо IoU рамок; метрика сегментации экземпляров COCO
Двоичная голова масок	«Маска на класс»	Предсказывает одну двоичную маску для каждого класса и предложения; сохраняется только канал предсказанного класса
Класс фона	«Класс 0»	Общий класс «нет объекта»; индексы реальных классов начинаются с 1

Дополнительные материалы

• Mask R-CNN (He et al., 2017) — статья; критически важен раздел 3 о RoIAlign
• FPN: Feature Pyramid Networks (Lin et al., 2017) — статья об FPN; её использует каждый современный детектор
• Руководство torchvision по Mask R-CNN — справочный источник для цикла тонкой настройки
• Model zoo Detectron2 — продакшен-реализации с обученными весами почти для каждого варианта обнаружения и сегментации

---

Источник: Instance Segmentation — Mask R-CNN 04.07 — Семантическая сегментация — U-Net · Фаза 04 — Компьютерное зрение · 04.09 — Генерация изображений — GAN · Полный каталог