객체 탐지
객체 탐지 (Object detection)
객체 탐지는 이미지에서 인스턴스(사람, 건물, 자동차 등)를 감지하는 컴퓨터 비전 작업이에요. 객체 탐지 모델은 이미지를 입력으로 받아 감지된 객체들의 바운딩 박스(bounding box) 좌표와 관련 라벨을 출력하죠. 이미지에는 객체가 여러 개 있을 수 있고(예: 자동차와 건물), 각각 자신만의 바운딩 박스와 라벨을 가집니다. 또 각 객체는 이미지의 다른 부분에 존재할 수 있어요(예: 이미지에 자동차가 여러 대 있을 수 있죠). 이 작업은 자율주행에서 보행자, 도로 표지판, 신호등 같은 것을 감지하는 데 흔히 쓰여요. 그 외에도 이미지 속 객체 개수 세기, 이미지 검색 등 다양한 응용이 있습니다.
이 가이드에서 배울 내용은 이렇습니다:
- RF-DETR을 mobile-ui-design 데이터셋으로 파인튜닝해 모바일 앱 스크린샷의 UI 요소를 탐지합니다.
- 파인튜닝한 모델로 추론(inference)을 수행합니다.
이 작업과 호환되는 모든 아키텍처와 체크포인트를 보려면 task-page를 확인해 보세요.
시작하기 전에 필요한 라이브러리가 모두 설치돼 있는지 확인하세요:
pip install -q datasets transformers accelerate timm trackio torchmetrics pycocotools
🤗 Datasets로 Hugging Face Hub에서 데이터셋을 불러오고, 🤗 Transformers로 모델을 훈련할 거예요.
모델을 커뮤니티와 공유하길 권장해요. Hub에 업로드하려면 Hugging Face 계정에 로그인하세요. 요청이 오면 토큰을 입력해 로그인합니다:
>>> from huggingface_hub import notebook_login
>>> notebook_login()
모델 이름과 이미지 크기 같은 전역 상수를 정의합니다. 이 튜토리얼은 RF-DETR을 쓰지만, Transformers의 어떤 객체 탐지 모델이든 선택할 수 있어요.
>>> MODEL_NAME = "Roboflow/rf-detr-medium"
mobile-ui-design 데이터셋 불러오기
mobile-ui-design 데이터셋은 텍스트, 이미지, 사각형, 그룹 같은 UI 요소를 탐지하기 위한 주석이 달린 모바일 앱 스크린샷을 담고 있어요.
먼저 데이터셋을 불러오고 카테고리 라벨을 추출합니다. 이 데이터셋은 이미 스플릿이 나뉘어 있어요.
>>> from datasets import load_dataset
>>> ds = load_dataset("merve/mobile-ui-design")
>>> CATEGORIES = sorted(set(
... cat for split in ds.values() for example in split for cat in example["objects"]["category"]
... ))
>>> label2id = {label: i for i, label in enumerate(CATEGORIES)}
>>> id2label = {i: label for label, i in label2id.items()}
>>> print(f"Categories ({len(CATEGORIES)}): {CATEGORIES}")
Categories (4): ['group', 'image', 'rectangle', 'text']
데이터셋은 문자열 카테고리 이름과 COCO 형식 (x, y, w, h)의 바운딩 박스를 써요. 카테고리를 정수 id로 변환하고, 면적을 계산하고, 훈련 전에 퇴화된 바운딩 박스를 걸러냅니다:
>>> def prepare_example(example, idx):
... objects = example["objects"]
... bboxes = objects["bbox"]
... categories = objects["category"]
... img_w, img_h = example["width"], example["height"]
... bboxes, cats, areas, ids = [], [], [], []
... for i, (bbox, cat) in enumerate(zip(bboxes, categories)):
... x, y, w, h = bbox
... if w <= 0 or h <= 0:
... continue
... x = max(0.0, min(x, img_w))
... y = max(0.0, min(y, img_h))
... w = min(w, img_w - x)
... h = min(h, img_h - y)
... if w <= 0 or h <= 0:
... continue
... bboxes.append([x, y, w, h])
... cats.append(label2id[cat])
... areas.append(w * h)
... ids.append(i)
... return {
... "image_id": idx, "image": example["image"],
... "width": example["width"], "height": example["height"],
... "objects": {"id": ids, "bbox": bboxes, "category": cats, "area": areas},
... }
>>> ds_prepared = ds["train"].map(prepare_example, with_indices=True, remove_columns=ds["train"].column_names)
>>> ds_prepared = ds_prepared.filter(lambda x: len(x["objects"]["bbox"]) > 0)
>>> split = ds_prepared.train_test_split(test_size=0.15, seed=1337)
>>> train_ds = split["train"]
>>> val_ds = split["test"]
>>> print(f"Train: {len(train_ds)}, Validation: {len(val_ds)}")
Train: 6669, Validation: 1177
데이터 전처리
[AutoImageProcessor]는 모델이 훈련할 수 있는 pixel_values, pixel_mask, labels를 만들도록 이미지 데이터를 처리해 줘요. 이미지 프로세서는 리사이즈, 패딩, 정규화를 담당하죠. 여기에 더해 일반화를 높이기 위해 무작위 데이터 증강을 선택적으로 추가할 수도 있습니다 (아래 데이터 증강 참고).
>>> import numpy as np
>>> from functools import partial
>>> from transformers import AutoImageProcessor
>>> image_processor = AutoImageProcessor.from_pretrained(MODEL_NAME)
image_processor는 주석을 COCO 형식({'image_id': int, 'annotations': list[Dict]})으로 기대해요. 각 예시의 주석을 형식에 맞추고 나머지는 프로세서가 처리하게 맡기면 됩니다:
>>> def format_image_annotations_as_coco(image_id, categories, areas, bboxes):
... annotations = []
... for category, area, bbox in zip(categories, areas, bboxes):
... annotations.append({
... "image_id": image_id,
... "category_id": category,
... "iscrowd": 0,
... "area": area,
... "bbox": list(bbox),
... })
... return {"image_id": image_id, "annotations": annotations}
>>> def transform_batch(examples, image_processor):
... images = []
... annotations = []
... for image_id, image, objects in zip(examples["image_id"], examples["image"], examples["objects"]):
... images.append(np.array(image.convert("RGB")))
... formatted = format_image_annotations_as_coco(
... image_id, objects["category"], objects["area"], objects["bbox"]
... )
... annotations.append(formatted)
... result = image_processor(images=images, annotations=annotations, return_tensors="pt")
... result.pop("pixel_mask", None)
... return result
>>> transform_fn = partial(transform_batch, image_processor=image_processor)
>>> train_ds = train_ds.with_transform(transform_fn)
>>> val_ds = val_ds.with_transform(transform_fn)
데이터 증강
위 변환은 이미지를 리사이즈하고 정규화할 뿐이에요. 훈련 스플릿에 적용하는 무작위 증강은 보통 일반화를 높여 주지만, 검증 스플릿은 평가를 결정적으로 유지하기 위해 증강 없이 두어야 해요. 흔히 쓰는 선택은 이미지와 그 바운딩 박스를 함께 증강하는 Albumentations입니다. bbox_params로 파이프라인을 정의해서 박스가 이미지와 일관되게 변환되게 하고, 증강된 박스에서 면적을 다시 계산해 봅니다:
>>> import albumentations as A
>>> train_augment = A.Compose(
... [
... A.Perspective(p=0.1),
... A.HorizontalFlip(p=0.5),
... A.RandomBrightnessContrast(p=0.5),
... A.HueSaturationValue(p=0.1),
... ],
... bbox_params=A.BboxParams(format="coco", label_fields=["category"], clip=True, min_area=25),
... )
>>> def augment_and_transform_batch(examples, image_processor, transform):
... images = []
... annotations = []
... for image_id, image, objects in zip(examples["image_id"], examples["image"], examples["objects"]):
... image = np.array(image.convert("RGB"))
... output = transform(image=image, bboxes=objects["bbox"], category=objects["category"])
... images.append(output["image"])
... areas = [w * h for (_, _, w, h) in output["bboxes"]]
... formatted = format_image_annotations_as_coco(
... image_id, output["category"], areas, output["bboxes"]
... )
... annotations.append(formatted)
... result = image_processor(images=images, annotations=annotations, return_tensors="pt")
... result.pop("pixel_mask", None)
... return result
증강 변환을 훈련 스플릿에만 적용하고, 검증에는 일반 transform_fn을 유지합니다:
>>> train_augment_fn = partial(augment_and_transform_batch, image_processor=image_processor, transform=train_augment)
>>> train_ds = train_ds.with_transform(train_augment_fn)
>>> val_ds = val_ds.with_transform(transform_fn)
이미지를 함께 배치하기 위해 커스텀 collate_fn을 만듭니다:
>>> import torch
>>> def collate_fn(batch):
... data = {}
... data["pixel_values"] = torch.stack([x["pixel_values"] for x in batch])
... data["labels"] = [x["labels"] for x in batch]
... if "pixel_mask" in batch[0]:
... data["pixel_mask"] = torch.stack([x["pixel_mask"] for x in batch])
... return data
mAP를 계산하는 함수 준비하기
객체 탐지 모델은 보통 COCO 스타일 지표 세트로 평가돼요. torchmetrics로 mAP(mean average precision)와 mAR(mean average recall) 지표를 계산하고, [Trainer]의 평가에 쓸 수 있도록 compute_metrics 함수로 감쌀 거예요.
훈련에 쓰이는 박스의 중간 형식은 YOLO(정규화)지만, 박스 면적을 제대로 처리하려면 Pascal VOC(절대값) 형식으로 지표를 계산합니다. 바운딩 박스를 Pascal VOC 형식으로 변환하는 함수를 정의해 볼게요:
>>> from transformers.image_transforms import center_to_corners_format
>>> def convert_bbox_yolo_to_pascal(boxes, image_size):
... """
... Convert bounding boxes from YOLO format (x_center, y_center, width, height) in range [0, 1]
... to Pascal VOC format (x_min, y_min, x_max, y_max) in absolute coordinates.
... Args:
... boxes (torch.Tensor): Bounding boxes in YOLO format
... image_size (tuple[int, int]): Image size in format (height, width)
... Returns:
... torch.Tensor: Bounding boxes in Pascal VOC format (x_min, y_min, x_max, y_max)
... """
... # convert center to corners format
... boxes = center_to_corners_format(boxes)
... # convert to absolute coordinates
... height, width = image_size
... boxes = boxes * torch.tensor([[width, height, width, height]])
... return boxes
그다음 compute_metrics 함수에서 평가 루프 결과로부터 예측·타깃 바운딩 박스, 점수, 라벨을 모아 스코어링 함수에 넘깁니다.
>>> import numpy as np
>>> from dataclasses import dataclass
>>> from torchmetrics.detection.mean_ap import MeanAveragePrecision
>>> @dataclass
>>> class ModelOutput:
... logits: torch.Tensor
... pred_boxes: torch.Tensor
>>> def _get_orig_size(image_target):
... """Robust orig_size extraction - Trainer serialization can truncate to 1 element."""
... orig = np.atleast_1d(np.asarray(image_target["orig_size"])).flatten()
... if len(orig) >= 2:
... return (int(orig[0]), int(orig[1]))
... return (int(orig[0]), int(orig[0]))
>>> @torch.no_grad()
>>> def compute_metrics(evaluation_results, image_processor, threshold=0.0, id2label=None):
... predictions, targets = evaluation_results.predictions, evaluation_results.label_ids
... image_sizes = []
... post_processed_targets = []
... post_processed_predictions = []
... for batch in targets:
... batch_sizes = []
... for image_target in batch:
... h, w = _get_orig_size(image_target)
... batch_sizes.append([h, w])
... boxes = torch.tensor(image_target["boxes"])
... boxes = convert_bbox_yolo_to_pascal(boxes, (h, w))
... labels = torch.tensor(image_target["class_labels"])
... post_processed_targets.append({"boxes": boxes, "labels": labels})
... image_sizes.append(torch.tensor(batch_sizes))
... for batch, target_sizes in zip(predictions, image_sizes):
... batch_logits, batch_boxes = batch[1], batch[2]
... output = ModelOutput(logits=torch.tensor(batch_logits), pred_boxes=torch.tensor(batch_boxes))
... post_processed_output = image_processor.post_process_object_detection(
... output, threshold=threshold, target_sizes=target_sizes
... )
... post_processed_predictions.extend(post_processed_output)
... metric = MeanAveragePrecision(box_format="xyxy", class_metrics=True)
... metric.update(post_processed_predictions, post_processed_targets)
... metrics = metric.compute()
... classes = metrics.pop("classes")
... map_per_class = metrics.pop("map_per_class")
... mar_100_per_class = metrics.pop("mar_100_per_class")
... for class_id, class_map, class_mar in zip(classes, map_per_class, mar_100_per_class):
... class_name = id2label[class_id.item()] if id2label is not None else class_id.item()
... metrics[f"map_{class_name}"] = class_map
... metrics[f"mar_100_{class_name}"] = class_mar
... metrics = {k: round(v.item(), 4) for k, v in metrics.items()}
... return metrics
>>> eval_compute_metrics_fn = partial(
... compute_metrics, image_processor=image_processor, id2label=id2label, threshold=0.0
... )
탐지 모델 훈련
이전 섹션에서 힘든 일을 대부분 끝냈으니, 이제 모델을 훈련할 준비가 됐어요! 이 데이터셋의 이미지는 리사이즈 후에도 여전히 꽤 큽니다. 즉 이 모델을 파인튜닝하려면 최소한 GPU 하나가 필요해요.
훈련은 다음 단계로 진행됩니다:
- 전처리에서 쓴 것과 같은 체크포인트로 [
AutoModelForObjectDetection]을 사용해 모델을 로드합니다. - [
TrainingArguments]에 훈련 하이퍼파라미터를 정의합니다. - 훈련 인자를 모델, 데이터셋, 이미지 프로세서, 데이터 콜레이터와 함께 [
Trainer]로 넘깁니다. - [
~Trainer.train]을 호출해 모델을 파인튜닝합니다.
전처리에 쓴 것과 같은 체크포인트에서 모델을 로드할 때, 앞서 데이터셋 메타데이터로 만든 label2id와 id2label 맵을 넘기는 걸 잊지 마세요. 또 기존 분류 헤드를 새 것으로 교체하기 위해 ignore_mismatched_sizes=True를 지정합니다.
>>> from transformers import AutoModelForObjectDetection
>>> model = AutoModelForObjectDetection.from_pretrained(
... MODEL_NAME,
... id2label=id2label,
... label2id=label2id,
... ignore_mismatched_sizes=True,
... )
[TrainingArguments]에서 output_dir로 모델을 저장할 위치를 지정하고, 하이퍼파라미터를 원하는 대로 설정합니다. num_train_epochs=5면 A100 GPU에서 훈련에 약 35분이 걸려요. 에폭 수를 늘리면 더 좋은 결과를 얻을 수 있습니다.
중요한 참고 사항:
- 사용하지 않는 컬럼을 제거하면 이미지 컬럼이 사라지므로 제거하지 마세요. 이미지 컬럼이 없으면
pixel_values를 만들 수 없어요. 그래서remove_unused_columns를False로 설정합니다. - 올바른 평가 결과를 얻으려면
eval_do_concat_batches=False를 설정하세요. 이미지마다 타깃 박스 개수가 다른데, 배치를 이어 붙이면 어떤 박스가 특정 이미지에 속하는지 알 수 없게 돼요.
모델을 Hub에 push해서 공유하고 싶다면 push_to_hub를 True로 설정하세요 (모델을 업로드하려면 Hugging Face에 로그인해야 해요).
>>> from transformers import TrainingArguments
>>> training_args = TrainingArguments(
... output_dir="rf_detr_finetuned_mobile_ui",
... num_train_epochs=5,
... bf16=True,
... per_device_train_batch_size=8,
... dataloader_num_workers=4,
... learning_rate=5e-5,
... lr_scheduler_type="cosine",
... weight_decay=1e-4,
... max_grad_norm=0.01,
... metric_for_best_model="eval_map",
... greater_is_better=True,
... load_best_model_at_end=True,
... eval_strategy="epoch",
... save_strategy="epoch",
... save_total_limit=2,
... remove_unused_columns=False,
... report_to="trackio",
... run_name="mobile-ui-detection",
... eval_do_concat_batches=False,
... push_to_hub=True,
... )
마지막으로 모든 것을 한데 묶고 [~transformers.Trainer.train]을 호출합니다:
>>> from transformers import Trainer
>>> trainer = Trainer(
... model=model,
... args=training_args,
... train_dataset=train_ds,
... eval_dataset=val_ds,
... processing_class=image_processor,
... data_collator=collate_fn,
... compute_metrics=eval_compute_metrics_fn,
... )
>>> trainer.train()
[2085/2085 38:39, Epoch 5/5]
| Epoch | Training Loss | Validation Loss | Map | Map 50 | Map 75 | Map Small | Map Medium | Map Large | Mar 1 | Mar 10 | Mar 100 | Mar Small | Mar Medium | Mar Large | Map Group | Mar 100 Group | Map Image | Mar 100 Image | Map Rectangle | Mar 100 Rectangle | Map Text | Mar 100 Text |
|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|
| 1 | No log | 9.9234 | 0.1303 | 0.2236 | 0.1478 | 0.0909 | 0.2030 | 0.2524 | 0.0421 | 0.2520 | 0.4683 | 0.3113 | 0.5607 | 0.6782 | 0.1244 | 0.5122 | 0.0958 | 0.5035 | 0.1285 | 0.4328 | 0.1725 | 0.4413 |
| 2 | No log | 9.8472 | 0.1893 | 0.3017 | 0.2124 | 0.1347 | 0.2789 | 0.3038 | 0.0549 | 0.2961 | 0.5140 | 0.3433 | 0.5941 | 0.7406 | 0.1305 | 0.5423 | 0.1979 | 0.5578 | 0.1964 | 0.4648 | 0.2324 | 0.4437 |
| 3 | No log | 9.6401 | 0.2275 | 0.3547 | 0.2657 | 0.1698 | 0.3336 | 0.3892 | 0.0611 | 0.3204 | 0.5270 | 0.3625 | 0.6143 | 0.7496 | 0.1602 | 0.5684 | 0.2617 | 0.5763 | 0.2249 | 0.4684 | 0.2631 | 0.4692 |
| 4 | No log | 9.5770 | 0.2733 | 0.4068 | 0.3133 | 0.2100 | 0.3867 | 0.4343 | 0.0668 | 0.3456 | 0.5593 | 0.3875 | 0.6393 | 0.7725 | 0.2013 | 0.5941 | 0.3158 | 0.6065 | 0.2733 | 0.4998 | 0.3028 | 0.4756 |
| 5 | 10.3700 | 11.0500 | 0.2827 | 0.4193 | 0.2913 | 0.2021 | 0.2814 | 0.3763 | 0.0609 | 0.3403 | 0.5668 | 0.4138 | 0.5669 | 0.7317 | 0.2092 | 0.5979 | 0.3334 | 0.6295 | 0.2793 | 0.5245 | 0.3089 | 0.5151 |
training_args에서 push_to_hub를 True로 설정했다면 훈련 체크포인트가 Hugging Face Hub에 push돼요. 훈련이 끝나면 [~transformers.Trainer.push_to_hub] 메서드를 호출해 최종 모델도 Hub에 올립니다.
>>> trainer.push_to_hub()
평가
>>> from pprint import pprint
>>> metrics = trainer.evaluate(eval_dataset=val_ds, metric_key_prefix="test")
>>> pprint(metrics)
{'test_loss': 11.05,
'test_map': 0.2827,
'test_map_50': 0.4193,
'test_map_75': 0.2913,
'test_map_group': 0.2092,
'test_map_image': 0.3334,
'test_map_large': 0.3763,
'test_map_medium': 0.2814,
'test_map_rectangle': 0.2793,
'test_map_small': 0.2021,
'test_map_text': 0.3089,
'test_mar_1': 0.0609,
'test_mar_10': 0.3403,
'test_mar_100': 0.5668,
'test_mar_100_group': 0.5979,
'test_mar_100_image': 0.6295,
'test_mar_100_rectangle': 0.5245,
'test_mar_100_text': 0.5151,
'test_mar_large': 0.7317,
'test_mar_medium': 0.5669,
'test_mar_small': 0.4138}
이 결과는 에폭 수를 늘리거나 [TrainingArguments]의 다른 하이퍼파라미터를 조정하면 더 개선할 수 있어요. 한번 해 보세요!
추론
이제 모델을 파인튜닝하고, 평가하고, Hugging Face Hub에 업로드했으니 추론에 쓸 수 있어요.
>>> import torch
>>> from PIL import Image, ImageDraw
>>> from transformers import AutoImageProcessor, AutoModelForObjectDetection
>>> from datasets import load_dataset
>>> ds = load_dataset("merve/mobile-ui-design", split="train")
>>> image = ds[5]["image"].convert("RGB")
Hugging Face Hub에서 모델과 이미지 프로세서를 로드합니다 (이번 세션에서 이미 훈련한 걸 쓰려면 건너뛰어도 됩니다):
>>> model_repo = "merve/rf_detr_finetuned_mobile_ui"
>>> image_processor = AutoImageProcessor.from_pretrained(model_repo)
>>> model = AutoModelForObjectDetection.from_pretrained(model_repo)
>>> model.eval()
그리고 바운딩 박스를 탐지합니다:
>>> with torch.no_grad():
... inputs = image_processor(images=[image], return_tensors="pt")
... outputs = model(**inputs)
... target_sizes = torch.tensor([[image.size[1], image.size[0]]])
... results = image_processor.post_process_object_detection(outputs, threshold=0.5, target_sizes=target_sizes)[0]
>>> for score, label, box in zip(results["scores"], results["labels"], results["boxes"]):
... box = [round(i, 2) for i in box.tolist()]
... print(
... f"Detected {model.config.id2label[label.item()]} with confidence "
... f"{round(score.item(), 3)} at location {box}"
... )
Detected text with confidence 0.727 at location [324.02, 340.55, 339.52, 359.12]
Detected rectangle with confidence 0.717 at location [39.97, 705.14, 335.93, 753.54]
Detected text with confidence 0.702 at location [199.94, 473.66, 213.41, 490.6]
Detected text with confidence 0.678 at location [153.14, 474.81, 165.33, 491.0]
Detected text with confidence 0.675 at location [262.67, 718.28, 281.44, 740.81]
Detected rectangle with confidence 0.655 at location [143.57, 242.51, 214.32, 274.26]
Detected text with confidence 0.653 at location [298.68, 637.77, 345.68, 656.26]
결과를 그려 볼게요:
>>> draw = ImageDraw.Draw(image)
>>> colors = {"group": "blue", "image": "green", "rectangle": "red", "text": "orange"}
>>> for score, label, box in zip(results["scores"], results["labels"], results["boxes"]):
... box = [round(i, 2) for i in box.tolist()]
... x, y, x2, y2 = tuple(box)
... label_name = model.config.id2label[label.item()]
... color = colors.get(label_name, "red")
... draw.rectangle((x, y, x2, y2), outline=color, width=2)
... draw.text((x, y), f"{label_name} {score:.2f}", fill=color)
>>> image
더 알아보기 (Learn more)
- task-page: 객체 탐지와 호환되는 모든 아키텍처·체크포인트 목록.
- COCO 스타일 지표: 객체 탐지 평가 지표.
- Albumentations: 이미지와 바운딩 박스를 함께 증강하는 라이브러리.