인스턴스 분할

인스턴스 분할 (Instance segmentation)

인스턴스 분할은 이미지에서 객체를 탐지하고 각 객체를 픽셀 단위로 분할하는 컴퓨터 비전 작업입니다. 바운딩 박스를 출력하는 객체 탐지와 달리, 인스턴스 분할은 탐지된 모든 객체에 대해 정밀한 마스크를 생성해 객체가 겹쳐 있어도 개별 인스턴스를 구분할 수 있게 해 줍니다.

출처: 문서

본문

이 가이드에서 다룰 내용은 다음과 같습니다.

  1. Hugging Face Hub에서 인스턴스 분할 데이터셋을 로드합니다.
  2. Transformer 기반 인스턴스 분할 모델인 RF-DETR-Seg를 Transformers Trainer로 파인튜닝합니다.
  3. mean IoU로 모델을 평가합니다.
  4. 추론을 실행하고 예측을 시각화합니다.

건물 인스턴스 마스크로 주석이 달린 약 9.6k개의 위성 이미지를 포함하는 satellite-building-segmentation 데이터셋을 사용하겠습니다.

[!TIP] 이 작업과 호환되는 모든 아키텍처와 체크포인트를 보려면 task-page를 확인하는 것을 권장합니다.

시작하기 전에 필요한 라이브러리를 모두 설치하세요.

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

모델을 커뮤니티와 공유하는 것을 권장합니다. 훈련이 끝나면 모델을 Hub에 업로드하기 위해 Hugging Face 계정에 로그인하세요.

>>> 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 딕셔너리가 있습니다. 각 인스턴스에는 [x, y, width, height] 형식의 bbox와 다각형 좌표가 있는 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 통계로 정규화, 균일한 크기로 패딩, 그리고 인스턴스 분할에 결정적으로 중요한 다각형 주석의 이진 마스크 변환, 마스크 리사이즈, 바운딩 박스를 모델이 기대하는 [0, 1] 범위의 [cx, cy, w, h] 형식으로 정규화하는 작업을 수행합니다.

>>> 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는 return_segmentation_masks=True와 함께 이미지와 COCO 형식 주석을 전달하면 모든 무거운 작업을 처리합니다.

  1. 다각형 분할을 이진 마스크로 래스터화
  2. 이미지, 바운딩 박스, 마스크를 모델의 입력 크기로 리사이즈
  3. ImageNet mean/std로 픽셀 값 정규화
  4. 바운딩 박스를 [x, y, w, h]에서 정규화된 [cx, cy, w, h]로 변환
  5. 이미지를 균일한 크기로 패딩하고 pixel_mask 생성

트랜스폼은 데이터셋의 objects 컬럼에서 이미지 프로세서가 기대하는 COCO 스타일 주석 딕셔너리를 재구성합니다.

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(실제 픽셀 vs 패딩), 그리고 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])

데이터 콜레이터

각 이미지마다 객체 인스턴스 수가 다르므로 라벨은 가변 길이 딕셔너리이고 단일 텐서로 쌓을 수 없습니다. 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

평가 메트릭 정의하기

훈련 중 분할 품질을 추적하기 위해 각 평가 에폭에서 union 기반 mean IoU를 계산합니다. 각 이미지에 대해:

  1. 모든 예측 인스턴스 마스크를 단일 이진 "buildings" 맵으로 병합합니다(각 쿼리의 마스크 logits을 sigmoid > 0.5로 이진화).
  2. 모든 그라운드 트루스 인스턴스 마스크를 단일 이진 맵으로 병합합니다.
  3. 두 맵 사이의 Intersection-over-Union을 계산합니다.

RF-DETR-Seg의 쿼리 마스크는 클래스 점수 임계값으로 게이팅되지 않으므로, 메트릭은 클래스 점수를 무시하고 쿼리별 마스크 logits을 직접 합집합합니다. 매칭되지 않은 쿼리는 거의 빈 마스크를 생성하므로 합집합을 오염시키지 않습니다.

이것은 "모델이 건물을 얼마나 잘 덮는가"를 나타내는 이미지별 메트릭이며, 전체 검증 세트에 걸쳐 평균을 냅니다.

[!TIP] 인스턴스 분할 벤치마크(예: COCO)는 보통 마스크 mean average precision(mAP)을 보고하는데, 이는 다양한 IoU 임계값 범위에 걸쳐 각 예측 인스턴스 마스크를 그라운드 트루스와 비교하므로 개별 객체를 올바르게 분리하는 것을 보상합니다. 여기서 사용하는 union 기반 mean IoU는 더 간단하고 빠른 대리자입니다. 인스턴스별 품질보다 전체적인 픽셀 커버리지를 측정하므로 훈련 중 진행 상황을 추적하기에 편리합니다. 표준적이고 인스턴스 인지적인 평가를 위해서는 대신 마스크 mAP를 계산하세요. 예를 들어 torchmetrics의 MeanAveragePrecision(iou_type="segm")을 사용할 수 있습니다.

이것을 트레이너를 서브클래싱하는 대신 compute_metrics 함수로 Trainer에 전달합니다. 아래 TrainingArguments에서 설정하는 eval_do_concat_batches=False를 사용하면 표준 평가 패스가 만든 예측값과 라벨이 배치별 출력 리스트로 compute_metrics에 전달되므로, 메트릭이 해당 예측값을 재사용하고 검증 세트에 대한 두 번째 forward 패스가 필요 없습니다. 모델 출력 튜플에서 인덱스 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 메트릭이 생성됩니다.

훈련 (Training)

데이터, 모델, 메트릭이 준비되었으니 훈련을 설정합니다. TrainingArguments에 대한 몇 가지 중요한 참고 사항입니다.

  • remove_unused_columns=False: 필수입니다. 기본 동작은 트랜스폼이 실행되기 전에 컬럼을 삭제하기 때문입니다.
  • eval_do_concat_batches=False: 인스턴스 분할 라벨은 가변 길이 딕셔너리이므로 배치 간에 연결할 수 없습니다. 또한 예측값을 배치별로 그룹화해 compute_metrics가 라벨과 일치시킬 수 있게 합니다.
  • metric_for_best_model="eval_mean_iou": loss뿐 아니라 분할 품질로 최상의 체크포인트를 선택합니다.
  • 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로 푸시됩니다. 훈련이 완료되면 push_to_hub() 메서드를 호출해 최종 모델도 Hub로 푸시하세요.

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

추론 (Inference)

이제 파인튜닝된 모델이 있으니 새로운 위성 이미지에 대한 추론에 사용해 보세요. 워크플로는 다음과 같습니다.

  1. 이미지 프로세서로 이미지를 전처리합니다.
  2. 모델에 forward 패스를 실행합니다.
  3. post_process_instance_segmentation으로 출력을 후처리해 픽셀 단위 마스크를 얻습니다.

후처리 단계는 원시 쿼리 출력(logits + 저해상도 마스크)을 전체 해상도 인스턴스 분할 맵으로 변환하고, 점수 임계값 처리와 마스크 이진화를 적용합니다.

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

예측을 시각화합니다. 분할 맵은 각 픽셀에 세그먼트 ID를 할당합니다(-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

더 알아보기 (Learn more)