examples/dataset_mnist.py

# hf_custom_proj/examples/dataset_mnist.py

import numpy as np
from PIL import Image

import torch
from torchvision.datasets import MNIST
from torchvision.datasets.vision import VisionDataset

class MNISTWithProcessor(VisionDataset):
    """
    Hugging Face Trainer에서 바로 사용할 수 있는 MNIST Dataset 예제.

    핵심 목표
    ----------
    - __getitem__에서 HF Trainer가 기대하는 dict 반환:
        {"pixel_values": Tensor(C,H,W), "labels": Tensor(long)}
    - torchvision 스타일 변환 훅 지원:
        - transforms      : (img, y) -> (img, y)  (VisionDataset 관례)
        - transform       : img -> img
        - target_transform: y -> y

    transforms(v2 포함)와 tv_tensors.Image 관련
    ------------------------------------------
    - torchvision.transforms.v2는 변환 결과로 tv_tensors.Image를 반환할 수 있음.
      tv_tensors.Image는 torch.Tensor의 서브클래스이므로, "Dataset 단계" 자체는 보통 문제 없음.
    - 문제는 processor가 입력 타입을 무엇까지 지원하느냐임.
        * processor가 torch.Tensor 입력을 지원하면: tv_tensors.Image도 그대로 처리 가능(권장)
        * processor가 PIL / np.ndarray만 지원하면: tv_tensors.Image에서 TypeError 가능
    - 본 예제는 processor 호출 직전에 image를 "안전하게" torch.Tensor로 고정하여
      입력 타입이 흔들리지 않도록 한다.
      (PIL -> np.ndarray -> torch.from_numpy 경로를 사용하여 확실히 변환)
    - 반드시, 이 구현은 processor가 torch.Tensor 입력을 지원해야 안전하게 처리 가능함

    전처리(Processor) 정책
    -----------------------
    - 최종 모델 입력 규약(크기/정규화/채널/배치 형태)은 ImageProcessor에서 표준화하는 것을 전제로 함.
    - processor(..., return_tensors="pt")는 단일 입력에도 (1,C,H,W)로 반환할 수 있으므로,
      Dataset에서는 더미 배치 차원을 제거해 항상 (C,H,W)만 반환하도록 강제함.
    - 이같이 해야 DataLoader에서 (B,C,H,W)로 쌓이게 됨.
    """

    def __init__(
        self,
        root: str,
        train: bool,
        processor,
        transforms=None,         # (img,y)->(img,y) 또는 img->img 모두 가능
        transform=None,          # img->img
        target_transform=None,   # y->y
        download: bool = True,
    ):
        # VisionDataset 훅 설정
        # (자체적으로 self.transforms/self.transform/self.target_transform을 관리)
        super().__init__(
            root=root,
            transforms=transforms,
            transform=transform,
            target_transform=target_transform,
        )

        # 내부 MNIST (PIL.Image.Image, int) 반환
        self.ds = MNIST(root=root, train=train, download=download)

        # HF ImageProcessor(또는 커스텀 Processor)
        self.processor = processor

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

    def _apply_transforms(self, image, label):
        """
        torchvision 스타일 변환 적용 유틸.

        - self.transforms가 있으면 (img, y)로 먼저 호출(정석).
        - 만약 사용자가 img->img만 처리하는 callable(v2.Compose 등)을 넣었다면 TypeError가 날 수 있어
          그 경우 image만 변환하고 label은 통과시키는 방어 구현.
        1) self.transforms가 있으면 우선 사용:
           - 원칙적으로 (img, y)로 호출을 시도
           - img->img 형태면 TypeError가 날 수 있으므로, 그 경우 image만 변환하고 label은 통과
        2) self.transforms가 없으면 transform / target_transform을 각각 적용
        """
        if self.transforms is not None:
            try:
                image, label = self.transforms(image, label)
            except TypeError:
                image = self.transforms(image)
            return image, label

        # transforms가 없으면 각각 적용
        if self.transform is not None:
            image = self.transform(image)
        if self.target_transform is not None:
            label = self.target_transform(label)

        return image, label

    @staticmethod
    def _to_torch_tensor_image(image) -> torch.Tensor:
        """
        image를 "확실하게" torch.Tensor로 변환.

        지원 입력(대표)
        -------------
        - torch.Tensor (tv_tensors.Image 포함)
        - PIL.Image.Image
        - np.ndarray

        PIL -> np.array -> torch.from_numpy 경로로 확실히 변환함.

        반환
        ----
        - torch.Tensor (CPU)
        - shape는 입력에 따라 (H,W) 또는 (H,W,C) 또는 (C,H,W)일 수 있음
          (최종 (C,H,W)로의 통일은 processor가 담당하는 전제)

        왜 이렇게 하나?
        --------------
        - torch.as_tensor(PIL)은 환경에 따라 동작이 애매할 수 있으므로,
          PIL -> np.ndarray -> torch.from_numpy 경로를 사용해 변환을 확실히 함.
        """
        # 1) 이미 Tensor면 그대로( tv_tensors.Image도 여기로 들어옴 )
        if torch.is_tensor(image):
            # 안전을 위해 CPU로
            return image.detach().to("cpu")

        # 2) PIL.Image.Image -> np.ndarray -> torch.Tensor
        if isinstance(image, Image.Image):
            arr = np.array(image)  # (H,W) or (H,W,C)
            # np.array(PIL)는 보통 uint8이지만, 모드에 따라 다를 수 있음
            return torch.from_numpy(arr)

        # 3) np.ndarray -> torch.Tensor
        if isinstance(image, np.ndarray):
            return torch.from_numpy(image)

        raise TypeError(f"Unsupported image type for tensor conversion: {type(image)}")

    def __getitem__(self, idx: int):
        """
        HF Trainer 호환 dict 반환.

        반환 형식:
          {
            "pixel_values": Tensor(C,H,W),
            "labels": Tensor(long),
          }

        처리 단계
        --------
        1) MNIST에서 (PIL, int) 로드
        2) torchvision transforms 적용 (v2면 tv_tensors.Image가 될 수 있음)
        3) processor 호출 직전 image를 "안전하게" torch.Tensor로 고정
           - tv_tensors.Image: torch.Tensor 서브클래스라 그대로 통과
           - PIL: np.array -> torch.from_numpy 로 확실히 변환
           - np.ndarray: torch.from_numpy
        4) processor로 pixel_values 생성
        5) pixel_values가 (1,C,H,W)이면 더미 배치 차원 제거 -> (C,H,W)
        6) labels를 torch.long으로 변환
        """
        # 1) 원본 로드
        image, label = self.ds[idx]  # (PIL.Image.Image, int)

        # 2) 변환 적용 (여기서 image가 PIL/Tensor/tv_tensor/np.ndarray가 될 수 있음)
        image, label = self._apply_transforms(image, label)

        # 3) 이후에 "추가 변환이 없다면" 타입을 명시적으로 고정하는 편이 안전함
        #    - 특히 v2(tv_tensors.Image) 경로에서도 processor 입력 타입을 단일화할 수 있음
        #    - (전제) processor가 Tensor 입력을 확실히 지원한다면을 가정함.
        image = self._to_torch_tensor_image(image)

        # 4) processor로 모델 입력 생성
        #    - processor가 torch.Tensor 입력을 지원해야 함(권장)
        out = self.processor(image, return_tensors="pt")
        pixel_values = out["pixel_values"]

        # 5) Dataset 반환 규약 통일: 항상 (C,H,W)
        #    - DataLoader default collate로 (B,C,H,W)가 쌓이도록 만들기 위함
        #    - processor가 단일 입력에도 (1,C,H,W)를 반환할 수 있으므로 더미 배치 차원 제거
        if pixel_values.ndim == 4:
            # (1,C,H,W) -> (C,H,W)
            if pixel_values.shape[0] == 1:
                pixel_values = pixel_values[0]
            else:
                raise ValueError(
                    f"Dataset received batched pixel_values with shape {tuple(pixel_values.shape)}"
                )
        elif pixel_values.ndim == 3:
            pass
        else:
            raise ValueError(f"Unexpected pixel_values shape: {tuple(pixel_values.shape)}")

        # 6) labels 텐서화(CE loss 기준 long)
        labels = torch.as_tensor(label, dtype=torch.long)

        return {
            "pixel_values": pixel_values,  # (C,H,W)
            "labels": labels,              # torch.long
        }