reworked iop_corr, added explainability tools and plotting tools, cleanup codebase

classes/v2/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/v2/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/v2/bridges.py

@@ -3,7 +3,7 @@ from __future__ import annotations
 import torch
 import torch.nn as nn
 
-from classes.SE_attention import SEBlock, SEGateLogger
+from classes.v2.SE_attention import SEBlock, SEGateLogger
 
 
 class Bridge(nn.Module):

classes/v2/croppers.py

@@ -10,9 +10,30 @@ import torch
 from PIL import Image, ImageDraw
 from torchvision import transforms
 
-from classes.geometry_features import compute_geometry_features, disc_cup_from_mask_image
-from classes.refuge_classification import _geometry_from_mask
-from classes.unet_segmenter import UNetSegmenter
+from classes.v2.geometry_features import compute_geometry_features, disc_cup_from_mask_image
+from classes.v2.unet_segmenter import UNetSegmenter
+
+
+def _geometry_from_mask(mask: np.ndarray, scale: float) -> Dict:
+    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,
+    }
 
 
 class UNetImageCropper:

classes/v2/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/v2/models.py

@@ -219,6 +219,15 @@ def _set_single_phase(model: SingleEyeHT, phase: str) -> None:
     # Ablation modes have no fusion bridge; fused_warmup is meaningless — treat as tower_warmup
     if bridge_mode in ("image_only", "metadata_only") and phase == "fused_warmup":
         phase = "tower_warmup"
+    if phase == "md_warmup":
+        _set_requires_grad(model.img_tower, False)
+        _set_requires_grad(model.md_tower, True)
+        _set_requires_grad(model.bridge.classifier_img, False)
+        _set_requires_grad(model.bridge.classifier_md, True)
+        _set_requires_grad(model.bridge.W_img, False)
+        _set_requires_grad(model.bridge.W_md, False)
+        _set_requires_grad(model.bridge.classifier_fused, False)
+        return
     if phase == "tower_warmup":
         _set_requires_grad(model.img_tower, bridge_mode != "metadata_only")
         _set_requires_grad(model.md_tower, bridge_mode != "image_only")
@@ -282,12 +291,27 @@ def train_single_epoch(
         x = batch.get("image_1")
         m = batch.get("matrix_1")
         y = batch.get("label_1")
+        if phase == "md_warmup":
+            if not torch.is_tensor(m):
+                continue
+            m = m.to(device)
+            y = _to_label_tensor(y, device)
+            md_feats = model.md_tower(m)
+            logits = model.bridge.classifier_md(md_feats)
+            loss = F.cross_entropy(logits, y)
+            opt.zero_grad(); loss.backward(); opt.step()
+            bs = y.shape[0]
+            total_loss += float(loss.item()) * bs
+            total_correct += int((logits.argmax(1) == y).sum())
+            total_n += bs
+            continue
         if not torch.is_tensor(x) or not torch.is_tensor(m):
             continue
         x = x.to(device)
         m = m.to(device)
         y = _to_label_tensor(y, device)
         bridge_mode = model.bridge.mode
+
         img_feats = None if bridge_mode == "metadata_only" else model.img_tower(x)
         md_feats  = None if bridge_mode == "image_only"    else model.md_tower(m)
 

classes/v2/papila_builders.py

@@ -33,10 +33,21 @@ def _nearest_pachy_key(x: float) -> int:
     return int(_PACHY_KEYS[idx])
 
 
+# Ratio derived from patients with both Pneumatic and Perkins readings (n=41, OD+OS combined).
+# Pneumatic / Perkins mean ratio = 1.158; applied to Perkins-only rows to put them on the
+# Pneumatic scale before IOP_corr is computed.
+_PERKINS_TO_PNEUMATIC_RATIO: float = 1.158
+
+
 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
+    """Prefer Pneumatic; scale Perkins to Pneumatic scale if Pneumatic is absent."""
+    pneumatic = row.get("Pneumatic", np.nan)
+    if not pd.isna(pneumatic):
+        return float(pneumatic)
+    perkins = row.get("Perkins", np.nan)
+    if not pd.isna(perkins):
+        return float(perkins) * _PERKINS_TO_PNEUMATIC_RATIO
+    return np.nan
 
 
 def _correct_iop(raw_iop: float, pachy: float) -> float:

classes/v2/towers.py

@@ -7,8 +7,8 @@ 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.backbones import BACKBONES, list_names, load_backbone_weights
+from classes.v2.SE_attention import SEBlock
 from classes.v2.data_bundle import DataBundle
 
 

classes/v2/transforms.py

@@ -7,7 +7,7 @@ from PIL import Image
 
 from torchvision import transforms
 
-from classes.backbones import BACKBONES
+from classes.v2.backbones import BACKBONES
 
 
 IMAGENET_MEAN: Tuple[float, float, float] = (0.485, 0.456, 0.406)

classes/v2/unet_segmenter.py

@@ -0,0 +1,894 @@
+"""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")
+            # Disc = ALL non-background pixels (full optic disc: rim + cup combined).
+            # Previously this was rim-only, which caused the cup structural prior
+            # (cup & disc) to produce empty cup masks since cup and rim don't overlap.
+            disc_mask = (arr != bg_val).astype(np.uint8)
+            # Cup = the darkest non-background value (0 in REFUGE = inner cup region).
+            # Using min-value rather than frequency avoids swapping when cup area > rim area.
+            cup_val = int(np.min(vals)) if vals.size > 1 else None
+            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/v2_hypertower.py

@@ -206,6 +206,9 @@ class V2HyperTower:
                         help="Single-eye model tower warmup (overrides --warmup-tower-epochs).")
         ap.add_argument("--single-warmup-fused-epochs", type=int, default=None,
                         help="Single-eye model fused warmup (overrides --warmup-fused-epochs).")
+        ap.add_argument("--warmup-md-epochs", type=int, default=0,
+                        help="MD-only warmup epochs before tower warmup. Trains only md_tower + "
+                             "classifier_md (no CNN forward pass, so 50-100 epochs is cheap).")
         ap.add_argument("--bilat-warmup-tower-epochs", type=int, default=None,
                         help="Bilateral model tower warmup (overrides --warmup-tower-epochs).")
         ap.add_argument("--bilat-warmup-fused-epochs", type=int, default=None,
@@ -424,7 +427,8 @@ class V2HyperTower:
             if getattr(args, "single_warmup_fused_epochs", None) is not None
             else int(_global_warmup_fused) if _global_warmup_fused is not None else 2
         )
-        _total_epochs = _warmup_tower + _warmup_fused + int(args.epochs) + fusion_epochs
+        _warmup_md    = int(getattr(args, "warmup_md_epochs", 0))
+        _total_epochs = _warmup_md + _warmup_tower + _warmup_fused + int(args.epochs) + fusion_epochs
 
         # sample IDs depend on mode: single uses eye IDs, others use patient IDs
         if tower_mode in ("single", "classic"):
@@ -572,6 +576,7 @@ class V2HyperTower:
             "run_id": run_name,
             "backbone": args.backbone,
             "epochs": args.epochs,
+            "warmup_md_epochs": getattr(args, "warmup_md_epochs", 0),
             "warmup_tower_epochs": args.warmup_tower_epochs,
             "warmup_fused_epochs": args.warmup_fused_epochs,
             "single_warmup_tower_epochs": args.single_warmup_tower_epochs,
@@ -659,6 +664,7 @@ class V2HyperTower:
             if getattr(args, "bilat_warmup_fused_epochs", None) is not None
             else int(global_warmup_fused) if global_warmup_fused is not None else 3
         )
+        single_warmup_md = int(getattr(args, "warmup_md_epochs", 0)) if run_single else 0
         if not run_single:
             single_warmup_tower = 0
             single_warmup_fused = 0
@@ -666,7 +672,7 @@ class V2HyperTower:
             bilat_warmup_tower = 0
             bilat_warmup_fused = 0
         main_epochs = int(args.epochs)
-        total_single_epochs = (single_warmup_tower + single_warmup_fused + main_epochs) if run_single else 0
+        total_single_epochs = (single_warmup_md + single_warmup_tower + single_warmup_fused + main_epochs) if run_single else 0
         total_bilat_epochs = (bilat_warmup_tower + bilat_warmup_fused + main_epochs) if run_bilat else 0
         total_epochs = max(total_single_epochs, total_bilat_epochs)
 
@@ -744,6 +750,7 @@ class V2HyperTower:
         train_single_loader = None
         train_eval_loader   = None   # non-shuffled, no sampler — for per-epoch train logging
         train_bilat_loader  = None
+        md_only_loader      = None   # image-free loader for md_warmup phase
         if run_single:
             single_sampler = build_balanced_sampler(eye_train) if use_balanced else None
             train_single_loader = make_loader(
@@ -761,6 +768,20 @@ class V2HyperTower:
                 shuffle=False,
                 **loader_kw,
             )
+            if single_warmup_md > 0:
+                # MD-only loader: drop image_1 so PIL never opens files during md_warmup.
+                # Always use balanced sampling for md_warmup — MD features alone are weaker
+                # than images and collapse to majority class without class balancing.
+                slots_md_only = {k: v for k, v in slots_eye.items() if k != "image_1"}
+                md_warmup_sampler = single_sampler if single_sampler is not None else build_balanced_sampler(eye_train)
+                md_only_loader = make_loader(
+                    eye_train, slots_md_only,
+                    image_transform=None,
+                    image_preprocessor=None,
+                    shuffle=True,
+                    sampler=md_warmup_sampler,
+                    **loader_kw,
+                )
         if run_bilat:
             bilat_sampler = build_balanced_sampler(bilat_train) if use_balanced else None
             train_bilat_loader = make_loader(
@@ -907,7 +928,7 @@ class V2HyperTower:
             print(
                 f"  [fold {fold+1}]  single_train_n={len(eye_train)} (eye-level)  "
                 f"val_n={len(bilat_val)}  "
-                f"single_warmup={single_warmup_tower}+{single_warmup_fused} total={total_single_epochs}",
+                f"single_warmup=md{single_warmup_md}+twr{single_warmup_tower}+fus{single_warmup_fused} total={total_single_epochs}",
                 flush=True,
             )
         else:
@@ -919,17 +940,20 @@ class V2HyperTower:
             )
 
         # ---- epoch loop ------------------------------------------------
+        _prev_phase_single = "inactive"  # used to detect md_warmup → next phase transition
         for epoch in range(total_epochs):
             if not run_single:
                 phase_single, main_epoch_single, single_active = "inactive", 0, False
-            elif epoch < single_warmup_tower:
+            elif epoch < single_warmup_md:
+                phase_single, main_epoch_single, single_active = "md_warmup", 0, True
+            elif epoch < (single_warmup_md + single_warmup_tower):
                 phase_single, main_epoch_single, single_active = "tower_warmup", 0, True
-            elif epoch < (single_warmup_tower + single_warmup_fused):
+            elif epoch < (single_warmup_md + single_warmup_tower + single_warmup_fused):
                 phase_single, main_epoch_single, single_active = "fused_warmup", 0, True
             elif epoch < total_single_epochs:
                 phase_single, main_epoch_single, single_active = (
                     "main",
-                    epoch - single_warmup_tower - single_warmup_fused + 1,
+                    epoch - single_warmup_md - single_warmup_tower - single_warmup_fused + 1,
                     True,
                 )
             else:
@@ -951,8 +975,9 @@ class V2HyperTower:
                 phase_bilat, main_epoch_bilat, bilat_active = "done", main_epochs, False
 
             if run_single and single_active:
+                _active_loader = md_only_loader if phase_single == "md_warmup" else train_single_loader
                 sl_loss, sl_acc = train_single_epoch(
-                    single, train_single_loader, opt_single, device,
+                    single, _active_loader, opt_single, device,
                     phase=phase_single, bcd_prob=float(args.bcd_prob),
                 )
             else:
@@ -966,7 +991,9 @@ class V2HyperTower:
             else:
                 bl_loss, bl_acc = nan, nan
 
-            if run_single and tower_mode == "single":
+            _skip_val_eval = (phase_single == "md_warmup")
+
+            if run_single and tower_mode == "single" and not _skip_val_eval:
                 y_cl, p_cl, p_cl_img, p_cl_md = collect_probs_single_components(
                     single, val_loader, device, aggregate_patient=False
                 )
@@ -980,7 +1007,7 @@ class V2HyperTower:
                 en_acc = en_auc = nan
                 en_n = 0
                 en_acc_img = en_acc_md = en_auc_img = en_auc_md = nan
-            elif run_single and tower_mode == "ensemble":
+            elif run_single and tower_mode == "ensemble" and not _skip_val_eval:
                 (y_en,
                  _p_en_f_od, _p_en_i_od, _p_en_m_od,
                  _p_en_f_os, _p_en_i_os, _p_en_m_os,
@@ -1010,7 +1037,7 @@ class V2HyperTower:
                 cl_acc_img = cl_acc_md = en_acc_img = en_acc_md = nan
                 cl_auc_img = cl_auc_md = en_auc_img = en_auc_md = nan
 
-            if run_bilat:
+            if run_bilat and not _skip_val_eval:
                 y_bi, p_bi, p_bi_img, p_bi_md = collect_probs_bilateral_components(
                     bilateral, val_loader, device
                 )
@@ -1034,7 +1061,7 @@ class V2HyperTower:
             p_cl_h = p_cl_h_img = p_cl_h_md = _z2
             p_en_h = p_en_h_img = p_en_h_md = _z2
 
-            if holdout_loader is not None:
+            if holdout_loader is not None and not _skip_val_eval:
                 if run_single and tower_mode == "single":
                     y_cl_h, p_cl_h, p_cl_h_img, p_cl_h_md = collect_probs_single_components(
                         single, holdout_loader, device, aggregate_patient=False
@@ -1092,7 +1119,7 @@ class V2HyperTower:
             tr_fe_corr = tr_fe_err = tr_n = 0
             y_tr = np.array([], dtype=np.int64)
             p_tr_f = p_tr_i = p_tr_m = np.zeros((0, num_classes), dtype=np.float32)
-            if run_single and train_eval_loader is not None:
+            if run_single and train_eval_loader is not None and not _skip_val_eval:
                 y_tr, p_tr_f, p_tr_i, p_tr_m, tr_ids = collect_probs_eye_level(
                     single, train_eval_loader, device, return_ids=True
                 )
@@ -1284,16 +1311,33 @@ class V2HyperTower:
                 **cm_row,
             }, optional_cols=epoch_fields)
 
+            # ---- md_warmup progress bar (replaces per-epoch print) --------
+            if phase_single == "md_warmup":
+                _bar_w = 30
+                _filled = int(_bar_w * (epoch + 1) / single_warmup_md)
+                _bar = "#" * _filled + "-" * (_bar_w - _filled)
+                _bar_msg = (
+                    f"  [fold {fold+1}] md_warmup [{_bar}] "
+                    f"{epoch + 1}/{single_warmup_md}  loss={sl_loss:.4f}"
+                )
+                print(f"\r{_bar_msg}", end="", flush=True)
+                fold_logger.info(_bar_msg)
+                _prev_phase_single = phase_single
+                continue  # skip normal log block entirely
+
+            if _prev_phase_single == "md_warmup":
+                print()  # seal the progress bar line
+
             if args.log_every > 0 and (epoch + 1) % args.log_every == 0:
                 hld_auc = target_holdout_single_auc if run_single else bi_auc_h
                 hld_suffix = f"  hld_auc={hld_auc:.4f}" if holdout_loader is not None else ""
 
                 # Human-readable phase progress for console logs.
                 if phase_single == "tower_warmup":
-                    single_phase_epoch = epoch + 1
+                    single_phase_epoch = epoch - single_warmup_md + 1
                     single_phase_total = single_warmup_tower
                 elif phase_single == "fused_warmup":
-                    single_phase_epoch = epoch - single_warmup_tower + 1
+                    single_phase_epoch = epoch - single_warmup_md - single_warmup_tower + 1
                     single_phase_total = single_warmup_fused
                 else:
                     single_phase_epoch = main_epoch_single
@@ -1343,6 +1387,8 @@ class V2HyperTower:
                 print(msg, flush=True)
                 fold_logger.info(msg)
 
+            _prev_phase_single = phase_single
+
         fold_logger.close()
 
         if args.save_checkpoints: