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)