dataset.py

"""
Faster R-CNN 밑바닥 구현 — [1/5] 데이터셋 (Pascal VOC)
=========================================================
Pascal VOC 2007/2012 데이터를 읽어 (이미지, 박스, 라벨)을 반환한다.

VOC 어노테이션은 XML 형식이며, 각 객체마다 다음을 담는다:
  - name       : 클래스 이름 (예: 'person', 'car')
  - bndbox     : xmin, ymin, xmax, ymax (좌상단·우하단 픽셀 좌표)
  - difficult  : 판별 어려운 객체 표시(학습 시 보통 제외)

핵심 개념:
  - 박스 좌표는 [x1, y1, x2, y2] 절대 픽셀 좌표로 통일한다.
  - 이미지를 리사이즈하면 박스도 같은 비율로 스케일해야 한다.
"""

import os
import xml.etree.ElementTree as ET

import torch
from torch.utils.data import Dataset
from PIL import Image
import torchvision.transforms.functional as F


# Pascal VOC 20개 클래스 (인덱스 0은 배경으로 예약 → 클래스는 1부터)
VOC_CLASSES = [
    "aeroplane", "bicycle", "bird", "boat", "bottle",
    "bus", "car", "cat", "chair", "cow",
    "diningtable", "dog", "horse", "motorbike", "person",
    "pottedplant", "sheep", "sofa", "train", "tvmonitor",
]
# 이름 → 인덱스 (배경=0 이므로 +1)
CLASS_TO_IDX = {name: i + 1 for i, name in enumerate(VOC_CLASSES)}
NUM_CLASSES = len(VOC_CLASSES) + 1  # +1 = 배경(background)


class VOCDataset(Dataset):
    """
    Pascal VOC 객체탐지 데이터셋.

    Args:
        root:        VOCdevkit/VOC2007 (또는 VOC2012) 경로
        split:       'train' | 'val' | 'trainval' | 'test'
        min_size:    리사이즈 시 이미지 짧은 변의 목표 길이
        max_size:    긴 변의 최대 길이(비율 유지하며 상한 적용)
        keep_difficult: difficult=1 객체를 포함할지 여부(학습 시 False 권장)
    """

    def __init__(self, root, split="trainval", min_size=600, max_size=1000,
                 keep_difficult=False):
        self.root = root
        self.min_size = min_size
        self.max_size = max_size
        self.keep_difficult = keep_difficult

        # ImageSets/Main/<split>.txt 에 이미지 ID 목록이 있다.
        split_file = os.path.join(root, "ImageSets", "Main", f"{split}.txt")
        with open(split_file) as f:
            self.ids = [line.strip() for line in f if line.strip()]

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

    def _load_annotation(self, img_id):
        """XML을 파싱해 박스와 라벨을 뽑는다."""
        ann_path = os.path.join(self.root, "Annotations", f"{img_id}.xml")
        tree = ET.parse(ann_path)
        boxes, labels = [], []

        for obj in tree.findall("object"):
            difficult = int(obj.findtext("difficult", "0"))
            if difficult and not self.keep_difficult:
                continue

            name = obj.findtext("name").strip().lower()
            if name not in CLASS_TO_IDX:
                continue

            bnd = obj.find("bndbox")
            # VOC 좌표는 1부터 시작 → 0-기반으로 보정(-1)
            x1 = float(bnd.findtext("xmin")) - 1
            y1 = float(bnd.findtext("ymin")) - 1
            x2 = float(bnd.findtext("xmax")) - 1
            y2 = float(bnd.findtext("ymax")) - 1
            boxes.append([x1, y1, x2, y2])
            labels.append(CLASS_TO_IDX[name])

        boxes = torch.as_tensor(boxes, dtype=torch.float32).reshape(-1, 4)
        labels = torch.as_tensor(labels, dtype=torch.int64)
        return boxes, labels

    def _resize(self, img, boxes):
        """
        짧은 변을 min_size로 맞추되, 긴 변이 max_size를 넘지 않도록 스케일.
        박스도 같은 비율로 조정한다. (Faster R-CNN 원논문 방식)
        """
        w, h = img.size
        short, long = min(w, h), max(w, h)
        scale = self.min_size / short
        if long * scale > self.max_size:
            scale = self.max_size / long

        new_w, new_h = int(round(w * scale)), int(round(h * scale))
        img = img.resize((new_w, new_h), Image.BILINEAR)
        if boxes.numel() > 0:
            boxes = boxes * scale  # 박스도 동일 배율 적용
        return img, boxes, scale

    def __getitem__(self, idx):
        img_id = self.ids[idx]
        img_path = os.path.join(self.root, "JPEGImages", f"{img_id}.jpg")
        img = Image.open(img_path).convert("RGB")

        boxes, labels = self._load_annotation(img_id)
        img, boxes, scale = self._resize(img, boxes)

        # 텐서 변환 + ImageNet 정규화(백본이 ImageNet 사전학습이므로)
        img = F.to_tensor(img)
        img = F.normalize(img,
                          mean=[0.485, 0.456, 0.406],
                          std=[0.229, 0.224, 0.225])

        target = {
            "boxes": boxes,          # [N,4] 절대 픽셀 (리사이즈 후)
            "labels": labels,        # [N]   1..20 (0=배경)
            "image_id": img_id,
            "scale": scale,          # 평가 시 원본 좌표로 되돌릴 때 사용
        }
        return img, target


def collate_fn(batch):
    """
    이미지마다 크기가 달라 기본 collate로 못 묶는다.
    리스트 형태로 그대로 넘기고, 모델 내부에서 처리한다.
    (간단화를 위해 batch_size=1 사용을 권장)
    """
    imgs, targets = list(zip(*batch))
    return list(imgs), list(targets)