[실전 예제/객체 탐지/PyTorch] 객체 검출 튜토리얼: COCO 데이터셋으로 PyTorch 데이터셋 만들기

객체 검출(Object Detection)이란?

  객체 검출(Object Detection)은 이미지 속의 객체의 종류(class)와 위치(bounding box)를 동시에 예측하는 비전 태스크입니다. COCO는 객체 검출 학습을 위한 가장 널리 사용되는 대표 데이터셋입니다. 이번 시간에는 PyTorch를 이용하여 객체 탐지 데이터셋을 만드는 방법에 대해 알아보도록 하겠습니다.

 

PyTorch로COCO데이터셋 만들기

COCO 데이터셋 특징

• COCO (Common Objects in Context)
• 80개 클래스
• 바운딩 박스 + 클래스 ID + 세그멘테이션 마스크 형태의 라벨로 구성
• JSON (MS COCO format)
• 객체 검출, 인스턴스 분할, 키포인트 검출 등에 사용

디렉토리 구조 예시

coco/
  train2017/
  val2017/
  annotations/
    instances_train2017.json
    instances_val2017.json

 

PyTorch 코드 예제

import os
import json
from typing import Dict, List, Tuple, Any

import torch
from torch.utils.data import Dataset

from PIL import Image
from torchvision.transforms import functional as F


class CocoDetectionDataset(Dataset):
    """
    COCO 객체 검출(Instances) 어노테이션을 사용한 Dataset
      - instances_train2017.json
      - instances_val2017.json
    
    Reads:
      - images_dir: .../train2017 or .../val2017
      - ann_file  : .../annotations/instances_train2017.json (or val)

    Returns:
      image: FloatTensor (3,H,W) in [0,1]
      target: dict with boxes (xyxy), labels (contiguous), image_id, area, iscrowd
    """
    def __init__(self, images_dir: str, ann_file: str, cat_id_to_contiguous: Dict[int, int]):
        self.images_dir = images_dir
        self.ann_file = ann_file
        self.cat_map = cat_id_to_contiguous

        if not os.path.isdir(self.images_dir):
            raise FileNotFoundError(f"[ERR] images_dir not found: {self.images_dir}")
        if not os.path.isfile(self.ann_file):
            raise FileNotFoundError(f"[ERR] ann_file not found: {self.ann_file}")

        with open(self.ann_file, "r", encoding="utf-8") as f:
            coco = json.load(f)

        # image_id -> image_info
        self.id_to_img = {img["id"]: img for img in coco.get("images", [])}
        self.img_ids = sorted(self.id_to_img.keys())

        # group annotations by image_id
        self.anns_by_img: Dict[int, List[dict]] = {}
        for ann in coco.get("annotations", []):
            img_id = ann["image_id"]
            self.anns_by_img.setdefault(img_id, []).append(ann)

    def __len__(self) -> int:
        return len(self.img_ids)

    def __getitem__(self, idx: int) -> Tuple[torch.Tensor, Dict[str, Any]]:
        img_id = self.img_ids[idx]
        info = self.id_to_img[img_id]
        file_name = info["file_name"]

        img_path = os.path.join(self.images_dir, file_name)
        img = Image.open(img_path).convert("RGB")
        image = F.to_tensor(img)  # float32, [0,1]

        anns = self.anns_by_img.get(img_id, [])

        boxes = []
        labels = []
        areas = []
        iscrowd = []

        for ann in anns:
            # COCO bbox: [x,y,w,h] -> xyxy
            x, y, w, h = ann["bbox"]
            if w <= 1 or h <= 1:
                continue

            x1, y1 = float(x), float(y)
            x2, y2 = float(x + w), float(y + h)

            cat_id = int(ann["category_id"])
            if cat_id not in self.cat_map:
                continue

            boxes.append([x1, y1, x2, y2])
            labels.append(self.cat_map[cat_id])
            areas.append(float(ann.get("area", w * h)))
            iscrowd.append(int(ann.get("iscrowd", 0)))

        if len(boxes) == 0:
            boxes_t = torch.zeros((0, 4), dtype=torch.float32)
            labels_t = torch.zeros((0,), dtype=torch.int64)
            areas_t = torch.zeros((0,), dtype=torch.float32)
            iscrowd_t = torch.zeros((0,), dtype=torch.int64)
        else:
            boxes_t = torch.tensor(boxes, dtype=torch.float32)
            labels_t = torch.tensor(labels, dtype=torch.int64)
            areas_t = torch.tensor(areas, dtype=torch.float32)
            iscrowd_t = torch.tensor(iscrowd, dtype=torch.int64)

        target = {
            "boxes": boxes_t,
            "labels": labels_t,
            "image_id": torch.tensor([img_id], dtype=torch.int64),
            "area": areas_t,
            "iscrowd": iscrowd_t,
        }
        return image, target

  COCO annotation에서 category_id는 반드시 1~80로 연속적이지 않을 수 있습니다. Detection 모델에서 클래스 레이블은 0(background)부터 연속적인 정수가 되도록 재매핑(mapping)하는 것이 안전합니다.

  또한 transform이 이미지만 변환하고 있기 때문에, 학습 시에는 이미지와 label(boxes) 모두를 동기화(transform 대상)에 포함시키는 것이 중요합니다.

 

출력 형태

• 이미지 (Tensor): [3, H, W]
• target (dict):

target = {
    'boxes':    Tensor[N, 4],   # [xmin, ymin, xmax, ymax] (float32)
    'labels':   Tensor[N],      # 클래스 인덱스 (int64)
    'image_id': Tensor[1],      # 이미지 고유 ID (int64)
    'area':     Tensor[N],      # 박스 면적 (float32)
    'iscrowd':  Tensor[N],      # crowd 여부 (int64)
}

 

마무리

  PyTorch를 이용하여 객체 검출 데이터셋을 어떻게 만드는지 살펴보았습니다. 다음 시간에는 모델 구성 및 학습 방법을 PyTorch로 작성하는 방법을 알아보도록 하겠습니다.

 

관련 내용

• [PyTorch] 맞춤형 데이터셋 만들기: torch.utils.data.Dataset() 사용 가이드
• [PyTorch] 효율적인 데이터 배치: torch.utils.data.DataLoader() 사용 가이드