بخش‌بندی نمونه‌ای با Mask R-CNN و پیاده‌سازی RoIAlign

در پروژه‌های واقعی بینایی ماشین، تمایز میان اشیاء هم‌کلاس اهمیتی حیاتی دارد. بخش‌بندی معنایی تفاوتی میان دو فرد در کنار هم قائل نمی‌شود، اما با بخش‌بندی نمونه‌ای (Instance Segmentation) می‌توان برای هر شیء ماسکی مستقل با مرزهای دقیق استخراج کرد. مدل Mask R-CNN استاندارد تثبیت‌شده و پایداری است که این فرآیند را از طریق ترکیب آشکارسازی و پیش‌بینی ماسک حل می‌کند.

چالش مهندسی اساسی در پیاده‌سازی این شبکه‌ها، انحراف مرزهای پیکسلی هنگام برش ویژگی‌ها بود. با معرفی لایه‌ی RoIAlign و حذف خطای گرد کردن مختصات، دقت ماسک‌ها جهش چشمگیری پیدا کرد. شناخت این لایه و نحوه‌ی عملکرد سرهای مجزای شبکه، مهارت کلیدی برای توسعه مدل‌های صنعتی است.

برای اجرای کدهای این آموزش و دریافت وزن‌های پیش‌آموزش‌دیده‌ی torchvision نیازی به کارت گرافیک‌های بسیار سنگین ندارید؛ اما با توجه به اینکه دانلود وزن‌های آماده یا بسته‌ها از PyPI و مخازن خارجی از داخل ایران ممکن است گاهی با کندی یا اختلال همراه باشد، استفاده از کش محلی یا میرورهای پایدار روند کار را روان‌تر می‌کند.

یک شاخه‌ی (branch) کوچک ماسک به آشکارساز Faster R-CNN اضافه کنید تا به بخش‌بندی نمونه‌ای (instance segmentation) برسید. بخش دشوار کار RoIAlign است، و از آنچه به نظر می‌رسد هم چالش‌برانگیزتر است.

اهداف یادگیری

• دنبال کردن معماری Mask R-CNN به‌صورت سرتاسری (end-to-end): ستون فقرات (backbone)، شبکه‌ی هرم ویژگی (FPN)، شبکه‌ی پیشنهاد ناحیه (RPN)، RoIAlign، سر جعبه (box head) و سر ماسک (mask head)
• پیاده‌سازی RoIAlign از صفر و توضیح اینکه چرا دیگر از RoIPool استفاده نمی‌شود
• استفاده از مدل پیش‌آموزش‌دیده‌ی maskrcnn_resnet50_fpn_v2 در torchvision برای تولید ماسک‌های نمونه‌ای با کیفیت پروداکشن و خواندن درست قالب خروجی آن
• فاین‌تیونینگ Mask R-CNN روی یک مجموعه‌داده‌ی کوچک و سفارشی با تعویض سرهای جعبه و ماسک و منجمد نگه‌داشتن ستون فقرات

صورت مسئله

بخش‌بندی معنایی (semantic segmentation) به ازای هر کلاس یک ماسک به شما می‌دهد. بخش‌بندی نمونه‌ای به ازای هر شیء یک ماسک مجزا تولید می‌کند، حتی زمانی که دو شیء از یک کلاس مشترک باشند. شمارش تک‌تک افراد، ردیابی اشیاء در طول فریم‌ها و اندازه‌گیری ابعاد پدیده‌ها (باکس دور هر آجر در یک دیوار، یا تک‌تک سلول‌ها در یک تصویر میکروسکوپی) همگی نیازمند بخش‌بندی نمونه‌ای هستند.

مدل Mask R-CNN (ارائه‌شده توسط He و همکاران، 2017) با بازتعریف بخش‌بندی نمونه‌ای به‌صورت «آشکارسازی به‌علاوه‌ی یک ماسک»، این مسئله را حل کرد. این طراحی آن‌قدر ترتمیز بود که تا پنج سال پس از آن، تقریباً تمام مقالات بخش‌بندی نمونه‌ای گونه‌ای از Mask R-CNN بودند؛ پیاده‌سازی torchvision آن نیز هنوز گزینه‌ی پیش‌فرض پروداکشن برای مجموعه‌داده‌های کوچک تا متوسط است.

مسئله‌ی مهندسی دشوار در اینجا نمونه‌برداری (sampling) است: چگونه یک ناحیه‌ی ویژگی با اندازه‌ی ثابت را از یک باکس پیشنهادی (proposal box) برش دهیم (crop کنیم)، در حالی که گوشه‌های آن دقیقاً روی مرزهای پیکسلی قرار نگرفته‌اند؟ اشتباه در این بخش، در همه‌جا به قیمت از دست رفتن چند دهم امتیاز mAP تمام می‌شود. راهکار این مسئله RoIAlign است.

مفاهیم پایه

معماری شبکه

پنج بخش اصلی که باید بشناسید:

1. ستون فقرات (Backbone) — شبکه‌ی ResNet-50 یا ResNet-101 آموزش‌دیده روی ImageNet. سلسله‌مراتبی از نقشه‌های ویژگی (feature maps) با گام‌ها (strides) برابر با 4، 8، 16 و 32 تولید می‌کند.
2. شبکه‌ی هرم ویژگی (FPN) — اتصالات بالا به پایین و جانبی که در هر سطح، C کانال از ویژگی‌های غنی معنایی فراهم می‌کنند. در فرآیند آشکارسازی، سطحی از FPN فراخوانی می‌شود که با اندازه‌ی شیء متناسب باشد.
3. شبکه‌ی پیشنهاد ناحیه (RPN) — یک سر کانولوشنی کوچک که در موقعیت هر انکر (anchor)، دو چیز را پیش‌بینی می‌کند: «آیا اینجا شیئی وجود دارد؟» و «چگونه این باکس را اصلاح کنم؟». این شبکه حدود 1000 پیشنهاد در هر تصویر تولید می‌کند.
4. لایه‌ی RoIAlign — یک وصله‌ی ویژگی (feature patch) با اندازه‌ی ثابت (مثلاً 7x7) را از هر باکسی روی هر سطحی از FPN نمونه‌برداری می‌کند؛ با نمونه‌برداری دوخطی (bilinear sampling) و بدون کوانتیزه‌سازی (quantisation).
5. سرها (Heads) — یک سر جعبه‌ی دو لایه که ابعاد باکس را دقیق‌تر کرده و کلاس آن را مشخص می‌کند، به‌همراه یک سر کانولوشنی کوچک که برای هر پیشنهاد، یک ماسک دودویی 28x28 تحویل می‌دهد.

چرا RoIAlign و نه RoIPool؟

مدل اولیه‌ی Fast R-CNN از RoIPool استفاده می‌کرد؛ روشی که باکس پیشنهادی را به یک شبکه (grid) تقسیم می‌کند، بیشترین مقدار ویژگی را در هر سلول برمی‌دارد و تمام مختصات را به نزدیک‌ترین عدد صحیح گرد می‌کند. این گرد کردن باعث انحراف نقشه‌ی ویژگی از مختصات پیکسلی ورودی می‌شود؛ انحرافی که گاهی تا یک پیکسل کامل از نقشه‌ی ویژگی می‌رسد. چنین خطایی روی یک تصویر 224x224 شاید کوچک باشد، اما وقتی گام نقشه‌ی ویژگی 32 است، فاجعه‌بار خواهد بود.

RoIPool:
  box (34.7, 51.3, 98.2, 142.9)
  round -> (34, 51, 98, 142)
  split grid -> round each cell boundary
  misalignment accumulates at every step

RoIAlign:
  box (34.7, 51.3, 98.2, 142.9)
  sample at exact float coordinates using bilinear interpolation
  no rounding anywhere

لایه‌ی RoIAlign معیار mask AP را روی مجموعه‌داده‌ی COCO بدون هیچ هزینه‌ای، 3 تا 4 پوینت بالا می‌برد. امروزه هر آشکارسازی که مکان‌یابی دقیق برایش اهمیت دارد، از آن استفاده می‌کند؛ از نسخه‌ی بخش‌بندی YOLOv7 گرفته تا RT-DETR و Mask2Former.

شبکه‌ی RPN در یک بند

در هر موقعیت از نقشه‌ی ویژگی، K باکس لنگر (anchor box) با ابعاد و شکل‌های مختلف قرار دهید. برای هر انکر، یک امتیاز شیء‌بودن (objectness score) و یک جابه‌جایی رگرسیونی برای تبدیل آن به باکسی دقیق‌تر پیش‌بینی کنید. حدود 1000 باکس برتر را بر اساس امتیاز نگه دارید، الگوریتم NMS را با آستانه‌ی IoU برابر با 0.7 روی آن‌ها اعمال کنید و بازمانده‌ها را تحویل سرها دهید. RPN با یک تابع زیان (loss) کوچک و اختصاصی آموزش می‌بیند؛ ساختاری درست مثل تابع زیان YOLO در درس 6، فقط با دو کلاس (شیء / بدون شیء).

سر ماسک

برای هر پیشنهاد (پس از اعمال RoIAlign)، سر ماسک در واقع یک FCN بسیار کوچک است: چهار لایه‌ی کانولوشن 3x3، یک لایه‌ی دیکانولوشن 2x، و یک لایه‌ی کانولوشن پایانی 1x1 که num_classes کانال خروجی با ابعاد 28x28 تولید می‌کند. فقط کانال مربوط به کلاس پیش‌بینی‌شده نگه داشته می‌شود و بقیه نادیده گرفته می‌شوند. این کار پیش‌بینی ماسک را از کلاسبندی مستقل می‌کند.

با بیش‌نمونه‌برداری (upsample) ماسک 28x28 به ابعاد پیکسلی اولیه‌ی باکس پیشنهادی، ماسک دودویی نهایی به دست می‌آید.

توابع زیان

مدل Mask R-CNN پنج تابع زیان (loss) دارد که با هم جمع می‌شوند:

$$L = L_{\text{rpn\_cls}} + L_{\text{rpn\_box}} + L_{\text{box\_cls}} + L_{\text{box\_reg}} + L_{\text{mask}}$$

• L_rpn_cls و L_rpn_box — امتیاز شیء‌بودن به‌علاوه‌ی رگرسیون باکس برای باکس‌های پیشنهادی RPN.
• L_box_cls — آنتروپی متقاطع (Cross-Entropy) روی (C+1) کلاس (شامل پس‌زمینه) در کلاسیفایر سر جعبه.
• L_box_reg — تابع زیان Smooth L1 برای اصلاح نهایی باکس در سر جعبه.
• L_mask — آنتروپی متقاطع دودویی (Binary Cross-Entropy) نقطه‌به‌نقطه روی ماسک خروجی 28x28.

هر تابع زیان یک وزن پیش‌فرض دارد؛ پیاده‌سازی torchvision این وزن‌ها را به‌صورت ورودی‌های تابع سازنده در اختیارتان می‌گذارد.

قالب خروجی

مدل torchvision.models.detection.maskrcnn_resnet50_fpn_v2 لیستی از دیکشنری‌ها را برمی‌گرداند؛ یک دیکشنری برای هر تصویر:

{
    "boxes":  (N, 4) in (x1, y1, x2, y2) pixel coordinates,
    "labels": (N,) class IDs, 0 = background so indices are 1-based,
    "scores": (N,) confidence scores,
    "masks":  (N, 1, H, W) float masks in [0, 1] — threshold at 0.5 for binary,
}

ماسک همین حالا وضوح کامل تصویر را دارد؛ چرا که خروجی 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)

تک‌تک اعداد در موقعیت‌هایی قرار دارند که با نمونه‌برداری دوخطی به دست آمده‌اند. نه گرد کردنی در کار است، نه کوانتیزه‌سازی، و نه حذف گرادیان‌ها (gradient).

گام 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])}")

46 میلیون پارامتر، 91 کلاس (مربوط به COCO). اولین کلاس (شناسه‌ی 0) پس‌زمینه است؛ همه‌ی کلاس‌هایی که مدل در عمل آشکار می‌کند، از شناسه‌ی 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)}")

ابعاد تنسور (tensor) ماسک به‌صورت (N, 1, H, W) است. با اعمال آستانه (threshold) روی مقدار 0.5، به ازای هر شیء یک ماسک دودویی به دست می‌آورید:

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

گام 5: تعویض سرها برای تعداد کلاس‌های سفارشی

روش رایج فاین‌تیونینگ: از ستون فقرات، 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: منجمد کردن بخش‌هایی که نیازی به آموزش ندارند

روی مجموعه‌داده‌های کوچک، ستون فقرات و FPN را منجمد کنید. در این حالت فقط رگرسیون و امتیاز شیء‌بودن RPN به‌همراه دو سر انتهایی آموزش می‌بینند.

def freeze_backbone_and_fpn(model):
    # مدل Mask R-CNN در torchvision بخش FPN را درون `model.backbone` (به صورت
    # `model.backbone.fpn`) قرار می‌دهد، بنابراین پیمایش `model.backbone.parameters()`
    # هم لایه‌های ویژگی ResNet و هم لایه‌های کانولوشن جانبی/خروجی FPN را پوشش می‌دهد.
    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 تصویری، همین نکته مرز میان همگرایی و بیش‌برازش (overfitting) است.

استفاده در عمل

حلقه‌ی کامل آموزش 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 — پرامپتی (prompt) که با پرسیدن سه سؤال، بین بخش‌بندی نمونه‌ای، معنایی یا پان‌پتیک دست به انتخاب می‌زند و مدل دقیق مناسب برای شروع را تعیین می‌کند.
• outputs/skill-mask-rcnn-head-swapper.md — مهارتی که 10 خط کد لازم برای تعویض سرها را روی هر مدلی از آشکارسازهای torchvision بر اساس num_classes جدید تولید می‌کند.

تمرین‌ها

1. (ساده) صحت پیاده‌سازی RoIAlign خود را در برابر torchvision.ops.roi_align روی 100 باکس تصادفی بسنجید. بیشترین اختلاف قدرمطلق را گزارش کنید. همچنین RoIPool (رفتار قبل از سال 2017) را اجرا کنید و نشان دهید که روی باکس‌های نزدیک به حاشیه، تا حدود 1 الی 2 پیکسل از نقشه‌ی ویژگی دچار انحراف می‌شود.
2. (متوسط) مدل maskrcnn_resnet50_fpn_v2 را روی یک مجموعه‌داده‌ی سفارشی 50 تصویری (شامل هر دو کلاسی؛ مثل بادکنک‌ها، ماهی‌ها، چاله یا لوگوها) فاین‌تیون کنید. ستون فقرات را منجمد کرده، مدل را برای 20 ایپاک آموزش دهید و شاخص mask AP@0.5 را گزارش کنید.
3. (سخت) سر ماسک Mask R-CNN را با سری جایگزین کنید که به جای 28x28، در ابعاد 56x56 پیش‌بینی انجام دهد. شاخص mAP@IoU=0.75 را قبل و بعد از این تغییر بسنجید. توضیح دهید که چرا این بهبود (یا عدم وجود آن) با موازنه‌ی مورد انتظار میان دقت مرزها و حافظه‌ی مصرفی همخوانی دارد.

اصطلاحات کلیدی

برای مطالعه‌ی بیشتر

• Mask R-CNN (He et al., 2017) — مقاله‌ی اصلی؛ مطالعه‌ی بخش 3 درباره‌ی RoIAlign اهمیت ویژه‌ای دارد
• FPN: Feature Pyramid Networks (Lin et al., 2017) — مقاله‌ی FPN؛ تقریباً تمام آشکارسازهای مدرن از آن استفاده می‌کنند
• آموزش Mask R-CNN در torchvision — مرجع اصلی برای پیاده‌سازی حلقه‌ی فاین‌تیونینگ
• مخزن مدل‌های Detectron2 — پیاده‌سازی‌های سطح پروداکشن همراه با وزن‌های آموزش‌دیده برای تقریباً تمامی گونه‌های آشکارسازی و بخش‌بندی