Added code

code/src/data.py

@@ -0,0 +1,247 @@
+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,
+    )