moved_repo_first_update

classes/SE_attention.py

@@ -0,0 +1,123 @@
+# se_block.py
+import torch
+import torch.nn as nn
+
+class SEGateLogger:
+    """
+    Lightweight stats over SE gates.
+    Use: logger.accumulate(gates) each batch; logger.get() at epoch end.
+    """
+    def __init__(self, enabled: bool = True, track_channels: bool = False, dim: int | None = None):
+        self.enabled = enabled
+        self.track_channels = track_channels
+        self.dim = dim
+        self.reset()
+
+    def reset(self):
+        self._n = 0
+        self._sum = 0.0
+        self._sum2 = 0.0
+        self._lt02 = 0
+        self._gt08 = 0
+        # optional per-channel
+        self._ch_sum = None
+        self._ch_count = 0
+        if self.track_channels and self.dim is not None:
+            self._ch_sum = torch.zeros(self.dim, dtype=torch.float32)
+
+    @torch.no_grad()
+    def accumulate(self, gates: torch.Tensor):
+        if not self.enabled:
+            return
+        # gates expected shape [N, C]; if a map/sequence gate is passed, reduce to [N, C]
+        if gates.dim() == 4:   # [N,C,H,W] gates (uncommon)
+            g = gates.mean(dim=(2,3))
+        elif gates.dim() == 3: # [N,T,C] gates (sequence)
+            g = gates.mean(dim=1)
+        elif gates.dim() == 2: # [N,C]
+            g = gates
+        else:
+            g = gates.view(gates.size(0), -1)
+
+        g = g.detach()
+        self._n    += g.numel()
+        self._sum  += g.sum().item()
+        self._sum2 += (g*g).sum().item()
+        self._lt02 += (g < 0.2).sum().item()
+        self._gt08 += (g > 0.8).sum().item()
+
+        if self._ch_sum is not None:
+            self._ch_sum += g.sum(dim=0).cpu()
+            self._ch_count += g.size(0)
+
+    def get(self, reset: bool = True):
+        if self._n == 0:
+            return None
+        mean = self._sum / self._n
+        var  = max(0.0, self._sum2 / self._n - mean * mean)
+        out = {
+            "mean": mean,
+            "std": var ** 0.5,
+            "pct_lt_0.2": self._lt02 / self._n,
+            "pct_gt_0.8": self._gt08 / self._n,
+        }
+        if self._ch_sum is not None and self._ch_count > 0:
+            out["channel_mean"] = (self._ch_sum / float(self._ch_count)).tolist()
+        if reset:
+            self.reset()
+        return out
+
+class SEBlock(nn.Module):
+    """
+    SE-style channel gating that works for vectors and maps.
+
+    Input:
+      - [N, C]         (vector)      -> squeeze = identity
+      - [N, C, H, W]   (image map)   -> squeeze over H,W
+      - [N, T, C]      (sequence)    -> squeeze over T
+
+    Gate modes:
+      - residual (default): gate = 1 + tanh(MLP(s)) in (0, 2)  [identity at init]
+      - plain:              gate = sigmoid(MLP(s)) in (0, 1)
+    """
+    def __init__(self, dim: int, reduction: int = 16, residual: bool = True, identity_init: bool = True):
+        super().__init__()
+        hid = max(1, dim // max(1, reduction))
+        self.fc1 = nn.Linear(dim, hid, bias=True)
+        self.act = nn.ReLU(inplace=True)
+        self.fc2 = nn.Linear(hid, dim, bias=True)
+        self.residual = residual
+
+        if residual and identity_init:
+            # make MLP output ~0 at start → gate ≈ 1.0
+            nn.init.zeros_(self.fc2.weight)
+            nn.init.zeros_(self.fc2.bias)
+
+    def _squeeze(self, x: torch.Tensor) -> torch.Tensor:
+        if x.dim() == 2:      # [N,C]
+            return x
+        if x.dim() == 4:      # [N,C,H,W]
+            return x.mean(dim=(2,3))
+        if x.dim() == 3:      # [N,T,C]
+            return x.mean(dim=1)
+        # fallback: flatten non-batch dims into channels
+        return x.view(x.size(0), -1)
+
+    def _broadcast(self, gate: torch.Tensor, like: torch.Tensor) -> torch.Tensor:
+        if like.dim() == 2:
+            return gate
+        if like.dim() == 3:
+            return gate.unsqueeze(1)            # [N,1,C]
+        if like.dim() == 4:
+            return gate.unsqueeze(-1).unsqueeze(-1)  # [N,C,1,1]
+        return gate.view_as(like)
+
+    def forward(self, x: torch.Tensor):
+        s = self._squeeze(x)                   # [N,C]
+        u = self.fc2(self.act(self.fc1(s)))    # [N,C]
+        if self.residual:
+            gate = 1.0 + torch.tanh(u)         # (0, 2) with identity at 1.0
+        else:
+            gate = torch.sigmoid(u)            # (0, 1)
+        y = x * self._broadcast(gate, x)
+        return y, gate  # return both the reweighted tensor and the gate for logging

classes/__init__.py

@@ -0,0 +1,19 @@
+from .clinical_data import ClinicalData
+from .dataset import ClinicalDataset
+from .image_tower import ImageTower
+from .md_tower import MDTower
+from .bridge import Bridge, VoteBridge
+# from .hypertower import HyperTower
+from .backbones import list_names, BackboneSpec, BACKBONES
+from .papila_builders import build_papila_clinical
+from .SE_attention import SEBlock, SEGateLogger
+from .early_stop import EarlyStopper
+__all__ = [
+    "ClinicalData",
+    "ClinicalDataset",
+    "ImageTower",
+    "MDTower",
+    "Bridge",
+    "VoteBridge",
+#    "HyperTower",
+]

classes/backbones.py

@@ -0,0 +1,178 @@
+# classes/backbones.py
+from __future__ import annotations
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Callable, Dict, List
+
+import torch
+from torch import nn
+from torchvision import models
+
+@dataclass(frozen=True)
+class BackboneSpec:
+    ctor: Callable                  # torchvision constructor
+    weights_default: object         # torchvision Weights enum DEFAULT member
+    strip: Callable[[nn.Module], tuple]   # fn(model)->(out_dim, model_no_head)
+    blocks: Callable[[nn.Module], List[nn.Module]]  # fn(model)->ordered blocks for freezing
+
+REFUGELIKE_BACKBONE_PATH = Path("models/refuge/classifier/refugelike_backbone.pt")
+REFUGE_DENSENET_PATH = Path("models/refuge/classifier/refuge_densenet_backbone.pt")
+REFUGE_EFFICIENT_B0_PATH = Path("models/refuge/classifier/refuge_efficient_b0_backbone.pt")
+REFUGE_EFFICIENT_B7_PATH = Path("models/refuge/classifier/refuge_efficient_b7_backbone.pt")
+
+# --- strip fns ---
+def _strip_efficientnet_b0(m: models.EfficientNet):
+    from torch import nn as _nn
+    out_dim = m.classifier[1].in_features
+    m.classifier = _nn.Identity()
+    return out_dim, m
+
+def _strip_resnet(m: models.ResNet):
+    out_dim = m.fc.in_features
+    m.fc = nn.Identity()
+    return out_dim, m
+
+def _strip_densenet(m: models.DenseNet):
+    out_dim = m.classifier.in_features
+    m.classifier = nn.Identity()
+    return out_dim, m
+
+def _strip_vgg(m: models.VGG):
+    out_dim = m.classifier[0].in_features   # 25088 for VGG16 at 224×224
+    m.classifier = nn.Identity()
+    return out_dim, m
+
+def _strip_mobilenet_v2(m: models.MobileNetV2):
+    out_dim = m.classifier[1].in_features
+    m.classifier = nn.Identity()
+    return out_dim, m
+
+def _strip_inception_v3(m: models.Inception3):
+    out_dim = m.fc.in_features
+    m.fc = nn.Identity()
+    if hasattr(m, "AuxLogits"):
+        m.aux_logits = False
+    return out_dim, m
+
+# --- block splitters for ratio-based freezing ---
+def _blocks_efficientnet_b0(m: models.EfficientNet):
+    return list(m.features)
+
+def _blocks_resnet(m: models.ResNet):
+    stem = nn.Sequential(m.conv1, m.bn1, m.relu, m.maxpool)
+    return [stem, m.layer1, m.layer2, m.layer3, m.layer4]
+
+def _blocks_densenet(m: models.DenseNet):
+    f = m.features
+    stem = nn.Sequential(f.conv0, f.norm0, f.relu0, f.pool0)
+    return [stem, f.denseblock1, f.transition1, f.denseblock2, f.transition2,
+            f.denseblock3, f.transition3, f.denseblock4, f.norm5]
+
+def _blocks_vgg(m: models.VGG):
+    stages, cur = [], []
+    for mod in m.features:
+        cur.append(mod)
+        if isinstance(mod, nn.MaxPool2d):
+            stages.append(nn.Sequential(*cur)); cur = []
+    if cur: stages.append(nn.Sequential(*cur))
+    return stages
+
+def _blocks_mobilenet_v2(m: models.MobileNetV2):
+    return list(m.features)
+
+def _blocks_inception_v3(m: models.Inception3):
+    blocks = []
+    for name, child in m.named_children():
+        if name in ("fc", "AuxLogits"):
+            continue
+        blocks.append(child)
+    return blocks
+
+# --- registry (covers paper models available in torchvision) ---
+BACKBONES: Dict[str, BackboneSpec] = {
+    "efficientnet_b0": BackboneSpec(
+        ctor=models.efficientnet_b0,
+        weights_default=models.EfficientNet_B0_Weights.DEFAULT,
+        strip=_strip_efficientnet_b0,
+        blocks=_blocks_efficientnet_b0,
+    ),
+    "resnet50": BackboneSpec(
+        ctor=models.resnet50,
+        weights_default=models.ResNet50_Weights.DEFAULT,
+        strip=_strip_resnet,
+        blocks=_blocks_resnet,
+    ),
+    "densenet121": BackboneSpec(
+        ctor=models.densenet121,
+        weights_default=models.DenseNet121_Weights.DEFAULT,
+        strip=_strip_densenet,
+        blocks=_blocks_densenet,
+    ),
+    "vgg16": BackboneSpec(
+        ctor=models.vgg16,
+        weights_default=models.VGG16_Weights.DEFAULT,
+        strip=_strip_vgg,
+        blocks=_blocks_vgg,
+    ),
+    "mobilenet_v2": BackboneSpec(
+        ctor=models.mobilenet_v2,
+        weights_default=models.MobileNet_V2_Weights.DEFAULT,
+        strip=_strip_mobilenet_v2,
+        blocks=_blocks_mobilenet_v2,
+    ),
+    "inception_v3": BackboneSpec(
+        ctor=models.inception_v3,
+        weights_default=models.Inception_V3_Weights.DEFAULT,
+        strip=_strip_inception_v3,
+        blocks=_blocks_inception_v3,
+    ),
+    "refugelike": BackboneSpec(
+        ctor=models.resnet50,
+        weights_default=None,
+        strip=_strip_resnet,
+        blocks=_blocks_resnet,
+    ),
+    "refuge_densenet": BackboneSpec(
+        ctor=models.densenet121,
+        weights_default=None,
+        strip=_strip_densenet,
+        blocks=_blocks_densenet,
+    ),
+    "refuge_efficient_b0": BackboneSpec(
+        ctor=models.efficientnet_b0,
+        weights_default=None,
+        strip=_strip_efficientnet_b0,
+        blocks=_blocks_efficientnet_b0,
+    ),
+    "refuge_efficient_b7": BackboneSpec(
+        ctor=models.efficientnet_b7,
+        weights_default=None,
+        strip=_strip_efficientnet_b0,
+        blocks=_blocks_efficientnet_b0,
+    ),
+    # Xception isn’t in torchvision 
+}
+
+def list_names() -> List[str]:
+    return list(BACKBONES.keys())
+    
+
+def load_backbone_weights(key: str, model: nn.Module) -> None:
+    if key == "refugelike":
+        path = REFUGELIKE_BACKBONE_PATH
+    elif key == "refuge_densenet":
+        path = REFUGE_DENSENET_PATH
+    elif key == "refuge_efficient_b0":
+        path = REFUGE_EFFICIENT_B0_PATH
+    elif key == "refuge_efficient_b7":
+        path = REFUGE_EFFICIENT_B7_PATH
+    else:
+        return
+
+    if not path.exists():
+        raise FileNotFoundError(
+            "Custom REFUGE backbone not found at "
+            f"{path}. Export it via refuge_build.py --export-backbone first."
+        )
+    state = torch.load(path, map_location="cpu")
+    model.load_state_dict(state, strict=False)

classes/bridge.py

@@ -0,0 +1,105 @@
+# bridge.py
+import torch
+import torch.nn as nn
+from classes.SE_attention import SEBlock, SEGateLogger
+
+# class SEBlock(nn.Module):
+#     def __init__(self, dim: int, reduction: int = 16):
+#         super().__init__()
+#         hidden = max(1, dim // max(1, reduction))
+#         self.net = nn.Sequential(
+#             nn.Linear(dim, hidden, bias=True),
+#             nn.ReLU(inplace=True),
+#             nn.Linear(hidden, dim, bias=True),
+#             nn.Sigmoid(),
+#         )
+
+#     def forward(self, x):
+#         return self.net(x)
+
+class Bridge(nn.Module):
+    def __init__(
+            self, 
+            img_dim, 
+            meta_dim, 
+            num_classes, 
+            fusion_dim=256, 
+            mode="fused",
+            use_se: bool = True,
+            se_reduction: int = 16,
+            se_pre_norm: bool = True,
+        ):
+
+        super().__init__()
+        self.mode = mode
+        self.use_se = use_se
+        # self.se_reduction = se_reduction
+        # self.se_pre_norm = se_pre_norm
+
+        #project towers to equal width
+        self.W_img = nn.Linear(img_dim, fusion_dim)
+        self.W_md = nn.Linear(meta_dim, fusion_dim)
+        
+        #(optional) : set layernorm for se so one tower doesn't dominate the other
+        self.ln_img = nn.LayerNorm(fusion_dim) if se_pre_norm else nn.Identity()
+        self.ln_md  = nn.LayerNorm(fusion_dim) if se_pre_norm else nn.Identity()
+        
+        #SE gate on the fused vector
+        self.se = SEBlock(fusion_dim, reduction=se_reduction, residual=True) if use_se else None
+        self.se_log = SEGateLogger(enabled=use_se, track_channels=False, dim=fusion_dim)
+
+        
+        #heads
+        self.classifier_fused = nn.Sequential(
+            nn.ReLU(), nn.Dropout(0.5), nn.Linear(fusion_dim, num_classes)
+        )
+        self.classifier_img = nn.Linear(img_dim, num_classes)
+        self.classifier_md = nn.Linear(meta_dim, num_classes)
+    def reset_se_stats(self):
+        """Call at epoch start."""
+        if getattr(self, "se_log", None):
+            self.se_log.reset()
+
+    def get_se_stats(self, reset: bool = True):
+        """Call after eval. Returns dict or None."""
+        if getattr(self, "se_log", None) and self.se_log.enabled:
+            return self.se_log.get(reset=reset)
+        return None
+    
+    def forward(self, img_feats, md_feats):
+        out_img = None if self.mode == "metadata_only" else self.classifier_img(img_feats)
+        out_md = None if self.mode == "image_only" else self.classifier_md(md_feats)
+ 
+        if self.mode == "fused":
+            hi = self.ln_img(self.W_img(img_feats)) #image features
+            hm = self.ln_md(self.W_md(md_feats)) #metadata features
+            fused = hi * hm  #elementwise product
+            #apply SE gates
+            if self.se is not None:
+                fused, gates = self.se(fused)
+                if self.se_log.enabled:
+                    self.se_log.accumulate(gates)     
+            
+            if self.se is not None and self.training and self.se_log.enabled:
+                if not hasattr(self, "_dbg_seen"):
+                    self._dbg_seen = 0
+                if self._dbg_seen < 3:                     # print only a few times
+                    print("[SE] gate mean this batch:", gates.mean().item())
+                    self._dbg_seen += 1
+            out_f = self.classifier_fused(fused)
+            return out_f, out_img, out_md
+        # if ablation modes:
+        if self.mode == "image_only":
+            return out_img, out_img, None
+        if self.mode == "metadata_only":
+            return out_md, None, out_md
+
+
+class VoteBridge(nn.Module):
+    def __init__(self, num_classes):
+        super().__init__()
+        self.vote_combiner = nn.Linear(num_classes * 2, num_classes)  # two sets of logits
+
+    def forward(self, out_img, out_md):
+        votes = torch.cat([out_img, out_md], dim=1)
+        return self.vote_combiner(votes)

classes/clinical_data.py

@@ -0,0 +1,264 @@
+# clinical_data.py
+from __future__ import annotations
+from pathlib import Path
+from typing import Iterable, Optional, Dict, List, Tuple
+import numpy as np
+import pandas as pd
+
+class ClinicalData:
+    """
+    Torch-free container for clinical metadata and file/label bookkeeping.
+    - Holds one or more dataframes (via add_df) and harmonizes columns
+    - Canonical IDs: 'Patient ID' must exist (or be specified and will be renamed)
+    - Canonical eye column: 'eyeID' recoded to 'OS'/'OD' if present; if absent, set to 0
+    - Feature typing (if cat_cols not provided):
+        * Categorical if (a) <= max_unique categorical threshold (default 4), or
+          (b) values cannot be coerced to float; otherwise numeric (scalar)
+    - Scaling/imputation:
+        * Numeric: min–max to [0,1], median imputation; + one missing flag per numeric feature
+        * Categorical: one-hot with '<UNK>' bucket at index 0
+    - Patient-level K-fold indices stored as dict: folds[k] -> {'train_ids': [...], 'test_ids': [...]}
+    """
+
+    def __init__(
+        self,
+        image_dir: str,
+        clinical_dir: Optional[str],
+        label_col: str,
+        # typing / detection
+        cat_cols: Optional[Iterable[str]] = None,
+        max_unique_for_cat: int = 4,
+        # splitting
+        n_splits: int = 5,
+        random_seed: int = 42,
+    ):
+        self.image_dir = Path(image_dir)
+        self.clinical_dir = Path(clinical_dir) if clinical_dir else None
+        self.label_col = label_col
+        self.max_unique_for_cat = max_unique_for_cat
+        self.n_splits = n_splits
+
+        # Internal state
+        self.frames: List[pd.DataFrame] = []         # raw frames as added
+        self.df: pd.DataFrame = pd.DataFrame()        # concatenated
+        self.scalar_cols: List[str] = []
+        self.cat_cols: List[str] = list(cat_cols) if cat_cols is not None else []
+        self.scalar_stats: Dict[str, Dict[str, float]] = {}
+        self.cat_maps: Dict[str, Dict[object, int]] = {}
+        self.feature_dim: int = 0
+        self.folds: Dict[int, Dict[str, List[object]]] = {}  # fold -> {'train_ids': [], 'test_ids': []}
+        self.random_seed = int(random_seed)
+
+    # ------------------- Public API -------------------
+    def add_df(
+        self,
+        df: pd.DataFrame,
+        id_column: Optional[str] = None,
+        eye_column: Optional[str] = None,
+        exclude_cols: Optional[Iterable[str]] = None,
+    ) -> None:
+        """
+        Add a dataframe and re-run harmonization, typing, stats, and K-fold indices.
+        QC rules:
+          - Must have patient ID column; if not provided under that name, specify id_column.
+          - eyeID, if present, must be binary; recoded to 'OS'/'OD'. If absent, create and set to 0.
+        """
+        df = df.copy()
+        # --- QC: Patient ID ---
+        pid_col = self._ensure_patient_id(df, id_column)
+        # --- QC: eyeID ---
+        self._canonicalize_eye_column(df, eye_column)
+        # --- Normalize label presence ---
+        if self.label_col not in df.columns:
+            raise ValueError(f"label_col '{self.label_col}' not found in added dataframe")
+
+        # append & refresh
+        self.frames.append(df)
+        self._refresh_master_df(exclude_cols=exclude_cols)
+        self._infer_or_validate_feature_types(exclude_cols=exclude_cols)
+        self._compute_numeric_stats()
+        self._build_cat_maps()
+        self._compute_feature_dim()
+        self._build_kfold_indices()
+
+    def get_split_ids(self, fold: int) -> Tuple[List[object], List[object]]:
+        rec = self.folds.get(fold)
+        if not rec: raise KeyError(f"Fold {fold} not available. Built folds: {sorted(self.folds.keys())}")
+        return rec['train_ids'], rec['test_ids']
+
+    def get_split_dfs(self, fold: int) -> Tuple[pd.DataFrame, pd.DataFrame]:
+        train_ids, test_ids = self.get_split_ids(fold)
+        train_df = self.df[self.df['Patient ID'].isin(train_ids)].reset_index(drop=True)
+        test_df  = self.df[self.df['Patient ID'].isin(test_ids)].reset_index(drop=True)
+        return train_df, test_df
+
+    def vectorize_row(self, row: pd.Series) -> np.ndarray:
+        """Return a numpy feature vector (torch-free)."""
+        feats: List[float] = []
+        miss: List[float] = []
+        # numeric
+        for col in self.scalar_cols:
+            v = pd.to_numeric(row.get(col), errors='coerce')
+            if pd.isna(v):
+                miss.append(1.0)
+                v = self.scalar_stats[col]['median']
+            else:
+                miss.append(0.0)
+            lo = self.scalar_stats[col]['min']; hi = self.scalar_stats[col]['max']
+            feats.append((float(v) - lo) / (hi - lo) if hi > lo else 0.0)
+        # categorical
+        for col in self.cat_cols:
+            mapping = self.cat_maps[col]
+            one = [0.0] * len(mapping)
+            key = row.get(col)
+            one[mapping.get(key, 0)] = 1.0  # 0 is <UNK>
+            feats.extend(one)
+        # numeric missing flags
+        feats.extend(miss)
+        return np.asarray(feats, dtype=np.float32)
+
+    def get_image_path(self, row: pd.Series, filename_template: str = "RET{pid:03d}{eye}.jpg") -> Path:
+        pid = int(row['Patient ID']); eye = row.get('eyeID', 0)
+        if eye in ("OS", "OD"):
+            eye_str = eye
+        else:
+            eye_str = str(eye)
+        return self.image_dir / filename_template.format(pid=pid, eye=eye_str)
+
+    # ------------------- Internal helpers -------------------
+    def _ensure_patient_id(self, df: pd.DataFrame, id_column: Optional[str]) -> str:
+        if 'Patient ID' in df.columns:
+            return 'Patient ID'
+        if id_column and id_column in df.columns:
+            df.rename(columns={id_column: 'Patient ID'}, inplace=True)
+            return 'Patient ID'
+        # try auto-detect common variants
+        candidates = [c for c in df.columns if c.lower().replace(" ", "") in {"patientid","patient","pid"}]
+        if len(candidates) == 1:
+            df.rename(columns={candidates[0]: 'Patient ID'}, inplace=True)
+            return 'Patient ID'
+        raise ValueError("A 'Patient ID' column is required; provide id_column=... if it has a different name.")
+
+    def _canonicalize_eye_column(self, df: pd.DataFrame, eye_column: Optional[str]) -> None:
+        # Find source
+        src = None
+        if 'eyeID' in df.columns: src = 'eyeID'
+        elif eye_column and eye_column in df.columns: src = eye_column
+        else:
+            # try auto detect
+            for c in df.columns:
+                if 'eye' in c.lower():
+                    src = c; break
+        if src is None:
+            df['eyeID'] = 0
+            return
+        # Map to OS/OD
+        s = df[src]
+        def norm(v):
+            if pd.isna(v): return None
+            x = str(v).strip().upper()
+            if x in {"OS","L","LEFT","0"}: return "OS"
+            if x in {"OD","R","RIGHT","1"}: return "OD"
+            # numbers like 2? fall back by parity
+            try:
+                num = int(float(x))
+                return "OD" if num % 2 == 1 else "OS"
+            except Exception:
+                return None
+        mapped = s.map(norm)
+        uniq = {u for u in mapped.dropna().unique().tolist()}
+        if not uniq.issubset({"OS","OD"}):
+            raise ValueError(f"eyeID must be binary; found values {sorted(uniq)}")
+        df['eyeID'] = mapped.fillna("OS")
+        if src != 'eyeID':
+            # keep original too if you want, but we standardize on 'eyeID'
+            pass
+
+    def _refresh_master_df(self, exclude_cols: Optional[Iterable[str]] = None) -> None:
+        self.df = pd.concat(self.frames, axis=0, ignore_index=True)
+        # drop columns explicitly excluded
+        if exclude_cols:
+            self.df = self.df.drop(columns=[c for c in exclude_cols if c in self.df.columns])
+
+    def _infer_or_validate_feature_types(self, exclude_cols: Optional[Iterable[str]] = None) -> None:
+        excluded = set(exclude_cols or []) | {self.label_col, 'Patient ID'}
+        # we keep canonical 'eyeID' as categorical if present
+        feature_candidates = [c for c in self.df.columns if c not in excluded]
+        # If user pre-specified cat_cols in __init__, respect them and infer the rest
+        cats = set(self.cat_cols) if self.cat_cols else set()
+        scalars = set()
+        for c in feature_candidates:
+            if c == 'eyeID':
+                cats.add('eyeID'); continue
+            if c in cats: continue
+            s = self.df[c]
+            # try numeric coercion
+            as_num = pd.to_numeric(s, errors='coerce')
+            num_missing = as_num.isna().mean()
+            num_unique = s.dropna().nunique()
+            if as_num.notna().any() and num_missing < 1.0 and num_unique > self.max_unique_for_cat:
+                scalars.add(c)
+            else:
+                # categorical if few uniques OR non-numeric
+                if num_unique <= self.max_unique_for_cat or as_num.isna().mean() > 0.0:
+                    cats.add(c)
+                else:
+                    scalars.add(c)
+        self.cat_cols = sorted(cats)
+        self.scalar_cols = sorted(scalars)
+
+    def _compute_numeric_stats(self) -> None:
+        self.scalar_stats.clear()
+        for col in self.scalar_cols:
+            s = pd.to_numeric(self.df[col], errors='coerce')
+            vals = s.dropna().astype(float).values
+            if vals.size == 0:
+                lo, hi, med = 0.0, 1.0, 0.0
+            else:
+                lo, hi = float(np.min(vals)), float(np.max(vals))
+                med = float(np.median(vals))
+                if hi <= lo: hi = lo + 1.0
+            self.scalar_stats[col] = {"min": lo, "max": hi, "median": med}
+
+    def _build_cat_maps(self) -> None:
+        self.cat_maps.clear()
+        for col in self.cat_cols:
+            cats = [v for v in self.df[col].dropna().unique().tolist()]
+            try: cats = sorted(cats)
+            except Exception: pass
+            mapping = {"<UNK>": 0}
+            for i, v in enumerate(cats, start=1): mapping[v] = i
+            self.cat_maps[col] = mapping
+
+    def _compute_feature_dim(self) -> None:
+        self.feature_dim = len(self.scalar_cols) + sum(len(m) for m in self.cat_maps.values()) + len(self.scalar_cols)
+
+    # ------------------- K-fold on unique patients -------------------
+    def _build_kfold_indices(self) -> None:
+        # unique patients and a per-patient label for stratification if possible
+        pats = self.df['Patient ID'].unique().tolist()
+        # Derive a patient label as the mode of their rows (fallback to first valid)
+        labels_by_pat = {}
+        for pid, grp in self.df.groupby('Patient ID'):
+            lab = grp[self.label_col].dropna()
+            if len(lab) == 0:
+                labels_by_pat[pid] = 0
+            else:
+                labels_by_pat[pid] = lab.mode().iloc[0]
+        y_pat = np.array([labels_by_pat[p] for p in pats])
+
+        # Try to use StratifiedGroupKFold if available, else fall back to StratifiedKFold on patient labels
+        try:
+            from sklearn.model_selection import StratifiedGroupKFold
+            sgkf = StratifiedGroupKFold(n_splits=self.n_splits, shuffle=True, random_state=self.random_seed)
+            split_iter = sgkf.split(X=pats, y=y_pat, groups=pats)
+        except Exception:
+            from sklearn.model_selection import StratifiedKFold
+            skf = StratifiedKFold(n_splits=self.n_splits, shuffle=True, random_state=self.random_seed)
+            split_iter = skf.split(X=np.zeros(len(pats)), y=y_pat)
+
+        self.folds.clear()
+        for i, (train_idx, test_idx) in enumerate(split_iter):
+            train_ids = [pats[j] for j in train_idx]
+            test_ids  = [pats[j] for j in test_idx]
+            self.folds[i] = {"train_ids": train_ids, "test_ids": test_ids}

classes/dataset.py

@@ -0,0 +1,58 @@
+# dataset.py
+from torch.utils.data import Dataset
+from PIL import Image
+import numpy as np
+import torch
+
+
+class ClinicalDataset(Dataset):
+    """Generic dataset wrapping a ClinicalData instance.
+    Returns (img_tensor, meta_tensor, label)."""
+
+    def __init__(
+        self,
+        clinical_data,
+        img_transform,
+        meta_transform=None,
+        image_preprocessor=None,
+        geometry_provider=None,
+        geometry_dim: int = 0,
+    ):
+        self.clinical = clinical_data
+        self.transform_image = img_transform
+        self.meta_transform = meta_transform or (lambda x: x)
+        self.image_preprocessor = image_preprocessor
+        self.geometry_provider = geometry_provider
+        self.geometry_dim = geometry_dim if geometry_provider is not None else 0
+
+    def __len__(self):
+        return len(self.clinical.df)
+
+    def __getitem__(self, idx: int):
+        row = self.clinical.df.iloc[idx]
+        # load & transform image
+        img_path = self.clinical.get_image_path(row)
+        orig_img = Image.open(img_path).convert("RGB")
+        img = orig_img
+        if self.image_preprocessor is not None:
+            img = self.image_preprocessor(img, img_path)
+        img_t = self.transform_image(img)
+        # encode & transform metadata
+        meta = self.clinical.encode_metadata(row)
+        meta_t = self.meta_transform(meta)
+        # label
+        label = self.clinical.get_label(row)
+        if self.geometry_dim > 0:
+            features = None
+            if self.geometry_provider is not None and hasattr(self.geometry_provider, "geometry_features"):
+                features = self.geometry_provider.geometry_features(orig_img, img_path)
+            if features is None:
+                geom_vec = torch.zeros(self.geometry_dim, dtype=torch.float32)
+            else:
+                features = np.asarray(features, dtype=np.float32)
+                if features.shape[0] != self.geometry_dim:
+                    geom_vec = torch.zeros(self.geometry_dim, dtype=torch.float32)
+                else:
+                    geom_vec = torch.from_numpy(features)
+            return img_t, meta_t, geom_vec, label
+        return img_t, meta_t, label

classes/early_stop.py

@@ -0,0 +1,89 @@
+import math, copy, torch
+
+class EarlyStopper:
+    def __init__(self, monitor: str, mode: str = "auto",
+                 patience: int = 5, min_delta: float = 0.0,
+                 save_path: str | None = None, restore_best: bool = True):
+        """
+        monitor: key in your epoch row, e.g. 'eval_loss', 'auc_fused', 'acc_fused'
+        mode: 'max' (higher is better), 'min', or 'auto' (min for '*loss*', else max)
+        patience: epochs without improvement before stopping
+        min_delta: required improvement magnitude
+        save_path: optional .pth file to save best weights each time it improves
+        restore_best: if True, load best weights back at the end
+        """
+        self.monitor = monitor
+        if mode == "auto":
+            mode = "min" if "loss" in monitor.lower() else "max"
+        self.mode = mode
+        self.patience = int(patience)
+        self.min_delta = float(min_delta)
+        self.save_path = save_path
+        self.restore_best = restore_best
+
+        self.best = -math.inf if mode == "max" else math.inf
+        self.bad_epochs = 0
+        self.best_state = None
+        self.best_epoch = -1
+        self.last_improved = False
+
+    def _is_better(self, val):
+        if val is None or (isinstance(val, float) and math.isnan(val)):
+            return False
+        if self.mode == "max":
+            return val > (self.best + self.min_delta)
+        else:
+            return val < (self.best - self.min_delta)
+
+    def step(self, metrics: dict, trainer, epoch: int) -> bool:
+        val = metrics.get(self.monitor, None)
+        improved = self._is_better(val)
+        self.last_improved = improved
+
+        if improved:
+            self.best = val
+            self.best_epoch = epoch
+            self.bad_epochs = 0
+            # snapshot + optional save
+            state = {
+                "img_tower": trainer.img_tower.state_dict(),
+                "md_tower":  trainer.md_tower.state_dict(),
+                "optimizer": trainer.optimizer.state_dict(),
+            }
+            if hasattr(trainer, "bridge"):   state["bridge"]   = trainer.bridge.state_dict()
+            if hasattr(trainer, "head_img"): state["head_img"] = trainer.head_img.state_dict()
+            if hasattr(trainer, "head_md"):  state["head_md"]  = trainer.head_md.state_dict()
+            # keep an in-memory copy for restore(); file save is optional
+            self.best_state = copy.deepcopy(state)
+            if self.save_path: torch.save(state, self.save_path)
+            print(f"[early] ↑ new best {self.monitor}={val:.5f} at epoch {epoch+1}")
+        else:
+            self.bad_epochs += 1
+
+        stop = self.bad_epochs >= self.patience
+        if stop:
+            print(f"[early] stopping: no improvement in {self.patience} epochs "
+                  f"(best {self.monitor}={self.best:.5f} @ epoch {self.best_epoch+1})")
+        return stop
+
+    def restore(self, trainer):
+        if not self.restore_best:
+            return
+        # Prefer in-memory best state; otherwise try loading from save_path
+        st = self.best_state
+        if st is None and self.save_path:
+            try:
+                st = torch.load(self.save_path, map_location="cpu")
+            except Exception:
+                st = None
+        if st is None:
+            return
+        trainer.img_tower.load_state_dict(st["img_tower"])
+        trainer.md_tower.load_state_dict(st["md_tower"])
+        if "bridge" in st and hasattr(trainer, "bridge"):
+            trainer.bridge.load_state_dict(st["bridge"])
+        if "head_img" in st and hasattr(trainer, "head_img"):
+            trainer.head_img.load_state_dict(st["head_img"])
+        if "head_md" in st and hasattr(trainer, "head_md"):
+            trainer.head_md.load_state_dict(st["head_md"])
+        trainer.optimizer.load_state_dict(st["optimizer"])

classes/geometry_features.py

@@ -0,0 +1,87 @@
+"""Shared helpers for deriving disc/cup geometry features."""
+
+from __future__ import annotations
+
+from collections import Counter
+from typing import Tuple
+
+import numpy as np
+from PIL import Image
+
+EPS = 1e-6
+FEATURE_DIM = 5
+
+
+def disc_cup_from_mask_image(mask_img: Image.Image) -> Tuple[np.ndarray, np.ndarray]:
+    """Return binary disc/cup masks from a REFUGE-style annotation image."""
+    arr = np.asarray(mask_img)
+    if arr.ndim == 3:
+        h, w, c = arr.shape
+        border = np.concatenate(
+            [arr[0, :, :], arr[-1, :, :], arr[:, 0, :], arr[:, -1, :]],
+            axis=0,
+        )
+        border_counts = Counter(map(tuple, border))
+        bg_color = border_counts.most_common(1)[0][0]
+        flat = arr.reshape(-1, c)
+        colors = Counter(map(tuple, flat))
+        colors.pop(bg_color, None)
+        disc = (~np.all(arr == bg_color, axis=-1)).astype(np.uint8)
+        if colors:
+            cup_color = min(colors.keys(), key=lambda col: sum(col))
+            cup = np.all(arr == cup_color, axis=-1).astype(np.uint8)
+        else:
+            cup = np.zeros((h, w), dtype=np.uint8)
+    else:
+        border = np.concatenate([arr[0, :], arr[-1, :], arr[:, 0], arr[:, -1]])
+        counts = Counter(border.tolist())
+        bg_value = counts.most_common(1)[0][0]
+        disc = (arr != bg_value).astype(np.uint8)
+        fg = arr[arr != bg_value]
+        if fg.size > 0:
+            cup_value = int(np.min(fg))
+            cup = (arr == cup_value).astype(np.uint8)
+        else:
+            cup = np.zeros_like(arr, dtype=np.uint8)
+    cup = (cup > 0) & (disc > 0)
+    return disc.astype(np.uint8), cup.astype(np.uint8)
+
+
+def compute_geometry_features(disc_mask: np.ndarray, cup_mask: np.ndarray) -> np.ndarray:
+    """Compute cup/disc geometry descriptors (area, rim, diameter ratios, centre shift)."""
+    disc = (disc_mask > 0).astype(np.float32)
+    cup = (cup_mask > 0).astype(np.float32)
+
+    disc_area = disc.sum()
+    cup_area = cup.sum()
+    area_ratio = cup_area / (disc_area + EPS)
+    rim_ratio = (disc_area - cup_area) / (disc_area + EPS)
+
+    disc_rows = np.any(disc > 0, axis=1)
+    cup_rows = np.any(cup > 0, axis=1)
+    disc_cols = np.any(disc > 0, axis=0)
+    cup_cols = np.any(cup > 0, axis=0)
+
+    disc_height = float(disc_rows.sum())
+    cup_height = float(cup_rows.sum())
+    disc_width = float(disc_cols.sum())
+    cup_width = float(cup_cols.sum())
+
+    vertical_ratio = cup_height / (disc_height + EPS)
+    horizontal_ratio = cup_width / (disc_width + EPS)
+
+    def _centre(mask: np.ndarray) -> Tuple[float, float]:
+        coords = np.argwhere(mask > 0)
+        if coords.size == 0:
+            return 0.5, 0.5
+        ys, xs = coords[:, 0], coords[:, 1]
+        return float(xs.mean()) / mask.shape[1], float(ys.mean()) / mask.shape[0]
+
+    disc_cx, disc_cy = _centre(disc)
+    cup_cx, cup_cy = _centre(cup)
+    centre_shift = float(np.hypot(cup_cx - disc_cx, cup_cy - disc_cy))
+
+    return np.array(
+        [area_ratio, rim_ratio, vertical_ratio, horizontal_ratio, centre_shift],
+        dtype=np.float32,
+    )

classes/image_tower.py

@@ -0,0 +1,125 @@
+# classes/image_tower.py
+from __future__ import annotations
+import math
+from typing import Optional
+
+import torch
+from torch import nn
+from torchvision import transforms
+from classes.backbones import BACKBONES, list_names, load_backbone_weights
+from classes.SE_attention import SEBlock
+
+def build_backbone(name: str, freeze_ratio: float = 0.0, augment: bool = True):
+    """
+    Operational builder:
+      - instantiate with DEFAULT weights
+      - strip classifier → features
+      - apply ratio-based freezing over coarse blocks
+      - return (model, out_dim, transform)
+    """
+    key = (name or "").lower()
+    if key not in BACKBONES:
+        raise ValueError(f"Unsupported backbone '{name}'. Valid options: {list_names()}")
+
+    spec = BACKBONES[key]
+    m = spec.ctor(weights=spec.weights_default)
+    out_dim, m = spec.strip(m)
+    load_backbone_weights(key, m)
+
+    # transforms: use the weights’ mean/std, but keep your augmentation pipeline
+    mean = getattr(spec.weights_default, "meta", {}).get("mean", (0.485, 0.456, 0.406))
+    std  = getattr(spec.weights_default, "meta", {}).get("std",  (0.229, 0.224, 0.225))
+    crop = 299 if key == "inception_v3" else 224
+
+    if augment:
+        transform = transforms.Compose([
+            transforms.Resize(256),
+            transforms.CenterCrop(crop),
+            transforms.RandomHorizontalFlip(),
+            transforms.RandomVerticalFlip(),
+            transforms.RandomRotation(15),
+            transforms.ColorJitter(0.1, 0.1, 0.1, 0.05),
+            transforms.ToTensor(),
+            transforms.Normalize(mean=mean, std=std),
+        ])
+    else:
+        transform = transforms.Compose([
+            transforms.Resize(256),
+            transforms.CenterCrop(crop),
+            transforms.ToTensor(),
+            transforms.Normalize(mean=mean, std=std),
+        ])
+
+    # ratio-based freezing: freeze earliest floor(N * freeze_ratio) blocks
+    fr = max(0.0, min(1.0, float(freeze_ratio)))
+    blocks = spec.blocks(m)
+    n = len(blocks)
+    freeze_n = int(math.floor(n * fr))
+    for b in blocks[:freeze_n]:
+        for p in b.parameters():
+            p.requires_grad = False
+
+    return m, out_dim, transform
+
+class ImageTower(nn.Module):
+    """
+    Vision backbone → pooled features.
+    - backbone: one of list_names() (default 'efficientnet_b0')
+    - always DEFAULT torchvision weights
+    - freeze_ratio ∈ [0,1] freezes earliest floor(N*freeze_ratio) blocks
+    - returns [N, out_dim] features from backbone forward
+    """
+    def __init__(self, backbone: str = "efficientnet_b0", freeze_ratio: float = 0.0,
+                 use_se: bool = False, se_reduction: int = 16, se_pre_norm: bool = True,
+                 augment: bool = True, geometry_dim: int = 0):
+        super().__init__()
+        self.backbone, base_dim, self.transform = build_backbone(backbone, freeze_ratio, augment=augment)
+        self._name = backbone
+        # Keep ordered blocks for dynamic freezing/thawing
+        key = (self._name or "").lower()
+        self._spec = BACKBONES[key]
+        self._blocks = self._spec.blocks(self.backbone)
+        # Optional tower-level SE over the final feature vector
+        self.base_dim = base_dim
+        self.geometry_dim = max(0, int(geometry_dim))
+        self.out_dim = self.base_dim + self.geometry_dim
+        self.tower_ln = nn.LayerNorm(self.base_dim) if se_pre_norm else nn.Identity()
+        self.tower_se = SEBlock(self.base_dim, reduction=se_reduction, residual=True) if use_se else None
+
+    def forward(self, x: torch.Tensor, geometry: Optional[torch.Tensor] = None) -> torch.Tensor:
+        y = self.backbone(x)
+        # sanity: pooled features, not logits
+        assert y.dim() == 2 and y.size(1) == self.base_dim, \
+            f"Expected features [N,{self.base_dim}], got {tuple(y.shape)}"
+        if self.tower_se is not None:
+            y, _ = self.tower_se(self.tower_ln(y))
+        if self.geometry_dim > 0:
+            if geometry is None or geometry.numel() == 0:
+                geom = torch.zeros(y.size(0), self.geometry_dim, device=y.device, dtype=y.dtype)
+            else:
+                if geometry.dim() == 1:
+                    geom = geometry.unsqueeze(0)
+                else:
+                    geom = geometry
+                geom = geom.to(device=y.device, dtype=y.dtype)
+                if geom.size(0) != y.size(0):
+                    raise ValueError(f"Geometry batch size mismatch: {geom.size(0)} vs {y.size(0)}")
+                if geom.size(1) != self.geometry_dim:
+                    raise ValueError(f"Expected geometry dim {self.geometry_dim}, got {geom.size(1)}")
+            y = torch.cat([y, geom], dim=1)
+        return y
+
+    def set_freeze_ratio(self, ratio: float):
+        """Dynamically freeze earliest floor(N*ratio) backbone blocks.
+        ratio in [0,1]."""
+        r = max(0.0, min(1.0, float(ratio)))
+        n = len(self._blocks)
+        freeze_n = int(math.floor(n * r))
+        # Unfreeze all first
+        for b in self._blocks:
+            for p in b.parameters():
+                p.requires_grad = True
+        # Freeze earliest blocks
+        for b in self._blocks[:freeze_n]:
+            for p in b.parameters():
+                p.requires_grad = False

classes/md_tower.py

@@ -0,0 +1,54 @@
+# md_tower.py
+import torch
+import torch.nn as nn
+from classes import ClinicalData
+from classes.SE_attention import SEBlock
+
+class MDTower(nn.Module):
+    """MLP over ClinicalData.vectorize_row outputs (convert to torch inside tower)."""
+    def __init__(self, clinical_data: ClinicalData, hidden_dim: int = 128, dropout: float = 0.1,
+                 use_se: bool = False, se_reduction: int = 16, se_pre_norm: bool = True):
+        super().__init__()
+        self.feature_dim = clinical_data.feature_dim
+        self.out_dim = hidden_dim
+        # two-block MLP so we can optionally freeze/thaw per block
+        self.block0 = nn.Sequential(
+            nn.Linear(self.feature_dim, hidden_dim),
+            nn.LayerNorm(hidden_dim),
+            nn.ReLU(inplace=True),
+            nn.Dropout(dropout),
+        )
+        self.block1 = nn.Sequential(
+            nn.Linear(hidden_dim, hidden_dim),
+            nn.ReLU(inplace=True),
+        )
+        self.net = nn.Sequential(self.block0, self.block1)
+        self.tower_ln = nn.LayerNorm(hidden_dim) if se_pre_norm else nn.Identity()
+        self.tower_se = SEBlock(hidden_dim, reduction=se_reduction, residual=True) if use_se else None
+
+    def forward(self, meta_np_or_torch) -> torch.Tensor:
+        if isinstance(meta_np_or_torch, torch.Tensor):
+            x = meta_np_or_torch
+        else:
+            x = torch.as_tensor(meta_np_or_torch, dtype=torch.float32)
+        h = self.net(x)
+        if self.tower_se is not None:
+            h, _ = self.tower_se(self.tower_ln(h))
+        return h
+
+    def set_freeze_ratio(self, ratio: float):
+        """Optionally freeze earliest blocks of the MLP.
+        With two blocks, ratio≥0.5 freezes block0; ratio≥1.0 freezes both."""
+        r = max(0.0, min(1.0, float(ratio)))
+        # Unfreeze all
+        for p in self.block0.parameters():
+            p.requires_grad = True
+        for p in self.block1.parameters():
+            p.requires_grad = True
+        # Freeze earliest blocks based on ratio threshold
+        if r >= 0.5:
+            for p in self.block0.parameters():
+                p.requires_grad = False
+        if r >= 1.0:
+            for p in self.block1.parameters():
+                p.requires_grad = False

classes/papila_builders.py

@@ -0,0 +1,99 @@
+# papila_builders.py
+from typing import List, Dict
+import numpy as np
+import pandas as pd
+
+from classes import ClinicalData  # adjust import path if needed
+
+# ---- Pachymetry → IOP correction (per PAPILA Table 3) ----
+_PACHY_TABLE: Dict[int, int] = {
+    475:+5, 485:+4, 495:+4, 505:+3, 515:+2, 525:+1, 535:+1,
+    545: 0, 555:-1, 565:-1, 575:-2, 585:-3, 595:-4, 605:-4, 615:-5,
+}
+_PACHY_KEYS = np.array(sorted(_PACHY_TABLE.keys()))
+
+def _nearest_pachy_key(x: float) -> int:
+    idx = int(np.argmin(np.abs(_PACHY_KEYS - float(x))))
+    return int(_PACHY_KEYS[idx])
+
+def _pick_iop(row: pd.Series) -> float:
+    """Prefer Pneumatic, else Perkins; may return NaN."""
+    raw = row["Pneumatic"] if not pd.isna(row.get("Pneumatic", np.nan)) else row.get("Perkins", np.nan)
+    return float(raw) if not pd.isna(raw) else np.nan
+
+def _correct_iop(raw_iop: float, pachy: float) -> float:
+    """Return corrected IOP using nearest pachymetry bin; if pachy missing, return raw."""
+    if pd.isna(raw_iop):
+        return np.nan
+    if pd.isna(pachy):
+        return float(raw_iop)
+    key = _nearest_pachy_key(float(pachy))
+    return float(raw_iop) + float(_PACHY_TABLE[key])
+
+def _apply_iop_and_drop_md(df: pd.DataFrame) -> pd.DataFrame:
+    """Add IOP_raw/IOP_corr and drop VF_MD if present (in-place safe)."""
+    # IOP_raw
+    df["IOP_raw"] = df.apply(_pick_iop, axis=1)
+
+    # IOP_corr
+    pachy = df.get("Pachymetry", pd.Series(np.nan, index=df.index))
+    df["IOP_corr"] = [
+        _correct_iop(r, p) for r, p in zip(df["IOP_raw"].values, pachy.values)
+    ]
+
+    # Drop VF_MD if present
+    if "VF_MD" in df.columns:
+        df.drop(columns=["VF_MD"], inplace=True)
+    return df
+
+
+def build_papila_clinical(
+    image_dir: str,
+    clinical_dir: str,
+    label_col: str,
+    cat_cols: List[str],
+    n_splits: int = 5,
+    random_seed: int = 42,
+) -> ClinicalData:
+    """
+    Build ClinicalData exactly like the user's original build_clinical:
+      - add_df(OD), set eyeID='OD'
+      - add_df(OS), set eyeID='OS'
+      - normalize 'Patient ID' on frames
+    THEN:
+      - compute IOP_raw / IOP_corr on each frame
+      - drop VF_MD
+      - refresh master df + kfold indices
+    """
+    clinical = ClinicalData(
+        image_dir=image_dir,
+        clinical_dir=clinical_dir,
+        label_col=label_col,
+        cat_cols=cat_cols,
+        n_splits=n_splits,
+        random_seed=random_seed,
+    )
+
+    # --- Load exactly like original build_clinical ---
+    clinical.add_df(pd.read_excel(f"{clinical_dir}/patient_data_od.xlsx", header=1), id_column="ID")
+    clinical.frames[0]["eyeID"] = "OD"
+
+    clinical.add_df(pd.read_excel(f"{clinical_dir}/patient_data_os.xlsx", header=1), id_column="ID")
+    clinical.frames[1]["eyeID"] = "OS"
+
+    # Normalize 'Patient ID' on the per-eye frames (string → int)
+    for frame in clinical.frames:
+        frame["Patient ID"] = frame["Patient ID"].astype(str).str.extract(r"(\d+)")[0].astype(int)
+
+    # Build initial master as in original
+    clinical._refresh_master_df()
+
+    # --- Post-processing ON THE FRAMES (so everything stays consistent) ---
+    for i in range(len(clinical.frames)):
+        clinical.frames[i] = _apply_iop_and_drop_md(clinical.frames[i])
+
+    # Refresh master again so IOP_raw/IOP_corr & MD removal propagate
+    clinical._refresh_master_df()
+    clinical._build_kfold_indices()
+
+    return clinical

classes/refuge_classification.py

@@ -0,0 +1,819 @@
+"""REFUGE glaucoma classification with rotation-based TTT."""
+
+from __future__ import annotations
+
+import math
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Callable, Dict, Iterable, List, Optional, Sequence, Tuple
+import random
+
+import numpy as np
+from PIL import Image
+import torch
+from torch import nn
+from torch.utils.data import DataLoader, Dataset
+from torchvision import models, transforms
+from torchvision.transforms import functional as TF
+import torch.nn.functional as F
+from sklearn.metrics import roc_auc_score
+from skimage.transform import warp_polar
+from tqdm import tqdm
+
+from classes.geometry_features import (
+    FEATURE_DIM,
+    EPS,
+    compute_geometry_features,
+    disc_cup_from_mask_image,
+)
+from classes.refuge_preprocessing import RefugePreprocessing, RefugeSample
+from classes.refuge_segmentation import RefugeSegmentation
+from classes.unet_segmenter import UNetSegmenter
+
+
+# ---------------------------------------------------------------------------
+# Dataset utilities
+# ---------------------------------------------------------------------------
+
+
+def _default_image_transform(size: int = 256) -> transforms.Compose:
+    return transforms.Compose(
+        [
+            transforms.Resize((size, size)),
+            transforms.ToTensor(),
+            transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
+        ]
+    )
+
+
+def _augment_image_transform(size: int = 256) -> transforms.Compose:
+    return transforms.Compose(
+        [
+            transforms.Resize((size, size)),
+            transforms.RandomHorizontalFlip(),
+            transforms.RandomRotation(10),
+            transforms.ColorJitter(0.1, 0.1, 0.1, 0.05),
+            transforms.ToTensor(),
+            transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
+        ]
+    )
+
+
+def _crop_from_geometry(image: Image.Image, geometry: Dict[str, float], size: int = 256) -> Image.Image:
+    cx, cy = geometry["centre_x"], geometry["centre_y"]
+    r = geometry["crop_radius"]
+    left = max(0.0, cx - r)
+    upper = max(0.0, cy - r)
+    right = min(image.width, cx + r)
+    lower = min(image.height, cy + r)
+    crop = image.crop((left, upper, right, lower))
+    return crop.resize((size, size), Image.BILINEAR)
+
+
+def _geometry_from_mask(mask: np.ndarray, scale: float) -> Dict[str, float]:
+    mask = np.asarray(mask) > 0
+    coords = np.argwhere(mask)
+    if coords.size == 0:
+        raise RuntimeError("Empty mask; cannot derive geometry")
+    ys, xs = coords[:, 0], coords[:, 1]
+    centre_x = float(xs.mean())
+    centre_y = float(ys.mean())
+    width = float(xs.max() - xs.min())
+    height = float(ys.max() - ys.min())
+    diameter = max(width, height)
+    radius = diameter / 2.0
+    crop_radius = radius * scale
+    return {
+        "centre_x": centre_x,
+        "centre_y": centre_y,
+        "radius": radius,
+        "crop_radius": crop_radius,
+        "crop_size": crop_radius * 2.0,
+    }
+
+
+def _compute_feature_vector(disc_mask: np.ndarray, cup_mask: np.ndarray) -> np.ndarray:
+    return compute_geometry_features(disc_mask, cup_mask)
+
+
+def _compute_polar_image(crop: Image.Image, size: int) -> Image.Image:
+    arr = np.asarray(crop).astype(np.float32) / 255.0
+    radius = min(arr.shape[0], arr.shape[1]) / 2.0
+    polar = warp_polar(
+        arr,
+        radius=radius,
+        scaling="linear",
+        channel_axis=-1,
+    )
+    polar = np.clip(polar, 0.0, 1.0)
+    polar_img = Image.fromarray((polar * 255).astype(np.uint8))
+    return polar_img.resize((size, size), Image.BILINEAR)
+
+
+def _crop_mask_from_geometry(mask: np.ndarray, geometry: Dict[str, float], size: int) -> np.ndarray:
+    mask_img = Image.fromarray((mask > 0).astype(np.uint8) * 255)
+    cx, cy = geometry["centre_x"], geometry["centre_y"]
+    r = geometry["crop_radius"]
+    left = max(0.0, cx - r)
+    upper = max(0.0, cy - r)
+    right = min(mask_img.width, cx + r)
+    lower = min(mask_img.height, cy + r)
+    crop = mask_img.crop((left, upper, right, lower)).resize((size, size), Image.NEAREST)
+    return (np.asarray(crop) > 0).astype(np.uint8)
+
+
+@dataclass
+class RefugeClassificationRecord:
+    sample: RefugeSample
+    geometry: Dict[str, float]
+    disc_mask: Optional[np.ndarray] = None
+    cup_mask: Optional[np.ndarray] = None
+
+
+class RefugeClassificationDataset(Dataset):
+    def __init__(
+        self,
+        records: Sequence[RefugeClassificationRecord],
+        transform: transforms.Compose,
+        polar_transform: transforms.Compose,
+        size: int = 256,
+    ) -> None:
+        self.records = list(records)
+        self.transform = transform
+        self.polar_transform = polar_transform
+        self.size = size
+
+    def __len__(self) -> int:
+        return len(self.records)
+
+    def __getitem__(self, idx: int) -> Dict[str, torch.Tensor]:
+        rec = self.records[idx]
+        image = Image.open(rec.sample.image_path).convert("RGB")
+        crop = _crop_from_geometry(image, rec.geometry, size=self.size)
+        polar_image = _compute_polar_image(crop, size=self.size)
+        tensor = self.transform(crop)
+        polar_tensor = self.polar_transform(polar_image)
+
+        features = np.zeros((FEATURE_DIM,), dtype=np.float32)
+        if rec.disc_mask is not None and rec.cup_mask is not None:
+            disc_crop = _crop_mask_from_geometry(rec.disc_mask, rec.geometry, self.size)
+            cup_crop = _crop_mask_from_geometry(rec.cup_mask, rec.geometry, self.size)
+            features = _compute_feature_vector(disc_crop, cup_crop)
+
+        feature_tensor = torch.from_numpy(features).float()
+        label = rec.sample.label
+        if label is None:
+            raise ValueError(f"Sample {rec.sample.sample_id} is missing glaucoma label")
+        return {
+            "image": tensor,
+            "polar": polar_tensor,
+            "features": feature_tensor,
+            "label": torch.tensor(label, dtype=torch.long),
+            "sample_id": rec.sample.sample_id,
+        }
+
+
+class RefugeTTTDataset(Dataset):
+    """Dataset providing unlabeled crops for test-time training."""
+
+    def __init__(self, records: Sequence[RefugeClassificationRecord], transform: transforms.Compose, size: int = 256) -> None:
+        self.records = list(records)
+        self.transform = transform
+        self.size = size
+
+    def __len__(self) -> int:
+        return len(self.records)
+
+    def __getitem__(self, idx: int) -> torch.Tensor:
+        rec = self.records[idx]
+        image = Image.open(rec.sample.image_path).convert("RGB")
+        crop = _crop_from_geometry(image, rec.geometry, size=self.size)
+        return self.transform(crop)
+
+
+class UNetGeometryProvider:
+    """Callable wrapper that derives disc geometry using a trained UNetSegmenter."""
+
+    def __init__(
+        self,
+        segmenter: UNetSegmenter,
+        threshold: float = 0.5,
+        tta: bool = False,
+    ) -> None:
+        self.segmenter = segmenter
+        self.threshold = threshold
+        self.tta = tta
+        self.segmenter.model.eval()
+
+    def __call__(self, sample: RefugeSample, scale: float) -> Tuple[Dict[str, float], np.ndarray, np.ndarray]:
+        image = Image.open(sample.image_path).convert("RGB")
+        resized = self.segmenter.preprocess_image(image)
+        tensor = transforms.ToTensor()(resized)
+        tensor = self.segmenter._normalize_tensor(tensor)
+        tensor = tensor.unsqueeze(0).to(self.segmenter.device)
+        with torch.no_grad():
+            logits = self.segmenter.model(tensor)
+            if self.tta:
+                t_h = torch.flip(tensor, dims=[3])
+                log_h = self.segmenter.model(t_h)
+                log_h = torch.flip(log_h, dims=[3])
+                t_v = torch.flip(tensor, dims=[2])
+                log_v = self.segmenter.model(t_v)
+                log_v = torch.flip(log_v, dims=[2])
+                logits = (logits + log_h + log_v) / 3.0
+            probs = torch.sigmoid(logits)[0].cpu().numpy()
+
+        disc_pred = (probs[0] > self.threshold).astype(np.uint8) * 255
+        cup_pred = (probs[1] > self.threshold).astype(np.uint8) * 255
+        disc_img = Image.fromarray(disc_pred, mode="L").resize(image.size, Image.NEAREST)
+        cup_img = Image.fromarray(cup_pred, mode="L").resize(image.size, Image.NEAREST)
+        disc_mask = (np.array(disc_img, dtype=np.uint8) > 0).astype(np.uint8)
+        cup_mask = (np.array(cup_img, dtype=np.uint8) > 0).astype(np.uint8)
+        cup_mask = (cup_mask > 0) & (disc_mask > 0)
+        cup_mask = cup_mask.astype(np.uint8)
+        geom = _geometry_from_mask(disc_mask, scale)
+        return geom, disc_mask, cup_mask
+
+
+# ---------------------------------------------------------------------------
+# Classification module
+# ---------------------------------------------------------------------------
+
+
+class ArcMarginProduct(nn.Module):
+    """Additive angular margin (ArcFace) head."""
+
+    def __init__(
+        self,
+        in_features: int,
+        out_features: int,
+        s: float = 30.0,
+        m: float = 0.5,
+        easy_margin: bool = False,
+    ) -> None:
+        super().__init__()
+        self.in_features = in_features
+        self.out_features = out_features
+        self.s = float(s)
+        self.m = float(m)
+        self.easy_margin = easy_margin
+        self.weight = nn.Parameter(torch.empty(out_features, in_features))
+        nn.init.xavier_uniform_(self.weight)
+
+        self.cos_m = math.cos(m)
+        self.sin_m = math.sin(m)
+        self.th = math.cos(math.pi - m)
+        self.mm = math.sin(math.pi - m) * m
+
+    def forward(self, input: torch.Tensor, label: Optional[torch.Tensor] = None) -> torch.Tensor:
+        cosine = F.linear(F.normalize(input), F.normalize(self.weight))
+        if label is None:
+            return cosine * self.s
+
+        sine = torch.sqrt(torch.clamp(1.0 - cosine.pow(2), min=0.0))
+        phi = cosine * self.cos_m - sine * self.sin_m
+        if self.easy_margin:
+            phi = torch.where(cosine > 0, phi, cosine)
+        else:
+            phi = torch.where(cosine > self.th, phi, cosine - self.mm)
+
+        one_hot = torch.zeros_like(cosine)
+        one_hot.scatter_(1, label.view(-1, 1), 1.0)
+        logits = (one_hot * phi) + ((1.0 - one_hot) * cosine)
+        logits *= self.s
+        return logits
+
+
+class RefugeClassification:
+    """Train and evaluate REFUGE glaucoma classifiers with TTT support."""
+
+    def __init__(
+        self,
+        preprocessing: RefugePreprocessing,
+        segmentation: RefugeSegmentation,
+        backbone: Optional[nn.Module] = None,
+        geometry_fn: Optional[
+            Callable[
+                [RefugeSample, float],
+                Tuple[Dict[str, float], Optional[np.ndarray], Optional[np.ndarray]],
+            ]
+        ] = None,
+        cache_dir: Optional[Path] = None,
+        use_all_labeled: bool = False,
+        auto_val_ratio: float = 0.1,
+        use_margin: bool = False,
+        margin_s: float = 30.0,
+        margin_m: float = 0.5,
+    ) -> None:
+        self.preprocessing = preprocessing
+        self.segmentation = segmentation
+        if backbone is not None:
+            self.backbone = backbone
+            in_features = getattr(self.backbone, "_feature_dim", None)
+            if in_features is None:
+                if hasattr(self.backbone, "fc") and hasattr(self.backbone.fc, "in_features"):
+                    in_features = self.backbone.fc.in_features  # type: ignore[attr-defined]
+                    self.backbone.fc = nn.Identity()  # type: ignore[attr-defined]
+                else:
+                    raise ValueError(
+                        "Provided backbone must have '_feature_dim' or expose fc.in_features"
+                    )
+        else:
+            self.backbone = self._default_backbone()
+            in_features = getattr(self.backbone, "_feature_dim", None)
+            if in_features is None:
+                in_features = self.backbone.fc.in_features  # type: ignore[attr-defined]
+                self.backbone.fc = nn.Identity()  # type: ignore[attr-defined]
+        self.feature_dim = in_features
+        self.use_polar = True
+        self.extra_feature_dim = FEATURE_DIM
+        combined_dim = self.feature_dim * (1 + int(self.use_polar)) + self.extra_feature_dim
+        self.margin_s = float(margin_s)
+        self.margin_m = float(margin_m)
+        self.use_margin = bool(use_margin)
+        if self.use_margin:
+            self.classifier_head = ArcMarginProduct(
+                combined_dim, 2, s=self.margin_s, m=self.margin_m
+            )
+        else:
+            self.classifier_head = nn.Linear(combined_dim, 2)
+        self.rotation_head = nn.Linear(self.feature_dim, 4)
+
+        self.train_dataset: Optional[RefugeClassificationDataset] = None
+        self.val_dataset: Optional[RefugeClassificationDataset] = None
+        self.train_loader: Optional[DataLoader] = None
+        self.val_loader: Optional[DataLoader] = None
+        self.ttt_transform = _default_image_transform()
+        self.train_transform = _augment_image_transform()
+        self.eval_transform = _default_image_transform()
+        self.polar_transform = _default_image_transform()
+        self.crop_scale = 2.5
+        self.crop_size = 256
+        self.geometry_cache: Dict[
+            str, Tuple[Dict[str, float], Optional[np.ndarray], Optional[np.ndarray]]
+        ] = {}
+        self.train_records: List[RefugeClassificationRecord] = []
+        self.val_records: List[RefugeClassificationRecord] = []
+        self._geometry_fn = geometry_fn
+        self.cache_dir = cache_dir
+        if self.cache_dir is not None:
+            self.cache_dir.mkdir(parents=True, exist_ok=True)
+        self.use_all_labeled = use_all_labeled
+        self.auto_val_ratio = auto_val_ratio
+
+    # ------------------------------------------------------------------
+    @staticmethod
+    def _default_backbone() -> nn.Module:
+        weights = models.ResNet50_Weights.IMAGENET1K_V2
+        model = models.resnet50(weights=weights)
+        in_features = model.fc.in_features
+        model.fc = nn.Identity()
+        setattr(model, "_feature_dim", in_features)
+        return model
+
+    # ------------------------------------------------------------------
+    def build_datasets(
+        self,
+        crop_scale: float = 2.5,
+        crop_size: int = 256,
+        batch_size: int = 16,
+        num_workers: int = 4,
+    ) -> None:
+        self.crop_scale = crop_scale
+        self.crop_size = crop_size
+        self.train_transform = _augment_image_transform(crop_size)
+        self.eval_transform = _default_image_transform(crop_size)
+        self.ttt_transform = _default_image_transform(crop_size)
+        self.polar_transform = _default_image_transform(crop_size)
+
+        manifest = list(self.preprocessing.build_manifest())
+        train_records: List[RefugeClassificationRecord] = []
+        val_records: List[RefugeClassificationRecord] = []
+
+        allowed_splits = {"train", "val"}
+        candidates = [
+            sample
+            for sample in manifest
+            if sample.label is not None and sample.split in allowed_splits
+        ]
+
+        print(
+            f"[classifier] Building datasets from {len(candidates)} labelled samples (train/val)"
+        )
+
+        for sample in tqdm(
+            candidates,
+            desc="Preparing records",
+            unit="sample",
+            leave=False,
+        ):
+            try:
+                geom, disc_mask, cup_mask = self._resolve_geometry(sample, crop_scale)
+            except RuntimeError:
+                continue
+            record = RefugeClassificationRecord(
+                sample=sample,
+                geometry=geom,
+                disc_mask=disc_mask,
+                cup_mask=cup_mask,
+            )
+            if sample.split == "train" or (
+                self.use_all_labeled and sample.split == "val"
+            ):
+                train_records.append(record)
+            else:
+                val_records.append(record)
+
+        if (not val_records or self.use_all_labeled) and train_records and self.auto_val_ratio > 0.0:
+            rng = random.Random(42)
+            label_groups: Dict[int, List[RefugeClassificationRecord]] = {}
+            for rec in train_records:
+                label = int(rec.sample.label or 0)
+                label_groups.setdefault(label, []).append(rec)
+
+            new_train: List[RefugeClassificationRecord] = []
+            new_val: List[RefugeClassificationRecord] = []
+            for recs in label_groups.values():
+                rng.shuffle(recs)
+                if len(recs) <= 1:
+                    new_train.extend(recs)
+                    continue
+                val_count = max(1, int(round(len(recs) * self.auto_val_ratio)))
+                if val_count >= len(recs):
+                    val_count = len(recs) - 1
+                new_val.extend(recs[:val_count])
+                new_train.extend(recs[val_count:])
+
+            if not new_val:
+                # Fallback: ensure at least one validation sample if possible
+                if len(new_train) > 1:
+                    new_val.append(new_train.pop())
+
+            if new_val:
+                val_records = new_val
+                train_records = new_train
+
+        self.train_records = train_records
+        self.val_records = val_records
+
+        print(
+            f"[classifier] Records ready → train: {len(train_records)}, val: {len(val_records)}"
+        )
+
+        self.train_dataset = RefugeClassificationDataset(
+            train_records,
+            transform=self.train_transform,
+            polar_transform=self.polar_transform,
+            size=crop_size,
+        )
+        self.val_dataset = RefugeClassificationDataset(
+            val_records,
+            transform=self.eval_transform,
+            polar_transform=self.polar_transform,
+            size=crop_size,
+        )
+
+        self.train_loader = DataLoader(
+            self.train_dataset,
+            batch_size=batch_size,
+            shuffle=True,
+            num_workers=num_workers,
+            pin_memory=True,
+        )
+        self.val_loader = DataLoader(
+            self.val_dataset,
+            batch_size=batch_size,
+            shuffle=False,
+            num_workers=num_workers,
+            pin_memory=True,
+        )
+
+        print(
+            "[classifier] DataLoaders prepared — training batches will start shortly"
+        )
+
+    # ------------------------------------------------------------------
+    def _resolve_geometry(
+        self, sample: RefugeSample, scale: float
+    ) -> Tuple[Dict[str, float], Optional[np.ndarray], Optional[np.ndarray]]:
+        key = self._cache_key(sample.sample_id, scale)
+        cached = self.geometry_cache.get(key)
+        if cached is not None:
+            return cached
+
+        cache_path = self._cache_path(sample.sample_id, scale)
+        if cache_path is not None and cache_path.exists():
+            data = np.load(cache_path, allow_pickle=False)
+            geom = {
+                "centre_x": float(data["centre_x"]),
+                "centre_y": float(data["centre_y"]),
+                "radius": float(data["radius"]),
+                "crop_radius": float(data["crop_radius"]),
+                "crop_size": float(data["crop_size"]),
+            }
+            disc_mask = None
+            cup_mask = None
+            if int(data["has_disc"]):
+                disc_mask = data["disc_mask"].astype(np.uint8)
+            if int(data["has_cup"]):
+                cup_mask = data["cup_mask"].astype(np.uint8)
+            self.geometry_cache[key] = (geom, disc_mask, cup_mask)
+            return geom, disc_mask, cup_mask
+
+        disc_mask: Optional[np.ndarray] = None
+        cup_mask: Optional[np.ndarray] = None
+
+        if sample.mask_path and sample.mask_path.exists():
+            mask_img = Image.open(sample.mask_path).convert("RGB")
+            disc_mask, cup_mask = disc_cup_from_mask_image(mask_img)
+            geom = _geometry_from_mask(disc_mask, scale)
+        elif self._geometry_fn is not None:
+            geom, disc_mask, cup_mask = self._geometry_fn(sample, scale)
+        else:
+            geom = self.segmentation.infer_disc_geometry(sample, scale=scale)
+            try:
+                pred_mask = self.segmentation.predict_mask(sample).numpy()
+                disc_mask = pred_mask.astype(np.uint8)
+            except Exception:
+                disc_mask = None
+            cup_mask = None
+
+        if cache_path is not None:
+            try:
+                np.savez_compressed(
+                    cache_path,
+                    centre_x=geom["centre_x"],
+                    centre_y=geom["centre_y"],
+                    radius=geom["radius"],
+                    crop_radius=geom["crop_radius"],
+                    crop_size=geom.get("crop_size", geom["crop_radius"] * 2.0),
+                    disc_mask=disc_mask if disc_mask is not None else np.array([], dtype=np.uint8),
+                    cup_mask=cup_mask if cup_mask is not None else np.array([], dtype=np.uint8),
+                    has_disc=int(disc_mask is not None),
+                    has_cup=int(cup_mask is not None),
+                )
+            except Exception:
+                pass
+
+        self.geometry_cache[key] = (geom, disc_mask, cup_mask)
+        return geom, disc_mask, cup_mask
+
+    def set_geometry_fn(
+        self,
+        geometry_fn: Optional[
+            Callable[
+                [RefugeSample, float],
+                Tuple[Dict[str, float], Optional[np.ndarray], Optional[np.ndarray]],
+            ]
+        ],
+    ) -> None:
+        self._geometry_fn = geometry_fn
+        self.geometry_cache.clear()
+
+    def build_records_for_samples(
+        self,
+        samples: Sequence[RefugeSample],
+        crop_scale: Optional[float] = None,
+        progress_prefix: Optional[str] = None,
+    ) -> List[RefugeClassificationRecord]:
+        scale = crop_scale if crop_scale is not None else self.crop_scale
+        records: List[RefugeClassificationRecord] = []
+        iterator: Iterable[RefugeSample]
+        if progress_prefix is not None:
+            iterator = tqdm(samples, desc=progress_prefix, unit="sample", leave=False)
+        else:
+            iterator = samples
+        for sample in iterator:
+            if sample.label is None:
+                continue
+            try:
+                geom, disc_mask, cup_mask = self._resolve_geometry(sample, scale)
+            except RuntimeError:
+                continue
+            records.append(
+                RefugeClassificationRecord(
+                    sample=sample,
+                    geometry=geom,
+                    disc_mask=disc_mask,
+                    cup_mask=cup_mask,
+                )
+            )
+        return records
+
+    def _cache_key(self, sample_id: str, scale: float) -> str:
+        scale_tag = int(round(scale * 100))
+        return f"{sample_id}_s{scale_tag}"
+
+    def _cache_path(self, sample_id: str, scale: float) -> Optional[Path]:
+        if self.cache_dir is None:
+            return None
+        return self.cache_dir / f"{self._cache_key(sample_id, scale)}.npz"
+
+    # ------------------------------------------------------------------
+    def train(
+        self,
+        epochs: int = 30,
+        lr: float = 1e-4,
+        weight_decay: float = 1e-4,
+        device: Optional[str] = None,
+        rotation_weight: float = 0.5,
+        checkpoint_dir: Optional[Path] = None,
+    ) -> Dict[str, float]:
+        if self.train_loader is None or self.val_loader is None:
+            self.build_datasets()
+
+        device = device or ("cuda" if torch.cuda.is_available() else "cpu")
+        self.backbone.to(device)
+        self.classifier_head.to(device)
+        self.rotation_head.to(device)
+
+        params = list(self.backbone.parameters()) + list(self.classifier_head.parameters()) + list(self.rotation_head.parameters())
+        optimizer = torch.optim.Adam(params, lr=lr, weight_decay=weight_decay)
+        clf_loss = nn.CrossEntropyLoss()
+        rot_loss = nn.CrossEntropyLoss()
+
+        best_auc = 0.0
+        history: Dict[str, float] = {}
+
+        epoch_iter = tqdm(range(1, epochs + 1), desc="Epochs", unit="epoch")
+
+        print(
+            f"[classifier] Starting training for {epochs} epochs with batch size {self.train_loader.batch_size}"
+        )
+
+        for epoch in epoch_iter:
+            self.backbone.train()
+            self.classifier_head.train()
+            self.rotation_head.train()
+            running_loss = 0.0
+
+            batch_iter = tqdm(
+                self.train_loader,  # type: ignore[arg-type]
+                desc=f"Train {epoch}/{epochs}",
+                leave=False,
+                unit="batch",
+            )
+
+            for batch in batch_iter:
+                images = batch["image"].to(device)
+                polars = batch["polar"].to(device)
+                extra_feats = batch["features"].to(device)
+                labels = batch["label"].to(device)
+                optimizer.zero_grad()
+
+                feats_img = self.backbone(images)
+                feats = feats_img
+                if self.use_polar:
+                    feats_polar = self.backbone(polars)
+                    feats = torch.cat([feats, feats_polar], dim=1)
+                if self.extra_feature_dim > 0:
+                    feats = torch.cat([feats, extra_feats], dim=1)
+                if self.use_margin:
+                    logits = self.classifier_head(feats, labels)
+                else:
+                    logits = self.classifier_head(feats)
+                loss_cls = clf_loss(logits, labels)
+
+                rot_imgs, rot_labels = self._build_rotation_batch(images)
+                feats_rot = self.backbone(rot_imgs)
+                logits_rot = self.rotation_head(feats_rot)
+                loss_rot = rot_loss(logits_rot, rot_labels)
+
+                loss = loss_cls + rotation_weight * loss_rot
+                loss.backward()
+                optimizer.step()
+                running_loss += loss.item() * images.size(0)
+
+            train_loss = running_loss / len(self.train_loader.dataset)  # type: ignore[arg-type]
+            metrics = self.evaluate(device=device)
+            history[f"epoch_{epoch}_loss"] = train_loss
+            history[f"epoch_{epoch}_auc"] = metrics.get("auc", float("nan"))
+
+            auc_val = metrics.get("auc", 0.0)
+            epoch_iter.set_postfix(loss=f"{train_loss:.4f}", auc=f"{auc_val:.4f}")
+
+            if auc_val > best_auc:
+                best_auc = metrics["auc"]
+                if checkpoint_dir is not None:
+                    checkpoint_dir.mkdir(parents=True, exist_ok=True)
+                    torch.save({
+                        "backbone": self.backbone.state_dict(),
+                        "classifier": self.classifier_head.state_dict(),
+                        "rotation": self.rotation_head.state_dict(),
+                    }, checkpoint_dir / "refuge_classifier_best.pt")
+
+        return {"best_auc": best_auc, **history}
+
+    # ------------------------------------------------------------------
+    def evaluate(
+        self,
+        split: str = "val",
+        apply_ttt: bool = False,
+        device: Optional[str] = None,
+    ) -> Dict[str, float]:
+        if split != "val":
+            raise ValueError("Only validation split supported currently")
+        if self.val_loader is None:
+            self.build_datasets()
+
+        device = device or ("cuda" if torch.cuda.is_available() else "cpu")
+        self.backbone.to(device)
+        self.classifier_head.to(device)
+        self.rotation_head.to(device)
+
+        if apply_ttt:
+            ttt_ds = RefugeTTTDataset(self.val_records, transform=self.ttt_transform, size=self.crop_size)
+            ttt_loader = DataLoader(ttt_ds, batch_size=32, shuffle=False)
+            self.apply_ttt(ttt_loader, device=device)
+
+        self.backbone.eval()
+        self.classifier_head.eval()
+        preds: List[float] = []
+        targets: List[int] = []
+
+        with torch.no_grad():
+            val_iter = tqdm(self.val_loader, desc="Validate", leave=False, unit="batch")
+            for batch in val_iter:  # type: ignore[arg-type]
+                images = batch["image"].to(device)
+                labels = batch["label"].to(device)
+                polars = batch["polar"].to(device)
+                extra_feats = batch["features"].to(device)
+                feats_img = self.backbone(images)
+                feats = feats_img
+                if self.use_polar:
+                    feats_polar = self.backbone(polars)
+                    feats = torch.cat([feats, feats_polar], dim=1)
+                if self.extra_feature_dim > 0:
+                    feats = torch.cat([feats, extra_feats], dim=1)
+                if self.use_margin:
+                    logits = self.classifier_head(feats)
+                else:
+                    logits = self.classifier_head(feats)
+                probs = torch.softmax(logits, dim=1)[:, 1]
+                preds.extend(probs.cpu().numpy().tolist())
+                targets.extend(labels.cpu().numpy().tolist())
+
+        auc = 0.0
+        try:
+            if len(set(targets)) > 1:
+                auc = float(roc_auc_score(targets, preds))
+        except ValueError:
+            auc = 0.0
+
+        return {"auc": auc}
+
+    # ------------------------------------------------------------------
+    def apply_ttt(self, loader: DataLoader, device: Optional[str] = None, steps: int = 1, lr: float = 1e-5) -> None:
+        device = device or ("cuda" if torch.cuda.is_available() else "cpu")
+        self.backbone.to(device)
+        self.rotation_head.to(device)
+        self.backbone.train()
+        self.rotation_head.train()
+
+        optimizer = torch.optim.Adam(list(self.backbone.parameters()) + list(self.rotation_head.parameters()), lr=lr)
+        criterion = nn.CrossEntropyLoss()
+
+        for _ in range(steps):
+            for batch in tqdm(loader, desc="TTT adapt", leave=False, unit="batch"):
+                if isinstance(batch, dict):
+                    images = batch["image"].to(device)
+                else:
+                    images = batch.to(device)
+                optimizer.zero_grad()
+                rot_imgs, rot_labels = self._build_rotation_batch(images)
+                feats = self.backbone(rot_imgs)
+                logits = self.rotation_head(feats)
+                loss = criterion(logits, rot_labels)
+                loss.backward()
+                optimizer.step()
+
+    # ------------------------------------------------------------------
+    def extract_backbone(self) -> nn.Module:
+        return self.backbone
+
+    def save_checkpoint(self, output_dir: Path) -> None:
+        output_dir.mkdir(parents=True, exist_ok=True)
+        torch.save({
+            "backbone": self.backbone.state_dict(),
+            "classifier": self.classifier_head.state_dict(),
+            "rotation": self.rotation_head.state_dict(),
+        }, output_dir / "refuge_classifier.pt")
+
+    def load_checkpoint(self, checkpoint_path: Path) -> None:
+        payload = torch.load(checkpoint_path, map_location="cpu")
+        self.backbone.load_state_dict(payload["backbone"])
+        self.classifier_head.load_state_dict(payload["classifier"])
+        self.rotation_head.load_state_dict(payload["rotation"])
+
+    # ------------------------------------------------------------------
+    def _build_rotation_batch(self, images: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
+        rotations = [0, 90, 180, 270]
+        rotated = []
+        labels = []
+        for idx, angle in enumerate(rotations):
+            rot = TF.rotate(images, angle)
+            rotated.append(rot)
+            labels.append(torch.full((images.size(0),), idx, dtype=torch.long, device=images.device))
+        batch = torch.cat(rotated, dim=0)
+        batch_labels = torch.cat(labels, dim=0)
+        return batch, batch_labels

classes/refuge_preprocessing.py

@@ -0,0 +1,306 @@
+"""Utilities for preparing REFUGE (REFUGE1/REFUGE2) datasets.
+
+Builds a unified manifest across all provided splits (REFUGE1 train/val/test
+and REFUGE2 validation/test), exposing image paths, glaucoma labels, disc/cup
+masks, and fovea coordinates so downstream segmentation/classification modules
+can operate without additional bookkeeping.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Dict, Iterable, List, Optional, Tuple
+
+import pandas as pd
+
+
+@dataclass
+class RefugeSample:
+    """Lightweight container describing a REFUGE sample."""
+
+    sample_id: str
+    dataset: str
+    split: str
+    image_path: Path
+    label: Optional[int]
+    device: Optional[str]
+    mask_path: Optional[Path]
+    fovea_coord: Optional[Tuple[float, float]]
+
+
+class RefugePreprocessing:
+    """Builds manifests and provides shared helpers for REFUGE workflows.
+
+    Responsibilities:
+      * scan the REFUGE directory structure and build a consistent manifest
+        (train/val/test, device vendor, ground-truth labels)
+      * expose convenience loaders for raw RGB frames, OD/OC masks, and
+        optional fovea landmarks
+      * compute geometric metadata (disc centres, diameters) so downstream
+        stages can crop ROIs lazily instead of storing pre-rendered tiles
+    """
+
+    def __init__(self, root_dir: Path | str) -> None:
+        self.root_dir = Path(root_dir)
+        self._manifest = None  # populated by build_manifest()
+
+    # ------------------------------------------------------------------
+    # Manifest handling
+    # ------------------------------------------------------------------
+    def build_manifest(self, refresh: bool = False) -> Iterable[RefugeSample]:
+        """Return an iterable of :class:`RefugeSample` records.
+
+        Parameters
+        ----------
+        refresh:
+            when True, force a rescan of the filesystem instead of reusing the
+            cached manifest.
+
+        Returns
+        -------
+        Iterable[RefugeSample]
+            A sequence containing one entry per sample in the REFUGE datasets.
+
+        Notes
+        -----
+        The actual manifest-building logic will live here: parsing the
+        directory structure, reading any provided CSV/Excel metadata, and
+        aligning masks/labels. For now, this method raises ``NotImplementedError``
+        so callers are reminded to hook it up before use.
+        """
+
+        if self._manifest is not None and not refresh:
+            return self._manifest
+
+        manifest: List[RefugeSample] = []
+
+        manifest.extend(self._collect_refuge1_train())
+        manifest.extend(self._collect_refuge1_val())
+        manifest.extend(self._collect_refuge1_test())
+        manifest.extend(self._collect_refuge2_val())
+        manifest.extend(self._collect_refuge2_test())
+
+        self._manifest = manifest
+        return self._manifest
+
+    # ------------------------------------------------------------------
+    # Accessors for downstream modules
+    # ------------------------------------------------------------------
+    def load_image(self, sample: RefugeSample):
+        """Return the RGB fundus image for ``sample``.
+
+        Implementors should handle color-space consistency (e.g., ensure RGB vs
+        BGR) and any global normalisation desired across devices.
+        """
+
+        raise NotImplementedError("Image loading to be implemented")
+
+    def load_mask(self, sample: RefugeSample):
+        """Return the optic disc/cup mask for ``sample`` if available."""
+
+        raise NotImplementedError("Mask loading to be implemented")
+
+    def disc_geometry(self, sample: RefugeSample) -> Dict[str, float]:
+        """Compute disc centre and diameter from the mask.
+
+        The segmentation module will rely on this to crop 2.5–3× disc-diameter
+        ROIs at training time.
+        """
+
+        raise NotImplementedError("Disc geometry helper to be implemented")
+
+    # ------------------------------------------------------------------
+    # Internal helpers
+    # ------------------------------------------------------------------
+    def _collect_refuge1_train(self) -> List[RefugeSample]:
+        base = self.root_dir / "Train" / "REFUGE1-train"
+        if not base.exists():
+            return []
+
+        fovea_path = base / "Fovea_location.xlsx"
+        fovea_map = self._read_fovea_table(fovea_path, img_col="ImgName")
+
+        samples: List[RefugeSample] = []
+        image_root = base / "Training400"
+        mask_root = base / "Disc_Cup_Masks"
+
+        for label_name, label_val in ("Glaucoma", 1), ("Non-Glaucoma", 0):
+            img_dir = image_root / label_name
+            mask_dir = mask_root / label_name
+            if not img_dir.exists():
+                continue
+            for image_path in sorted(img_dir.glob("*.jpg")):
+                img_name = image_path.name
+                mask_path = (mask_dir / image_path.with_suffix(".bmp").name)
+                fovea = fovea_map.get(img_name)
+                sample_id = f"refuge1_train_{image_path.stem}"
+                samples.append(
+                    RefugeSample(
+                        sample_id=sample_id,
+                        dataset="refuge1",
+                        split="train",
+                        image_path=image_path,
+                        label=label_val,
+                        device=None,
+                        mask_path=mask_path if mask_path.exists() else None,
+                        fovea_coord=fovea,
+                    )
+                )
+        return samples
+
+    def _collect_refuge1_val(self) -> List[RefugeSample]:
+        base = self.root_dir / "Train" / "REFUGE1-val"
+        if not base.exists():
+            return []
+
+        fovea_path = base / "Fovea_locations.xlsx"
+        df = pd.read_excel(fovea_path)
+        samples: List[RefugeSample] = []
+        image_root = base / "REFUGE-Validation400"
+        mask_root = base / "Disc_Cup_Masks"
+
+        for _, row in df.iterrows():
+            img_name = row["ImgName"]
+            image_path = image_root / img_name
+            mask_path = mask_root / Path(img_name).with_suffix(".bmp").name
+            fovea = self._extract_fovea(row, x_key="Fovea_X", y_key="Fovea_Y")
+            label = int(row.get("Glaucoma Label", 0)) if not pd.isna(row.get("Glaucoma Label", 0)) else None
+            sample_id = f"refuge1_val_{Path(img_name).stem}"
+            samples.append(
+                RefugeSample(
+                    sample_id=sample_id,
+                    dataset="refuge1",
+                    split="val",
+                    image_path=image_path,
+                    label=label,
+                    device=None,
+                    mask_path=mask_path if mask_path.exists() else None,
+                    fovea_coord=fovea,
+                )
+            )
+        return samples
+
+    def _collect_refuge1_test(self) -> List[RefugeSample]:
+        base = self.root_dir / "Train" / "REFUGE1-test"
+        if not base.exists():
+            return []
+
+        df = pd.read_excel(base / "Glaucoma_label_and_Fovea_location.xlsx")
+        image_root = base / "Test400"
+        mask_root = base / "Disc_Cup_Masks"
+        samples: List[RefugeSample] = []
+
+        for _, row in df.iterrows():
+            img_name = row["ImgName"]
+            image_path = image_root / img_name
+            mask_path = mask_root / Path(img_name).with_suffix(".bmp").name
+            fovea = self._extract_fovea(row, x_key="Fovea_X", y_key="Fovea_Y")
+            label = int(row.get("Label(Glaucoma=1)", 0)) if not pd.isna(row.get("Label(Glaucoma=1)", 0)) else None
+            sample_id = f"refuge1_test_{Path(img_name).stem}"
+            samples.append(
+                RefugeSample(
+                    sample_id=sample_id,
+                    dataset="refuge1",
+                    split="test",
+                    image_path=image_path,
+                    label=label,
+                    device=None,
+                    mask_path=mask_path if mask_path.exists() else None,
+                    fovea_coord=fovea,
+                )
+            )
+        return samples
+
+    def _collect_refuge2_val(self) -> List[RefugeSample]:
+        base = self.root_dir / "Validation"
+        if not base.exists():
+            return []
+
+        label_df = pd.read_csv(base / "glaucoma.csv")
+        fovea_df = pd.read_csv(base / "fovea.csv")
+        fovea_map = {
+            row["ImageName"]: (float(row["Fovea_X"]), float(row["Fovea_Y"]))
+            for _, row in fovea_df.iterrows()
+        }
+        samples: List[RefugeSample] = []
+        image_root = base / "Images"
+        mask_root = base / "Disc_Masks"
+
+        for _, row in label_df.iterrows():
+            img_name = row["FileName"]
+            image_path = image_root / img_name
+            mask_path = mask_root / Path(img_name).with_suffix(".png").name
+            label = row.get("Glaucoma Risk")
+            label = int(label) if label == label else None
+            sample_id = f"refuge2_val_{Path(img_name).stem}"
+            samples.append(
+                RefugeSample(
+                    sample_id=sample_id,
+                    dataset="refuge2",
+                    split="val",
+                    image_path=image_path,
+                    label=label,
+                    device=None,
+                    mask_path=mask_path if mask_path.exists() else None,
+                    fovea_coord=fovea_map.get(img_name),
+                )
+            )
+        return samples
+
+    def _collect_refuge2_test(self) -> List[RefugeSample]:
+        base = self.root_dir / "Test"
+        if not base.exists():
+            return []
+
+        label_df = pd.read_excel(base / "task1.xls", header=None, names=["ImgName", "Glaucoma"])
+        fovea_df = pd.read_excel(base / "fovea.xlsx")
+        fovea_map = {
+            row["ImageName"]: (float(row["Fovea_X"]), float(row["Fovea_Y"]))
+            for _, row in fovea_df.iterrows()
+        }
+        samples: List[RefugeSample] = []
+        image_root = base / "refuge2-test"
+        mask_root = base / "Disc_Mask"
+
+        for _, row in label_df.iterrows():
+            img_name = row["ImgName"]
+            image_path = image_root / img_name
+            mask_path = mask_root / Path(img_name).with_suffix(".png").name
+            label = row.get("Glaucoma")
+            label = int(label) if label == label else None
+            sample_id = f"refuge2_test_{Path(img_name).stem}"
+            samples.append(
+                RefugeSample(
+                    sample_id=sample_id,
+                    dataset="refuge2",
+                    split="test",
+                    image_path=image_path,
+                    label=label,
+                    device=None,
+                    mask_path=mask_path if mask_path.exists() else None,
+            fovea_coord=fovea_map.get(img_name),
+                )
+            )
+        return samples
+
+    @staticmethod
+    def _read_fovea_table(path: Path, img_col: str) -> Dict[str, Tuple[float, float]]:
+        if not path.exists():
+            return {}
+        df = pd.read_excel(path)
+        mapping: Dict[str, Tuple[float, float]] = {}
+        for _, row in df.iterrows():
+            mapping[row[img_col]] = (
+                float(row.get("Fovea_X", float("nan"))),
+                float(row.get("Fovea_Y", float("nan"))),
+            )
+        return mapping
+
+    @staticmethod
+    def _extract_fovea(row: pd.Series, x_key: str, y_key: str) -> Optional[Tuple[float, float]]:
+        x_val = row.get(x_key)
+        y_val = row.get(y_key)
+        if pd.isna(x_val) or pd.isna(y_val):
+            return None
+        return float(x_val), float(y_val)

classes/refuge_segmentation.py

@@ -0,0 +1,383 @@
+"""REFUGE optic disc / cup segmentation utilities."""
+
+from __future__ import annotations
+
+import math
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Dict, Iterable, List, Optional, Sequence, Tuple
+
+import numpy as np
+from PIL import Image
+import torch
+from torch import nn
+from torch.utils.data import DataLoader, Dataset
+from torchvision import transforms
+
+from classes.refuge_preprocessing import RefugePreprocessing, RefugeSample
+
+
+# ---------------------------------------------------------------------------
+# Dataset helpers
+# ---------------------------------------------------------------------------
+
+
+def _load_rgb(path: Path) -> Image.Image:
+    img = Image.open(path)
+    if img.mode != "RGB":
+        img = img.convert("RGB")
+    return img
+
+
+def _load_mask_array(path: Path) -> np.ndarray:
+    mask_img = Image.open(path).convert("L")
+    mask = np.array(mask_img, dtype=np.float32)
+    # REFUGE masks encode disc/cup with different intensities; treat any
+    # positive value as disc for coarse localisation.
+    mask = np.where(mask > 0, 1.0, 0.0)
+    return mask
+
+
+@dataclass
+class RefugeSegmentationSample:
+    sample: RefugeSample
+    image_path: Path
+    mask_path: Path
+
+
+class RefugeSegmentationDataset(Dataset):
+    """Simple segmentation dataset returning tensors."""
+
+    def __init__(
+        self,
+        samples: Sequence[RefugeSegmentationSample],
+        image_size: int = 512,
+    ) -> None:
+        self.samples = list(samples)
+        self.image_size = image_size
+        self.to_tensor = transforms.ToTensor()
+
+    def __len__(self) -> int:
+        return len(self.samples)
+
+    def __getitem__(self, idx: int) -> Dict[str, torch.Tensor]:
+        rec = self.samples[idx]
+        image = _load_rgb(rec.image_path)
+        mask_arr = _load_mask_array(rec.mask_path)
+
+        if self.image_size is not None:
+            image = image.resize((self.image_size, self.image_size), Image.BILINEAR)
+            mask_img = Image.fromarray(mask_arr).resize(
+                (self.image_size, self.image_size), Image.NEAREST
+            )
+            mask_arr = np.array(mask_img, dtype=np.float32)
+
+        image_tensor = self.to_tensor(image)
+        mask_tensor = torch.from_numpy(mask_arr).unsqueeze(0)  # [1,H,W]
+        return {
+            "image": image_tensor,
+            "mask": mask_tensor,
+            "sample_id": rec.sample.sample_id,
+        }
+
+
+# ---------------------------------------------------------------------------
+# Model definition (lightweight U-Net)
+# ---------------------------------------------------------------------------
+
+
+class DoubleConv(nn.Module):
+    def __init__(self, in_channels: int, out_channels: int):
+        super().__init__()
+        self.net = nn.Sequential(
+            nn.Conv2d(in_channels, out_channels, 3, padding=1, bias=False),
+            nn.BatchNorm2d(out_channels),
+            nn.ReLU(inplace=True),
+            nn.Conv2d(out_channels, out_channels, 3, padding=1, bias=False),
+            nn.BatchNorm2d(out_channels),
+            nn.ReLU(inplace=True),
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.net(x)
+
+
+class UNet(nn.Module):
+    def __init__(self, in_channels: int = 3, base_channels: int = 64):
+        super().__init__()
+        self.enc1 = DoubleConv(in_channels, base_channels)
+        self.enc2 = DoubleConv(base_channels, base_channels * 2)
+        self.enc3 = DoubleConv(base_channels * 2, base_channels * 4)
+        self.enc4 = DoubleConv(base_channels * 4, base_channels * 8)
+
+        self.pool = nn.MaxPool2d(2)
+        self.bottleneck = DoubleConv(base_channels * 8, base_channels * 16)
+
+        self.up4 = nn.ConvTranspose2d(base_channels * 16, base_channels * 8, 2, stride=2)
+        self.dec4 = DoubleConv(base_channels * 16, base_channels * 8)
+        self.up3 = nn.ConvTranspose2d(base_channels * 8, base_channels * 4, 2, stride=2)
+        self.dec3 = DoubleConv(base_channels * 8, base_channels * 4)
+        self.up2 = nn.ConvTranspose2d(base_channels * 4, base_channels * 2, 2, stride=2)
+        self.dec2 = DoubleConv(base_channels * 4, base_channels * 2)
+        self.up1 = nn.ConvTranspose2d(base_channels * 2, base_channels, 2, stride=2)
+        self.dec1 = DoubleConv(base_channels * 2, base_channels)
+
+        self.out = nn.Conv2d(base_channels, 1, 1)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        e1 = self.enc1(x)
+        e2 = self.enc2(self.pool(e1))
+        e3 = self.enc3(self.pool(e2))
+        e4 = self.enc4(self.pool(e3))
+        b = self.bottleneck(self.pool(e4))
+
+        d4 = self.up4(b)
+        d4 = torch.cat([d4, e4], dim=1)
+        d4 = self.dec4(d4)
+        d3 = self.up3(d4)
+        d3 = torch.cat([d3, e3], dim=1)
+        d3 = self.dec3(d3)
+        d2 = self.up2(d3)
+        d2 = torch.cat([d2, e2], dim=1)
+        d2 = self.dec2(d2)
+        d1 = self.up1(d2)
+        d1 = torch.cat([d1, e1], dim=1)
+        d1 = self.dec1(d1)
+        return self.out(d1)
+
+
+# ---------------------------------------------------------------------------
+# Segmentation manager
+# ---------------------------------------------------------------------------
+
+
+class RefugeSegmentation:
+    """Train and run coarse-to-fine OD/OC segmentation for REFUGE."""
+
+    def __init__(
+        self,
+        preprocessing: RefugePreprocessing,
+        model: Optional[nn.Module] = None,
+    ) -> None:
+        self.preprocessing = preprocessing
+        self.model = model or UNet()
+        self.train_dataset: Optional[RefugeSegmentationDataset] = None
+        self.val_dataset: Optional[RefugeSegmentationDataset] = None
+        self.train_loader: Optional[DataLoader] = None
+        self.val_loader: Optional[DataLoader] = None
+
+    # ------------------------------------------------------------------
+    def build_datasets(
+        self,
+        image_size: int = 512,
+        batch_size: int = 8,
+        num_workers: int = 4,
+    ) -> None:
+        manifest = self.preprocessing.build_manifest()
+
+        train_samples: List[RefugeSegmentationSample] = []
+        val_samples: List[RefugeSegmentationSample] = []
+
+        for sample in manifest:
+            if not sample.mask_path or not sample.mask_path.exists():
+                continue
+            rec = RefugeSegmentationSample(sample=sample, image_path=sample.image_path, mask_path=sample.mask_path)
+            if sample.split == "train":
+                train_samples.append(rec)
+            elif sample.split in {"val", "validation"}:
+                val_samples.append(rec)
+
+        if not val_samples:
+            # Fall back to using a subset of training data for validation
+            split = max(1, int(0.1 * len(train_samples)))
+            val_samples = train_samples[:split]
+            train_samples = train_samples[split:]
+
+        self.train_dataset = RefugeSegmentationDataset(train_samples, image_size=image_size)
+        self.val_dataset = RefugeSegmentationDataset(val_samples, image_size=image_size)
+        self.train_loader = DataLoader(
+            self.train_dataset,
+            batch_size=batch_size,
+            shuffle=True,
+            num_workers=num_workers,
+            pin_memory=True,
+        )
+        self.val_loader = DataLoader(
+            self.val_dataset,
+            batch_size=batch_size,
+            shuffle=False,
+            num_workers=num_workers,
+            pin_memory=True,
+        )
+
+    # ------------------------------------------------------------------
+    def train(
+        self,
+        epochs: int = 40,
+        lr: float = 1e-3,
+        weight_decay: float = 1e-5,
+        device: Optional[str] = None,
+        checkpoint_dir: Optional[Path] = None,
+    ) -> Dict[str, float]:
+        if self.train_loader is None or self.val_loader is None:
+            self.build_datasets()
+
+        device = device or ("cuda" if torch.cuda.is_available() else "cpu")
+        self.model.to(device)
+        criterion = nn.BCEWithLogitsLoss()
+        optimizer = torch.optim.Adam(self.model.parameters(), lr=lr, weight_decay=weight_decay)
+
+        best_dice = 0.0
+        history: Dict[str, float] = {}
+
+        for epoch in range(1, epochs + 1):
+            print(f"[Seg] Processing epoch {epoch}/{epochs}")
+            self.model.train()
+            running_loss = 0.0
+            for batch in self.train_loader:  # type: ignore[arg-type]
+                images = batch["image"].to(device)
+                masks = batch["mask"].to(device)
+                optimizer.zero_grad()
+                logits = self.model(images)
+                loss = criterion(logits, masks)
+                loss.backward()
+                optimizer.step()
+                running_loss += loss.item() * images.size(0)
+
+            train_loss = running_loss / len(self.train_loader.dataset)  # type: ignore[arg-type]
+            val_metrics = self.evaluate(device=device)
+            history[f"epoch_{epoch}_loss"] = train_loss
+            history[f"epoch_{epoch}_dice"] = val_metrics.get("dice", float("nan"))
+
+            if val_metrics.get("dice", 0.0) > best_dice:
+                best_dice = val_metrics["dice"]
+                if checkpoint_dir is not None:
+                    checkpoint_dir.mkdir(parents=True, exist_ok=True)
+                    torch.save(self.model.state_dict(), checkpoint_dir / "refuge_segmentation_best.pt")
+
+        return {"best_dice": best_dice, **history}
+
+    # ------------------------------------------------------------------
+    def evaluate(self, split: str = "val", device: Optional[str] = None) -> Dict[str, float]:
+        if split != "val":
+            raise ValueError("Only validation split supported currently")
+        if self.val_loader is None:
+            self.build_datasets()
+
+        device = device or ("cuda" if torch.cuda.is_available() else "cpu")
+        self.model.to(device)
+        self.model.eval()
+
+        dices: List[float] = []
+        criterion = nn.BCEWithLogitsLoss()
+        losses: List[float] = []
+
+        with torch.no_grad():
+            for batch in self.val_loader:  # type: ignore[arg-type]
+                images = batch["image"].to(device)
+                masks = batch["mask"].to(device)
+                logits = self.model(images)
+                loss = criterion(logits, masks)
+                losses.append(loss.item() * images.size(0))
+                probs = torch.sigmoid(logits)
+                preds = (probs > 0.5).float()
+                dice = self._dice_coefficient(preds, masks)
+                dices.extend(dice)
+
+        mean_dice = float(np.mean(dices)) if dices else 0.0
+        mean_loss = float(np.sum(losses) / len(self.val_loader.dataset))  # type: ignore[arg-type]
+        return {"dice": mean_dice, "loss": mean_loss}
+
+    # ------------------------------------------------------------------
+    def predict_mask(self, sample: RefugeSample, device: Optional[str] = None) -> torch.Tensor:
+        if self.train_dataset is None:
+            self.build_datasets()
+        device = device or ("cuda" if torch.cuda.is_available() else "cpu")
+        self.model.to(device)
+        self.model.eval()
+
+        image = _load_rgb(sample.image_path)
+        original_size = image.size  # (width, height)
+        image_resized = image.resize((self.train_dataset.image_size, self.train_dataset.image_size), Image.BILINEAR)  # type: ignore[union-attr]
+        tensor = transforms.ToTensor()(image_resized).unsqueeze(0).to(device)
+
+        with torch.no_grad():
+            logits = self.model(tensor)
+            mask_resized = torch.sigmoid(logits)[0, 0]
+
+        mask_np = mask_resized.cpu().numpy()
+        mask_np = (mask_np > 0.5).astype(np.float32)
+        mask_img = Image.fromarray(mask_np)
+        mask_img = mask_img.resize(original_size, Image.NEAREST)
+        return torch.from_numpy(np.array(mask_img, dtype=np.float32))
+
+    def infer_disc_geometry(
+        self,
+        sample: RefugeSample,
+        scale: float = 2.5,
+    ) -> Dict[str, float]:
+        if sample.mask_path and sample.mask_path.exists():
+            mask = _load_mask_array(sample.mask_path)
+        else:
+            mask = self.predict_mask(sample).numpy()
+
+        coords = np.argwhere(mask > 0.5)
+        if coords.size == 0:
+            raise RuntimeError(f"Unable to locate disc for sample {sample.sample_id}")
+
+        ys, xs = coords[:, 0], coords[:, 1]
+        centre_x = float(xs.mean())
+        centre_y = float(ys.mean())
+        width = float(xs.max() - xs.min())
+        height = float(ys.max() - ys.min())
+        diameter = max(width, height)
+        radius = diameter / 2.0
+        crop_radius = radius * scale
+        return {
+            "centre_x": centre_x,
+            "centre_y": centre_y,
+            "radius": radius,
+            "crop_radius": crop_radius,
+            "crop_size": crop_radius * 2.0,
+        }
+
+    def batch_crops(
+        self,
+        samples: Iterable[RefugeSample],
+        scale: float = 2.5,
+        output_dir: Optional[Path] = None,
+        size: int = 256,
+    ) -> Dict[str, Path]:
+        output_paths: Dict[str, Path] = {}
+        if output_dir is not None:
+            output_dir.mkdir(parents=True, exist_ok=True)
+
+        for sample in samples:
+            geom = self.infer_disc_geometry(sample, scale=scale)
+            image = _load_rgb(sample.image_path)
+            cx, cy = geom["centre_x"], geom["centre_y"]
+            r = geom["crop_radius"]
+            left = max(0.0, cx - r)
+            upper = max(0.0, cy - r)
+            right = min(image.width, cx + r)
+            lower = min(image.height, cy + r)
+            crop = image.crop((left, upper, right, lower)).resize((size, size), Image.BILINEAR)
+            if output_dir is not None:
+                out_path = output_dir / f"{sample.sample_id}_crop.png"
+                crop.save(out_path)
+                output_paths[sample.sample_id] = out_path
+        return output_paths
+
+    # ------------------------------------------------------------------
+    @staticmethod
+    def _dice_coefficient(preds: torch.Tensor, targets: torch.Tensor) -> List[float]:
+        eps = 1e-6
+        dices = []
+        preds = preds.view(preds.size(0), -1)
+        targets = targets.view(targets.size(0), -1)
+        for p, t in zip(preds, targets):
+            intersection = float((p * t).sum().item())
+            union = float(p.sum().item() + t.sum().item())
+            dice = (2.0 * intersection + eps) / (union + eps)
+            dices.append(dice)
+        return dices

classes/tower_watcher.py

@@ -0,0 +1,149 @@
+# tower_watcher.py
+import matplotlib.pyplot as plt
+
+# tower_watcher.py
+import matplotlib.pyplot as plt
+
+
+class TowerWatcher:
+    """
+    Live monitor:
+    - Cumulative batch-level: loss & accuracy per batch across all epochs.
+    - Epoch batch-level: loss & accuracy per batch within the current epoch (resets each epoch).
+    - TP/FP/TN/FN bar charts per tower, one chart each, new group each epoch.
+    """
+
+    def __init__(self):
+        plt.ion()
+        # 2 line plots (cum loss, cum acc), 2 line plots (epoch loss, epoch acc), 3 bar plots
+        self.fig, self.axs = plt.subplots(7, 1, figsize=(10, 28))
+        self.reset()
+
+    def reset(self):
+        # Cumulative batch-level
+        self.global_batches = []
+        self.loss_cum = {"fusion": [], "image": [], "meta": []}
+        self.acc_cum = {"fusion": [], "image": [], "meta": []}
+        # Epoch batch-level
+        self.epoch_batches = []
+        self.loss_epoch_batch = {"fusion": [], "image": [], "meta": []}
+        self.acc_epoch_batch = {"fusion": [], "image": [], "meta": []}
+        # Epoch markers for cum plots
+        self.epoch_markers = []
+        # Stats per epoch for bars
+        self.epoch_stats = {"fusion": [], "image": [], "meta": []}
+        # Track current epoch
+        self.current_epoch = -1
+
+    def on_epoch_start(self, epoch):
+        # mark epoch boundary in cumulative
+        x = self.global_batches[-1] + 1 if self.global_batches else 0
+        self.epoch_markers.append(x)
+        # reset epoch batch-level data
+        self.epoch_batches = []
+        for d in [self.loss_epoch_batch, self.acc_epoch_batch]:
+            for k in d:
+                d[k].clear()
+        self.current_epoch = epoch
+
+    def on_batch_end(self, idx, stats: dict):
+        # Cumulative
+        self.global_batches.append(len(self.global_batches) + 1)
+        for key, lk, ak in [
+            ("fusion", "loss_f", "acc_f"),
+            ("image", "loss_i", "acc_i"),
+            ("meta", "loss_m", "acc_m"),
+        ]:
+            self.loss_cum[key].append(stats.get(lk, 0))
+            self.acc_cum[key].append(stats.get(ak, 0))
+        # Epoch-level
+        self.epoch_batches.append(len(self.epoch_batches) + 1)
+        for key, lk, ak in [
+            ("fusion", "loss_f", "acc_f"),
+            ("image", "loss_i", "acc_i"),
+            ("meta", "loss_m", "acc_m"),
+        ]:
+            self.loss_epoch_batch[key].append(stats.get(lk, 0))
+            self.acc_epoch_batch[key].append(stats.get(ak, 0))
+        # redraw
+        self._draw_batch_plots()
+
+    def on_epoch_end(self, epoch, stats: dict):
+        # record per-epoch TP/FP/TN/FN
+        for key in ["fusion", "image", "meta"]:
+            self.epoch_stats[key].append(
+                {
+                    "tp": stats.get("tp", 0),
+                    "fp": stats.get("fp", 0),
+                    "tn": stats.get("tn", 0),
+                    "fn": stats.get("fn", 0),
+                }
+            )
+        self._draw_epoch_bars()
+
+    def _draw_batch_plots(self):
+        # Cumulative Loss
+        ax = self.axs[0]
+        ax.clear()
+        ax.plot(self.global_batches, self.loss_cum["fusion"], label="Fusion")
+        ax.plot(self.global_batches, self.loss_cum["image"], label="Image Tower")
+        ax.plot(self.global_batches, self.loss_cum["meta"], label="MD Tower")
+        for x in self.epoch_markers:
+            ax.axvline(x=x, color="gray", linestyle="--")
+        ax.set_ylabel("Cumulative Loss")
+        ax.legend()
+
+        # Epoch Loss
+        ax = self.axs[1]
+        ax.clear()
+        ax.plot(self.epoch_batches, self.loss_epoch_batch["fusion"], label="Fusion")
+        ax.plot(self.epoch_batches, self.loss_epoch_batch["image"], label="Image Tower")
+        ax.plot(self.epoch_batches, self.loss_epoch_batch["meta"], label="MD Tower")
+        ax.set_ylabel(f"Epoch {self.current_epoch+1} Loss")
+        ax.set_xlabel("Batch (Epoch)")
+        ax.legend()
+
+        # Cumulative Accuracy
+        ax = self.axs[2]
+        ax.clear()
+        ax.plot(self.global_batches, self.acc_cum["fusion"], label="Fusion")
+        ax.plot(self.global_batches, self.acc_cum["image"], label="Image Tower")
+        ax.plot(self.global_batches, self.acc_cum["meta"], label="MD Tower")
+        for x in self.epoch_markers:
+            ax.axvline(x=x, color="gray", linestyle="--")
+        ax.set_ylabel("Cumulative Accuracy")
+        ax.legend()
+
+        # Epoch Accuracy
+        ax = self.axs[3]
+        ax.clear()
+        ax.plot(self.epoch_batches, self.acc_epoch_batch["fusion"], label="Fusion")
+        ax.plot(self.epoch_batches, self.acc_epoch_batch["image"], label="Image Tower")
+        ax.plot(self.epoch_batches, self.acc_epoch_batch["meta"], label="MD Tower")
+        ax.set_ylabel(f"Epoch {self.current_epoch+1} Accuracy")
+        ax.set_xlabel("Batch (Epoch)")
+        ax.legend()
+
+        plt.pause(0.01)
+
+    def _draw_epoch_bars(self):
+        # Bar charts per tower
+        for i, key in enumerate(["fusion", "image", "meta"]):
+            ax = self.axs[4 + i]
+            ax.clear()
+            data = self.epoch_stats[key]
+            epochs = list(range(1, len(data) + 1))
+            tp = [d["tp"] for d in data]
+            fp = [d["fp"] for d in data]
+            tn = [d["tn"] for d in data]
+            fn = [d["fn"] for d in data]
+            width = 0.2
+            ax.bar([e - width for e in epochs], tp, width, label="TP")
+            ax.bar(epochs, fp, width, label="FP")
+            ax.bar([e + width for e in epochs], tn, width, label="TN")
+            ax.bar([e + 2 * width for e in epochs], fn, width, label="FN")
+            ax.set_title(f"{key.title()} Tower Stats")
+            ax.set_xlabel("Epoch")
+            ax.set_ylabel("Count")
+            ax.legend()
+        plt.pause(0.01)

classes/unet_segmenter.py

@@ -0,0 +1,891 @@
+"""U-Net based optic disc/cup segmenter for REFUGE + Papila."""
+
+from __future__ import annotations
+
+import math
+import os
+from concurrent.futures import ThreadPoolExecutor, as_completed
+from contextlib import suppress
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Iterable, List, Optional, Set, Tuple
+
+import numpy as np
+import pandas as pd
+from PIL import Image, ImageDraw, ImageOps
+from PIL.Image import Resampling
+from skimage import measure
+import torch
+from torch import nn
+from torch.utils.data import DataLoader, Dataset
+from torchvision import transforms
+from tqdm import tqdm
+
+
+@dataclass
+class ManifestEntry:
+    sample_id: str
+    dataset: str
+    image_path: Path
+    annotation_disc: Path
+    annotation_cup: Path
+    annotation_type_disc: str
+    annotation_type_cup: str
+    split: str  # train / holdout / etc.
+
+
+class UNet(nn.Module):
+    def __init__(
+        self, in_channels: int = 3, base_channels: int = 32, out_channels: int = 2
+    ):
+        super().__init__()
+        self.enc1 = self._block(in_channels, base_channels)
+        self.enc2 = self._block(base_channels, base_channels * 2)
+        self.enc3 = self._block(base_channels * 2, base_channels * 4)
+        self.enc4 = self._block(base_channels * 4, base_channels * 8)
+
+        self.pool = nn.MaxPool2d(2)
+        self.bottleneck = self._block(base_channels * 8, base_channels * 16)
+
+        self.up4 = nn.ConvTranspose2d(
+            base_channels * 16, base_channels * 8, 2, stride=2
+        )
+        self.dec4 = self._block(base_channels * 16, base_channels * 8)
+        self.up3 = nn.ConvTranspose2d(base_channels * 8, base_channels * 4, 2, stride=2)
+        self.dec3 = self._block(base_channels * 8, base_channels * 4)
+        self.up2 = nn.ConvTranspose2d(base_channels * 4, base_channels * 2, 2, stride=2)
+        self.dec2 = self._block(base_channels * 4, base_channels * 2)
+        self.up1 = nn.ConvTranspose2d(base_channels * 2, base_channels, 2, stride=2)
+        self.dec1 = self._block(base_channels * 2, base_channels)
+
+        self.out_conv = nn.Conv2d(base_channels, out_channels, kernel_size=1)
+
+    @staticmethod
+    def _block(in_ch: int, out_ch: int) -> nn.Module:
+        return nn.Sequential(
+            nn.Conv2d(in_ch, out_ch, kernel_size=3, padding=1, bias=False),
+            nn.BatchNorm2d(out_ch),
+            nn.ReLU(inplace=True),
+            nn.Conv2d(out_ch, out_ch, kernel_size=3, padding=1, bias=False),
+            nn.BatchNorm2d(out_ch),
+            nn.ReLU(inplace=True),
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        e1 = self.enc1(x)
+        e2 = self.enc2(self.pool(e1))
+        e3 = self.enc3(self.pool(e2))
+        e4 = self.enc4(self.pool(e3))
+        b = self.bottleneck(self.pool(e4))
+
+        d4 = self.up4(b)
+        d4 = torch.cat([d4, e4], dim=1)
+        d4 = self.dec4(d4)
+        d3 = self.up3(d4)
+        d3 = torch.cat([d3, e3], dim=1)
+        d3 = self.dec3(d3)
+        d2 = self.up2(d3)
+        d2 = torch.cat([d2, e2], dim=1)
+        d2 = self.dec2(d2)
+        d1 = self.up1(d2)
+        d1 = torch.cat([d1, e1], dim=1)
+        d1 = self.dec1(d1)
+        return self.out_conv(d1)
+
+
+class SegmentationDataset(Dataset):
+    def __init__(
+        self,
+        entries: List[ManifestEntry],
+        segmenter: "UNetSegmenter",
+        augment: bool,
+    ) -> None:
+        self.entries = entries
+        self.segmenter = segmenter
+        self.augment = augment
+
+    def __len__(self) -> int:
+        return len(self.entries)
+
+    def __getitem__(self, idx: int):
+        entry = self.entries[idx]
+        image = self.segmenter.load_preprocessed_image(entry)
+        disc_mask, cup_mask = self.segmenter.load_masks(entry)
+
+        if self.augment:
+            image = self.segmenter.jitter_image(image)
+            image, disc_mask, cup_mask = self.segmenter.augment_geometric(
+                image, disc_mask, cup_mask
+            )
+        image_tensor = transforms.ToTensor()(image)
+        image_tensor = self.segmenter._normalize_tensor(image_tensor)
+
+        mask = np.stack([disc_mask, cup_mask], axis=0).astype(np.float32)
+        mask_tensor = torch.from_numpy(mask)
+        return image_tensor, mask_tensor
+
+
+class UNetSegmenter:
+    def __init__(
+        self,
+        manifest_path: Path,
+        device: Optional[str] = None,
+        cup_weight: float = 1.0,
+        disc_weight: float = 1.0,
+        target_size: int = 512,
+        val_ratio: float = 0.1,
+        train_datasets: Optional[Iterable[str]] = None,
+        val_datasets: Optional[Iterable[str]] = None,
+        holdout_datasets: Optional[Iterable[str]] = None,
+        normalize: str = "none",
+        use_stronger_aug: bool = False,
+        mask_cache_dir: Optional[Path] = None,
+        image_cache_dir: Optional[Path] = None,
+        in_memory_cache: bool = False,
+        loader_workers: int = 0,
+    ) -> None:
+        self.manifest_path = manifest_path
+        self.device = device or ("cuda" if torch.cuda.is_available() else "cpu")
+        self.cup_weight = cup_weight
+        self.disc_weight = disc_weight
+        self.target_size = target_size
+        self.val_ratio = val_ratio
+        self.normalize = (normalize or "none").lower()
+        self.use_stronger_aug = bool(use_stronger_aug)
+        self.mask_cache_dir = Path(mask_cache_dir).resolve() if mask_cache_dir else None
+        if self.mask_cache_dir:
+            self.mask_cache_dir.mkdir(parents=True, exist_ok=True)
+        self.image_cache_dir = Path(image_cache_dir).resolve() if image_cache_dir else None
+        if self.image_cache_dir:
+            self.image_cache_dir.mkdir(parents=True, exist_ok=True)
+        self.in_memory_cache = bool(in_memory_cache)
+        self._mem_image_cache: dict[str, np.ndarray] = {}
+        self._mem_mask_cache: dict[str, Tuple[np.ndarray, np.ndarray]] = {}
+        self.loader_workers = max(0, int(loader_workers))
+
+        self.train_dataset_filter = self._normalize_filter(train_datasets)
+        self.val_dataset_filter = self._normalize_filter(val_datasets)
+        self.holdout_dataset_filter = self._normalize_filter(holdout_datasets)
+
+        self.model = UNet().to(self.device)
+        self._manifest: List[ManifestEntry] = []
+        self.train_entries: List[ManifestEntry] = []
+        self.val_entries: List[ManifestEntry] = []
+        self.holdout_entries: List[ManifestEntry] = []
+        self.read_manifest()
+
+    def prebuild_in_memory_cache(
+        self,
+        *,
+        cache_workers: int = 0,
+        include_train: bool = True,
+        include_val: bool = True,
+        include_holdout: bool = False,
+    ) -> None:
+        if not self.in_memory_cache:
+            return
+        selected: List[ManifestEntry] = []
+        if include_train:
+            selected.extend(self.train_entries)
+        if include_val:
+            selected.extend(self.val_entries)
+        if include_holdout:
+            selected.extend(self.holdout_entries)
+        if not selected:
+            return
+
+        # Deduplicate by cache key.
+        dedup = {}
+        for entry in selected:
+            dedup[self._entry_cache_key(entry)] = entry
+        entries = list(dedup.values())
+        workers = max(0, int(cache_workers))
+        print(
+            f"[UNetSegmenter] prebuilding in-memory cache for {len(entries)} samples "
+            f"(cache_workers={workers})",
+            flush=True,
+        )
+
+        def _warm_one(entry: ManifestEntry) -> None:
+            self.load_preprocessed_image(entry)
+            self.load_masks(entry)
+
+        if workers <= 1:
+            for entry in tqdm(entries, desc="Warm cache", unit="sample"):
+                _warm_one(entry)
+        else:
+            with ThreadPoolExecutor(max_workers=workers) as ex:
+                futures = [ex.submit(_warm_one, entry) for entry in entries]
+                for fut in tqdm(as_completed(futures), total=len(futures), desc="Warm cache", unit="sample"):
+                    fut.result()
+
+    # ------------------------------------------------------------------
+    def read_manifest(self) -> None:
+        df = pd.read_csv(self.manifest_path)
+        entries: List[ManifestEntry] = []
+        for _, row in df.iterrows():
+            entry = ManifestEntry(
+                sample_id=row["sample_id"],
+                dataset=row["dataset"],
+                image_path=Path(row["image_path"]),
+                annotation_disc=Path(row["annotation_disc"]),
+                annotation_cup=Path(row["annotation_cup"]),
+                annotation_type_disc=row["annotation_type_disc"],
+                annotation_type_cup=row["annotation_type_cup"],
+                split=row["split"],
+            )
+            entries.append(entry)
+        self._manifest = entries
+        self.holdout_entries = [e for e in entries if e.split == "holdout"]
+        if self.holdout_dataset_filter is not None:
+            self.holdout_entries = [
+                e for e in self.holdout_entries if e.dataset in self.holdout_dataset_filter
+            ]
+
+        trainable = [e for e in entries if e.split != "holdout"]
+        if self.train_dataset_filter is not None:
+            trainable = [
+                e for e in trainable if e.dataset in self.train_dataset_filter
+            ]
+
+        if not trainable:
+            self.val_entries = []
+            self.train_entries = []
+            return
+
+        val_pool = trainable
+        if self.val_dataset_filter is not None:
+            filtered = [e for e in trainable if e.dataset in self.val_dataset_filter]
+            if filtered:
+                val_pool = filtered
+
+        if len(trainable) == 1:
+            val_count = 0
+        else:
+            val_count = max(1, int(len(trainable) * self.val_ratio))
+            val_count = min(val_count, len(val_pool), len(trainable) - 1)
+
+        selected_val: List[ManifestEntry] = []
+        if val_count > 0:
+            selected_val = list(val_pool[:val_count])
+        self.val_entries = selected_val
+        selected_ids = {id(item) for item in selected_val}
+        self.train_entries = [e for e in trainable if id(e) not in selected_ids]
+
+        if not self.train_entries and trainable:
+            # Fallback when filtering removed all train entries (e.g. val_count forced entire set)
+            self.train_entries = trainable
+            self.val_entries = []
+
+    # ------------------------------------------------------------------
+    def preprocess_image(self, image: Image.Image) -> Image.Image:
+        return image.resize((self.target_size, self.target_size), Resampling.BILINEAR)
+
+    def jitter_image(self, image: Image.Image) -> Image.Image:
+        # Photometric jitter only; geometric ops are applied jointly (image+mask)
+        return transforms.ColorJitter(0.1, 0.1, 0.1, 0.05)(image)
+
+    def augment_geometric(
+        self,
+        image: Image.Image,
+        disc_mask: np.ndarray,
+        cup_mask: np.ndarray,
+    ) -> tuple[Image.Image, np.ndarray, np.ndarray]:
+        if not self.use_stronger_aug:
+            return image, disc_mask, cup_mask
+
+        img = image
+        disc_pil = Image.fromarray((disc_mask > 0).astype(np.uint8) * 255)
+        cup_pil = Image.fromarray((cup_mask > 0).astype(np.uint8) * 255)
+
+        # Random horizontal flip
+        if np.random.rand() < 0.5:
+            img = ImageOps.mirror(img)
+            disc_pil = ImageOps.mirror(disc_pil)
+            cup_pil = ImageOps.mirror(cup_pil)
+        # Random vertical flip
+        if np.random.rand() < 0.5:
+            img = ImageOps.flip(img)
+            disc_pil = ImageOps.flip(disc_pil)
+            cup_pil = ImageOps.flip(cup_pil)
+        # Random rotation (multiples of 90° to keep masks aligned)
+        rotations = np.random.choice([0, 90, 180, 270])
+        if rotations:
+            img = img.rotate(rotations, expand=False)
+            disc_pil = disc_pil.rotate(rotations, expand=False)
+            cup_pil = cup_pil.rotate(rotations, expand=False)
+
+        disc_mask = (np.array(disc_pil) > 0).astype(np.float32)
+        cup_mask = (np.array(cup_pil) > 0).astype(np.float32)
+        return img, disc_mask, cup_mask
+
+    @staticmethod
+    def _slugify(text: str) -> str:
+        return "".join(ch if ch.isalnum() or ch in ("-", "_") else "_" for ch in text)
+
+    def _entry_cache_key(self, entry: ManifestEntry) -> str:
+        return self._slugify(f"{entry.dataset}_{entry.sample_id}_sz{self.target_size}")
+
+    def _mask_cache_path(self, entry: ManifestEntry) -> Optional[Path]:
+        if self.mask_cache_dir is None:
+            return None
+        slug = self._slugify(f"{entry.dataset}_{entry.sample_id}")
+        fname = f"{slug}_sz{self.target_size}.npz"
+        return self.mask_cache_dir / fname
+
+    def _image_cache_path(self, entry: ManifestEntry) -> Optional[Path]:
+        if self.image_cache_dir is None:
+            return None
+        slug = self._slugify(f"{entry.dataset}_{entry.sample_id}")
+        fname = f"{slug}_img_sz{self.target_size}.npz"
+        return self.image_cache_dir / fname
+
+    def _load_image_cache(self, cache_path: Path) -> Optional[Image.Image]:
+        try:
+            data = np.load(str(cache_path), allow_pickle=False)
+            arr = data["image"].astype(np.uint8, copy=False)
+            if arr.ndim != 3 or arr.shape[2] != 3:
+                return None
+            return Image.fromarray(arr, mode="RGB")
+        except Exception:
+            with suppress(OSError, FileNotFoundError):
+                cache_path.unlink()
+            return None
+
+    def _save_image_cache(self, cache_path: Optional[Path], image: Image.Image) -> None:
+        if cache_path is None:
+            return
+        cache_path.parent.mkdir(parents=True, exist_ok=True)
+        tmp_path = cache_path.with_suffix(cache_path.suffix + ".tmp.npz")
+        try:
+            arr = np.asarray(image, dtype=np.uint8)
+            np.savez_compressed(tmp_path, image=arr)
+            os.replace(tmp_path, cache_path)
+        except Exception:
+            with suppress(OSError, FileNotFoundError):
+                tmp_path.unlink()
+
+    def load_preprocessed_image(self, entry: ManifestEntry) -> Image.Image:
+        key = self._entry_cache_key(entry)
+        if self.in_memory_cache:
+            cached = self._mem_image_cache.get(key)
+            if cached is not None:
+                return Image.fromarray(cached, mode="RGB")
+        cache_path = self._image_cache_path(entry)
+        if cache_path and cache_path.exists():
+            cached = self._load_image_cache(cache_path)
+            if cached is not None:
+                if self.in_memory_cache:
+                    self._mem_image_cache[key] = np.asarray(cached, dtype=np.uint8)
+                return cached
+        image = Image.open(entry.image_path).convert("RGB")
+        image = self.preprocess_image(image)
+        if self.in_memory_cache:
+            self._mem_image_cache[key] = np.asarray(image, dtype=np.uint8)
+        self._save_image_cache(cache_path, image)
+        return image
+
+    def _load_mask_cache(self, cache_path: Path) -> Optional[Tuple[np.ndarray, np.ndarray]]:
+        try:
+            data = np.load(str(cache_path), allow_pickle=False)
+            disc = data["disc"].astype(np.float32)
+            cup = data["cup"].astype(np.float32)
+            return disc, cup
+        except Exception:
+            with suppress(OSError, FileNotFoundError):
+                cache_path.unlink()
+            return None
+
+    def _save_mask_cache(
+        self,
+        cache_path: Optional[Path],
+        disc_mask: np.ndarray,
+        cup_mask: np.ndarray,
+    ) -> None:
+        if cache_path is None:
+            return
+        cache_path.parent.mkdir(parents=True, exist_ok=True)
+        tmp_path = cache_path.with_suffix(cache_path.suffix + ".tmp.npz")
+        try:
+            np.savez_compressed(
+                tmp_path,
+                disc=disc_mask.astype(np.uint8),
+                cup=cup_mask.astype(np.uint8),
+            )
+            os.replace(tmp_path, cache_path)
+        except Exception:
+            with suppress(OSError, FileNotFoundError):
+                tmp_path.unlink()
+
+    def _normalize_tensor(self, tensor: torch.Tensor) -> torch.Tensor:
+        if self.normalize == "per_image":
+            mean = tensor.mean(dim=(1, 2), keepdim=True)
+            std = tensor.std(dim=(1, 2), keepdim=True).clamp(min=1e-6)
+            return (tensor - mean) / std
+        if self.normalize == "imagenet":
+            mean = torch.tensor([0.485, 0.456, 0.406]).view(-1, 1, 1)
+            std = torch.tensor([0.229, 0.224, 0.225]).view(-1, 1, 1)
+            return (tensor - mean) / std
+        return tensor
+
+    # ------------------------------------------------------------------
+    def extract_masks_from_image(
+        self,
+        mask_path: Path,
+        disc_color: Optional[tuple[int, int, int]] = None,
+        cup_color: Optional[tuple[int, int, int]] = None,
+    ) -> Tuple[np.ndarray, Optional[np.ndarray], Tuple[int, int]]:
+        raw = Image.open(mask_path)
+        arr = np.array(raw)
+        if arr.ndim == 2:
+            h, w = arr.shape
+            flat = arr.reshape(-1).astype(np.int64, copy=False)
+            edges = np.concatenate([arr[0, :], arr[-1, :], arr[:, 0], arr[:, -1]], axis=0).astype(np.int64, copy=False)
+            edge_counts = np.bincount(edges, minlength=256)
+            bg_val = int(np.argmax(edge_counts))
+            counts = np.bincount(flat, minlength=256)
+            counts[bg_val] = 0
+            vals = np.where(counts > 0)[0]
+            if vals.size < 1:
+                raise ValueError(f"Mask {mask_path} does not contain discernible labels")
+            vals = vals[np.argsort(-counts[vals])]
+            disc_val = int(vals[0])
+            cup_val = int(vals[1]) if vals.size > 1 else None
+            disc_mask = (arr == disc_val).astype(np.uint8)
+            cup_mask = (arr == cup_val).astype(np.uint8) if cup_val is not None else np.zeros_like(disc_mask, dtype=np.uint8)
+            return disc_mask, cup_mask if cup_mask.any() else None, (w, h)
+
+        image = raw.convert("RGB")
+        arr = np.array(image)
+        h, w, c = arr.shape
+
+        if disc_color is None or cup_color is None:
+            # Fast color discovery via NumPy (avoid Python-level per-pixel tuple counting).
+            edges = np.concatenate(
+                [arr[0, :, :], arr[-1, :, :], arr[:, 0, :], arr[:, -1, :]], axis=0
+            )
+            edge_colors, edge_counts = np.unique(edges.reshape(-1, c), axis=0, return_counts=True)
+            bg_color_np = edge_colors[int(np.argmax(edge_counts))]
+
+            colors_np, counts_np = np.unique(arr.reshape(-1, c), axis=0, return_counts=True)
+            keep = np.any(colors_np != bg_color_np.reshape(1, -1), axis=1)
+            colors_np = colors_np[keep]
+            counts_np = counts_np[keep]
+            if colors_np.shape[0] < 1:
+                raise ValueError(f"Mask {mask_path} does not contain discernible labels")
+            order = np.argsort(-counts_np)
+            colors_np = colors_np[order]
+            disc_color = tuple(int(v) for v in colors_np[0].tolist())
+            cup_color = (
+                tuple(int(v) for v in colors_np[1].tolist())
+                if colors_np.shape[0] > 1
+                else None
+            )
+
+        disc_mask = np.zeros((h, w), dtype=np.uint8)
+        cup_mask = np.zeros((h, w), dtype=np.uint8)
+
+        if disc_color is not None:
+            disc_mask[np.all(arr == disc_color, axis=-1)] = 1
+        if cup_color is not None:
+            cup_mask[np.all(arr == cup_color, axis=-1)] = 1
+
+        return disc_mask, cup_mask if cup_mask.any() else None, (w, h)
+
+    def load_contour_from_file(self, contour_path: Path) -> np.ndarray:
+        # Fast path: contour files are typically CSV or whitespace-delimited x,y pairs.
+        try:
+            arr = np.loadtxt(str(contour_path), delimiter=",", comments="#", dtype=np.float32)
+        except Exception:
+            try:
+                arr = np.loadtxt(str(contour_path), comments="#", dtype=np.float32)
+            except Exception:
+                return np.zeros((0, 2), dtype=np.float32)
+        if arr.size == 0:
+            return np.zeros((0, 2), dtype=np.float32)
+        if arr.ndim == 1:
+            if arr.shape[0] < 2:
+                return np.zeros((0, 2), dtype=np.float32)
+            arr = arr.reshape(1, -1)
+        if arr.shape[1] < 2:
+            return np.zeros((0, 2), dtype=np.float32)
+        return arr[:, :2].astype(np.float32, copy=False)
+
+    def coords_to_mask(
+        self,
+        coords: Optional[np.ndarray],
+        size: Tuple[int, int],
+    ) -> np.ndarray:
+        if coords is None or len(coords) == 0:
+            return np.zeros((self.target_size, self.target_size), dtype=np.float32)
+
+        width, height = map(int, size)
+        target_shape = (height, width)
+        arr = np.asarray(coords)
+        if arr.size == 0:
+            return np.zeros((self.target_size, self.target_size), dtype=np.float32)
+
+        if arr.ndim == 2 and arr.shape[-1] != 2:
+            mask = (arr > 0).astype(np.uint8)
+            return self._resize_mask(mask)
+
+        if arr.ndim > 2:
+            arr = arr.reshape(-1, arr.shape[-1])
+        arr = arr.astype(float, copy=False)
+        if arr.shape[-1] != 2:
+            raise ValueError(f"Expected coordinate pairs, got shape {arr.shape}")
+
+        points = [tuple(map(float, pt)) for pt in arr]
+        if len(points) < 3:
+            return np.zeros(target_shape, dtype=np.float32)
+
+        img = Image.new("L", size, 0)
+        draw = ImageDraw.Draw(img)
+        draw.polygon(points, outline=1, fill=1)
+        mask = np.array(img, dtype=np.uint8)
+        return self._resize_mask(mask)
+
+    def _resize_mask(self, mask: np.ndarray) -> np.ndarray:
+        img = Image.fromarray((mask > 0).astype(np.uint8) * 255)
+        img = img.resize((self.target_size, self.target_size), Resampling.NEAREST)
+        return (np.array(img, dtype=np.uint8) > 0).astype(np.float32)
+
+    def load_masks(self, entry: ManifestEntry) -> Tuple[np.ndarray, np.ndarray]:
+        key = self._entry_cache_key(entry)
+        if self.in_memory_cache:
+            cached = self._mem_mask_cache.get(key)
+            if cached is not None:
+                disc_u8, cup_u8 = cached
+                return disc_u8.astype(np.float32), cup_u8.astype(np.float32)
+        cache_path = self._mask_cache_path(entry)
+        if cache_path and cache_path.exists():
+            cached = self._load_mask_cache(cache_path)
+            if cached is not None:
+                if self.in_memory_cache:
+                    disc, cup = cached
+                    self._mem_mask_cache[key] = (
+                        disc.astype(np.uint8),
+                        cup.astype(np.uint8),
+                    )
+                return cached
+
+        image = Image.open(entry.image_path)
+        size = image.size
+
+        disc_coords = cup_coords = None
+        if entry.annotation_type_disc == "mask":
+            disc_coords, cup_coords_from_disc, size = self.extract_masks_from_image(
+                entry.annotation_disc
+            )
+            if cup_coords_from_disc is not None:
+                cup_coords = cup_coords_from_disc
+        else:
+            disc_coords = self.load_contour_from_file(entry.annotation_disc)
+
+        if entry.annotation_type_cup == "mask":
+            _, cup_coords_from_cup, size_cup = self.extract_masks_from_image(
+                entry.annotation_cup
+            )
+            if cup_coords_from_cup is not None:
+                cup_coords = cup_coords_from_cup
+            if disc_coords is None:
+                disc_coords, _, size = self.extract_masks_from_image(
+                    entry.annotation_cup
+                )
+            else:
+                size = size_cup
+        else:
+            cup_coords = self.load_contour_from_file(entry.annotation_cup)
+
+        disc_mask = self.coords_to_mask(disc_coords, size).astype(np.float32)
+        cup_mask = self.coords_to_mask(cup_coords, size).astype(np.float32)
+        if self.in_memory_cache:
+            self._mem_mask_cache[key] = (
+                disc_mask.astype(np.uint8),
+                cup_mask.astype(np.uint8),
+            )
+        self._save_mask_cache(cache_path, disc_mask, cup_mask)
+        return disc_mask, cup_mask
+
+    # ------------------------------------------------------------------
+    def build_loaders(self, batch_size: int = 4, num_workers: int = 0) -> Tuple[DataLoader, DataLoader]:
+        train_ds = SegmentationDataset(self.train_entries, self, augment=True)
+        val_ds = SegmentationDataset(self.val_entries, self, augment=False)
+        train_loader = DataLoader(
+            train_ds, batch_size=batch_size, shuffle=True, num_workers=num_workers, pin_memory=True
+        )
+        val_loader = DataLoader(
+            val_ds, batch_size=batch_size, shuffle=False, num_workers=num_workers, pin_memory=True
+        )
+        return train_loader, val_loader
+
+    # ------------------------------------------------------------------
+    def dice_score(self, preds: torch.Tensor, targets: torch.Tensor) -> torch.Tensor:
+        preds = (preds > 0.5).float()
+        intersection = (preds * targets).sum(dim=(2, 3))
+        union = preds.sum(dim=(2, 3)) + targets.sum(dim=(2, 3))
+        dice = (2 * intersection + 1e-6) / (union + 1e-6)
+        return dice.mean(dim=0)
+
+    def train(
+        self,
+        epochs: int = 40,
+        batch_size: int = 4,
+        lr: float = 1e-3,
+        weight_decay: float = 1e-5,
+        checkpoint_dir: Path = Path("models/unet_segmenter"),
+    ) -> None:
+        print(
+            f"[UNetSegmenter] training on device={self.device} "
+            f"(epochs={epochs}, batch_size={batch_size}, workers={self.loader_workers})"
+        )
+        train_loader, val_loader = self.build_loaders(batch_size=batch_size, num_workers=self.loader_workers)
+        optimizer = torch.optim.Adam(
+            self.model.parameters(), lr=lr, weight_decay=weight_decay
+        )
+        criterion = nn.BCEWithLogitsLoss()
+        best_dice = -math.inf
+        checkpoint_dir.mkdir(parents=True, exist_ok=True)
+        best_path = checkpoint_dir / "best.pt"
+
+        epoch_bar = tqdm(range(1, epochs + 1), desc="Epochs", unit="epoch")
+
+        for epoch in epoch_bar:
+            self.model.train()
+            batch_bar = tqdm(
+                train_loader,
+                desc=f"Train {epoch}/{epochs}",
+                leave=False,
+                unit="batch",
+                total=len(train_loader),
+            )
+            train_loss_total = 0.0
+            train_samples = 0
+            for images, masks in batch_bar:
+                images = images.to(self.device)
+                masks = masks.to(self.device)
+                optimizer.zero_grad()
+                logits = self.model(images)
+                loss_disc = criterion(logits[:, 0:1], masks[:, 0:1])
+                loss_cup = criterion(logits[:, 1:2], masks[:, 1:2])
+                loss = self.disc_weight * loss_disc + self.cup_weight * loss_cup
+                loss.backward()
+                optimizer.step()
+                batch_size = images.size(0)
+                train_loss_total += loss.item() * batch_size
+                train_samples += batch_size
+
+            train_loss = (
+                train_loss_total / train_samples if train_samples else float("nan")
+            )
+
+            self.model.eval()
+            dices = []
+            val_bar = tqdm(
+                val_loader,
+                desc="Validate",
+                leave=False,
+                unit="batch",
+                total=len(val_loader),
+            )
+            with torch.no_grad():
+                for images, masks in val_bar:
+                    images = images.to(self.device)
+                    masks = masks.to(self.device)
+                    logits = self.model(images)
+                    probs = torch.sigmoid(logits)
+                    dice = self.dice_score(probs, masks)
+                    dices.append(dice.cpu())
+            if dices:
+                mean_dice = torch.stack(dices).mean(dim=0)
+                disc_dice = mean_dice[0].item()
+                cup_dice = mean_dice[1].item()
+                weight_sum = self.disc_weight + self.cup_weight
+                score = (
+                    (self.disc_weight * disc_dice + self.cup_weight * cup_dice)
+                    / weight_sum
+                    if weight_sum
+                    else 0.0
+                )
+                epoch_bar.set_postfix(
+                    loss=f"{train_loss:.4f}",
+                    dice_disc=f"{disc_dice:.3f}",
+                    dice_cup=f"{cup_dice:.3f}",
+                    dice_w=f"{score:.3f}",
+                )
+            else:
+                disc_dice = cup_dice = 0.0
+                score = 0.0
+                epoch_bar.set_postfix(loss=f"{train_loss:.4f}")
+
+            if score > best_dice:
+                best_dice = score
+                torch.save({"model": self.model.state_dict()}, best_path)
+
+        if best_path.exists():
+            state = torch.load(best_path, map_location=self.device)
+            self.model.load_state_dict(state["model"])
+
+    # ------------------------------------------------------------------
+    def evaluate_holdout(
+        self, output_dir: Path = Path("analysis_data/segmenter_eval")
+    ) -> pd.DataFrame:
+        return self.evaluate_dataset(split_filter={"holdout"}, output_dir=output_dir)
+
+    @staticmethod
+    def overlay_masks(
+        image: Image.Image, disc: np.ndarray, cup: np.ndarray
+    ) -> Image.Image:
+        overlay = image.copy()
+        disc_img = Image.fromarray((disc * 255).astype(np.uint8))
+        cup_img = Image.fromarray((cup * 255).astype(np.uint8))
+        disc_color = Image.new("RGBA", image.size, (255, 0, 0, 0))
+        cup_color = Image.new("RGBA", image.size, (0, 255, 0, 0))
+        disc_color.paste((255, 0, 0, 100), mask=disc_img)
+        cup_color.paste((0, 255, 0, 100), mask=cup_img)
+        overlay = overlay.convert("RGBA")
+        overlay = Image.alpha_composite(overlay, disc_color)
+        overlay = Image.alpha_composite(overlay, cup_color)
+        return overlay.convert("RGB")
+
+    # ------------------------------------------------------------------
+    @staticmethod
+    def _normalize_filter(values: Optional[Iterable[str]]) -> Optional[Set[str]]:
+        if values is None:
+            return None
+        if isinstance(values, str):
+            return {values}
+        return {str(item) for item in values}
+
+    @staticmethod
+    def _dice_from_masks(pred: np.ndarray, target: np.ndarray) -> float:
+        pred = (pred > 0).astype(np.float32)
+        target = (target > 0).astype(np.float32)
+        intersection = float((pred * target).sum())
+        denom = float(pred.sum() + target.sum())
+        return (2.0 * intersection + 1e-6) / (denom + 1e-6)
+
+    def get_entries(
+        self,
+        dataset_filter: Optional[Iterable[str]] = None,
+        split_filter: Optional[Iterable[str]] = None,
+    ) -> List[ManifestEntry]:
+        dataset_set = self._normalize_filter(dataset_filter)
+        split_set = self._normalize_filter(split_filter)
+        entries = self._manifest
+        if dataset_set is not None:
+            entries = [e for e in entries if e.dataset in dataset_set]
+        if split_set is not None:
+            entries = [e for e in entries if e.split in split_set]
+        return list(entries)
+
+    def evaluate_dataset(
+        self,
+        dataset_filter: Optional[Iterable[str]] = None,
+        split_filter: Optional[Iterable[str]] = None,
+        output_dir: Path = Path("analysis_data/segmenter_eval"),
+        save_overlays: bool = True,
+        metrics_path: Optional[Path] = None,
+        threshold: float = 0.5,
+        tta: bool = False,
+    ) -> pd.DataFrame:
+        entries = self.get_entries(
+            dataset_filter=dataset_filter, split_filter=split_filter
+        )
+        if not entries:
+            return pd.DataFrame(
+                columns=[
+                    "sample_id",
+                    "dataset",
+                    "split",
+                    "dice_disc",
+                    "dice_cup",
+                ]
+            )
+
+        output_dir.mkdir(parents=True, exist_ok=True)
+        if metrics_path is None:
+            suffix_parts = []
+            if dataset_filter is not None:
+                suffix_parts.append("-".join(sorted(self._normalize_filter(dataset_filter))))
+            if split_filter is not None:
+                suffix_parts.append("-".join(sorted(self._normalize_filter(split_filter))))
+            suffix = "_".join(part for part in suffix_parts if part)
+            csv_name = f"metrics{'_' + suffix if suffix else ''}.csv"
+            metrics_path = output_dir / csv_name
+
+        records = []
+        self.model.eval()
+        progress = tqdm(
+            entries,
+            desc="Evaluate",
+            unit="sample",
+            leave=False,
+        )
+        for entry in progress:
+            orig_image = Image.open(entry.image_path).convert("RGB")
+            image = self.preprocess_image(orig_image)
+            tensor = transforms.ToTensor()(image)
+            tensor = self._normalize_tensor(tensor)
+            tensor = tensor.unsqueeze(0).to(self.device)
+            with torch.no_grad():
+                logits = self.model(tensor)
+                if tta:
+                    t_h = torch.flip(tensor, dims=[3])
+                    log_h = self.model(t_h)
+                    log_h = torch.flip(log_h, dims=[3])
+                    t_v = torch.flip(tensor, dims=[2])
+                    log_v = self.model(t_v)
+                    log_v = torch.flip(log_v, dims=[2])
+                    logits = (logits + log_h + log_v) / 3.0
+                probs = torch.sigmoid(logits)[0].cpu().numpy()
+
+            disc_pred = (probs[0] > threshold).astype(np.uint8)
+            cup_pred = (probs[1] > threshold).astype(np.uint8)
+            # Structural prior: cup within disc
+            cup_pred = (cup_pred > 0) & (disc_pred > 0)
+            cup_pred = cup_pred.astype(np.uint8)
+
+            disc_gt, cup_gt = self.load_masks(entry)
+            disc_gt = disc_gt.astype(np.uint8)
+            cup_gt = cup_gt.astype(np.uint8)
+
+            dice_disc = self._dice_from_masks(disc_pred, disc_gt)
+            dice_cup = self._dice_from_masks(cup_pred, cup_gt)
+
+            records.append(
+                {
+                    "sample_id": entry.sample_id,
+                    "dataset": entry.dataset,
+                    "split": entry.split,
+                    "dice_disc": dice_disc,
+                    "dice_cup": dice_cup,
+                }
+            )
+
+            progress.set_postfix(
+                dice_disc=f"{dice_disc:.3f}", dice_cup=f"{dice_cup:.3f}"
+            )
+
+            if save_overlays:
+                overlay_gt = self.overlay_masks(image, disc_gt, cup_gt)
+                overlay_pred = self.overlay_masks(image, disc_pred, cup_pred)
+                combined = Image.new("RGB", (image.width * 2, image.height))
+                combined.paste(overlay_gt, (0, 0))
+                combined.paste(overlay_pred, (image.width, 0))
+                combined.save(output_dir / f"{entry.sample_id}_eval.png")
+
+        metrics_df = pd.DataFrame(records)
+        summary = metrics_df[["dice_disc", "dice_cup"]].mean()
+        summary_row = {
+            "sample_id": "__mean__",
+            "dataset": "summary",
+            "split": "summary",
+            "dice_disc": summary["dice_disc"],
+            "dice_cup": summary["dice_cup"],
+        }
+        metrics_with_summary = pd.concat(
+            [metrics_df, pd.DataFrame([summary_row])], ignore_index=True
+        )
+        metrics_with_summary.to_csv(metrics_path, index=False)
+        return metrics_with_summary

classes/v2/__init__.py

@@ -0,0 +1,102 @@
+from .network_manager import (
+    FoldResult,
+    LoaderBundle,
+    NetworkManager,
+    PatientSplit,
+)
+from .split_manager import (
+    PatientFirstSplitManager,
+    SplitPlan,
+    build_patient_split_plans,
+)
+from .profiles import (
+    DatasetProfile,
+    SimpleDatasetProfile,
+    SlotDescriptor,
+    PapilaProfile,
+    build_papila_profile,
+)
+from .loader_factory import SlotLoaderFactory
+from .slot_dataset import SlotDataset, slot_collate
+from .papila_data import PapilaData
+from .papila_builders import build_papila_data
+from .data_bundle import DataBundle
+from .dataset import ClinicalDataset
+from .config_builder import (
+    ConfigAssembly,
+    assemble_config,
+    load_config,
+    resolve_imports,
+)
+from .filters import RegexFilter, ColumnFilter, apply_regex_filters, apply_column_filters
+from .transforms import (
+    ImageTransformConfig,
+    backbone_transform_config,
+    build_backbone_transform,
+    build_imagenet_transform,
+    ResizeTransform,
+    CenterCropTransform,
+    ROICropTransform,
+    JitterBundleTransform,
+    UnetMaskProvider,
+    TRANSFORM_REGISTRY,
+    build_transform_chain,
+)
+from .model_builder import V2ModelBundle, build_model_bundle
+from .towers import ImageTower, MDTower, SiameseImageTower, build_backbone
+from .bridges import Bridge, VoteBridge
+from .v2_hypertower import V2HyperTower, V2ModeComparisonOps, V2ModeComparator
+from .hypertower_logger import HypertowerLogger
+
+__all__ = [
+    "NetworkManager",
+    "PatientSplit",
+    "LoaderBundle",
+    "FoldResult",
+    "PatientFirstSplitManager",
+    "SplitPlan",
+    "build_patient_split_plans",
+    "DatasetProfile",
+    "SimpleDatasetProfile",
+    "SlotDescriptor",
+    "PapilaProfile",
+    "build_papila_profile",
+    "PapilaData",
+    "build_papila_data",
+    "DataBundle",
+    "ClinicalDataset",
+    "SlotLoaderFactory",
+    "SlotDataset",
+    "slot_collate",
+    "ConfigAssembly",
+    "assemble_config",
+    "load_config",
+    "resolve_imports",
+    "RegexFilter",
+    "ColumnFilter",
+    "apply_regex_filters",
+    "apply_column_filters",
+    "ImageTransformConfig",
+    "backbone_transform_config",
+    "build_backbone_transform",
+    "build_imagenet_transform",
+    "ResizeTransform",
+    "CenterCropTransform",
+    "ROICropTransform",
+    "JitterBundleTransform",
+    "UnetMaskProvider",
+    "TRANSFORM_REGISTRY",
+    "build_transform_chain",
+    "V2ModelBundle",
+    "build_model_bundle",
+    "ImageTower",
+    "MDTower",
+    "SiameseImageTower",
+    "build_backbone",
+    "Bridge",
+    "VoteBridge",
+    "V2HyperTower",
+    "V2ModeComparisonOps",
+    "V2ModeComparator",
+    "HypertowerLogger",
+]

classes/v2/bridges.py

@@ -0,0 +1,93 @@
+from __future__ import annotations
+
+import torch
+import torch.nn as nn
+
+from classes.SE_attention import SEBlock, SEGateLogger
+
+
+class Bridge(nn.Module):
+    def __init__(
+        self,
+        img_dim,
+        meta_dim,
+        num_classes,
+        fusion_dim=256,
+        mode="fused",
+        use_se: bool = True,
+        se_reduction: int = 16,
+        se_pre_norm: bool = True,
+    ):
+        super().__init__()
+        self.mode = mode
+        self.use_se = use_se
+
+        # project towers to equal width
+        self.W_img = nn.Linear(img_dim, fusion_dim)
+        self.W_md = nn.Linear(meta_dim, fusion_dim)
+
+        # optional: layernorm before SE
+        self.ln_img = nn.LayerNorm(fusion_dim) if se_pre_norm else nn.Identity()
+        self.ln_md = nn.LayerNorm(fusion_dim) if se_pre_norm else nn.Identity()
+
+        # SE gate on the fused vector
+        self.se = SEBlock(fusion_dim, reduction=se_reduction, residual=True) if use_se else None
+        self.se_log = SEGateLogger(enabled=use_se, track_channels=False, dim=fusion_dim)
+
+        # heads
+        self.classifier_fused = nn.Sequential(
+            nn.ReLU(),
+            nn.Dropout(0.5),
+            nn.Linear(fusion_dim, num_classes),
+        )
+        self.classifier_img = nn.Linear(img_dim, num_classes)
+        self.classifier_md = nn.Linear(meta_dim, num_classes)
+
+    def reset_se_stats(self):
+        """Call at epoch start."""
+        if getattr(self, "se_log", None):
+            self.se_log.reset()
+
+    def get_se_stats(self, reset: bool = True):
+        """Call after eval. Returns dict or None."""
+        if getattr(self, "se_log", None) and self.se_log.enabled:
+            return self.se_log.get(reset=reset)
+        return None
+
+    def forward(self, img_feats, md_feats):
+        out_img = None if self.mode == "metadata_only" else self.classifier_img(img_feats)
+        out_md = None if self.mode == "image_only" else self.classifier_md(md_feats)
+
+        if self.mode == "fused":
+            hi = self.ln_img(self.W_img(img_feats))  # image features
+            hm = self.ln_md(self.W_md(md_feats))  # metadata features
+            fused = hi * hm  # elementwise product
+            # apply SE gates
+            if self.se is not None:
+                fused, gates = self.se(fused)
+                if self.se_log.enabled:
+                    self.se_log.accumulate(gates)
+
+            if self.se is not None and self.training and self.se_log.enabled:
+                if not hasattr(self, "_dbg_seen"):
+                    self._dbg_seen = 0
+                if self._dbg_seen < 3:  # print only a few times
+                    print("[SE] gate mean this batch:", gates.mean().item())
+                    self._dbg_seen += 1
+            out_f = self.classifier_fused(fused)
+            return out_f, out_img, out_md
+        # if ablation modes:
+        if self.mode == "image_only":
+            return out_img, out_img, None
+        if self.mode == "metadata_only":
+            return out_md, None, out_md
+
+
+class VoteBridge(nn.Module):
+    def __init__(self, num_classes):
+        super().__init__()
+        self.vote_combiner = nn.Linear(num_classes * 2, num_classes)  # two sets of logits
+
+    def forward(self, out_img, out_md):
+        votes = torch.cat([out_img, out_md], dim=1)
+        return self.vote_combiner(votes)

classes/v2/config_builder.py

@@ -0,0 +1,276 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any, Dict, Iterable, List, Optional
+
+import json
+
+from classes.v2.papila_data import PapilaData
+
+
+@dataclass
+class ImportSpec:
+    id: str
+    class_name: str
+    params: Dict[str, Any]
+
+
+@dataclass
+class DataSourceSpec:
+    node_id: str
+    label: str
+    output_type: str
+    source: Optional[Dict[str, Any]]
+    source_ref: Optional[Dict[str, Any]]
+
+
+@dataclass
+class TransformSpec:
+    node_id: str
+    label: str
+    transform_type: str
+    params: Dict[str, Any]
+
+
+@dataclass
+class LoaderSpec:
+    node_id: str
+    label: str
+    input_type: str
+    input_index: str
+    input_key: str
+    output_key: str
+    transforms: List[TransformSpec]
+    data_source: Optional[DataSourceSpec]
+
+
+@dataclass
+class TowerSpec:
+    node_id: str
+    label: str
+    tower_type: str
+    params: Dict[str, Any]
+
+
+@dataclass
+class BridgeSpec:
+    node_id: str
+    label: str
+    method: str
+    params: Dict[str, Any]
+
+
+@dataclass
+class ClassifierSpec:
+    node_id: str
+    label: str
+
+
+@dataclass
+class ConfigAssembly:
+    raw: Dict[str, Any]
+    imports: Dict[str, ImportSpec]
+    data_sources: Dict[str, DataSourceSpec]
+    transforms: Dict[str, TransformSpec]
+    loaders: Dict[str, LoaderSpec]
+    towers: Dict[str, TowerSpec]
+    bridges: Dict[str, BridgeSpec]
+    classifiers: Dict[str, ClassifierSpec]
+
+
+def load_config(path: Path) -> Dict[str, Any]:
+    payload = json.loads(Path(path).read_text())
+    if not isinstance(payload, dict):
+        raise ValueError("Config JSON must be an object.")
+    return payload
+
+
+def assemble_config(path: Path) -> ConfigAssembly:
+    config = load_config(path)
+    meta = config.get("meta", {})
+    imports = _build_imports(meta.get("imports", []))
+    nodes = {node["id"]: node for node in config.get("nodes", [])}
+    edges = config.get("edges", [])
+
+    data_sources: Dict[str, DataSourceSpec] = {}
+    transforms: Dict[str, TransformSpec] = {}
+    loaders: Dict[str, LoaderSpec] = {}
+    towers: Dict[str, TowerSpec] = {}
+    bridges: Dict[str, BridgeSpec] = {}
+    classifiers: Dict[str, ClassifierSpec] = {}
+
+    for node in nodes.values():
+        ntype = node.get("type")
+        if ntype == "data":
+            data_sources[node["id"]] = DataSourceSpec(
+                node_id=node["id"],
+                label=node.get("label", ""),
+                output_type=node.get("outputType", ""),
+                source=node.get("source"),
+                source_ref=node.get("sourceRef"),
+            )
+        elif ntype == "transform":
+            transforms[node["id"]] = TransformSpec(
+                node_id=node["id"],
+                label=node.get("label", ""),
+                transform_type=node.get("transformType", ""),
+                params=_extract_transform_params(node),
+            )
+        elif ntype == "loader":
+            loaders[node["id"]] = LoaderSpec(
+                node_id=node["id"],
+                label=node.get("label", ""),
+                input_type=node.get("inputType", ""),
+                input_index=node.get("inputIndex", ""),
+                input_key=node.get("inputKey", ""),
+                output_key=node.get("outputKey", ""),
+                transforms=[],
+                data_source=None,
+            )
+        elif ntype in ("image_tower", "metadata_tower"):
+            towers[node["id"]] = TowerSpec(
+                node_id=node["id"],
+                label=node.get("label", ""),
+                tower_type=node.get("towerType", "image" if ntype == "image_tower" else "metadata"),
+                params=_extract_tower_params(node),
+            )
+        elif ntype == "bridge":
+            bridges[node["id"]] = BridgeSpec(
+                node_id=node["id"],
+                label=node.get("label", ""),
+                method=node.get("bridgeMethod", "fusion"),
+                params=_extract_bridge_params(node),
+            )
+        elif ntype == "classifier":
+            classifiers[node["id"]] = ClassifierSpec(
+                node_id=node["id"],
+                label=node.get("label", ""),
+            )
+
+    # attach transforms + data sources to loaders by walking upstream
+    for loader_id, loader in loaders.items():
+        chain = _upstream_chain(loader_id, nodes, edges)
+        for node_id in reversed(chain):
+            if node_id in transforms:
+                loader.transforms.append(transforms[node_id])
+            if node_id in data_sources:
+                loader.data_source = data_sources[node_id]
+
+    return ConfigAssembly(
+        raw=config,
+        imports=imports,
+        data_sources=data_sources,
+        transforms=transforms,
+        loaders=loaders,
+        towers=towers,
+        bridges=bridges,
+        classifiers=classifiers,
+    )
+
+
+def resolve_imports(assembly: ConfigAssembly) -> Dict[str, Any]:
+    resolved: Dict[str, Any] = {}
+    for import_id, spec in assembly.imports.items():
+        if spec.class_name == "PapilaData":
+            params = spec.params
+            resolved[import_id] = PapilaData.from_dirs(
+                image_dir=params.get("image_dir", "Papila/FundusImages"),
+                clinical_dir=params.get("clinical_dir", "Papila/ClinicalData"),
+                label_col=params.get("label_col", "Diagnosis"),
+                cat_cols=params.get("cat_cols", ["Gender", "Phakic/Pseudophakic"]),
+            )
+        else:
+            raise ValueError(f"Unsupported import class {spec.class_name!r}")
+    return resolved
+
+
+def _build_imports(entries: Iterable[Dict[str, Any]]) -> Dict[str, ImportSpec]:
+    specs: Dict[str, ImportSpec] = {}
+    for entry in entries or []:
+        import_id = entry.get("id")
+        if not import_id:
+            continue
+        specs[import_id] = ImportSpec(
+            id=import_id,
+            class_name=entry.get("className", ""),
+            params=entry.get("params", {}) or {},
+        )
+    return specs
+
+
+def _extract_transform_params(node: Dict[str, Any]) -> Dict[str, Any]:
+    return {
+        "transformType": node.get("transformType"),
+        "roiMaskSource": node.get("roiMaskSource"),
+        "roiScale": node.get("roiScale"),
+        "roiTargetSize": node.get("roiTargetSize"),
+        "roiFallback": node.get("roiFallback"),
+        "centerCropSize": node.get("centerCropSize"),
+        "jitterHFlip": node.get("jitterHFlip"),
+        "jitterVFlip": node.get("jitterVFlip"),
+        "jitterRotation": node.get("jitterRotation"),
+        "jitterColorEnabled": node.get("jitterColorEnabled"),
+        "jitterColor": node.get("jitterColor"),
+        "resizeSize": node.get("resizeSize"),
+    }
+
+
+def _extract_tower_params(node: Dict[str, Any]) -> Dict[str, Any]:
+    if node.get("towerType") == "metadata":
+        return {
+            "hidden_dim": node.get("mdHiddenDim"),
+            "dropout": node.get("mdDropout"),
+            "use_se": node.get("mdUseSe"),
+            "se_reduction": node.get("mdSeReduction"),
+            "se_pre_norm": node.get("mdSePreNorm"),
+            "freeze_ratio": node.get("mdFreezeRatio"),
+        }
+    return {
+        "backbone": node.get("imageBackbone"),
+        "freeze_ratio": node.get("imageFreezeRatio"),
+        "augment": node.get("imageAugment"),
+        "geometry_dim": node.get("imageGeometryDim"),
+        "use_se": node.get("imageUseSe"),
+        "se_reduction": node.get("imageSeReduction"),
+        "se_pre_norm": node.get("imageSePreNorm"),
+    }
+
+
+def _extract_bridge_params(node: Dict[str, Any]) -> Dict[str, Any]:
+    return {
+        "fusion_dim": node.get("bridgeFusionDim"),
+        "use_se": node.get("bridgeUseSe"),
+        "se_reduction": node.get("bridgeSeReduction"),
+        "se_pre_norm": node.get("bridgeSePreNorm"),
+    }
+
+
+def _edge_from(edge: Dict[str, Any]) -> Optional[str]:
+    return edge.get("from") or edge.get("source")
+
+
+def _edge_to(edge: Dict[str, Any]) -> Optional[str]:
+    return edge.get("to") or edge.get("target")
+
+
+def _upstream_chain(start_id: str, nodes: Dict[str, Dict[str, Any]], edges: List[Dict[str, Any]]) -> List[str]:
+    chain: List[str] = []
+    visited = set()
+    current = start_id
+    while True:
+        if current in visited:
+            break
+        visited.add(current)
+        incoming = [edge for edge in edges if _edge_to(edge) == current]
+        if not incoming:
+            break
+        # prefer first incoming edge for now
+        current = _edge_from(incoming[0])
+        if not current:
+            break
+        chain.append(current)
+        node = nodes.get(current)
+        if node and node.get("type") == "data":
+            break
+    return chain

classes/v2/data_bundle.py

@@ -0,0 +1,241 @@
+from __future__ import annotations
+
+from pathlib import Path
+from typing import Callable, Dict, Iterable, List, Optional, Tuple
+
+import numpy as np
+import pandas as pd
+
+
+class DataBundle:
+    """
+    Generic, torch-free container for metadata and file/label bookkeeping.
+
+    Keeps feature typing, vectorization, and patient-level splits generic.
+    Dataset-specific preprocessing (e.g., eye canonicalization) should live
+    in the dataset builder (e.g., papila_builders in v2).
+    """
+
+    def __init__(
+        self,
+        *,
+        image_dir: str,
+        clinical_dir: Optional[str] = None,
+        label_col: str,
+        patient_col: str = "Patient ID",
+        cat_cols: Optional[Iterable[str]] = None,
+        max_unique_for_cat: int = 4,
+        n_splits: int = 5,
+        random_seed: int = 42,
+        filename_template: str = "RET{pid:03d}{eye}.jpg",
+        image_path_fn: Optional[Callable[[pd.Series], Path]] = None,
+    ) -> None:
+        self.image_dir = Path(image_dir)
+        self.label_col = label_col
+        self.patient_col = patient_col
+        self.max_unique_for_cat = max_unique_for_cat
+        self.n_splits = n_splits
+        self.filename_template = filename_template
+        self.image_path_fn = image_path_fn
+        self.clinical_dir = Path(clinical_dir) if clinical_dir else None
+
+        # Internal state
+        self.frames: List[pd.DataFrame] = []
+        self.df: pd.DataFrame = pd.DataFrame()
+        self.scalar_cols: List[str] = []
+        self.cat_cols: List[str] = list(cat_cols) if cat_cols is not None else []
+        self.scalar_stats: Dict[str, Dict[str, float]] = {}
+        self.cat_maps: Dict[str, Dict[object, int]] = {}
+        self.feature_dim: int = 0
+        self.folds: Dict[int, Dict[str, List[object]]] = {}
+        self.random_seed = int(random_seed)
+
+    # ------------------- Public API -------------------
+    def add_df(
+        self,
+        df: pd.DataFrame,
+        *,
+        id_column: Optional[str] = None,
+        exclude_cols: Optional[Iterable[str]] = None,
+    ) -> None:
+        """
+        Add a dataframe and re-run typing, stats, and K-fold indices.
+        QC rules:
+          - Must have patient ID column; if not provided under that name, specify id_column.
+        """
+        df = df.copy()
+        self._ensure_patient_id(df, id_column)
+        if self.label_col not in df.columns:
+            raise ValueError(f"label_col '{self.label_col}' not found in added dataframe")
+
+        self.frames.append(df)
+        self._refresh_master_df(exclude_cols=exclude_cols)
+        self._infer_or_validate_feature_types(exclude_cols=exclude_cols)
+        self._compute_numeric_stats()
+        self._build_cat_maps()
+        self._compute_feature_dim()
+        self._build_kfold_indices()
+
+    def get_split_ids(self, fold: int) -> Tuple[List[object], List[object]]:
+        rec = self.folds.get(fold)
+        if not rec:
+            raise KeyError(f"Fold {fold} not available. Built folds: {sorted(self.folds.keys())}")
+        return rec["train_ids"], rec["test_ids"]
+
+    def get_split_dfs(self, fold: int) -> Tuple[pd.DataFrame, pd.DataFrame]:
+        train_ids, test_ids = self.get_split_ids(fold)
+        train_df = self.df[self.df[self.patient_col].isin(train_ids)].reset_index(drop=True)
+        test_df = self.df[self.df[self.patient_col].isin(test_ids)].reset_index(drop=True)
+        return train_df, test_df
+
+    def vectorize_row(self, row: pd.Series) -> np.ndarray:
+        """Return a numpy feature vector (torch-free)."""
+        feats: List[float] = []
+        miss: List[float] = []
+        # numeric
+        for col in self.scalar_cols:
+            v = pd.to_numeric(row.get(col), errors="coerce")
+            if pd.isna(v):
+                miss.append(1.0)
+                v = self.scalar_stats[col]["median"]
+            else:
+                miss.append(0.0)
+            lo = self.scalar_stats[col]["min"]
+            hi = self.scalar_stats[col]["max"]
+            feats.append((float(v) - lo) / (hi - lo) if hi > lo else 0.0)
+        # categorical
+        for col in self.cat_cols:
+            mapping = self.cat_maps[col]
+            one = [0.0] * len(mapping)
+            key = row.get(col)
+            one[mapping.get(key, 0)] = 1.0  # 0 is <UNK>
+            feats.extend(one)
+        # numeric missing flags
+        feats.extend(miss)
+        return np.asarray(feats, dtype=np.float32)
+
+    def get_image_path(self, row: pd.Series) -> Path:
+        if self.image_path_fn is not None:
+            return Path(self.image_path_fn(row))
+        pid = int(row[self.patient_col])
+        eye = row.get("eyeID", "")
+        if eye in ("OS", "OD"):
+            eye_str = eye
+        else:
+            eye_str = str(eye)
+        return self.image_dir / self.filename_template.format(pid=pid, eye=eye_str)
+
+    def encode_metadata(self, row: pd.Series) -> np.ndarray:
+        return self.vectorize_row(row)
+
+    def get_label(self, row: pd.Series) -> int:
+        return int(row[self.label_col])
+
+    # ------------------- Internal helpers -------------------
+    def _ensure_patient_id(self, df: pd.DataFrame, id_column: Optional[str]) -> None:
+        if self.patient_col in df.columns:
+            return
+        if id_column and id_column in df.columns:
+            df.rename(columns={id_column: self.patient_col}, inplace=True)
+            return
+        candidates = [
+            c
+            for c in df.columns
+            if c.lower().replace(" ", "") in {"patientid", "patient", "pid"}
+        ]
+        if len(candidates) == 1:
+            df.rename(columns={candidates[0]: self.patient_col}, inplace=True)
+            return
+        raise ValueError(
+            f"A '{self.patient_col}' column is required; provide id_column=... if it has a different name."
+        )
+
+    def _refresh_master_df(self, exclude_cols: Optional[Iterable[str]] = None) -> None:
+        self.df = pd.concat(self.frames, axis=0, ignore_index=True)
+        if exclude_cols:
+            self.df = self.df.drop(columns=[c for c in exclude_cols if c in self.df.columns])
+
+    def _infer_or_validate_feature_types(self, exclude_cols: Optional[Iterable[str]] = None) -> None:
+        excluded = set(exclude_cols or []) | {self.label_col, self.patient_col}
+        feature_candidates = [c for c in self.df.columns if c not in excluded]
+        cats = set(self.cat_cols) if self.cat_cols else set()
+        scalars = set()
+        for c in feature_candidates:
+            if c in cats:
+                continue
+            s = self.df[c]
+            as_num = pd.to_numeric(s, errors="coerce")
+            num_missing = as_num.isna().mean()
+            num_unique = s.dropna().nunique()
+            if as_num.notna().any() and num_missing < 1.0 and num_unique > self.max_unique_for_cat:
+                scalars.add(c)
+            else:
+                if num_unique <= self.max_unique_for_cat or as_num.isna().mean() > 0.0:
+                    cats.add(c)
+                else:
+                    scalars.add(c)
+        self.cat_cols = sorted(cats)
+        self.scalar_cols = sorted(scalars)
+
+    def _compute_numeric_stats(self) -> None:
+        self.scalar_stats.clear()
+        for col in self.scalar_cols:
+            s = pd.to_numeric(self.df[col], errors="coerce")
+            vals = s.dropna().astype(float).values
+            if vals.size == 0:
+                lo, hi, med = 0.0, 1.0, 0.0
+            else:
+                lo, hi = float(np.min(vals)), float(np.max(vals))
+                med = float(np.median(vals))
+                if hi <= lo:
+                    hi = lo + 1.0
+            self.scalar_stats[col] = {"min": lo, "max": hi, "median": med}
+
+    def _build_cat_maps(self) -> None:
+        self.cat_maps.clear()
+        for col in self.cat_cols:
+            cats = [v for v in self.df[col].dropna().unique().tolist()]
+            try:
+                cats = sorted(cats)
+            except Exception:
+                pass
+            mapping = {"<UNK>": 0}
+            for i, v in enumerate(cats, start=1):
+                mapping[v] = i
+            self.cat_maps[col] = mapping
+
+    def _compute_feature_dim(self) -> None:
+        self.feature_dim = len(self.scalar_cols) + sum(len(m) for m in self.cat_maps.values()) + len(self.scalar_cols)
+
+    # ------------------- K-fold on unique patients -------------------
+    def _build_kfold_indices(self) -> None:
+        pats = self.df[self.patient_col].unique().tolist()
+        labels_by_pat: Dict[object, object] = {}
+        for pid, grp in self.df.groupby(self.patient_col):
+            lab = grp[self.label_col].dropna()
+            if len(lab) == 0:
+                labels_by_pat[pid] = 0
+            else:
+                labels_by_pat[pid] = lab.mode().iloc[0]
+        y_pat = np.array([labels_by_pat[p] for p in pats])
+
+        try:
+            from sklearn.model_selection import StratifiedGroupKFold
+
+            sgkf = StratifiedGroupKFold(
+                n_splits=self.n_splits, shuffle=True, random_state=self.random_seed
+            )
+            split_iter = sgkf.split(X=pats, y=y_pat, groups=pats)
+        except Exception:
+            from sklearn.model_selection import StratifiedKFold
+
+            skf = StratifiedKFold(
+                n_splits=self.n_splits, shuffle=True, random_state=self.random_seed
+            )
+            split_iter = skf.split(X=np.zeros(len(pats)), y=y_pat)
+
+        self.folds.clear()
+        for i, (train_idx, test_idx) in enumerate(split_iter):
+            train_ids = [pats[j] for j in train_idx]
+            test_ids = [pats[j] for j in test_idx]
+            self.folds[i] = {"train_ids": train_ids, "test_ids": test_ids}

classes/v2/dataset.py

@@ -0,0 +1,57 @@
+from torch.utils.data import Dataset
+from PIL import Image
+import numpy as np
+import torch
+
+
+class ClinicalDataset(Dataset):
+    """Generic dataset wrapping a DataBundle-like instance.
+    Returns (img_tensor, meta_tensor, label)."""
+
+    def __init__(
+        self,
+        clinical_data,
+        img_transform,
+        meta_transform=None,
+        image_preprocessor=None,
+        geometry_provider=None,
+        geometry_dim: int = 0,
+    ):
+        self.clinical = clinical_data
+        self.transform_image = img_transform
+        self.meta_transform = meta_transform or (lambda x: x)
+        self.image_preprocessor = image_preprocessor
+        self.geometry_provider = geometry_provider
+        self.geometry_dim = geometry_dim if geometry_provider is not None else 0
+
+    def __len__(self):
+        return len(self.clinical.df)
+
+    def __getitem__(self, idx: int):
+        row = self.clinical.df.iloc[idx]
+        # load & transform image
+        img_path = self.clinical.get_image_path(row)
+        orig_img = Image.open(img_path).convert("RGB")
+        img = orig_img
+        if self.image_preprocessor is not None:
+            img = self.image_preprocessor(img, img_path)
+        img_t = self.transform_image(img)
+        # encode & transform metadata
+        meta = self.clinical.encode_metadata(row)
+        meta_t = self.meta_transform(meta)
+        # label
+        label = self.clinical.get_label(row)
+        if self.geometry_dim > 0:
+            features = None
+            if self.geometry_provider is not None and hasattr(self.geometry_provider, "geometry_features"):
+                features = self.geometry_provider.geometry_features(orig_img, img_path)
+            if features is None:
+                geom_vec = torch.zeros(self.geometry_dim, dtype=torch.float32)
+            else:
+                features = np.asarray(features, dtype=np.float32)
+                if features.shape[0] != self.geometry_dim:
+                    geom_vec = torch.zeros(self.geometry_dim, dtype=torch.float32)
+                else:
+                    geom_vec = torch.from_numpy(features)
+            return img_t, meta_t, geom_vec, label
+        return img_t, meta_t, label

classes/v2/filters.py

@@ -0,0 +1,119 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Iterable, List, Sequence, Tuple, Union
+import re
+
+import pandas as pd
+
+
+@dataclass
+class RegexFilter:
+    pattern: str
+    flags: int = 0
+
+    def apply_paths(self, paths: Sequence[str]) -> Tuple[List[str], List[str]]:
+        if not self.pattern:
+            return list(paths), []
+        try:
+            regex = re.compile(self.pattern, self.flags)
+        except re.error as err:
+            return list(paths), [f'Invalid regex "{self.pattern}": {err}']
+        filtered = [p for p in paths if regex.search(p)]
+        return filtered, []
+
+
+@dataclass
+class ColumnFilter:
+    column: str
+    operator: str
+    value: str
+    case_insensitive: bool = True
+
+    def apply_df(self, df: pd.DataFrame) -> Tuple[pd.DataFrame, List[str]]:
+        warnings: List[str] = []
+        if not self.column:
+            return df, ["Column filter missing column name."]
+        columns = list(df.columns)
+        col_index = _resolve_column_index(columns, self.column, warnings)
+        if col_index is None:
+            return df, warnings
+        col_name = columns[col_index]
+        if self.value is None or self.value == "":
+            return df, [f'Column filter "{self.column}" missing value.']
+        series = df[col_name]
+        mask = series.apply(
+            lambda cell: compare_cell(
+                cell, self.value, self.operator, case_insensitive=self.case_insensitive
+            )
+        )
+        return df[mask], warnings
+
+
+FilterSpec = Union[RegexFilter, ColumnFilter]
+
+
+def apply_regex_filters(paths: Sequence[str], filters: Iterable[RegexFilter]) -> Tuple[List[str], List[str]]:
+    filtered = list(paths)
+    warnings: List[str] = []
+    for filt in filters:
+        filtered, warn = filt.apply_paths(filtered)
+        warnings.extend(warn)
+    return filtered, warnings
+
+
+def apply_column_filters(df: pd.DataFrame, filters: Iterable[ColumnFilter]) -> Tuple[pd.DataFrame, List[str]]:
+    filtered = df
+    warnings: List[str] = []
+    for filt in filters:
+        filtered, warn = filt.apply_df(filtered)
+        warnings.extend(warn)
+    return filtered, warnings
+
+
+def compare_cell(cell, raw_value: str, operator: str, case_insensitive: bool = True) -> bool:
+    cell_str = "" if cell is None else str(cell).strip()
+    value_str = "" if raw_value is None else str(raw_value).strip()
+    if case_insensitive:
+        cell_str = cell_str.lower()
+        value_str = value_str.lower()
+    if operator == "=":
+        return cell_str == value_str
+    if operator == "!=":
+        return cell_str != value_str
+    cell_num = _to_float(cell_str)
+    value_num = _to_float(value_str)
+    if cell_num is None or value_num is None:
+        return False
+    if operator == ">":
+        return cell_num > value_num
+    if operator == ">=":
+        return cell_num >= value_num
+    if operator == "<":
+        return cell_num < value_num
+    if operator == "<=":
+        return cell_num <= value_num
+    return False
+
+
+def _resolve_column_index(columns: Sequence[str], column: str, warnings: List[str]) -> int | None:
+    try:
+        return columns.index(column)
+    except ValueError:
+        lower = column.lower()
+        matches = [idx for idx, col in enumerate(columns) if str(col).lower() == lower]
+        if matches:
+            if len(matches) > 1:
+                warnings.append(
+                    f'Column "{column}" matched multiple headers; using "{columns[matches[0]]}".'
+                )
+            return matches[0]
+    warnings.append(f'Column "{column}" not found.')
+    return None
+
+
+def _to_float(value: str) -> float | None:
+    try:
+        return float(value)
+    except (TypeError, ValueError):
+        return None

classes/v2/hypertower_logger.py

@@ -0,0 +1,128 @@
+from __future__ import annotations
+
+import csv
+import json
+import logging
+from pathlib import Path
+from typing import Optional
+
+
+DEFAULT_OPTIONAL_EPOCH_COLS = [
+    "pct_fused",
+    "pct_img",
+    "pct_md",
+    "phase",
+    "se_mean",
+    "se_std",
+    "se_pct_lt_0.2",
+    "se_pct_gt_0.8",
+    "holdout_loss",
+    "holdout_acc_fused",
+    "holdout_acc_img",
+    "holdout_acc_md",
+    "holdout_auc_fused",
+    "holdout_auc_img",
+    "holdout_auc_md",
+    "best_monitor",
+    "best_so_far",
+    "best_epoch",
+    "early_best_so_far",
+    "early_bad_epochs",
+    "early_improved",
+    "early_monitor",
+    "holdout_best_monitor",
+    "holdout_best_so_far",
+    "holdout_best_epoch",
+]
+
+
+class HypertowerLogger:
+    """
+    Shared logging utility for V2 tower workflows.
+    - train.log line logging
+    - epoch_log.csv row logging with stable header
+    - lightweight JSON/array artifact helpers
+    """
+
+    def __init__(
+        self,
+        *,
+        run_dir: Path,
+        train_log_path: Optional[Path] = None,
+        epoch_log_path: Optional[Path] = None,
+        logger_name: Optional[str] = None,
+    ) -> None:
+        self.run_dir = Path(run_dir).resolve()
+        self.run_dir.mkdir(parents=True, exist_ok=True)
+        self.train_log_path = Path(train_log_path) if train_log_path else (self.run_dir / "train.log")
+        self.epoch_log_path = Path(epoch_log_path) if epoch_log_path else (self.run_dir / "epoch_log.csv")
+
+        self._logger_name = logger_name or f"hypertower.{id(self)}"
+        self.logger = logging.getLogger(self._logger_name)
+        self.logger.setLevel(logging.INFO)
+        self.logger.handlers = []
+        fh = logging.FileHandler(str(self.train_log_path))
+        fh.setFormatter(logging.Formatter("%(asctime)s - %(message)s"))
+        self.logger.addHandler(fh)
+        self.logger.propagate = False
+
+        self._epoch_log_fp = None
+        self._epoch_log_writer = None
+        self._epoch_log_fields: list[str] | None = None
+
+    def info(self, msg: str) -> None:
+        self.logger.info(msg)
+
+    def warning(self, msg: str) -> None:
+        self.logger.warning(msg)
+
+    def error(self, msg: str) -> None:
+        self.logger.error(msg)
+
+    def write_epoch_row(
+        self,
+        row: dict,
+        *,
+        path: str | Path | None = None,
+        optional_cols: Optional[list[str]] = None,
+    ) -> None:
+        optional = optional_cols if optional_cols is not None else DEFAULT_OPTIONAL_EPOCH_COLS
+        if self._epoch_log_writer is None:
+            fieldnames = list(dict.fromkeys([*row.keys(), *optional]))
+            target_path = Path(path) if path is not None else self.epoch_log_path
+            target_path.parent.mkdir(parents=True, exist_ok=True)
+            self._epoch_log_fp = open(target_path, "w", newline="", encoding="utf-8")
+            self._epoch_log_writer = csv.DictWriter(self._epoch_log_fp, fieldnames=fieldnames)
+            self._epoch_log_writer.writeheader()
+            self._epoch_log_fields = fieldnames
+
+        assert self._epoch_log_fields is not None
+        assert self._epoch_log_writer is not None
+        assert self._epoch_log_fp is not None
+        for key in self._epoch_log_fields:
+            row.setdefault(key, None)
+        self._epoch_log_writer.writerow({k: row.get(k) for k in self._epoch_log_fields})
+        self._epoch_log_fp.flush()
+
+    def write_json(self, path: str | Path, payload: dict) -> None:
+        target = Path(path)
+        if not target.is_absolute():
+            target = self.run_dir / target
+        target.parent.mkdir(parents=True, exist_ok=True)
+        target.write_text(json.dumps(payload, indent=2), encoding="utf-8")
+
+    def close(self) -> None:
+        if self._epoch_log_fp is not None:
+            try:
+                self._epoch_log_fp.close()
+            except Exception:
+                pass
+            self._epoch_log_fp = None
+            self._epoch_log_writer = None
+            self._epoch_log_fields = None
+        for handler in list(self.logger.handlers):
+            try:
+                handler.close()
+            except Exception:
+                pass
+            self.logger.removeHandler(handler)

classes/v2/loader_factory.py

@@ -0,0 +1,165 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Any, Callable, Optional
+
+from torch.utils.data import DataLoader
+
+from .network_manager import LoaderBundle, PatientSplit
+from .slot_dataset import SlotDataset, slot_collate
+from .profiles.base import SlotDescriptor, SimpleDatasetProfile
+
+
+def _default_slot_descriptors(patient_col: str, label_col: str) -> dict[str, SlotDescriptor]:
+    return {
+        "id_1": SlotDescriptor(
+            key="id_1",
+            kind="id",
+            description=f"Patient identifier column ({patient_col})",
+            required=True,
+            shape_hint="scalar",
+        ),
+        "label_1": SlotDescriptor(
+            key="label_1",
+            kind="label",
+            description=f"Label column ({label_col})",
+            required=True,
+            shape_hint="scalar",
+        ),
+        "image_1": SlotDescriptor(
+            key="image_1",
+            kind="image",
+            description="Primary image slot",
+            required=False,
+            shape_hint="HWC or CHW",
+        ),
+        "matrix_1": SlotDescriptor(
+            key="matrix_1",
+            kind="matrix",
+            description="Primary matrix slot",
+            required=False,
+            shape_hint="[feature_dim]",
+        ),
+    }
+
+
+def _row_to_sample(
+    row: Any,
+    *,
+    clinical: Any,
+    patient_col: str,
+    label_col: str,
+) -> dict[str, Any]:
+    return {
+        "id_1": row[patient_col],
+        "label_1": row[label_col],
+        "image_1": clinical.get_image_path(row) if hasattr(clinical, "get_image_path") else None,
+        "matrix_1": clinical.vectorize_row(row) if hasattr(clinical, "vectorize_row") else None,
+    }
+
+
+@dataclass
+class SlotLoaderFactory:
+    """
+    Generic loader factory that emits dict batches keyed by slot names.
+    """
+
+    image_transform: Optional[Callable] = None
+    matrix_transform: Optional[Callable] = None
+    num_workers: int = 0
+
+    def build(
+        self,
+        *,
+        clinical: Any,
+        split: PatientSplit,
+        args: Any,
+        fold: int,
+        profile: Optional[Any] = None,
+    ) -> LoaderBundle:
+        batch_size = int(getattr(args, "batch_size", 8))
+        slot_desc = self._resolve_slot_descriptors(clinical=clinical, profile=profile)
+
+        train_samples = self._build_samples(split.train, clinical, profile, slot_desc)
+        val_samples = self._build_samples(split.val, clinical, profile, slot_desc)
+        holdout_samples = (
+            self._build_samples(split.holdout, clinical, profile, slot_desc)
+            if split.holdout is not None
+            else None
+        )
+
+        train_loader = DataLoader(
+            SlotDataset(
+                train_samples,
+                slot_desc,
+                image_transform=self.image_transform,
+                matrix_transform=self.matrix_transform,
+            ),
+            batch_size=batch_size,
+            shuffle=True,
+            num_workers=self.num_workers,
+            collate_fn=slot_collate,
+        )
+        val_loader = DataLoader(
+            SlotDataset(
+                val_samples,
+                slot_desc,
+                image_transform=self.image_transform,
+                matrix_transform=self.matrix_transform,
+            ),
+            batch_size=batch_size,
+            shuffle=False,
+            num_workers=self.num_workers,
+            collate_fn=slot_collate,
+        )
+        holdout_loader = None
+        if holdout_samples is not None:
+            holdout_loader = DataLoader(
+                SlotDataset(
+                    holdout_samples,
+                    slot_desc,
+                    image_transform=self.image_transform,
+                    matrix_transform=self.matrix_transform,
+                ),
+                batch_size=batch_size,
+                shuffle=False,
+                num_workers=self.num_workers,
+                collate_fn=slot_collate,
+            )
+        return LoaderBundle(train=train_loader, val=val_loader, holdout=holdout_loader)
+
+    @staticmethod
+    def _resolve_slot_descriptors(
+        *,
+        clinical: Any,
+        profile: Optional[Any],
+    ) -> dict[str, SlotDescriptor]:
+        if profile is not None and hasattr(profile, "slot_descriptors"):
+            return profile.slot_descriptors()
+        patient_col = getattr(clinical, "patient_col", "Patient ID")
+        label_col = getattr(clinical, "label_col", "Diagnosis")
+        return _default_slot_descriptors(patient_col, label_col)
+
+    @staticmethod
+    def _build_samples(
+        df,
+        clinical: Any,
+        profile: Optional[Any],
+        slot_desc: dict[str, SlotDescriptor],
+    ) -> list[dict[str, Any]]:
+        if df is None or df.empty:
+            return []
+        if profile is not None and hasattr(profile, "build_samples"):
+            return profile.build_samples(df=df, clinical=clinical)
+
+        patient_col = getattr(profile, "patient_col", None) if profile is not None else None
+        label_col = getattr(profile, "label_col", None) if profile is not None else None
+        pcol = patient_col or "Patient ID"
+        lcol = label_col or getattr(clinical, "label_col", "Diagnosis")
+        samples = []
+        for _, row in df.iterrows():
+            sample = _row_to_sample(row, clinical=clinical, patient_col=pcol, label_col=lcol)
+            for key in slot_desc.keys():
+                sample.setdefault(key, None)
+            samples.append(sample)
+        return samples

classes/v2/model_builder.py

@@ -0,0 +1,148 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Any, Callable, Optional
+
+import torch
+from torch import nn
+
+from classes.v2.bridges import Bridge, VoteBridge
+from classes.v2.towers import ImageTower, MDTower
+
+from .config_builder import ConfigAssembly
+from .transforms import build_transform_chain
+
+
+@dataclass
+class V2ModelBundle:
+    image_tower: Optional[ImageTower]
+    metadata_tower: Optional[MDTower]
+    bridge: Optional[nn.Module]
+    classifier: Optional[nn.Module]
+    image_transform: Optional[Callable]
+    matrix_transform: Optional[Callable]
+
+
+def build_model_bundle(
+    assembly: ConfigAssembly,
+    clinical: Any,
+    *,
+    device: Optional[torch.device] = None,
+    strict: bool = True,
+) -> V2ModelBundle:
+    """
+    Build torch modules and input transforms from a V2 config assembly.
+    """
+    image_tower_spec = _pick_tower(assembly, "image")
+    md_tower_spec = _pick_tower(assembly, "metadata")
+    bridge_spec = _pick_bridge(assembly)
+    image_loader = _pick_loader(assembly, input_type="image")
+
+    clinical_core = getattr(clinical, "clinical", clinical)
+    num_classes = _infer_num_classes(clinical)
+
+    img_tower = None
+    if image_tower_spec is not None:
+        img_tower = ImageTower(
+            backbone=image_tower_spec.params.get("backbone", "efficientnet_b0"),
+            freeze_ratio=float(image_tower_spec.params.get("freeze_ratio", 0.0) or 0.0),
+            use_se=bool(image_tower_spec.params.get("use_se", False)),
+            se_reduction=int(image_tower_spec.params.get("se_reduction", 16) or 16),
+            se_pre_norm=bool(image_tower_spec.params.get("se_pre_norm", True)),
+            augment=bool(image_tower_spec.params.get("augment", True)),
+            geometry_dim=int(image_tower_spec.params.get("geometry_dim", 0) or 0),
+        )
+        if device is not None:
+            img_tower = img_tower.to(device)
+
+    md_tower = None
+    if md_tower_spec is not None:
+        md_tower = MDTower(
+            clinical_core,
+            hidden_dim=int(md_tower_spec.params.get("hidden_dim", 128) or 128),
+            dropout=float(md_tower_spec.params.get("dropout", 0.1) or 0.1),
+            use_se=bool(md_tower_spec.params.get("use_se", False)),
+            se_reduction=int(md_tower_spec.params.get("se_reduction", 16) or 16),
+            se_pre_norm=bool(md_tower_spec.params.get("se_pre_norm", True)),
+        )
+        if device is not None:
+            md_tower = md_tower.to(device)
+
+    bridge = None
+    if bridge_spec is not None and img_tower is not None and md_tower is not None:
+        if bridge_spec.method == "consensus":
+            bridge = VoteBridge(num_classes=num_classes)
+        else:
+            bridge = Bridge(
+                img_dim=img_tower.out_dim,
+                meta_dim=md_tower.out_dim,
+                num_classes=num_classes,
+                fusion_dim=int(bridge_spec.params.get("fusion_dim", 256) or 256),
+                mode="fused",
+                use_se=bool(bridge_spec.params.get("use_se", True)),
+                se_reduction=int(bridge_spec.params.get("se_reduction", 16) or 16),
+                se_pre_norm=bool(bridge_spec.params.get("se_pre_norm", True)),
+            )
+        if device is not None:
+            bridge = bridge.to(device)
+
+    classifier = None
+    if assembly.classifiers:
+        classifier = nn.Identity()
+        if device is not None:
+            classifier = classifier.to(device)
+
+    image_transform = None
+    if image_loader is not None and image_tower_spec is not None:
+        image_transform = build_transform_chain(
+            image_loader.transforms,
+            backbone_name=image_tower_spec.params.get("backbone", "efficientnet_b0"),
+            augment=bool(image_tower_spec.params.get("augment", True)),
+            strict=strict,
+        )
+
+    return V2ModelBundle(
+        image_tower=img_tower,
+        metadata_tower=md_tower,
+        bridge=bridge,
+        classifier=classifier,
+        image_transform=image_transform,
+        matrix_transform=None,
+    )
+
+
+def _pick_tower(assembly: ConfigAssembly, tower_type: str):
+    matches = [tower for tower in assembly.towers.values() if tower.tower_type == tower_type]
+    if not matches:
+        return None
+    if len(matches) > 1:
+        raise ValueError(f"Multiple {tower_type} towers found; only one is supported for now.")
+    return matches[0]
+
+
+def _pick_bridge(assembly: ConfigAssembly):
+    if not assembly.bridges:
+        return None
+    if len(assembly.bridges) > 1:
+        raise ValueError("Multiple bridges found; only one is supported for now.")
+    return next(iter(assembly.bridges.values()))
+
+
+def _pick_loader(assembly: ConfigAssembly, input_type: str):
+    matches = [loader for loader in assembly.loaders.values() if loader.input_type == input_type]
+    if not matches:
+        return None
+    if len(matches) > 1:
+        raise ValueError(f"Multiple loaders with input_type={input_type!r} found.")
+    return matches[0]
+
+
+def _infer_num_classes(clinical: Any) -> int:
+    df = getattr(clinical, "df", None)
+    label_col = getattr(clinical, "label_col", None)
+    if df is None and hasattr(clinical, "clinical"):
+        df = clinical.clinical.df
+        label_col = clinical.clinical.label_col
+    if df is None or label_col is None or label_col not in df.columns:
+        return 2
+    return int(df[label_col].dropna().nunique())

classes/v2/network_manager.py

@@ -0,0 +1,200 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Any, Optional, Protocol
+
+import pandas as pd
+
+
+@dataclass
+class PatientSplit:
+    """Patient-disjoint split definition for a fold."""
+
+    train: pd.DataFrame
+    val: pd.DataFrame
+    holdout: Optional[pd.DataFrame] = None
+
+
+@dataclass
+class LoaderBundle:
+    """All loaders needed by a training run."""
+
+    train: Any
+    val: Any
+    holdout: Optional[Any] = None
+
+
+@dataclass
+class FoldResult:
+    """Normalized fold output from trainer implementations."""
+
+    fold: int
+    metrics: dict[str, Any]
+    artifacts: dict[str, Any]
+
+
+class SplitManager(Protocol):
+    def build_plans(
+        self,
+        *,
+        clinical: Any,
+        args: Any,
+        profile: Optional[Any] = None,
+    ) -> list[PatientSplit]:
+        ...
+
+
+class GraphFactory(Protocol):
+    def build(
+        self,
+        *,
+        clinical: Any,
+        args: Any,
+        fold: int,
+        profile: Optional[Any] = None,
+    ) -> Any:
+        ...
+
+
+class LoaderFactory(Protocol):
+    def build(
+        self,
+        *,
+        clinical: Any,
+        split: PatientSplit,
+        args: Any,
+        fold: int,
+        profile: Optional[Any] = None,
+    ) -> LoaderBundle:
+        ...
+
+
+class Trainer(Protocol):
+    def fit(
+        self,
+        *,
+        graph: Any,
+        loaders: LoaderBundle,
+        args: Any,
+        fold: int,
+        profile: Optional[Any] = None,
+    ) -> FoldResult:
+        ...
+
+
+class NetworkManager:
+    """
+    V2 orchestration entrypoint.
+
+    This class is intentionally small and modular:
+    - split policy is delegated to a SplitManager
+    - graph assembly is delegated to a GraphFactory
+    - dataloaders are delegated to a LoaderFactory
+    - train/eval/checkpoint lifecycle is delegated to a Trainer
+    """
+
+    def __init__(
+        self,
+        *,
+        clinical: Any,
+        args: Any,
+        split_manager: SplitManager,
+        graph_factory: GraphFactory,
+        loader_factory: LoaderFactory,
+        trainer: Trainer,
+        profile: Optional[Any] = None,
+    ) -> None:
+        self.clinical = clinical
+        self.args = args
+        self.split_manager = split_manager
+        self.graph_factory = graph_factory
+        self.loader_factory = loader_factory
+        self.trainer = trainer
+        self.profile = profile
+        self._split_plans: Optional[list[PatientSplit]] = None
+
+    def run_fold(self, fold: int) -> FoldResult:
+        plans = self._get_split_plans()
+        if fold < 0 or fold >= len(plans):
+            raise IndexError(f"Requested fold {fold} but only {len(plans)} fold plans are available")
+        split = plans[fold]
+        self._validate_patient_disjointness(split)
+        self._validate_labels(split)
+
+        graph = self.graph_factory.build(
+            clinical=self.clinical,
+            args=self.args,
+            fold=fold,
+            profile=self.profile,
+        )
+        loaders = self.loader_factory.build(
+            clinical=self.clinical,
+            split=split,
+            args=self.args,
+            fold=fold,
+            profile=self.profile,
+        )
+        return self.trainer.fit(
+            graph=graph,
+            loaders=loaders,
+            args=self.args,
+            fold=fold,
+            profile=self.profile,
+        )
+
+    def run_all_folds(self, n_splits: Optional[int] = None) -> list[FoldResult]:
+        plans = self._get_split_plans()
+        max_folds = len(plans)
+        if n_splits is None:
+            n = max_folds
+        else:
+            n = int(n_splits)
+            if n < 1:
+                raise ValueError("n_splits must be >= 1")
+            if n > max_folds:
+                raise ValueError(f"Requested {n} folds but only {max_folds} fold plans are available")
+        return [self.run_fold(fold) for fold in range(n)]
+
+    def _get_split_plans(self) -> list[PatientSplit]:
+        if self._split_plans is None:
+            self._split_plans = self.split_manager.build_plans(
+                clinical=self.clinical,
+                args=self.args,
+                profile=self.profile,
+            )
+            if not self._split_plans:
+                raise ValueError("SplitManager returned no fold plans")
+        return self._split_plans
+
+    def _validate_patient_disjointness(self, split: PatientSplit) -> None:
+        train_ids = self._patient_ids(split.train)
+        val_ids = self._patient_ids(split.val)
+        holdout_ids = self._patient_ids(split.holdout) if split.holdout is not None else set()
+
+        if train_ids & val_ids:
+            overlap = sorted(train_ids & val_ids)[:10]
+            raise ValueError(f"Patient leakage between train/val: {overlap}")
+        if train_ids & holdout_ids:
+            overlap = sorted(train_ids & holdout_ids)[:10]
+            raise ValueError(f"Patient leakage between train/holdout: {overlap}")
+        if val_ids & holdout_ids:
+            overlap = sorted(val_ids & holdout_ids)[:10]
+            raise ValueError(f"Patient leakage between val/holdout: {overlap}")
+
+    def _validate_labels(self, split: PatientSplit) -> None:
+        label_col = getattr(self.clinical, "label_col", None)
+        if not label_col:
+            return
+        for name, df in (("train", split.train), ("val", split.val), ("holdout", split.holdout)):
+            if df is None:
+                continue
+            if label_col not in df.columns:
+                raise ValueError(f"{name} split is missing label column {label_col!r}")
+
+    @staticmethod
+    def _patient_ids(df: Optional[pd.DataFrame]) -> set[Any]:
+        if df is None or df.empty:
+            return set()
+        if "Patient ID" not in df.columns:
+            raise ValueError("Split dataframes must include 'Patient ID'")
+        return set(df["Patient ID"].tolist())

classes/v2/papila_builders.py

@@ -0,0 +1,152 @@
+from __future__ import annotations
+
+from typing import Dict, List
+
+import numpy as np
+import pandas as pd
+
+from classes.v2.data_bundle import DataBundle
+
+# ---- Pachymetry → IOP correction (per PAPILA Table 3) ----
+_PACHY_TABLE: Dict[int, int] = {
+    475: +5,
+    485: +4,
+    495: +4,
+    505: +3,
+    515: +2,
+    525: +1,
+    535: +1,
+    545: 0,
+    555: -1,
+    565: -1,
+    575: -2,
+    585: -3,
+    595: -4,
+    605: -4,
+    615: -5,
+}
+_PACHY_KEYS = np.array(sorted(_PACHY_TABLE.keys()))
+
+
+def _nearest_pachy_key(x: float) -> int:
+    idx = int(np.argmin(np.abs(_PACHY_KEYS - float(x))))
+    return int(_PACHY_KEYS[idx])
+
+
+def _pick_iop(row: pd.Series) -> float:
+    """Prefer Pneumatic, else Perkins; may return NaN."""
+    raw = row["Pneumatic"] if not pd.isna(row.get("Pneumatic", np.nan)) else row.get("Perkins", np.nan)
+    return float(raw) if not pd.isna(raw) else np.nan
+
+
+def _correct_iop(raw_iop: float, pachy: float) -> float:
+    """Return corrected IOP using nearest pachymetry bin; if pachy missing, return raw."""
+    if pd.isna(raw_iop):
+        return np.nan
+    if pd.isna(pachy):
+        return float(raw_iop)
+    key = _nearest_pachy_key(float(pachy))
+    return float(raw_iop) + float(_PACHY_TABLE[key])
+
+
+def _apply_iop_and_drop_md(df: pd.DataFrame) -> pd.DataFrame:
+    """Add IOP_raw/IOP_corr and drop VF_MD if present (in-place safe)."""
+    df["IOP_raw"] = df.apply(_pick_iop, axis=1)
+    pachy = df.get("Pachymetry", pd.Series(np.nan, index=df.index))
+    df["IOP_corr"] = [
+        _correct_iop(r, p) for r, p in zip(df["IOP_raw"].values, pachy.values)
+    ]
+    if "VF_MD" in df.columns:
+        df.drop(columns=["VF_MD"], inplace=True)
+    return df
+
+
+def _canonicalize_eye_column(df: pd.DataFrame) -> None:
+    if "eyeID" in df.columns:
+        src = "eyeID"
+    else:
+        src = None
+        for c in df.columns:
+            if "eye" in c.lower():
+                src = c
+                break
+        if src is None:
+            df["eyeID"] = "OS"
+            return
+
+    s = df[src]
+
+    def norm(v):
+        if pd.isna(v):
+            return None
+        x = str(v).strip().upper()
+        if x in {"OS", "L", "LEFT", "0"}:
+            return "OS"
+        if x in {"OD", "R", "RIGHT", "1"}:
+            return "OD"
+        try:
+            num = int(float(x))
+            return "OD" if num % 2 == 1 else "OS"
+        except Exception:
+            return None
+
+    mapped = s.map(norm)
+    uniq = {u for u in mapped.dropna().unique().tolist()}
+    if not uniq.issubset({"OS", "OD"}):
+        raise ValueError(f"eyeID must be binary; found values {sorted(uniq)}")
+    df["eyeID"] = mapped.fillna("OS")
+
+
+def build_papila_data(
+    *,
+    image_dir: str,
+    clinical_dir: str,
+    label_col: str,
+    cat_cols: List[str],
+    n_splits: int = 5,
+    random_seed: int = 42,
+) -> DataBundle:
+    """
+    Build a DataBundle for PAPILA with dataset-specific preprocessing:
+      - load OD/OS Excel sheets
+      - normalize Patient ID
+      - canonicalize eyeID
+      - compute IOP_raw / IOP_corr, drop VF_MD
+      - build feature typing & folds
+    """
+    bundle = DataBundle(
+        image_dir=image_dir,
+        clinical_dir=clinical_dir,
+        label_col=label_col,
+        patient_col="Patient ID",
+        cat_cols=cat_cols,
+        n_splits=n_splits,
+        random_seed=random_seed,
+        filename_template="RET{pid:03d}{eye}.jpg",
+    )
+
+    od = pd.read_excel(f"{clinical_dir}/patient_data_od.xlsx", header=1)
+    od["eyeID"] = "OD"
+    os = pd.read_excel(f"{clinical_dir}/patient_data_os.xlsx", header=1)
+    os["eyeID"] = "OS"
+
+    for frame in (od, os):
+        if "Patient ID" not in frame.columns and "ID" in frame.columns:
+            frame.rename(columns={"ID": "Patient ID"}, inplace=True)
+        frame["Patient ID"] = frame["Patient ID"].astype(str).str.extract(r"(\d+)")[0].astype(int)
+        _canonicalize_eye_column(frame)
+
+    bundle.add_df(od, id_column="ID")
+    bundle.add_df(os, id_column="ID")
+
+    for i in range(len(bundle.frames)):
+        bundle.frames[i] = _apply_iop_and_drop_md(bundle.frames[i])
+
+    bundle._refresh_master_df()
+    bundle._infer_or_validate_feature_types()
+    bundle._compute_numeric_stats()
+    bundle._build_cat_maps()
+    bundle._compute_feature_dim()
+    bundle._build_kfold_indices()
+
+    return bundle

classes/v2/papila_data.py

@@ -0,0 +1,61 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Iterable, Optional
+
+import pandas as pd
+
+from classes.v2.data_bundle import DataBundle
+from classes.v2.papila_builders import build_papila_data
+
+
+@dataclass
+class PapilaData:
+    """
+    V2-friendly wrapper around the DataBundle pipeline.
+
+    Keeps all formatting/normalization behavior from build_papila_clinical,
+    but exposes a minimal surface area for the V2 engine.
+    """
+
+    clinical: DataBundle
+    patient_col: str = "Patient ID"
+
+    @property
+    def df(self) -> pd.DataFrame:
+        return self.clinical.df
+
+    @property
+    def label_col(self) -> str:
+        return self.clinical.label_col
+
+    @property
+    def feature_dim(self) -> int:
+        return self.clinical.feature_dim
+
+    def get_image_path(self, row: pd.Series):
+        return self.clinical.get_image_path(row)
+
+    def vectorize_row(self, row: pd.Series):
+        return self.clinical.vectorize_row(row)
+
+    @classmethod
+    def from_dirs(
+        cls,
+        *,
+        image_dir: str,
+        clinical_dir: str,
+        label_col: str,
+        cat_cols: Iterable[str],
+        n_splits: int = 5,
+        random_seed: int = 42,
+    ) -> "PapilaData":
+        clinical = build_papila_data(
+            image_dir=image_dir,
+            clinical_dir=clinical_dir,
+            label_col=label_col,
+            cat_cols=list(cat_cols),
+            n_splits=n_splits,
+            random_seed=random_seed,
+        )
+        return cls(clinical=clinical)

classes/v2/profiles/__init__.py

@@ -0,0 +1,10 @@
+from .base import DatasetProfile, SimpleDatasetProfile, SlotDescriptor
+from .papila import PapilaProfile, build_papila_profile
+
+__all__ = [
+    "DatasetProfile",
+    "SimpleDatasetProfile",
+    "SlotDescriptor",
+    "PapilaProfile",
+    "build_papila_profile",
+]

classes/v2/profiles/base.py

@@ -0,0 +1,53 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Protocol, Any
+
+import pandas as pd
+
+
+@dataclass(frozen=True)
+class SlotDescriptor:
+    """
+    Metadata for a generic batch slot key (e.g., image_1, matrix_1).
+    """
+
+    key: str
+    kind: str
+    description: str
+    required: bool = True
+    shape_hint: str | None = None
+
+
+class DatasetProfile(Protocol):
+    """
+    Dataset-specific wiring that stays outside the generic V2 engine.
+    """
+
+    name: str
+    patient_col: str
+    label_col: str
+
+    def slot_descriptors(self) -> dict[str, SlotDescriptor]:
+        ...
+
+    def semantic_aliases(self) -> dict[str, str]:
+        ...
+
+    def build_samples(self, *, df: pd.DataFrame, clinical: Any) -> list[dict[str, Any]]:
+        ...
+
+
+@dataclass(frozen=True)
+class SimpleDatasetProfile:
+    name: str
+    patient_col: str
+    label_col: str
+    slots: dict[str, SlotDescriptor]
+    aliases: dict[str, str]
+
+    def slot_descriptors(self) -> dict[str, SlotDescriptor]:
+        return dict(self.slots)
+
+    def semantic_aliases(self) -> dict[str, str]:
+        return dict(self.aliases)

classes/v2/profiles/papila.py

@@ -0,0 +1,155 @@
+from __future__ import annotations
+
+import pandas as pd
+
+from dataclasses import dataclass
+
+from .base import SimpleDatasetProfile, SlotDescriptor
+
+
+@dataclass(frozen=True)
+class PapilaProfile(SimpleDatasetProfile):
+    sample_mode: str = "patient"  # "patient" | "eye"
+
+    def build_samples(self, *, df: pd.DataFrame, clinical) -> list[dict[str, object]]:
+        samples: list[dict[str, object]] = []
+        patient_col = self.patient_col
+        label_col = self.label_col
+
+        mode = (self.sample_mode or "patient").lower()
+        if mode not in {"patient", "eye"}:
+            raise ValueError(f"Unsupported sample_mode '{self.sample_mode}'. Expected 'patient' or 'eye'.")
+
+        if mode == "eye":
+            for _, row in df.iterrows():
+                pid = row[patient_col]
+                label = row[label_col]
+                image_1 = clinical.get_image_path(row) if hasattr(clinical, "get_image_path") else None
+                matrix_1 = clinical.vectorize_row(row) if hasattr(clinical, "vectorize_row") else None
+                samples.append(
+                    {
+                        "id_1": pid,
+                        "label_1": label,
+                        "image_1": image_1,
+                        "matrix_1": matrix_1,
+                    }
+                )
+            return samples
+
+        for pid, grp in df.groupby(patient_col):
+            label_series = grp[label_col]
+            if label_series.empty:
+                continue
+            mode_vals = label_series.mode()
+            label = mode_vals.iloc[0] if not mode_vals.empty else label_series.iloc[0]
+
+            def _row_for_eye(eye: str):
+                if "eyeID" not in grp.columns:
+                    return None
+                match = grp[grp["eyeID"].astype(str).str.upper() == eye]
+                if match.empty:
+                    return None
+                return match.iloc[0]
+
+            row_od = _row_for_eye("OD")
+            row_os = _row_for_eye("OS")
+            row_any = grp.iloc[0]
+
+            image_1 = clinical.get_image_path(row_od) if row_od is not None else None
+            image_2 = clinical.get_image_path(row_os) if row_os is not None else None
+            matrix_1 = clinical.vectorize_row(row_od) if row_od is not None else None
+            matrix_2 = clinical.vectorize_row(row_os) if row_os is not None else None
+
+            if image_1 is None and hasattr(clinical, "get_image_path"):
+                image_1 = clinical.get_image_path(row_any)
+            if matrix_1 is None and hasattr(clinical, "vectorize_row"):
+                matrix_1 = clinical.vectorize_row(row_any)
+
+            samples.append(
+                {
+                    "id_1": pid,
+                    "label_1": label,
+                    "image_1": image_1,
+                    "image_2": image_2,
+                    "matrix_1": matrix_1,
+                    "matrix_2": matrix_2,
+                }
+            )
+        return samples
+
+
+def build_papila_profile(
+    *,
+    patient_col: str = "Patient ID",
+    label_col: str = "Diagnosis",
+    sample_mode: str = "patient",
+) -> PapilaProfile:
+    """
+    PAPILA-specific semantic map for generic V2 slot keys.
+
+    The engine remains slot-based (image_1/image_2/matrix_1/...).
+    PAPILA meaning is captured here so run config stays dataset-local.
+    """
+
+    slots = {
+        "id_1": SlotDescriptor(
+            key="id_1",
+            kind="id",
+            description=f"Patient identifier column ({patient_col})",
+            required=True,
+            shape_hint="scalar",
+        ),
+        "label_1": SlotDescriptor(
+            key="label_1",
+            kind="label",
+            description=f"Diagnosis label column ({label_col})",
+            required=True,
+            shape_hint="scalar",
+        ),
+        "image_1": SlotDescriptor(
+            key="image_1",
+            kind="image",
+            description="Fundus image slot 1 (PAPILA: OD / right eye)",
+            required=False,
+            shape_hint="HWC or CHW",
+        ),
+        "image_2": SlotDescriptor(
+            key="image_2",
+            kind="image",
+            description="Fundus image slot 2 (PAPILA: OS / left eye)",
+            required=False,
+            shape_hint="HWC or CHW",
+        ),
+        "matrix_1": SlotDescriptor(
+            key="matrix_1",
+            kind="matrix",
+            description="Clinical metadata feature vector",
+            required=False,
+            shape_hint="[feature_dim]",
+        ),
+        "matrix_2": SlotDescriptor(
+            key="matrix_2",
+            kind="matrix",
+            description="Optional auxiliary tabular vector (reserved for experiments)",
+            required=False,
+            shape_hint="[feature_dim_2]",
+        ),
+    }
+
+    aliases = {
+        "id_1": "patient_id",
+        "label_1": "diagnosis",
+        "image_1": "od_fundus",
+        "image_2": "os_fundus",
+        "matrix_1": "clinical_metadata",
+        "matrix_2": "aux_metadata",
+    }
+
+    return PapilaProfile(
+        name="papila",
+        patient_col=patient_col,
+        label_col=label_col,
+        slots=slots,
+        aliases=aliases,
+        sample_mode=sample_mode,
+    )

classes/v2/slot_dataset.py

@@ -0,0 +1,98 @@
+from __future__ import annotations
+
+from typing import Any, Callable, Optional
+
+from pathlib import Path
+from PIL import Image
+import numpy as np
+import torch
+from torch.utils.data import Dataset
+from torchvision import transforms
+
+from .profiles.base import SlotDescriptor
+
+
+def slot_collate(batch: list[dict[str, Any]]) -> dict[str, Any]:
+    if not batch:
+        return {}
+    keys = batch[0].keys()
+    out: dict[str, Any] = {}
+    for key in keys:
+        vals = [item.get(key) for item in batch]
+        if all(isinstance(v, torch.Tensor) for v in vals):
+            try:
+                out[key] = torch.stack(vals, dim=0)
+            except Exception:
+                out[key] = vals
+        else:
+            out[key] = vals
+    return out
+
+
+class SlotDataset(Dataset):
+    """
+    Dataset that yields dicts of slot-keyed values.
+
+    Sample records are expected to be dicts with keys matching slot descriptors.
+    Image slots accept filesystem paths; matrix slots accept array-like values.
+    """
+
+    def __init__(
+        self,
+        samples: list[dict[str, Any]],
+        slot_descriptors: dict[str, SlotDescriptor],
+        *,
+        image_transform: Optional[Callable[[Image.Image], torch.Tensor]] = None,
+        matrix_transform: Optional[Callable[[Any], torch.Tensor]] = None,
+        image_preprocessor: Optional[Callable[..., Image.Image]] = None,
+    ) -> None:
+        self.samples = samples
+        self.slot_descriptors = slot_descriptors
+        self.image_transform = image_transform or transforms.ToTensor()
+        self.matrix_transform = matrix_transform or self._default_matrix_transform
+        self.image_preprocessor = image_preprocessor
+
+    def __len__(self) -> int:
+        return len(self.samples)
+
+    def __getitem__(self, idx: int) -> dict[str, Any]:
+        record = self.samples[idx]
+        out: dict[str, Any] = {}
+        for key, desc in self.slot_descriptors.items():
+            val = record.get(key)
+            if desc.kind == "image":
+                out[key] = self._load_image(val, required=desc.required)
+            elif desc.kind == "matrix":
+                out[key] = self._load_matrix(val, required=desc.required)
+            else:
+                out[key] = val
+        return out
+
+    def _load_image(self, value: Any, *, required: bool) -> Optional[torch.Tensor]:
+        if value is None:
+            if required:
+                raise ValueError("Missing required image slot")
+            return None
+        path = Path(value)
+        img = Image.open(path).convert("RGB")
+        if self.image_preprocessor is not None:
+            try:
+                img = self.image_preprocessor(img, path)
+            except TypeError:
+                img = self.image_preprocessor(img)
+        return self.image_transform(img)
+
+    def _load_matrix(self, value: Any, *, required: bool) -> Optional[torch.Tensor]:
+        if value is None:
+            if required:
+                raise ValueError("Missing required matrix slot")
+            return None
+        return self.matrix_transform(value)
+
+    @staticmethod
+    def _default_matrix_transform(value: Any) -> torch.Tensor:
+        if isinstance(value, torch.Tensor):
+            return value.float()
+        if isinstance(value, np.ndarray):
+            return torch.from_numpy(value.astype(np.float32, copy=False))
+        return torch.as_tensor(value, dtype=torch.float32)

classes/v2/split_manager.py

@@ -0,0 +1,197 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Any, Iterable, Optional
+
+import numpy as np
+import pandas as pd
+from sklearn.model_selection import KFold, StratifiedKFold
+
+from .network_manager import PatientSplit
+
+
+@dataclass(frozen=True)
+class SplitPlan:
+    train_patient_ids: set[Any]
+    val_patient_ids: set[Any]
+    holdout_patient_ids: set[Any]
+
+
+def build_patient_split_plans(
+    patient_ids: Iterable[Any],
+    patient_labels: Iterable[Any],
+    *,
+    n_splits: int,
+    seed: int,
+    holdout_per_class: int = 0,
+    holdout_seed: int = 123,
+) -> list[SplitPlan]:
+    """
+    Core vector-based splitter.
+
+    Inputs are one row per patient:
+    - patient_ids: unique patient IDs
+    - patient_labels: one label per patient
+    """
+    ids = np.asarray(list(patient_ids))
+    labels = np.asarray(list(patient_labels))
+    if ids.ndim != 1 or labels.ndim != 1:
+        raise ValueError("patient_ids and patient_labels must be 1D arrays")
+    if ids.size != labels.size:
+        raise ValueError(f"Length mismatch: ids={ids.size}, labels={labels.size}")
+    if ids.size == 0:
+        raise ValueError("No patients available for splitting")
+    if len(set(ids.tolist())) != ids.size:
+        raise ValueError("patient_ids must be unique (one label per patient)")
+    if n_splits < 2:
+        raise ValueError("n_splits must be >= 2")
+
+    holdout_ids: set[Any] = set()
+    if holdout_per_class > 0:
+        rng = np.random.default_rng(holdout_seed)
+        for label in np.unique(labels):
+            idx = np.where(labels == label)[0]
+            if idx.size == 0:
+                continue
+            n = min(holdout_per_class, idx.size)
+            chosen = rng.choice(idx, size=n, replace=False)
+            holdout_ids.update(ids[chosen].tolist())
+
+    keep_mask = ~np.isin(ids, list(holdout_ids))
+    cv_ids = ids[keep_mask]
+    cv_labels = labels[keep_mask]
+    if cv_ids.size < n_splits:
+        raise ValueError(
+            f"Not enough patients ({cv_ids.size}) for n_splits={n_splits} after holdout removal"
+        )
+
+    use_stratified = _can_stratify(cv_labels, n_splits)
+    if use_stratified:
+        splitter = StratifiedKFold(n_splits=n_splits, shuffle=True, random_state=seed)
+        splits = list(splitter.split(cv_ids, cv_labels))
+    else:
+        splitter = KFold(n_splits=n_splits, shuffle=True, random_state=seed)
+        splits = list(splitter.split(cv_ids))
+
+    plans: list[SplitPlan] = []
+    for train_idx, val_idx in splits:
+        plans.append(
+            SplitPlan(
+                train_patient_ids=set(cv_ids[train_idx].tolist()),
+                val_patient_ids=set(cv_ids[val_idx].tolist()),
+                holdout_patient_ids=set(holdout_ids),
+            )
+        )
+    return plans
+
+
+class PatientFirstSplitManager:
+    """
+    Patient-level splitter for V2.
+
+    Behavior:
+    - Optional binary filtering happens first (labels in {0,1} only).
+    - Optional holdout is sampled at the patient level (never per-eye rows).
+    - K-fold split is built on remaining patients.
+    - Returned dataframes contain all rows for each selected patient.
+    """
+
+    def __init__(
+        self,
+        *,
+        patient_col: str = "Patient ID",
+        label_col: Optional[str] = None,
+    ) -> None:
+        self.patient_col = patient_col
+        self.label_col = label_col
+
+    def build_plans(
+        self,
+        *,
+        clinical: Any,
+        args: Any,
+        profile: Optional[Any] = None,
+    ) -> list[PatientSplit]:
+        profile_label_col = getattr(profile, "label_col", None) if profile is not None else None
+        profile_patient_col = getattr(profile, "patient_col", None) if profile is not None else None
+        patient_col = profile_patient_col or self.patient_col
+        label_col = self.label_col or profile_label_col or getattr(clinical, "label_col", None)
+        if label_col is None:
+            raise ValueError("Could not resolve label column from SplitManager or clinical.label_col")
+
+        if not hasattr(clinical, "df"):
+            raise ValueError("Clinical object must expose a dataframe at .df")
+        df_full = clinical.df.copy()
+        self._validate_columns(df_full, label_col, patient_col=patient_col)
+
+        eval_mode = str(getattr(args, "eval_mode", "multiclass")).lower()
+        if eval_mode == "binary":
+            df_full = df_full[df_full[label_col].isin([0, 1])].reset_index(drop=True)
+
+        holdout_per_class = int(getattr(args, "holdout_per_class", 0) or 0)
+        holdout_seed = int(getattr(args, "holdout_seed", 123))
+        n_splits = int(getattr(args, "n_splits", 5))
+        fold_seed = int(getattr(args, "fold_seed", 42))
+
+        patient_table = self._patient_label_table(df_full, label_col, patient_col=patient_col)
+        plans = build_patient_split_plans(
+            patient_ids=patient_table[patient_col].to_numpy(),
+            patient_labels=patient_table["_label"].to_numpy(),
+            n_splits=n_splits,
+            seed=fold_seed,
+            holdout_per_class=holdout_per_class,
+            holdout_seed=holdout_seed,
+        )
+
+        out: list[PatientSplit] = []
+        for plan in plans:
+            train_df = (
+                df_full[df_full[patient_col].isin(plan.train_patient_ids)]
+                .reset_index(drop=True)
+            )
+            val_df = (
+                df_full[df_full[patient_col].isin(plan.val_patient_ids)]
+                .reset_index(drop=True)
+            )
+            holdout_df = None
+            if plan.holdout_patient_ids:
+                holdout_df = (
+                    df_full[df_full[patient_col].isin(plan.holdout_patient_ids)]
+                    .reset_index(drop=True)
+                )
+            out.append(PatientSplit(train=train_df, val=val_df, holdout=holdout_df))
+        return out
+
+    def _validate_columns(self, df: pd.DataFrame, label_col: str, patient_col: Optional[str] = None) -> None:
+        pcol = patient_col or self.patient_col
+        if pcol not in df.columns:
+            raise ValueError(f"Missing required patient column: {pcol!r}")
+        if label_col not in df.columns:
+            raise ValueError(f"Missing required label column: {label_col!r}")
+
+    def _patient_label_table(
+        self,
+        df: pd.DataFrame,
+        label_col: str,
+        patient_col: Optional[str] = None,
+    ) -> pd.DataFrame:
+        pcol = patient_col or self.patient_col
+        grouped = (
+            df.groupby(pcol, as_index=False)[label_col]
+            .agg(lambda x: x.mode().iloc[0] if not x.mode().empty else x.iloc[0])
+            .rename(columns={label_col: "_label"})
+            .sort_values(pcol)
+            .reset_index(drop=True)
+        )
+        if grouped.empty:
+            raise ValueError("No patients available for splitting")
+        return grouped
+
+
+def _can_stratify(labels: np.ndarray, n_splits: int) -> bool:
+    if labels.size == 0:
+        return False
+    unique, counts = np.unique(labels, return_counts=True)
+    if len(unique) < 2:
+        return False
+    return bool(np.all(counts >= n_splits))

classes/v2/towers.py

@@ -0,0 +1,279 @@
+from __future__ import annotations
+
+import math
+from typing import Optional
+
+import torch
+from torch import nn
+from torchvision import transforms
+
+from classes.backbones import BACKBONES, list_names, load_backbone_weights
+from classes.SE_attention import SEBlock
+from classes.v2.data_bundle import DataBundle
+
+
+def build_backbone(name: str, freeze_ratio: float = 0.0, augment: bool = True):
+    """
+    Operational builder:
+      - instantiate with DEFAULT weights
+      - strip classifier → features
+      - apply ratio-based freezing over coarse blocks
+      - return (model, out_dim, transform)
+    """
+    key = (name or "").lower()
+    if key not in BACKBONES:
+        raise ValueError(f"Unsupported backbone '{name}'. Valid options: {list_names()}")
+
+    spec = BACKBONES[key]
+    m = spec.ctor(weights=spec.weights_default)
+    out_dim, m = spec.strip(m)
+    load_backbone_weights(key, m)
+
+    # transforms: use the weights’ mean/std, but keep your augmentation pipeline
+    mean = getattr(spec.weights_default, "meta", {}).get("mean", (0.485, 0.456, 0.406))
+    std = getattr(spec.weights_default, "meta", {}).get("std", (0.229, 0.224, 0.225))
+    crop = 299 if key == "inception_v3" else 224
+
+    if augment:
+        transform = transforms.Compose(
+            [
+                transforms.Resize(256),
+                transforms.CenterCrop(crop),
+                transforms.RandomHorizontalFlip(),
+                transforms.RandomVerticalFlip(),
+                transforms.RandomRotation(15),
+                transforms.ColorJitter(0.1, 0.1, 0.1, 0.05),
+                transforms.ToTensor(),
+                transforms.Normalize(mean=mean, std=std),
+            ]
+        )
+    else:
+        transform = transforms.Compose(
+            [
+                transforms.Resize(256),
+                transforms.CenterCrop(crop),
+                transforms.ToTensor(),
+                transforms.Normalize(mean=mean, std=std),
+            ]
+        )
+
+    # ratio-based freezing: freeze earliest floor(N * freeze_ratio) blocks
+    fr = max(0.0, min(1.0, float(freeze_ratio)))
+    blocks = spec.blocks(m)
+    n = len(blocks)
+    freeze_n = int(math.floor(n * fr))
+    for b in blocks[:freeze_n]:
+        for p in b.parameters():
+            p.requires_grad = False
+
+    return m, out_dim, transform
+
+
+class ImageTower(nn.Module):
+    """
+    Vision backbone → pooled features.
+    - backbone: one of list_names() (default 'efficientnet_b0')
+    - always DEFAULT torchvision weights
+    - freeze_ratio ∈ [0,1] freezes earliest floor(N*freeze_ratio) blocks
+    - returns [N, out_dim] features from backbone forward
+    """
+
+    def __init__(
+        self,
+        backbone: str = "efficientnet_b0",
+        freeze_ratio: float = 0.0,
+        use_se: bool = False,
+        se_reduction: int = 16,
+        se_pre_norm: bool = True,
+        augment: bool = True,
+        geometry_dim: int = 0,
+    ):
+        super().__init__()
+        self.backbone, base_dim, self.transform = build_backbone(
+            backbone, freeze_ratio, augment=augment
+        )
+        self._name = backbone
+        # Keep ordered blocks for dynamic freezing/thawing
+        key = (self._name or "").lower()
+        self._spec = BACKBONES[key]
+        self._blocks = self._spec.blocks(self.backbone)
+        # Optional tower-level SE over the final feature vector
+        self.base_dim = base_dim
+        self.geometry_dim = max(0, int(geometry_dim))
+        self.out_dim = self.base_dim + self.geometry_dim
+        self.tower_ln = nn.LayerNorm(self.base_dim) if se_pre_norm else nn.Identity()
+        self.tower_se = (
+            SEBlock(self.base_dim, reduction=se_reduction, residual=True)
+            if use_se
+            else None
+        )
+
+    def forward(
+        self, x: torch.Tensor, geometry: Optional[torch.Tensor] = None
+    ) -> torch.Tensor:
+        y = self.backbone(x)
+        # sanity: pooled features, not logits
+        assert y.dim() == 2 and y.size(1) == self.base_dim, (
+            f"Expected features [N,{self.base_dim}], got {tuple(y.shape)}"
+        )
+        if self.tower_se is not None:
+            y, _ = self.tower_se(self.tower_ln(y))
+        if self.geometry_dim > 0:
+            if geometry is None or geometry.numel() == 0:
+                geom = torch.zeros(
+                    y.size(0), self.geometry_dim, device=y.device, dtype=y.dtype
+                )
+            else:
+                if geometry.dim() == 1:
+                    geom = geometry.unsqueeze(0)
+                else:
+                    geom = geometry
+                geom = geom.to(device=y.device, dtype=y.dtype)
+                if geom.size(0) != y.size(0):
+                    raise ValueError(
+                        f"Geometry batch size mismatch: {geom.size(0)} vs {y.size(0)}"
+                    )
+                if geom.size(1) != self.geometry_dim:
+                    raise ValueError(
+                        f"Expected geometry dim {self.geometry_dim}, got {geom.size(1)}"
+                    )
+            y = torch.cat([y, geom], dim=1)
+        return y
+
+    def set_freeze_ratio(self, ratio: float):
+        """Dynamically freeze earliest floor(N*ratio) backbone blocks."""
+        r = max(0.0, min(1.0, float(ratio)))
+        n = len(self._blocks)
+        freeze_n = int(math.floor(n * r))
+        # Unfreeze all first
+        for b in self._blocks:
+            for p in b.parameters():
+                p.requires_grad = True
+        # Freeze earliest blocks
+        for b in self._blocks[:freeze_n]:
+            for p in b.parameters():
+                p.requires_grad = False
+
+
+class SiameseImageTower(nn.Module):
+    """
+    Shared-weight bilateral image tower.
+
+    Runs OD and OS images through a single shared backbone, then returns
+    cat([f_mean, f_delta]) where:
+        f_mean  = (f_od + f_os) / 2   -- shared bilateral representation
+        f_delta = f_od - f_os          -- asymmetry, signed OD-relative
+
+    out_dim = 2 * backbone_out_dim
+
+    When x_os is None (single-eye fallback):
+        f_mean  = f_od
+        f_delta = zeros
+    so the module degrades gracefully when only one eye is available.
+
+    The shared backbone means both eyes contribute to every gradient update,
+    effectively doubling the training signal for the visual pathway without
+    doubling parameters.
+    """
+
+    def __init__(
+        self,
+        backbone: str = "efficientnet_b0",
+        freeze_ratio: float = 0.0,
+        use_se: bool = False,
+        se_reduction: int = 16,
+        se_pre_norm: bool = True,
+        augment: bool = True,
+    ):
+        super().__init__()
+        self._tower = ImageTower(
+            backbone=backbone,
+            freeze_ratio=freeze_ratio,
+            use_se=use_se,
+            se_reduction=se_reduction,
+            se_pre_norm=se_pre_norm,
+            augment=augment,
+            geometry_dim=0,
+        )
+        self.out_dim = self._tower.out_dim * 2
+        self.transform = self._tower.transform
+
+    def forward(
+        self,
+        x_od: torch.Tensor,
+        x_os: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        f_od = self._tower(x_od)
+        if x_os is None:
+            f_mean = f_od
+            f_delta = torch.zeros_like(f_od)
+        else:
+            f_os = self._tower(x_os)
+            f_mean = (f_od + f_os) * 0.5
+            f_delta = f_od - f_os
+        return torch.cat([f_mean, f_delta], dim=1)
+
+    def set_freeze_ratio(self, ratio: float) -> None:
+        """Delegates to the shared inner tower."""
+        self._tower.set_freeze_ratio(ratio)
+
+
+class MDTower(nn.Module):
+    """MLP over DataBundle.vectorize_row outputs (convert to torch inside tower)."""
+
+    def __init__(
+        self,
+        clinical_data: DataBundle,
+        hidden_dim: int = 128,
+        dropout: float = 0.1,
+        use_se: bool = False,
+        se_reduction: int = 16,
+        se_pre_norm: bool = True,
+    ):
+        super().__init__()
+        self.feature_dim = clinical_data.feature_dim
+        self.out_dim = hidden_dim
+        # two-block MLP so we can optionally freeze/thaw per block
+        self.block0 = nn.Sequential(
+            nn.Linear(self.feature_dim, hidden_dim),
+            nn.LayerNorm(hidden_dim),
+            nn.ReLU(inplace=True),
+            nn.Dropout(dropout),
+        )
+        self.block1 = nn.Sequential(
+            nn.Linear(hidden_dim, hidden_dim),
+            nn.ReLU(inplace=True),
+        )
+        self.net = nn.Sequential(self.block0, self.block1)
+        self.tower_ln = nn.LayerNorm(hidden_dim) if se_pre_norm else nn.Identity()
+        self.tower_se = (
+            SEBlock(hidden_dim, reduction=se_reduction, residual=True)
+            if use_se
+            else None
+        )
+
+    def forward(self, meta_np_or_torch) -> torch.Tensor:
+        if isinstance(meta_np_or_torch, torch.Tensor):
+            x = meta_np_or_torch
+        else:
+            x = torch.as_tensor(meta_np_or_torch, dtype=torch.float32)
+        h = self.net(x)
+        if self.tower_se is not None:
+            h, _ = self.tower_se(self.tower_ln(h))
+        return h
+
+    def set_freeze_ratio(self, ratio: float):
+        """Optionally freeze earliest blocks of the MLP."""
+        r = max(0.0, min(1.0, float(ratio)))
+        # Unfreeze all
+        for p in self.block0.parameters():
+            p.requires_grad = True
+        for p in self.block1.parameters():
+            p.requires_grad = True
+        # Freeze earliest blocks based on ratio threshold
+        if r >= 0.5:
+            for p in self.block0.parameters():
+                p.requires_grad = False
+        if r >= 1.0:
+            for p in self.block1.parameters():
+                p.requires_grad = False

classes/v2/transforms.py

@@ -0,0 +1,315 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Callable, Iterable, Optional, Tuple, Union
+import numpy as np
+from PIL import Image
+
+from torchvision import transforms
+
+from classes.backbones import BACKBONES
+
+
+IMAGENET_MEAN: Tuple[float, float, float] = (0.485, 0.456, 0.406)
+IMAGENET_STD: Tuple[float, float, float] = (0.229, 0.224, 0.225)
+
+
+@dataclass
+class ImageTransformConfig:
+    """
+    Mirrors the hypertower v1 preprocessing:
+      - Resize(256)
+      - CenterCrop(crop)
+      - Optional augmentations (H/V flip, rotation, color jitter)
+      - ToTensor + Normalize(mean/std)
+    """
+
+    crop_size: int = 224
+    resize_size: int = 256
+    mean: Tuple[float, float, float] = IMAGENET_MEAN
+    std: Tuple[float, float, float] = IMAGENET_STD
+    augment: bool = True
+    rotation_deg: int = 15
+    color_jitter: Tuple[float, float, float, float] = (0.1, 0.1, 0.1, 0.05)
+    hflip: bool = True
+    vflip: bool = True
+
+    def build(self) -> transforms.Compose:
+        ops = [
+            transforms.Resize(self.resize_size),
+            transforms.CenterCrop(self.crop_size),
+        ]
+        if self.augment:
+            if self.hflip:
+                ops.append(transforms.RandomHorizontalFlip())
+            if self.vflip:
+                ops.append(transforms.RandomVerticalFlip())
+            if self.rotation_deg:
+                ops.append(transforms.RandomRotation(self.rotation_deg))
+            if self.color_jitter:
+                ops.append(transforms.ColorJitter(*self.color_jitter))
+        ops.extend(
+            [
+                transforms.ToTensor(),
+                transforms.Normalize(mean=self.mean, std=self.std),
+            ]
+        )
+        return transforms.Compose(ops)
+
+
+def backbone_transform_config(backbone_name: str, augment: bool = True) -> ImageTransformConfig:
+    """
+    Build a transform config that matches v1 ImageTower/backbone preprocessing.
+    Uses DEFAULT weights mean/std and InceptionV3 crop size when relevant.
+    """
+    key = (backbone_name or "").lower()
+    if key not in BACKBONES:
+        raise ValueError(f"Unsupported backbone '{backbone_name}'.")
+    spec = BACKBONES[key]
+    mean = getattr(spec.weights_default, "meta", {}).get("mean", IMAGENET_MEAN)
+    std = getattr(spec.weights_default, "meta", {}).get("std", IMAGENET_STD)
+    crop = 299 if key == "inception_v3" else 224
+    return ImageTransformConfig(crop_size=crop, mean=mean, std=std, augment=augment)
+
+
+def build_backbone_transform(backbone_name: str, augment: bool = True) -> transforms.Compose:
+    return backbone_transform_config(backbone_name, augment=augment).build()
+
+
+def build_imagenet_transform(augment: bool = True, crop_size: int = 224) -> transforms.Compose:
+    return ImageTransformConfig(crop_size=crop_size, augment=augment).build()
+
+
+@dataclass
+class ResizeTransform:
+    size: Union[int, Tuple[int, int]] = 256
+    interpolation: int = Image.BILINEAR
+
+    def __post_init__(self) -> None:
+        self._op = transforms.Resize(self.size, interpolation=self.interpolation)
+
+    def __call__(self, image: Image.Image) -> Image.Image:
+        return self._op(image)
+
+
+@dataclass
+class CenterCropTransform:
+    size: Union[int, Tuple[int, int]] = 224
+
+    def __post_init__(self) -> None:
+        self._op = transforms.CenterCrop(self.size)
+
+    def __call__(self, image: Image.Image) -> Image.Image:
+        return self._op(image)
+
+
+class UnetMaskProvider:
+    """
+    Placeholder for a UNet-powered mask provider.
+    This will be replaced once a UNet tower is wired in.
+    """
+
+    def __call__(self, image: Image.Image, image_path: Optional[str] = None):
+        raise NotImplementedError("UNet mask provider is not wired yet.")
+
+
+@dataclass
+class ROICropTransform:
+    """
+    Crop an image using a binary mask (GT or UNet).
+    Expects a mask of the same spatial size as the image; nonzero pixels are ROI.
+    """
+
+    mask_source: str = "gt"  # "gt" | "unet"
+    mask_provider: Optional[Callable[[Image.Image, Optional[str]], np.ndarray]] = None
+    scale: float = 2.5
+    target_size: Optional[Tuple[int, int]] = (224, 224)
+    fallback_to_original: bool = True
+
+    def __post_init__(self) -> None:
+        if self.mask_source not in {"gt", "unet"}:
+            raise ValueError(f"mask_source must be 'gt' or 'unet', got '{self.mask_source}'.")
+
+    def __call__(
+        self,
+        image: Image.Image,
+        mask: Optional[Union[np.ndarray, Image.Image]] = None,
+        image_path: Optional[str] = None,
+    ) -> Image.Image:
+        resolved_mask = mask
+        if resolved_mask is None and self.mask_provider is not None:
+            resolved_mask = self.mask_provider(image, image_path)
+        if resolved_mask is None:
+            if self.fallback_to_original:
+                return image
+            raise ValueError("ROI crop requested but no mask provided.")
+
+        mask_arr = (
+            np.asarray(resolved_mask)
+            if not isinstance(resolved_mask, Image.Image)
+            else np.array(resolved_mask)
+        )
+        if mask_arr.ndim == 3:
+            mask_arr = mask_arr[..., 0]
+        mask_arr = mask_arr > 0
+        if not np.any(mask_arr):
+            return image if self.fallback_to_original else image
+
+        ys, xs = np.where(mask_arr)
+        y_min, y_max = ys.min(), ys.max()
+        x_min, x_max = xs.min(), xs.max()
+        cx = (x_min + x_max) / 2.0
+        cy = (y_min + y_max) / 2.0
+        width = (x_max - x_min + 1)
+        height = (y_max - y_min + 1)
+        size = max(width, height) * float(self.scale)
+
+        left = int(round(cx - size / 2))
+        right = int(round(cx + size / 2))
+        upper = int(round(cy - size / 2))
+        lower = int(round(cy + size / 2))
+
+        left = max(0, left)
+        upper = max(0, upper)
+        right = min(image.width, right)
+        lower = min(image.height, lower)
+        crop = image.crop((left, upper, right, lower))
+        if self.target_size is not None:
+            crop = crop.resize(self.target_size, Image.BILINEAR)
+        return crop
+
+
+@dataclass
+class JitterBundleTransform:
+    """
+    Augmentations bundle: flips, rotation, color jitter.
+    """
+
+    hflip: bool = True
+    vflip: bool = True
+    rotation_deg: int = 15
+    color_jitter: Optional[Tuple[float, float, float, float]] = (0.1, 0.1, 0.1, 0.05)
+
+    def __post_init__(self) -> None:
+        ops = []
+        if self.hflip:
+            ops.append(transforms.RandomHorizontalFlip())
+        if self.vflip:
+            ops.append(transforms.RandomVerticalFlip())
+        if self.rotation_deg:
+            ops.append(transforms.RandomRotation(self.rotation_deg))
+        if self.color_jitter:
+            ops.append(transforms.ColorJitter(*self.color_jitter))
+        self._op = transforms.Compose(ops) if ops else None
+
+    def __call__(self, image: Image.Image) -> Image.Image:
+        if self._op is None:
+            return image
+        return self._op(image)
+
+
+TRANSFORM_REGISTRY = {
+    "resize": ResizeTransform,
+    "roi_crop": ROICropTransform,
+    "center_crop": CenterCropTransform,
+    "jitter_bundle": JitterBundleTransform,
+}
+
+
+def _parse_color_jitter(value: Optional[Union[str, Iterable[float]]]) -> Optional[Tuple[float, float, float, float]]:
+    if value is None:
+        return None
+    if isinstance(value, str):
+        parts = [p.strip() for p in value.split(",") if p.strip()]
+        if not parts:
+            return None
+        try:
+            nums = [float(p) for p in parts]
+        except ValueError:
+            return None
+        if len(nums) == 1:
+            return (nums[0], nums[0], nums[0], nums[0])
+        if len(nums) >= 4:
+            return (nums[0], nums[1], nums[2], nums[3])
+        return tuple(nums + [nums[-1]] * (4 - len(nums)))  # pad to length 4
+    try:
+        vals = list(value)
+    except TypeError:
+        return None
+    if not vals:
+        return None
+    vals = [float(v) for v in vals]
+    if len(vals) == 1:
+        return (vals[0], vals[0], vals[0], vals[0])
+    if len(vals) >= 4:
+        return (vals[0], vals[1], vals[2], vals[3])
+    return tuple(vals + [vals[-1]] * (4 - len(vals)))
+
+
+def build_transform_chain(
+    transform_specs: Iterable[object],
+    *,
+    backbone_name: str,
+    augment: bool = True,
+    mask_provider: Optional[Callable[[Image.Image, Optional[str]], np.ndarray]] = None,
+    strict: bool = True,
+) -> transforms.Compose:
+    """
+    Build an image transform pipeline from a list of transform specs plus the
+    standard ToTensor + Normalize steps. This mirrors the V1 preprocessing
+    but uses the explicit transform nodes from config.
+    """
+    ops: list[Callable[[Image.Image], Image.Image]] = []
+    for spec in transform_specs:
+        transform_type = getattr(spec, "transform_type", None)
+        params = getattr(spec, "params", None)
+        if transform_type is None and isinstance(spec, dict):
+            transform_type = spec.get("transformType") or spec.get("transform_type")
+            params = spec
+        params = params or {}
+
+        if transform_type == "resize":
+            size = params.get("resizeSize", 256)
+            ops.append(ResizeTransform(size=size))
+        elif transform_type == "center_crop":
+            size = params.get("centerCropSize", 224)
+            ops.append(CenterCropTransform(size=size))
+        elif transform_type == "jitter_bundle":
+            if not augment:
+                continue
+            jitter = JitterBundleTransform(
+                hflip=bool(params.get("jitterHFlip", True)),
+                vflip=bool(params.get("jitterVFlip", True)),
+                rotation_deg=int(params.get("jitterRotation", 15) or 0),
+                color_jitter=_parse_color_jitter(params.get("jitterColor"))
+                if params.get("jitterColorEnabled", True)
+                else None,
+            )
+            ops.append(jitter)
+        elif transform_type == "roi_crop":
+            roi = ROICropTransform(
+                mask_source=params.get("roiMaskSource", "gt"),
+                mask_provider=mask_provider,
+                scale=float(params.get("roiScale", 2.5)),
+                target_size=(int(params.get("roiTargetSize", 224)), int(params.get("roiTargetSize", 224)))
+                if params.get("roiTargetSize") is not None
+                else None,
+                fallback_to_original=bool(params.get("roiFallback", True)),
+            )
+            if roi.mask_provider is None and roi.mask_source == "unet":
+                if strict:
+                    raise ValueError("ROI crop requires a mask provider for 'unet' source.")
+            ops.append(roi)
+        else:
+            if strict:
+                raise ValueError(f"Unsupported transform type: {transform_type!r}")
+
+    # Always end with tensor + normalize, using backbone defaults
+    cfg = backbone_transform_config(backbone_name, augment=augment)
+    ops.extend(
+        [
+            transforms.ToTensor(),
+            transforms.Normalize(mean=cfg.mean, std=cfg.std),
+        ]
+    )
+    return transforms.Compose(ops)