[실전 예제/객체 탐지/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/
  images/
    train2017/
    val2017/
  annotations/
    instances_train2017.json
    instances_val2017.json

 

PyTorch 코드 예제

from pycocotools.coco import COCO
from PIL import Image
import os
import torch

class COCOObjectDetectionDataset(torch.utils.data.Dataset):
    def __init__(self, img_dir, ann_file, transform=None):
        self.coco = COCO(ann_file)
        self.img_ids = list(self.coco.imgs.keys())
        self.img_dir = img_dir
        self.transform = transform

    def __getitem__(self, idx):
        img_id = self.img_ids[idx]
        img_info = self.coco.loadImgs(img_id)[0]
        img_path = os.path.join(self.img_dir, img_info['file_name'])
        image = Image.open(img_path).convert('RGB')

        ann_ids = self.coco.getAnnIds(imgIds=img_id)
        anns = self.coco.loadAnns(ann_ids)

        boxes, labels = [], []
        for ann in anns:
            x, y, w, h = ann['bbox']
            boxes.append([x, y, x + w, y + h])
            labels.append(ann['category_id'])

        target = {
            'boxes': torch.tensor(boxes, dtype=torch.float32),
            'labels': torch.tensor(labels, dtype=torch.int64),
            'image_id': torch.tensor([img_id])
        }

        if self.transform:
            image = self.transform(image)

        return image, target

    def __len__(self):
        return len(self.img_ids)

 

출력 형태

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

  {
    'boxes': Tensor[N, 4],      # [xmin, ymin, xmax, ymax]
    'labels': Tensor[N],        # 클래스 인덱스
    'image_id': Tensor[1]
  }

 

마무리

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

 

관련 내용

• 준비중