from __future__ import annotations """Naruto Sign data pipeline: stratified split, transforms, imbalance-aware sampling. The dataset is expected as an ``ImageFolder`` tree:: data_root/ Bird/ img001.png ... Boar/ ... ... If the dataset already ships a ``train/`` and ``test/`` split, point ``data_root`` at the parent and set ``use_predefined_split=True``; otherwise a single folder is split here in a *stratified*, *seeded* way so the protocol stays reproducible. Design notes (see PROTOCOL §6, §10): * ``hflip`` is **off by default** — hand seals are left/right sensitive, a mirror flip can turn a valid seal into an invalid / different one. * Normalization mean/std come from the timm pretrained config so the frozen EdgeNeXt encoder sees inputs in the distribution it was trained on. """ import os from collections import Counter from dataclasses import dataclass, field import numpy as np import torch from torch.utils.data import DataLoader, Subset, WeightedRandomSampler from torchvision import transforms from torchvision.datasets import ImageFolder @dataclass class AugConfig: """Augmentation hyper-parameters (a subset is searched by Optuna).""" img_size: int = 256 use_hflip: bool = False # OFF by default: seals are chirality-sensitive. rotation_deg: float = 15.0 color_jitter: float = 0.2 rrc_scale_min: float = 0.7 # RandomResizedCrop lower scale bound. randaug_magnitude: int = 0 # 0 disables RandAugment. mean: tuple[float, float, float] = (0.485, 0.456, 0.406) std: tuple[float, float, float] = (0.229, 0.224, 0.225) @dataclass class DataConfig: """Everything needed to materialize the dataloaders.""" data_root: str use_predefined_split: bool = False val_frac: float = 0.15 test_frac: float = 0.15 batch_size: int = 32 num_workers: int = 4 seed: int = 42 use_weighted_sampler: bool = False aug: AugConfig = field(default_factory=AugConfig) def build_transforms(aug: AugConfig, *, train: bool) -> transforms.Compose: """Build train/eval transforms for EdgeNeXt-style input. Args: aug: Augmentation configuration. train: If ``True`` return the stochastic train pipeline, else the deterministic eval pipeline (resize -> center-crop). Returns: A composed torchvision transform. """ if train: ops: list = [ transforms.RandomResizedCrop( aug.img_size, scale=(aug.rrc_scale_min, 1.0) ) ] if aug.randaug_magnitude > 0: ops.append(transforms.RandAugment(magnitude=aug.randaug_magnitude)) if aug.rotation_deg > 0: ops.append(transforms.RandomRotation(aug.rotation_deg)) if aug.color_jitter > 0: ops.append( transforms.ColorJitter( brightness=aug.color_jitter, contrast=aug.color_jitter, saturation=aug.color_jitter, ) ) if aug.use_hflip: ops.append(transforms.RandomHorizontalFlip()) ops += [transforms.ToTensor(), transforms.Normalize(aug.mean, aug.std)] return transforms.Compose(ops) resize = int(round(aug.img_size * 1.14)) # standard resize->center-crop ratio return transforms.Compose( [ transforms.Resize(resize), transforms.CenterCrop(aug.img_size), transforms.ToTensor(), transforms.Normalize(aug.mean, aug.std), ] ) def _stratified_indices( targets: list[int], val_frac: float, test_frac: float, seed: int ) -> tuple[list[int], list[int], list[int]]: """Return stratified train/val/test index lists (per-class shuffle+slice).""" rng = np.random.default_rng(seed) by_class: dict[int, list[int]] = {} for idx, y in enumerate(targets): by_class.setdefault(y, []).append(idx) train_idx, val_idx, test_idx = [], [], [] for _, idxs in sorted(by_class.items()): idxs = np.array(idxs) rng.shuffle(idxs) n = len(idxs) # Force >=1 only when the fraction is positive (frac=0 -> exactly 0). n_test = max(1, int(round(n * test_frac))) if test_frac > 0 else 0 n_val = max(1, int(round(n * val_frac))) if val_frac > 0 else 0 test_idx += idxs[:n_test].tolist() val_idx += idxs[n_test : n_test + n_val].tolist() train_idx += idxs[n_test + n_val :].tolist() return train_idx, val_idx, test_idx def _make_sampler(targets: list[int], num_classes: int) -> WeightedRandomSampler: """Inverse-frequency WeightedRandomSampler (one sample weight per element).""" counts = Counter(targets) class_w = {c: 1.0 / max(1, counts.get(c, 0)) for c in range(num_classes)} weights = torch.tensor([class_w[t] for t in targets], dtype=torch.double) return WeightedRandomSampler(weights, num_samples=len(weights), replacement=True) def compute_class_weights(targets: list[int], num_classes: int) -> torch.Tensor: """Normalized inverse-frequency weights for weighted cross-entropy. Args: targets: Per-sample integer labels of the *train* split. num_classes: Total number of classes. Returns: Float tensor of shape ``[num_classes]`` averaging to ~1.0. """ counts = Counter(targets) freq = np.array([counts.get(c, 0) for c in range(num_classes)], dtype=np.float64) freq = np.clip(freq, 1.0, None) w = freq.sum() / (num_classes * freq) return torch.tensor(w, dtype=torch.float32) @dataclass class DataBundle: """Container returned by :func:`build_dataloaders`.""" train_loader: DataLoader val_loader: DataLoader test_loader: DataLoader class_names: list[str] train_targets: list[int] num_classes: int def build_dataloaders(cfg: DataConfig) -> DataBundle: """Build train/val/test dataloaders with reproducible stratified split. Args: cfg: Data configuration (paths, split fractions, batch size, sampler). Returns: A :class:`DataBundle` with loaders, class names and train targets. """ train_tf = build_transforms(cfg.aug, train=True) eval_tf = build_transforms(cfg.aug, train=False) if cfg.use_predefined_split: # The Kaggle Naruto Sign download ships train/ and test/ but NO val/. # Never alias val = test (that leaks test into model selection): if val/ is # absent, carve a stratified val OUT of train (seeded). test_ds = ImageFolder(os.path.join(cfg.data_root, "test"), eval_tf) class_names = test_ds.classes val_dir = os.path.join(cfg.data_root, "val") if os.path.isdir(val_dir): train_ds = ImageFolder(os.path.join(cfg.data_root, "train"), train_tf) val_ds = ImageFolder(val_dir, eval_tf) train_targets = list(train_ds.targets) else: base_train = ImageFolder(os.path.join(cfg.data_root, "train"), train_tf) base_eval = ImageFolder(os.path.join(cfg.data_root, "train"), eval_tf) val_frac = cfg.val_frac / (1.0 - cfg.test_frac) # val share of train tr, va, _ = _stratified_indices(base_train.targets, val_frac, 0.0, cfg.seed) train_ds = Subset(base_train, tr) val_ds = Subset(base_eval, va) train_targets = [base_train.targets[i] for i in tr] else: base_train = ImageFolder(cfg.data_root, train_tf) base_eval = ImageFolder(cfg.data_root, eval_tf) class_names = base_train.classes tr, va, te = _stratified_indices( base_train.targets, cfg.val_frac, cfg.test_frac, cfg.seed ) train_ds = Subset(base_train, tr) val_ds = Subset(base_eval, va) test_ds = Subset(base_eval, te) train_targets = [base_train.targets[i] for i in tr] num_classes = len(class_names) if cfg.use_weighted_sampler: sampler = _make_sampler(train_targets, num_classes) train_loader = DataLoader( train_ds, batch_size=cfg.batch_size, sampler=sampler, num_workers=cfg.num_workers, pin_memory=True, drop_last=False, ) else: train_loader = DataLoader( train_ds, batch_size=cfg.batch_size, shuffle=True, num_workers=cfg.num_workers, pin_memory=True, drop_last=False, ) val_loader = DataLoader( val_ds, batch_size=cfg.batch_size, shuffle=False, num_workers=cfg.num_workers ) test_loader = DataLoader( test_ds, batch_size=cfg.batch_size, shuffle=False, num_workers=cfg.num_workers ) return DataBundle( train_loader=train_loader, val_loader=val_loader, test_loader=test_loader, class_names=class_names, train_targets=train_targets, num_classes=num_classes, )