src/data/transforms.py

from __future__ import annotations

import numpy as np
import torch
import torch.nn.functional as F
from torchvision import transforms

IMAGENET_MEAN = [0.485, 0.456, 0.406]
IMAGENET_STD = [0.229, 0.224, 0.225]


def _as_tuple_pair(value, default: tuple[float, float]) -> tuple[float, float]:
    if value is None:
        return default
    if isinstance(value, (list, tuple)) and len(value) == 2:
        return float(value[0]), float(value[1])
    raise ValueError(f"Expected a 2-element list/tuple, got: {value!r}")


def resolve_normalization_stats(
    normalization_cfg: dict | None = None,
    default_mean: list[float] | tuple[float, ...] = IMAGENET_MEAN,
    default_std: list[float] | tuple[float, ...] = IMAGENET_STD,
) -> tuple[list[float], list[float]]:
    normalization_cfg = normalization_cfg or {}
    mean = normalization_cfg.get("mean", default_mean)
    std = normalization_cfg.get("std", default_std)

    if not isinstance(mean, (list, tuple)) or len(mean) != 3:
        raise ValueError(f"Expected normalization mean to be a 3-element list/tuple, got: {mean!r}")
    if not isinstance(std, (list, tuple)) or len(std) != 3:
        raise ValueError(f"Expected normalization std to be a 3-element list/tuple, got: {std!r}")

    return [float(v) for v in mean], [float(v) for v in std]


def get_train_transform(
    image_size: int,
    augmentation_cfg: dict | None = None,
    normalization_cfg: dict | None = None,
): #증강 포함 전처리
    augmentation_cfg = augmentation_cfg or {}
    crop_scale = _as_tuple_pair(augmentation_cfg.get("train_crop_scale"), (0.7, 1.0))
    horizontal_flip_prob = float(augmentation_cfg.get("horizontal_flip_prob", 0.5))
    rotation_degrees = float(augmentation_cfg.get("rotation_degrees", 15))
    color_jitter_cfg = augmentation_cfg.get("color_jitter", {})
    grayscale_prob = float(augmentation_cfg.get("grayscale_prob", 0.05))
    mean, std = resolve_normalization_stats(normalization_cfg)

    return transforms.Compose(
        [
            transforms.RandomResizedCrop(image_size, scale=crop_scale),     # 강화 증강: 넓은 크롭 범위
            transforms.RandomHorizontalFlip(p=horizontal_flip_prob),        # 좌우 반전
            transforms.RandomRotation(degrees=rotation_degrees),            # 다양한 촬영 각도 대응
            transforms.ColorJitter(
                brightness=float(color_jitter_cfg.get("brightness", 0.3)),
                contrast=float(color_jitter_cfg.get("contrast", 0.3)),
                saturation=float(color_jitter_cfg.get("saturation", 0.3)),
                hue=float(color_jitter_cfg.get("hue", 0.05)),
            ),                                                             # 조명·색상 변이 대응
            transforms.RandomGrayscale(p=grayscale_prob),                  # 색상 과의존 방지
            transforms.ToTensor(),
            transforms.Normalize(mean=mean, std=std),
        ]
    )


def get_valid_transform(
    image_size: int,
    augmentation_cfg: dict | None = None,
    normalization_cfg: dict | None = None,
): #증강 제외 전처리
    augmentation_cfg = augmentation_cfg or {}
    resize_size = int(augmentation_cfg.get("valid_resize_size", image_size + 32))
    mean, std = resolve_normalization_stats(normalization_cfg)
    return transforms.Compose(
        [
            transforms.Resize(resize_size),
            transforms.CenterCrop(image_size),
            transforms.ToTensor(),
            transforms.Normalize(mean=mean, std=std),
        ]
    )


def build_classification_transform(
    image_size: int,
    is_train: bool,
    augmentation_cfg: dict | None = None,
    normalization_cfg: dict | None = None,
): # 학습용,검증용 변환 선택
    if is_train:
        return get_train_transform(
            image_size,
            augmentation_cfg=augmentation_cfg,
            normalization_cfg=normalization_cfg,
        )
    return get_valid_transform(
        image_size,
        augmentation_cfg=augmentation_cfg,
        normalization_cfg=normalization_cfg,
    )


# CLIP 모델 사전학습 시 사용된 mean/std (OpenAI 공식 값)
CLIP_MEAN = [0.48145466, 0.4578275, 0.40821073]
CLIP_STD  = [0.26862954, 0.26130258, 0.27577711]


def get_clip_train_transform(image_size: int):
    """CLIP fine-tuning용 학습 증강 transform.

    기본 증강(크롭, 플립) 외에 ColorJitter와 RandomRotation을 추가해
    자동차 도메인의 다양한 촬영 조건에 대한 강건성을 높인다.
    Normalize는 CLIP 사전학습 mean/std를 사용한다.
    """
    return transforms.Compose(
        [
            transforms.RandomResizedCrop(image_size, scale=(0.8, 1.0)),  # 80~100% 랜덤 크롭 후 리사이즈
            transforms.RandomHorizontalFlip(),                            # 좌우 반전 (자동차는 대칭 구조)
            transforms.ColorJitter(                                       # 색상 변형: 밝기·대비·채도·색조 랜덤 조정
                brightness=0.1,
                contrast=0.1,
                saturation=0.4,
                hue=0.1,
            ),
            transforms.RandomRotation(degrees=15),                        # ±15도 랜덤 회전
            transforms.ToTensor(),
            transforms.Normalize(mean=CLIP_MEAN, std=CLIP_STD),          # CLIP 전용 정규화
        ]
    )


def get_clip_valid_transform(image_size: int):
    """CLIP fine-tuning용 검증/추론 transform.

    증강 없이 리사이즈·크롭·정규화만 적용한다.
    Normalize는 CLIP 사전학습 mean/std를 사용한다.
    """
    return transforms.Compose(
        [
            transforms.Resize(image_size + 32),          # 여백을 두고 리사이즈 후
            transforms.CenterCrop(image_size),            # 중앙 크롭으로 정보 손실 최소화
            transforms.ToTensor(),
            transforms.Normalize(mean=CLIP_MEAN, std=CLIP_STD),
        ]
    )


# ---------------------------------------------------------------------------
# Batch-level Augmentation: Mixup / CutMix
# ---------------------------------------------------------------------------
# Compose 파이프라인에 넣을 수 없고, 학습 루프에서 배치 단위로 호출한다.
#
# 사용 예시 (학습 루프):
#   mixup   = Mixup(alpha=0.4, num_classes=196)
#   cutmix  = CutMix(alpha=1.0, num_classes=196)
#   aug     = MixupCutMix(mixup=mixup, cutmix=cutmix, mixup_prob=0.5)
#
#   for images, labels in loader:
#       images, soft_labels = aug(images, labels)   # soft_labels: (B, num_classes)
#       logits = model(images)
#       loss = F.cross_entropy(logits, soft_labels)  # soft label CE 지원
# ---------------------------------------------------------------------------

class Mixup:
    """배치 내 두 이미지를 선형 보간해 새로운 이미지를 생성한다 (Mixup 논문, Zhang et al. 2018).

    λ ~ Beta(alpha, alpha) 로 샘플링.
    mixed_image = λ * image_a + (1-λ) * image_b
    mixed_label = λ * label_a + (1-λ) * label_b  (soft label)

    Args:
        alpha:       Beta 분포 파라미터. 클수록 λ가 0.5에 집중 (강한 혼합).
        num_classes: 소프트 레이블 생성에 필요한 클래스 수.
    """

    def __init__(self, alpha: float = 0.4, num_classes: int = 196):
        self.alpha = alpha
        self.num_classes = num_classes

    def __call__(
        self, images: torch.Tensor, labels: torch.Tensor
    ) -> tuple[torch.Tensor, torch.Tensor]:
        """
        Args:
            images: (B, C, H, W) float tensor
            labels: (B,) int tensor — 클래스 인덱스
        Returns:
            mixed_images: (B, C, H, W)
            soft_labels:  (B, num_classes) — 소프트 레이블
        """
        # Beta 분포에서 λ 샘플링 (0~1 사이 혼합 비율)
        lam = float(np.random.beta(self.alpha, self.alpha))

        B = images.size(0)
        # 배치 내 무작위 순서로 섞어 혼합 쌍을 만든다
        perm = torch.randperm(B, device=images.device)

        # 이미지 선형 보간
        mixed_images = lam * images + (1 - lam) * images[perm]

        # 원-핫 인코딩 후 소프트 레이블 계산
        labels_onehot = F.one_hot(labels, num_classes=self.num_classes).float()
        soft_labels = lam * labels_onehot + (1 - lam) * labels_onehot[perm]

        return mixed_images, soft_labels


class CutMix:
    """한 이미지의 사각형 영역을 잘라 다른 이미지에 붙여넣는다 (CutMix 논문, Yun et al. 2019).

    λ ~ Beta(alpha, alpha) 로 샘플링 → 박스 크기 결정.
    박스 면적 비율에 따라 소프트 레이블을 혼합한다.

    Args:
        alpha:       Beta 분포 파라미터.
        num_classes: 소프트 레이블 생성에 필요한 클래스 수.
    """

    def __init__(self, alpha: float = 1.0, num_classes: int = 196):
        self.alpha = alpha
        self.num_classes = num_classes

    @staticmethod
    def _rand_bbox(H: int, W: int, lam: float) -> tuple[int, int, int, int]:
        """λ에 비례하는 면적의 랜덤 박스 좌표를 반환한다.

        박스 변의 길이 = sqrt(1 - λ) * 이미지 변의 길이
        → λ가 클수록 박스가 작아져 원본 이미지가 더 많이 보존된다.
        """
        cut_ratio = np.sqrt(1.0 - lam)
        cut_h = int(H * cut_ratio)
        cut_w = int(W * cut_ratio)

        # 박스 중심점 랜덤 선택
        cx = np.random.randint(W)
        cy = np.random.randint(H)

        # 이미지 경계 내로 클리핑
        x1 = max(cx - cut_w // 2, 0)
        y1 = max(cy - cut_h // 2, 0)
        x2 = min(cx + cut_w // 2, W)
        y2 = min(cy + cut_h // 2, H)

        return x1, y1, x2, y2

    def __call__(
        self, images: torch.Tensor, labels: torch.Tensor
    ) -> tuple[torch.Tensor, torch.Tensor]:
        """
        Args:
            images: (B, C, H, W) float tensor
            labels: (B,) int tensor — 클래스 인덱스
        Returns:
            mixed_images: (B, C, H, W)
            soft_labels:  (B, num_classes) — 소프트 레이블
        """
        lam = float(np.random.beta(self.alpha, self.alpha))

        B, _, H, W = images.shape
        perm = torch.randperm(B, device=images.device)

        x1, y1, x2, y2 = self._rand_bbox(H, W, lam)

        # 박스 영역만 다른 이미지로 교체 (clone으로 원본 보존)
        mixed_images = images.clone()
        mixed_images[:, :, y1:y2, x1:x2] = images[perm, :, y1:y2, x1:x2]

        # 실제 박스 면적 비율로 λ 재계산 (경계 클리핑으로 인한 오차 보정)
        lam = 1 - (x2 - x1) * (y2 - y1) / (W * H)

        # 소프트 레이블: 원본 이미지 비율(lam) + 붙여넣은 이미지 비율(1-lam)
        labels_onehot = F.one_hot(labels, num_classes=self.num_classes).float()
        soft_labels = lam * labels_onehot + (1 - lam) * labels_onehot[perm]

        return mixed_images, soft_labels


class MixupCutMix:
    """Mixup과 CutMix를 확률적으로 선택해 적용하는 래퍼 클래스.

    Args:
        mixup:      Mixup 인스턴스
        cutmix:     CutMix 인스턴스
        mixup_prob: Mixup을 선택할 확률 (나머지 확률로 CutMix 선택)
    """

    def __init__(self, mixup: Mixup, cutmix: CutMix, mixup_prob: float = 0.5):
        self.mixup = mixup
        self.cutmix = cutmix
        self.mixup_prob = mixup_prob

    def __call__(
        self, images: torch.Tensor, labels: torch.Tensor
    ) -> tuple[torch.Tensor, torch.Tensor]:
        """확률에 따라 Mixup 또는 CutMix를 적용한다."""
        if np.random.rand() < self.mixup_prob:
            return self.mixup(images, labels)
        return self.cutmix(images, labels)