248 lines
8.8 KiB
Python
248 lines
8.8 KiB
Python
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,
|
|
)
|