Add distributed server implementation and protocol definitions

- Introduced `protocol.py` for shared data models used in server/client communication, including request and response schemas for registration, job submission, and status updates.
- Implemented `server.py` to manage a SQLite job queue and client registry, handling job polling, status updates, and job completion.
- Created a cheat sheet for server usage, detailing commands for starting the server, submitting jobs, and monitoring clients.
- Added several experiment configuration files for various training setups, including geometry vector injections and baseline ensembles.

v4/classes/accessory/backbones.py

@@ -109,6 +109,12 @@ BACKBONES: Dict[str, BackboneSpec] = {
         strip=_strip_efficientnet,
         blocks=_blocks_efficientnet,
     ),
+    "resnet18": BackboneSpec(
+        ctor=models.resnet18,
+        weights_default=models.ResNet18_Weights.DEFAULT,
+        strip=_strip_resnet,
+        blocks=_blocks_resnet,
+    ),
     "resnet50": BackboneSpec(
         ctor=models.resnet50,
         weights_default=models.ResNet50_Weights.DEFAULT,

v4/classes/accessory/transforms.py

@@ -44,6 +44,37 @@ class ImageTransformConfig:
         ]
         return transforms.Compose(ops)
 
+    def build_precache(self) -> transforms.Compose:
+        """Deterministic prefix: PIL → resized CHW float32 in [0, 1].
+
+        Output is suitable for caching; per-batch ``build_postcache`` finishes
+        the pipeline (augment + normalize) on tensors.
+        """
+        return transforms.Compose([
+            transforms.Resize(self.resize_size),
+            transforms.CenterCrop(self.crop_size),
+            transforms.ToTensor(),
+        ])
+
+    def build_postcache(self) -> transforms.Compose:
+        """Per-batch tail run on cached float32 [0, 1] CHW tensors.
+
+        Augmentations operate on tensors (torchvision v1 supports this for
+        Flip/Rotation/ColorJitter on tensor input).  Normalize is applied last.
+        """
+        ops = []
+        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.append(transforms.Normalize(mean=self.mean, std=self.std))
+        return transforms.Compose(ops)
+
 
 def backbone_transform_config(backbone_name: str, augment: bool = True) -> ImageTransformConfig:
     """Build an ImageTransformConfig using the backbone's default normalisation stats."""
@@ -64,3 +95,15 @@ def build_backbone_transform(backbone_name: str, augment: bool = True) -> transf
 def build_eval_transform(backbone_name: str) -> transforms.Compose:
     """Deterministic eval transform — no augmentation, backbone-matched normalisation."""
     return build_backbone_transform(backbone_name, augment=False)
+
+
+def build_split_transforms(
+    backbone_name: str, augment: bool = True
+) -> tuple[transforms.Compose, transforms.Compose]:
+    """Return (precache, postcache) transform pair for tensor-cached image towers.
+
+    precache  : PIL → CHW float32 in [0, 1]   (deterministic, run once at fill)
+    postcache : tensor → augmented + normalized tensor   (run per batch)
+    """
+    cfg = backbone_transform_config(backbone_name, augment=augment)
+    return cfg.build_precache(), cfg.build_postcache()

v4/classes/accessory/unet.py

@@ -0,0 +1,212 @@
+"""unet — REFUGE-trained UNet wrapper for v4.
+
+Lean accessory module: model definition + a thin segmenter wrapper that handles
+weight loading, preprocessing, inference, and fine-tuning.
+
+Used by:
+  - GeometrySegEncoder tower (produces disc/cup seg maps as CNN input)
+  - (future) ImageEncoder cropping (locates disc bbox for image cropping)
+
+The segmenter is intentionally domain-agnostic: it takes PIL images in and
+returns binary (disc, cup) numpy masks.  Fine-tuning consumes any DataLoader
+yielding (image_tensor, mask_tensor) pairs — mask preparation (parsing GT
+contour files, etc.) lives in the consumer.
+"""
+from __future__ import annotations
+
+from pathlib import Path
+
+import numpy as np
+import torch
+from PIL import Image
+from PIL.Image import Resampling
+from torch import nn
+from torchvision import transforms
+
+
+# ---------------------------------------------------------------------------
+# Model
+# ---------------------------------------------------------------------------
+
+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, 3, padding=1, bias=False),
+            nn.BatchNorm2d(out_ch),
+            nn.ReLU(inplace=True),
+            nn.Conv2d(out_ch, out_ch, 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.dec4(torch.cat([self.up4(b),  e4], dim=1))
+        d3 = self.dec3(torch.cat([self.up3(d4), e3], dim=1))
+        d2 = self.dec2(torch.cat([self.up2(d3), e2], dim=1))
+        d1 = self.dec1(torch.cat([self.up1(d2), e1], dim=1))
+        return self.out_conv(d1)
+
+
+# ---------------------------------------------------------------------------
+# Segmenter wrapper
+# ---------------------------------------------------------------------------
+
+_IMAGENET_MEAN = (0.485, 0.456, 0.406)
+_IMAGENET_STD  = (0.229, 0.224, 0.225)
+
+
+class UNetSegmenter:
+    """Wraps a UNet with preprocessing, weight loading, inference, and fine-tuning.
+
+    Parameters
+    ----------
+    target_size  : square resolution UNet operates at (default 512)
+    normalize    : "per_image" | "imagenet" | "none"
+    device       : torch device string; defaults to cuda if available
+    """
+
+    def __init__(
+        self,
+        *,
+        target_size: int = 512,
+        normalize:   str = "per_image",
+        device:      str | torch.device | None = None,
+        in_channels: int = 3,
+        base_channels: int = 32,
+        out_channels: int = 2,
+    ):
+        self.target_size = target_size
+        self.normalize   = normalize
+        self.device      = (
+            torch.device(device) if device is not None
+            else torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        )
+        self.model = UNet(in_channels, base_channels, out_channels).to(self.device)
+        self._to_tensor = transforms.ToTensor()
+
+    # ── lifecycle ────────────────────────────────────────────────────────────
+
+    def to(self, device: str | torch.device) -> "UNetSegmenter":
+        self.device = torch.device(device)
+        self.model.to(self.device)
+        return self
+
+    def load_weights(self, path: str | Path) -> "UNetSegmenter":
+        """Load a UNet checkpoint (raw state_dict or {'model': state_dict})."""
+        state = torch.load(Path(path), map_location=self.device, weights_only=False)
+        sd    = state["model"] if isinstance(state, dict) and "model" in state else state
+        self.model.load_state_dict(sd)
+        self.model.eval()
+        return self
+
+    # ── preprocessing ────────────────────────────────────────────────────────
+
+    def _normalize_tensor(self, t: torch.Tensor) -> torch.Tensor:
+        if self.normalize == "per_image":
+            mean = t.mean(dim=(-2, -1), keepdim=True)
+            std  = t.std (dim=(-2, -1), keepdim=True).clamp(min=1e-6)
+            return (t - mean) / std
+        if self.normalize == "imagenet":
+            mean = torch.tensor(_IMAGENET_MEAN, device=t.device).view(-1, 1, 1)
+            std  = torch.tensor(_IMAGENET_STD,  device=t.device).view(-1, 1, 1)
+            return (t - mean) / std
+        return t
+
+    def preprocess(self, image: Image.Image) -> torch.Tensor:
+        """PIL image → normalized (C, H, W) tensor on segmenter device."""
+        resized = image.convert("RGB").resize(
+            (self.target_size, self.target_size), Resampling.BILINEAR
+        )
+        return self._normalize_tensor(self._to_tensor(resized).to(self.device))
+
+    # ── inference ────────────────────────────────────────────────────────────
+
+    @torch.no_grad()
+    def predict(
+        self,
+        image:     Image.Image,
+        *,
+        threshold: float = 0.5,
+        tta:       bool  = False,
+    ) -> tuple[np.ndarray, np.ndarray]:
+        """Single image → (disc_mask, cup_mask) binary uint8 arrays at target_size.
+
+        cup_mask is restricted to disc area (cup ⊆ disc).
+        """
+        self.model.eval()
+        x      = self.preprocess(image).unsqueeze(0)
+        logits = self.model(x)
+        if tta:
+            log_h  = torch.flip(self.model(torch.flip(x, dims=[3])), dims=[3])
+            log_v  = torch.flip(self.model(torch.flip(x, dims=[2])), dims=[2])
+            logits = (logits + log_h + log_v) / 3.0
+        probs = torch.sigmoid(logits)[0].cpu().numpy()
+        disc  = (probs[0] > threshold).astype(np.uint8)
+        cup   = ((probs[1] > threshold) & (disc > 0)).astype(np.uint8)
+        return disc, cup
+
+    # ── fine-tuning ──────────────────────────────────────────────────────────
+
+    def finetune(
+        self,
+        dataloader,
+        *,
+        epochs:     int   = 10,
+        lr:         float = 1e-5,
+        log_prefix: str   = "[UNetSegmenter]",
+    ) -> "UNetSegmenter":
+        """Fine-tune on (image_tensor, mask_tensor) pairs.
+
+        image_tensor : (B, C, H, W) — already preprocessed (normalized)
+        mask_tensor  : (B, 2, H, W) float32 — channel 0 disc, channel 1 cup
+        """
+        import time
+        opt  = torch.optim.Adam(self.model.parameters(), lr=lr)
+        crit = nn.BCEWithLogitsLoss()
+        for ep in range(1, epochs + 1):
+            self.model.train()
+            running, n_batches, t0 = 0.0, 0, time.time()
+            for img, mask in dataloader:
+                img, mask = img.to(self.device), mask.to(self.device)
+                opt.zero_grad()
+                loss = crit(self.model(img), mask)
+                loss.backward()
+                opt.step()
+                running   += float(loss.item())
+                n_batches += 1
+            avg = running / max(n_batches, 1)
+            print(
+                f"  {log_prefix} ep{ep:03d}/{epochs:03d}  loss={avg:.4f}  "
+                f"({time.time() - t0:.1f}s)",
+                flush=True,
+            )
+        self.model.eval()
+        return self

v4/classes/bridges/mono_bridge.py

@@ -0,0 +1,46 @@
+"""mono_bridge — MonoBridge: passthrough for single-tower fusion stages.
+
+The v4 stage runner always expects tower → bridge → head.  For configs that
+have only one tower feeding a head, MonoBridge is the no-op bridge that lets
+the architecture be "head sits directly on tower" without any extra projection,
+SE, or fusion logic.
+
+Optional LayerNorm is exposed for consistency with FusionBridge but defaults
+off to keep the embedding numerically identical to the tower's output.
+"""
+from __future__ import annotations
+
+import torch
+from torch import nn
+
+
+class MonoBridge(nn.Module):
+    """Single-input passthrough bridge.
+
+    Parameters
+    ----------
+    input_dims : list[int] — must be length 1
+    use_ln     : if True, wrap the embedding in a LayerNorm
+    """
+
+    def __init__(self, input_dims: list[int], use_ln: bool = False):
+        super().__init__()
+        if len(input_dims) != 1:
+            raise ValueError(
+                f"MonoBridge expects exactly 1 input dim, got {len(input_dims)}"
+            )
+        self.out_dim = input_dims[0]
+        self.ln      = nn.LayerNorm(self.out_dim) if use_ln else nn.Identity()
+
+    def forward(self, embeddings: list[torch.Tensor]) -> torch.Tensor:
+        if len(embeddings) != 1:
+            raise ValueError(
+                f"MonoBridge forward expects 1 embedding, got {len(embeddings)}"
+            )
+        return self.ln(embeddings[0])
+
+    def set_phase(self, phase: str) -> None:
+        """Freeze during tower_warmup; trainable otherwise (matches FusionBridge)."""
+        enabled = phase not in ("tower_warmup", "cd_warmup")
+        for p in self.parameters():
+            p.requires_grad_(enabled)

v4/classes/logging/prediction_store.py

@@ -205,7 +205,7 @@ class PredictionStore:
                 grp.create_dataset("y_true",     data=buf.y_true)
                 grp.create_dataset("loss",       data=buf.loss)
                 grp.create_dataset("head_names", data=np.array(buf.head_names, dtype=object), dtype=_STR_DT)
-                grp.create_dataset("split",      data=buf.split.astype(str),                  dtype=_STR_DT)
+                grp.create_dataset("split",      data=buf.split,                              dtype=_STR_DT)
                 _write_entity_ids(grp, buf.entity_ids)
 
     @classmethod
@@ -355,7 +355,7 @@ class FeatureStore:
             for phase, buf in self._phases.items():
                 grp = f.create_group(phase)
                 grp.create_dataset("y_true", data=buf.y_true)
-                grp.create_dataset("split",  data=buf.split.astype(str), dtype=_STR_DT)
+                grp.create_dataset("split",  data=buf.split,             dtype=_STR_DT)
                 _write_entity_ids(grp, buf.entity_ids)
                 for head, (arr, _) in buf._heads.items():
                     grp.create_dataset(head, data=arr, compression="gzip", compression_opts=4)

v4/classes/profiles/fundus_images.py

@@ -0,0 +1,574 @@
+"""fundus_images — disc/cup geometry for fundus image profiles.
+
+Contains all fundus-specific geometry logic: mask parsing, feature computation,
+and source-specific loaders.  Profiles whose ImageDataView supports geometry
+should implement build_geometry_loader(source, **kwargs) and/or
+build_seg_map_loader(source, **kwargs) and delegate here.
+
+Geometry vectors (5 scalar CDR features per eye)
+-----------------------------------------------
+  build_geometry_loader("gt", contour_dir=...)         → GTGeometryLoader
+
+Seg maps (3-class disc/cup label map per eye, fed to a CNN tower)
+-----------------------------------------------------------------
+  build_seg_map_loader("gt",   contour_dir=..., ...)   → GTSegMapLoader
+  build_seg_map_loader("unet", weights_path=..., ...)  → UNetSegMapLoader
+"""
+from __future__ import annotations
+
+import copy
+from collections import Counter
+from pathlib import Path
+from typing import Iterable, Tuple
+
+import numpy as np
+import torch
+from PIL import Image, ImageDraw
+from PIL.Image import Resampling
+from torch.utils.data import DataLoader, Dataset
+
+EPS = 1e-6
+
+_FEATURE_DIM   = 5
+_FEATURE_NAMES = ["area_cdr", "rim_ratio", "vertical_cdr", "horizontal_cdr", "centre_shift"]
+
+_MASK_SIZE = (512, 512)  # canonical rasterisation size; CDR ratios are scale-invariant
+
+
+# ---------------------------------------------------------------------------
+# Mask utilities
+# ---------------------------------------------------------------------------
+
+def disc_cup_from_mask_image(mask_img: Image.Image) -> Tuple[np.ndarray, np.ndarray]:
+    """Return binary (disc, cup) masks from a REFUGE-style colour 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
+        )
+        bg_color = Counter(map(tuple, border)).most_common(1)[0][0]
+        colors   = Counter(map(tuple, arr.reshape(-1, c)))
+        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]])
+        bg_value = Counter(border.tolist()).most_common(1)[0][0]
+        disc = (arr != bg_value).astype(np.uint8)
+        fg   = arr[arr != bg_value]
+        cup  = (arr == int(np.min(fg))).astype(np.uint8) if fg.size > 0 else np.zeros_like(arr)
+    cup = (cup > 0) & (disc > 0)
+    return disc.astype(np.uint8), cup.astype(np.uint8)
+
+
+def _contour_to_mask(coords: np.ndarray, size: Tuple[int, int]) -> np.ndarray:
+    """Rasterize a polygon contour (Nx2 xy array) into a binary mask of (width, height)."""
+    from PIL import ImageDraw
+    if coords is None or coords.size == 0:
+        return np.zeros((size[1], size[0]), dtype=np.uint8)
+    img  = Image.new("L", size, 0)
+    draw = ImageDraw.Draw(img)
+    draw.polygon([tuple(map(float, pt)) for pt in coords], outline=1, fill=1)
+    return np.array(img, dtype=np.uint8)
+
+
+# ---------------------------------------------------------------------------
+# Feature computation
+# ---------------------------------------------------------------------------
+
+def compute_geometry_features(disc_mask: np.ndarray, cup_mask: np.ndarray) -> np.ndarray:
+    """Compute 5 cup/disc structural descriptors from binary masks.
+
+    Returns float32 [area_cdr, rim_ratio, vertical_cdr, horizontal_cdr, 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_cdr  = cup_area  / (disc_area + EPS)
+    rim_ratio = (disc_area - cup_area) / (disc_area + EPS)
+
+    disc_h = float(np.any(disc > 0, axis=1).sum())
+    cup_h  = float(np.any(cup  > 0, axis=1).sum())
+    disc_w = float(np.any(disc > 0, axis=0).sum())
+    cup_w  = float(np.any(cup  > 0, axis=0).sum())
+
+    vertical_cdr   = cup_h / (disc_h + EPS)
+    horizontal_cdr = cup_w / (disc_w + EPS)
+
+    def _centre(m: np.ndarray) -> Tuple[float, float]:
+        coords = np.argwhere(m > 0)
+        if coords.size == 0:
+            return 0.5, 0.5
+        ys, xs = coords[:, 0], coords[:, 1]
+        return float(xs.mean()) / m.shape[1], float(ys.mean()) / m.shape[0]
+
+    dcx, dcy = _centre(disc)
+    ccx, ccy = _centre(cup)
+    centre_shift = float(np.hypot(ccx - dcx, ccy - dcy))
+
+    return np.array(
+        [area_cdr, rim_ratio, vertical_cdr, horizontal_cdr, centre_shift],
+        dtype=np.float32,
+    )
+
+
+# ---------------------------------------------------------------------------
+# Loaders
+# ---------------------------------------------------------------------------
+
+class GTGeometryLoader:
+    """Pre-computes per-eye geometry vectors from PAPILA GT contour annotations.
+
+    File naming: RET{pid:03d}{eye}_{disc|cup}_exp{n}.txt
+    Averages exp1 and exp2 when both are present; zero vector for missing entries.
+
+    Usage:
+        loader = GTGeometryLoader(contour_dir)
+        loader.precompute(df, patient_col="Patient ID")
+        vecs = loader.all_vectors()   # {(pid, eye): ndarray}
+    """
+
+    _EXPERTS = (1, 2)
+
+    feature_dim   = _FEATURE_DIM
+    feature_names = _FEATURE_NAMES
+
+    def __init__(self, contour_dir: str | Path) -> None:
+        self._dir   = Path(contour_dir)
+        self._cache: dict[tuple, np.ndarray] = {}
+
+    def precompute(self, df, patient_col: str = "Patient ID") -> None:
+        n_ok = 0
+        for _, row in df.iterrows():
+            pid = int(row[patient_col])
+            eye = str(row.get("eyeID", "OD"))
+            key = (pid, eye)
+            if key in self._cache:
+                continue
+            vec = self._compute(pid, eye)
+            self._cache[key] = vec if vec is not None else np.zeros(self.feature_dim, dtype=np.float32)
+            if vec is not None:
+                n_ok += 1
+        print(f"[GTGeometryLoader] {n_ok}/{len(self._cache)} geometry vectors computed", flush=True)
+
+    def all_vectors(self) -> dict:
+        return dict(self._cache)
+
+    def _compute(self, pid: int, eye: str) -> "np.ndarray | None":
+        stem = f"RET{pid:03d}{eye}"
+        vecs: list[np.ndarray] = []
+        for exp in self._EXPERTS:
+            disc_path = self._dir / f"{stem}_disc_exp{exp}.txt"
+            cup_path  = self._dir / f"{stem}_cup_exp{exp}.txt"
+            if not disc_path.exists():
+                continue
+            try:
+                disc_c = np.loadtxt(disc_path)
+                if disc_c.ndim == 1:
+                    disc_c = disc_c.reshape(-1, 2)
+                disc_mask = _contour_to_mask(disc_c, _MASK_SIZE)
+                if cup_path.exists():
+                    cup_c = np.loadtxt(cup_path)
+                    if cup_c.ndim == 1:
+                        cup_c = cup_c.reshape(-1, 2)
+                    cup_mask = _contour_to_mask(cup_c, _MASK_SIZE)
+                else:
+                    cup_mask = np.zeros((_MASK_SIZE[1], _MASK_SIZE[0]), dtype=np.uint8)
+                cup_mask = ((cup_mask > 0) & (disc_mask > 0)).astype(np.uint8)
+                vecs.append(compute_geometry_features(disc_mask, cup_mask))
+            except Exception:
+                continue
+        if not vecs:
+            return None
+        return np.stack(vecs).mean(axis=0).astype(np.float32)
+
+
+# ---------------------------------------------------------------------------
+# Seg-map utilities (shared by GT and UNet seg-map loaders)
+# ---------------------------------------------------------------------------
+
+def _combine_disc_cup(disc: np.ndarray, cup: np.ndarray) -> np.ndarray:
+    """Merge binary disc + cup masks into a uint8 label map: 0=bg, 1=rim, 2=cup."""
+    disc = (disc > 0).astype(np.uint8)
+    cup  = ((cup > 0) & (disc > 0)).astype(np.uint8)
+    return (disc + cup).astype(np.uint8)
+
+
+def _crop_to_disc_bbox(seg_map: np.ndarray) -> np.ndarray:
+    """Crop a label map tightly to the disc bounding box (anywhere seg_map > 0)."""
+    rows = np.any(seg_map > 0, axis=1)
+    cols = np.any(seg_map > 0, axis=0)
+    if not rows.any():
+        return seg_map
+    r0, r1 = int(np.argmax(rows)), int(len(rows) - 1 - np.argmax(rows[::-1]))
+    c0, c1 = int(np.argmax(cols)), int(len(cols) - 1 - np.argmax(cols[::-1]))
+    return seg_map[r0:r1 + 1, c0:c1 + 1]
+
+
+def _seg_map_to_array(
+    seg_map: np.ndarray, channels: int, target_size: int
+) -> np.ndarray:
+    """Resize a {0,1,2} seg map and convert to a (C, H, W) float32 array.
+
+    channels=1  →  (1, H, W) values in {0, 0.5, 1.0}
+    channels=3  →  (3, H, W) one-hot [bg, rim, cup]
+    """
+    pil = Image.fromarray(seg_map.astype(np.uint8), mode="L").resize(
+        (target_size, target_size), Resampling.NEAREST
+    )
+    arr = np.array(pil, dtype=np.uint8)
+    if channels == 1:
+        return (arr.astype(np.float32) / 2.0)[None, :, :]
+    if channels == 3:
+        return np.stack([
+            (arr == 0).astype(np.float32),
+            (arr == 1).astype(np.float32),
+            (arr == 2).astype(np.float32),
+        ], axis=0)
+    raise ValueError(f"channels must be 1 or 3, got {channels}")
+
+
+def _load_papila_contour(path: Path) -> np.ndarray:
+    """Load (x, y) contour pairs from a PAPILA whitespace/comma-delimited text file."""
+    for delim in (",", None):
+        try:
+            arr = np.loadtxt(str(path), delimiter=delim, comments="#", dtype=np.float32)
+            if arr.size > 0 and arr.ndim >= 1:
+                if arr.ndim == 1:
+                    arr = arr.reshape(-1, 2)
+                if arr.shape[1] >= 2:
+                    return arr[:, :2]
+        except Exception:
+            continue
+    return np.zeros((0, 2), dtype=np.float32)
+
+
+def _papila_disc_cup_masks(
+    pid: int, eye: str, contour_dir: Path, image_size: Tuple[int, int],
+    mask_size: int, experts: Tuple[int, ...] = (1, 2),
+) -> Tuple[np.ndarray, np.ndarray] | None:
+    """Average masks across PAPILA experts. Returns (disc, cup) at mask_size, or None."""
+    stem = f"RET{pid:03d}{eye}"
+    discs, cups = [], []
+    for exp in experts:
+        disc_path = contour_dir / f"{stem}_disc_exp{exp}.txt"
+        cup_path  = contour_dir / f"{stem}_cup_exp{exp}.txt"
+        if not disc_path.exists():
+            continue
+        disc_c = _load_papila_contour(disc_path)
+        if len(disc_c) < 3:
+            continue
+        disc_m = _rasterise_polygon(disc_c, image_size, mask_size)
+        if cup_path.exists():
+            cup_c = _load_papila_contour(cup_path)
+            cup_m = (_rasterise_polygon(cup_c, image_size, mask_size)
+                     if len(cup_c) >= 3 else np.zeros_like(disc_m))
+        else:
+            cup_m = np.zeros_like(disc_m)
+        discs.append(disc_m)
+        cups.append(cup_m)
+    if not discs:
+        return None
+    disc = (np.mean(discs, axis=0) > 0.5).astype(np.uint8)
+    cup  = (np.mean(cups,  axis=0) > 0.5).astype(np.uint8)
+    return disc, cup
+
+
+def _rasterise_polygon(
+    coords: np.ndarray, image_size: Tuple[int, int], target_size: int
+) -> np.ndarray:
+    """Rasterise an (N, 2) polygon contour into a (target_size, target_size) binary mask.
+
+    image_size is (width, height) of the coord space (the original fundus image).
+    """
+    if coords is None or len(coords) < 3:
+        return np.zeros((target_size, target_size), dtype=np.uint8)
+    img  = Image.new("L", image_size, 0)
+    pts  = [tuple(map(float, p)) for p in coords]
+    ImageDraw.Draw(img).polygon(pts, outline=1, fill=1)
+    img  = img.resize((target_size, target_size), Resampling.NEAREST)
+    return (np.array(img, dtype=np.uint8) > 0).astype(np.uint8)
+
+
+# ---------------------------------------------------------------------------
+# Seg-map loaders
+# ---------------------------------------------------------------------------
+
+class GTSegMapLoader:
+    """Pre-computes per-eye disc/cup seg maps from PAPILA GT contour annotations.
+
+    Output: dict {(pid, eye): np.ndarray (C, H, W) float32} cached for the fold.
+    """
+
+    def __init__(
+        self,
+        contour_dir: str | Path,
+        *,
+        channels:     int  = 3,
+        mask_size:    int  = 512,
+        target_size:  int  = 224,
+        crop_to_disc: bool = True,
+    ) -> None:
+        self._dir         = Path(contour_dir)
+        self._channels    = channels
+        self._mask_size   = mask_size
+        self._target_size = target_size
+        self._crop        = crop_to_disc
+        self._cache: dict[tuple, np.ndarray] = {}
+
+    @property
+    def cache_dim(self) -> tuple[int, int, int]:
+        return (self._channels, self._target_size, self._target_size)
+
+    def precompute(self, samples: Iterable[Tuple[int, str, Path]]) -> None:
+        """samples: iterable of (pid, eye, image_path) tuples."""
+        n_ok = 0
+        blank = np.zeros((self._mask_size, self._mask_size), dtype=np.uint8)
+        for pid, eye, image_path in samples:
+            key = (pid, eye)
+            if key in self._cache:
+                continue
+            with Image.open(image_path) as _im:
+                im_size = _im.size  # (W, H)
+            res = _papila_disc_cup_masks(
+                pid, eye, self._dir, im_size, self._mask_size,
+            )
+            if res is None:
+                seg = blank
+            else:
+                disc, cup = res
+                seg = _combine_disc_cup(disc, cup)
+                n_ok += 1
+            if self._crop:
+                seg = _crop_to_disc_bbox(seg)
+            self._cache[key] = _seg_map_to_array(seg, self._channels, self._target_size)
+        print(
+            f"[GTSegMapLoader] {n_ok}/{len(self._cache)} GT seg maps computed "
+            f"(channels={self._channels}, target={self._target_size})",
+            flush=True,
+        )
+
+    def all_seg_maps(self) -> dict:
+        return self._cache
+
+
+class _UNetFTDataset(Dataset):
+    """Pre-cached (image_tensor, mask_tensor) pairs for fine-tuning a UNet.
+
+    Decode + resize + normalize + GT mask rasterisation are all deterministic,
+    so we do them once at construction and store float32 tensors on CPU.  This
+    drops per-batch cost to a tensor lookup + GPU transfer.
+    """
+
+    def __init__(
+        self,
+        records:     list,         # list of (pid, eye, image_path)
+        contour_dir: Path,
+        segmenter,                 # UNetSegmenter — used for image preprocessing
+    ):
+        import time
+        S = segmenter.target_size
+        self._imgs:  list[torch.Tensor] = []
+        self._masks: list[torch.Tensor] = []
+
+        print(
+            f"[UNetSegMapLoader] pre-caching {len(records)} (image, mask) pairs "
+            f"at {S}×{S}...",
+            flush=True,
+        )
+        t0     = time.time()
+        report = max(1, len(records) // 4)
+        for i, (pid, eye, image_path) in enumerate(records, 1):
+            with Image.open(image_path) as raw:
+                im_size = raw.size
+                img_t   = segmenter.preprocess(raw).detach().cpu()
+            res = _papila_disc_cup_masks(pid, eye, contour_dir, im_size, S)
+            if res is None:
+                disc = np.zeros((S, S), dtype=np.uint8)
+                cup  = np.zeros_like(disc)
+            else:
+                disc, cup = res
+            mask_t = torch.from_numpy(np.stack([disc, cup], axis=0).astype(np.float32))
+            self._imgs.append(img_t)
+            self._masks.append(mask_t)
+            if i % report == 0 or i == len(records):
+                print(
+                    f"  [UNet ft cache] {i}/{len(records)}  ({time.time() - t0:.1f}s)",
+                    flush=True,
+                )
+
+    def __len__(self) -> int:
+        return len(self._imgs)
+
+    def __getitem__(self, idx: int):
+        return self._imgs[idx], self._masks[idx]
+
+
+class UNetSegMapLoader:
+    """Pre-computes per-eye seg maps via a REFUGE-pretrained UNet.
+
+    Optionally fine-tunes the UNet per fold on the training split's GT contours.
+
+    Output: dict {(pid, eye): np.ndarray (C, H, W) float32} cached for the fold.
+    """
+
+    def __init__(
+        self,
+        weights_path: str | Path,
+        *,
+        contour_dir:         str | Path,
+        channels:            int   = 3,
+        target_size:         int   = 224,
+        unet_size:           int   = 512,
+        normalize:           str   = "per_image",
+        threshold:           float = 0.5,
+        crop_to_disc:        bool  = True,
+        finetune_epochs:     int   = 0,
+        finetune_lr:         float = 1e-5,
+        finetune_batch_size: int   = 4,
+        device:              str | None = None,
+    ) -> None:
+        from v4.classes.accessory.unet import UNetSegmenter
+
+        self._weights_path        = Path(weights_path)
+        self._contour_dir         = Path(contour_dir)
+        self._channels            = channels
+        self._target_size         = target_size
+        self._threshold           = threshold
+        self._crop                = crop_to_disc
+        self._ft_epochs           = finetune_epochs
+        self._ft_lr               = finetune_lr
+        self._ft_batch_size       = finetune_batch_size
+
+        self._segmenter = UNetSegmenter(
+            target_size=unet_size, normalize=normalize, device=device,
+        ).load_weights(self._weights_path)
+        self._base_state = copy.deepcopy(self._segmenter.model.state_dict())
+        self._cache: dict[tuple, np.ndarray] = {}
+
+    @property
+    def cache_dim(self) -> tuple[int, int, int]:
+        return (self._channels, self._target_size, self._target_size)
+
+    def reset_cache(self) -> None:
+        """Clear cached seg maps (call between folds)."""
+        self._cache.clear()
+
+    def reset_weights(self) -> None:
+        """Restore base REFUGE weights (undo any prior fine-tuning)."""
+        self._segmenter.model.load_state_dict(copy.deepcopy(self._base_state))
+
+    def finetune(self, train_samples: list) -> None:
+        """Fine-tune the UNet on the training fold's GT contours.
+
+        train_samples: list of (pid, eye, image_path) tuples — train split only.
+        """
+        if self._ft_epochs <= 0:
+            return
+        ds = _UNetFTDataset(train_samples, self._contour_dir, self._segmenter)
+        loader = DataLoader(
+            ds, batch_size=self._ft_batch_size, shuffle=True, num_workers=0,
+        )
+        print(
+            f"[UNetSegMapLoader] fine-tuning UNet for {self._ft_epochs} epochs "
+            f"on {len(train_samples)} samples (lr={self._ft_lr}, "
+            f"batch_size={self._ft_batch_size})",
+            flush=True,
+        )
+        self._segmenter.finetune(
+            loader, epochs=self._ft_epochs, lr=self._ft_lr,
+            log_prefix="[UNet ft]",
+        )
+
+    def precompute(self, samples: Iterable[Tuple[int, str, Path]]) -> None:
+        """Run UNet inference on every sample and cache the resulting seg map."""
+        import time
+        samples = list(samples)
+        todo    = [s for s in samples if (s[0], s[1]) not in self._cache]
+        if not todo:
+            return
+        print(
+            f"[UNetSegMapLoader] running UNet inference on {len(todo)} images...",
+            flush=True,
+        )
+        t0     = time.time()
+        report = max(1, len(todo) // 4)
+        for i, (pid, eye, image_path) in enumerate(todo, 1):
+            with Image.open(image_path) as raw:
+                disc, cup = self._segmenter.predict(raw, threshold=self._threshold)
+            seg = _combine_disc_cup(disc, cup)
+            if self._crop:
+                seg = _crop_to_disc_bbox(seg)
+            self._cache[(pid, eye)] = _seg_map_to_array(
+                seg, self._channels, self._target_size,
+            )
+            if i % report == 0 or i == len(todo):
+                print(
+                    f"  [UNet inf] {i}/{len(todo)}  ({time.time() - t0:.1f}s)",
+                    flush=True,
+                )
+        print(
+            f"[UNetSegMapLoader] {len(todo)} seg maps cached via UNet "
+            f"(channels={self._channels}, target={self._target_size})",
+            flush=True,
+        )
+
+    def all_seg_maps(self) -> dict:
+        return self._cache
+
+
+# ---------------------------------------------------------------------------
+# Factories
+# ---------------------------------------------------------------------------
+
+def build_geometry_loader(source: str, **kwargs):
+    """Return the appropriate geometry-vector loader for the given source string.
+
+    Parameters
+    ----------
+    source      : "gt" | "unet"
+    contour_dir : (gt) path to contour annotation directory
+    """
+    if source == "gt":
+        contour_dir = kwargs.get("contour_dir")
+        if contour_dir is None:
+            raise ValueError("build_geometry_loader source='gt' requires contour_dir")
+        return GTGeometryLoader(contour_dir)
+    raise NotImplementedError(f"build_geometry_loader: source={source!r} not implemented")
+
+
+def build_seg_map_loader(source: str, **kwargs):
+    """Return the appropriate seg-map loader for the given source string.
+
+    Parameters
+    ----------
+    source : "gt" | "unet"
+
+    GT kwargs:
+        contour_dir, channels=3, mask_size=512, target_size=224, crop_to_disc=True
+    UNet kwargs:
+        weights_path, contour_dir, channels=3, target_size=224, unet_size=512,
+        normalize="per_image", threshold=0.5, crop_to_disc=True,
+        finetune_epochs=0, finetune_lr=1e-5, finetune_batch_size=4, device=None
+    """
+    if source == "gt":
+        if "contour_dir" not in kwargs:
+            raise ValueError("build_seg_map_loader source='gt' requires contour_dir")
+        return GTSegMapLoader(**kwargs)
+    if source == "unet":
+        if "weights_path" not in kwargs:
+            raise ValueError("build_seg_map_loader source='unet' requires weights_path")
+        if "contour_dir" not in kwargs:
+            raise ValueError(
+                "build_seg_map_loader source='unet' requires contour_dir "
+                "(needed for per-fold fine-tuning, even if finetune_epochs=0)"
+            )
+        return UNetSegMapLoader(**kwargs)
+    raise NotImplementedError(f"build_seg_map_loader: source={source!r} not implemented")

v4/classes/profiles/v4papila.py

@@ -418,6 +418,30 @@ class ImageDataView:
             eye_filter=eye,
         )
 
+    # ── Geometry hook ─────────────────────────────────────────────────────────
+
+    def _resolve_paths(self, kwargs: dict) -> dict:
+        """Resolve any *_dir / *_path kwargs against the repo root."""
+        repo_root = Path(__file__).resolve().parents[3]
+        out = {}
+        for k, v in kwargs.items():
+            if (k.endswith("_dir") or k.endswith("_path")) and v is not None:
+                p = Path(v)
+                out[k] = str(repo_root / p) if not p.is_absolute() else v
+            else:
+                out[k] = v
+        return out
+
+    def build_geometry_loader(self, source: str, **kwargs):
+        """Return a geometry-vector loader (delegates to fundus_images)."""
+        from v4.classes.profiles.fundus_images import build_geometry_loader as _build
+        return _build(source, **self._resolve_paths(kwargs))
+
+    def build_seg_map_loader(self, source: str, **kwargs):
+        """Return a seg-map loader (delegates to fundus_images)."""
+        from v4.classes.profiles.fundus_images import build_seg_map_loader as _build
+        return _build(source, **self._resolve_paths(kwargs))
+
 
 # ---------------------------------------------------------------------------
 # PapilaBundle — the v4 DataBundle returned by build_data
@@ -515,6 +539,7 @@ class PapilaBundle:
         *,
         level: str = "eye",
         label_filter: list[int] | None = None,
+        eye_filter: str | None = None,
     ) -> LoaderShell:
         """Build a LoaderShell from a split DataFrame.
 
@@ -533,6 +558,8 @@ class PapilaBundle:
 
         if label_filter is not None:
             df = df[df[lc].isin(label_filter)]
+        if eye_filter is not None and "eyeID" in df.columns:
+            df = df[df["eyeID"] == eye_filter]
 
         entries: list[ShellEntry] = []
 

v4/classes/stages/fusion.py

@@ -11,7 +11,8 @@ import torch.nn.functional as F
 from v4.classes.dataset import LoaderShell, to_label_tensor
 from v4.classes.metrics import score_arrays, compute_extended_metrics, tune_binary_threshold
 from v4.classes.stages.helpers import (
-    encode_embedding, get_out_dim, resolve_input_dims, phase_for_epoch,
+    class_weights_from_shell, encode_embedding, get_out_dim, resolve_input_dims,
+    phase_for_epoch,
 )
 
 
@@ -25,15 +26,15 @@ def collect_probs(
     loader,
     device,
     num_classes:  int,
-) -> tuple[np.ndarray, np.ndarray]:
-    """Eval pass for one fusion stage; returns (y_true, softmax_probs)."""
+) -> tuple[np.ndarray, np.ndarray, list, np.ndarray]:
+    """Eval pass for one fusion stage; returns (y_true, softmax_probs, entity_ids, embeddings)."""
     from v4.classes.dataset import to_label_tensor
     bridge.eval(); primary_head.eval()
     for t in towers.values():
         t.eval()
     inputs       = stage_cfg["inputs"]
     is_bilateral = isinstance(inputs, dict)
-    y_all, p_all = [], []
+    y_all, p_all, ids_all, z_all = [], [], [], []
 
     with torch.no_grad():
         for batch in loader:
@@ -57,10 +58,14 @@ def collect_probs(
             logits = primary_head(z)
             y_all.append(to_label_tensor(y, device).cpu().numpy())
             p_all.append(F.softmax(logits, dim=1).cpu().numpy())
+            z_all.append(z.cpu().numpy())
+            ids_all.extend(batch.get("entity_id", []))
 
     if not y_all:
-        return np.zeros(0, dtype=np.int64), np.zeros((0, num_classes), dtype=np.float32)
-    return np.concatenate(y_all), np.concatenate(p_all, axis=0)
+        return (np.zeros(0, dtype=np.int64), np.zeros((0, num_classes), dtype=np.float32),
+                [], np.zeros((0, 0), dtype=np.float32))
+    return (np.concatenate(y_all), np.concatenate(p_all, axis=0),
+            ids_all, np.concatenate(z_all, axis=0))
 
 
 def run(
@@ -88,9 +93,13 @@ def run(
     inputs       = stage_cfg["inputs"]
     is_bilateral = isinstance(inputs, dict)
 
-    s_train = data.build_shells(split.train, level=level, label_filter=label_filter)
-    s_val   = data.build_shells(split.val,   level=level, label_filter=label_filter)
-    s_test  = (data.build_shells(split.test, level=level, label_filter=label_filter)
+    eye_filter = stage_cfg.get("eye_filter", None)
+    s_train = data.build_shells(split.train, level=level, label_filter=label_filter,
+                                eye_filter=eye_filter)
+    s_val   = data.build_shells(split.val,   level=level, label_filter=label_filter,
+                                eye_filter=eye_filter)
+    s_test  = (data.build_shells(split.test, level=level, label_filter=label_filter,
+                                 eye_filter=eye_filter)
                if split.test is not None else LoaderShell(entries=[]))
 
     if not s_val.entries:
@@ -116,7 +125,11 @@ def run(
         h_dim = get_out_dim(hs["input"], towers, {**stage_models, name: bridge})
         h_mod = importlib.import_module(hs.get("module", "v4.classes.heads.classifier"))
         h_cls = getattr(h_mod, hs.get("class", "ClassificationHead"))
-        head_models[hs["name"]] = h_cls(h_dim, num_classes).to(device)
+        existing = stage_models.get(hs["name"])
+        head_models[hs["name"]] = (
+            existing.to(device) if existing is not None
+            else h_cls(h_dim, num_classes, **hs.get("args", {})).to(device)
+        )
 
     primary_hs_cfg = next((hs for hs in head_stage_cfgs if not hs.get("bcd", False)), None)
     bcd_head_cfgs  = [hs for hs in head_stage_cfgs if hs.get("bcd", False)]
@@ -132,11 +145,23 @@ def run(
         for p in m.parameters():
             p.requires_grad_(False)
         m.eval()
+    for h in head_models.values():
+        for p in h.parameters():
+            p.requires_grad_(True)
 
     # ── Optimizer ────────────────────────────────────────────────────────────
     train_towers = stage_cfg.get("train_towers", False)
+    if train_towers:
+        # Only train towers that are direct inputs to this stage (not all towers globally).
+        # For nt_od with inputs ["img_od", "cd_od"] this trains only those two; other
+        # eye's towers remain untouched.
+        direct_inputs = list(inputs.values()) if isinstance(inputs, dict) else inputs
+        tower_params  = [p for n in direct_inputs if n in towers
+                         for p in towers[n].parameters()]
+    else:
+        tower_params = []
     opt_params = (
-        ([p for t in towers.values() for p in t.parameters()] if train_towers else []) +
+        tower_params +
         list(bridge.parameters()) +
         [p for h in head_models.values() for p in h.parameters()]
     )
@@ -147,6 +172,17 @@ def run(
     wf         = 0 if is_bilateral else warmup_cfg.get("fused_epochs",  0)
     bcd_prob   = cfg["training"].get("bcd_prob", 0.5)
 
+    cw = class_weights_from_shell(
+        s_train, num_classes, device,
+        enabled=cfg["training"].get("class_weighted", False),
+    )
+    if cw is not None:
+        print(
+            f"  fold{fold+1} [{name}] class weights: "
+            + ", ".join(f"{i}={w:.3f}" for i, w in enumerate(cw.tolist())),
+            flush=True,
+        )
+
     # ── Epoch loop ────────────────────────────────────────────────────────────
     for epoch in range(epochs):
         bridge.train()
@@ -200,7 +236,7 @@ def run(
             if is_bilateral or phase == "fused_warmup":
                 logits = head_logits.get(primary_hs_cfg["name"])
             elif phase == "tower_warmup" and bcd_head_cfgs:
-                losses = [F.cross_entropy(head_logits[hs["name"]], y_t)
+                losses = [F.cross_entropy(head_logits[hs["name"]], y_t, weight=cw)
                           for hs in bcd_head_cfgs if hs["name"] in head_logits]
                 if not losses:
                     continue
@@ -217,7 +253,7 @@ def run(
 
             if logits is None:
                 continue
-            loss = F.cross_entropy(logits, y_t)
+            loss = F.cross_entropy(logits, y_t, weight=cw)
             opt.zero_grad(); loss.backward(); opt.step()
             total_correct += int((logits.argmax(1) == y_t).sum())
             total_loss    += loss.item() * len(y_t)
@@ -226,8 +262,8 @@ def run(
         tr_loss = total_loss / total_n if total_n else nan
         tr_acc  = total_correct / total_n if total_n else nan
 
-        y_v, p_v = collect_probs(bridge, primary_head, stage_cfg, towers,
-                                  stage_models, cfg_stages, val_loader, device, num_classes)
+        y_v, p_v, _, _ = collect_probs(bridge, primary_head, stage_cfg, towers,
+                                       stage_models, cfg_stages, val_loader, device, num_classes)
         _, val_auc, _ = score_arrays(y_v, p_v, num_classes) if y_v.size else (nan, nan, nan)
         print(
             f"  fold{fold+1} [{name}] ep{epoch+1:03d}/{epochs} [{phase:14s}]"
@@ -236,8 +272,8 @@ def run(
         )
 
     # ── Final eval ────────────────────────────────────────────────────────────
-    y_val, p_val = collect_probs(bridge, primary_head, stage_cfg, towers,
-                                  stage_models, cfg_stages, val_loader, device, num_classes)
+    y_val, p_val, ids_val, z_val = collect_probs(bridge, primary_head, stage_cfg, towers,
+                                                  stage_models, cfg_stages, val_loader, device, num_classes)
     val_acc, val_auc, val_n = (score_arrays(y_val, p_val, num_classes)
                                if y_val.size else (nan, nan, nan))
     ext = compute_extended_metrics(y_val, p_val, num_classes) if y_val.size else {}
@@ -247,10 +283,11 @@ def run(
             and num_classes == 2 and y_val.size >= 2):
         val_threshold = tune_binary_threshold(y_val, p_val[:, 1])
 
+    y_te = p_te = ids_te = z_te = None
     test_auc = test_acc = test_n = nan
     if test_loader is not None:
-        y_te, p_te = collect_probs(bridge, primary_head, stage_cfg, towers,
-                                    stage_models, cfg_stages, test_loader, device, num_classes)
+        y_te, p_te, ids_te, z_te = collect_probs(bridge, primary_head, stage_cfg, towers,
+                                                  stage_models, cfg_stages, test_loader, device, num_classes)
         test_acc, test_auc, test_n = (score_arrays(y_te, p_te, num_classes)
                                       if y_te.size else (nan, nan, nan))
 
@@ -270,4 +307,10 @@ def run(
         f"{name}_test_acc":      test_acc,
         f"{name}_test_n":        test_n,
     }
-    return updated, metrics
+    pred_data = {
+        name: {
+            "val_y": y_val, "val_p": p_val, "val_ids": ids_val, "val_z": z_val,
+            "test_y": y_te, "test_p": p_te, "test_ids": ids_te, "test_z": z_te,
+        }
+    }
+    return updated, metrics, pred_data

v4/classes/stages/helpers.py

@@ -1,9 +1,32 @@
 """stages/helpers — shared utilities for stage runners."""
 from __future__ import annotations
 
+from collections import Counter
+
 import torch
 
 
+def class_weights_from_shell(
+    shell, num_classes: int, device, *, enabled: bool = True
+) -> torch.Tensor | None:
+    """Return inverse-frequency CE weights normalised to mean 1, or None.
+
+    weights[i] = (N_total / num_classes) / N_class_i  → rare classes weighted higher.
+    Mean(weights) ≈ 1 so overall loss magnitude is unchanged.
+
+    Classes absent from the shell get weight 1.0 (no division-by-zero).
+    """
+    if not enabled or shell is None or not shell.entries:
+        return None
+    counts  = Counter(int(e.label) for e in shell.entries)
+    n_total = sum(counts.values())
+    weights = []
+    for c in range(num_classes):
+        n_c = counts.get(c, 0)
+        weights.append(1.0 if n_c == 0 else n_total / (num_classes * n_c))
+    return torch.tensor(weights, dtype=torch.float32, device=device)
+
+
 def get_out_dim(name: str, towers: dict, stage_models: dict) -> int:
     if name in towers:
         return towers[name].out_dim

v4/classes/stages/parallel.py

@@ -0,0 +1,394 @@
+"""stages/parallel — runs multiple same-type sub-stages in a shared epoch loop.
+
+Used to train bilateral pairs (OD + OS) simultaneously rather than sequentially.
+Sub-stages must all be the same type: either all 'warm' or all 'fusion'.
+"""
+from __future__ import annotations
+
+import importlib
+from random import choice, random as _random
+
+import numpy as np
+import torch
+import torch.nn.functional as F
+
+from v4.classes.dataset import LoaderShell, to_label_tensor
+from v4.classes.metrics import score_arrays, compute_extended_metrics, tune_binary_threshold
+from v4.classes.stages.helpers import (
+    class_weights_from_shell, get_out_dim, resolve_input_dims, phase_for_epoch,
+)
+from v4.classes.stages.fusion import collect_probs
+
+
+def run(
+    stage_cfg:        dict,
+    cfg:              dict,
+    towers:           dict,
+    stage_models:     dict,
+    data,
+    split,
+    label_filter,
+    num_classes:      int,
+    device,
+    fold:             int,
+    cfg_stages:       list[dict],
+    _make_loader,
+    _balanced_sampler,
+) -> tuple[dict, dict, dict]:
+    sub_cfgs  = stage_cfg["stages"]
+    sub_types = {s["type"] for s in sub_cfgs}
+
+    if sub_types == {"warm"}:
+        return _parallel_warm(sub_cfgs, cfg, towers, stage_models, data, split,
+                              label_filter, num_classes, device, fold,
+                              cfg_stages, _make_loader, _balanced_sampler)
+    elif sub_types == {"fusion"}:
+        return _parallel_fusion(sub_cfgs, cfg, towers, stage_models, data, split,
+                                label_filter, num_classes, device, fold,
+                                cfg_stages, _make_loader)
+    else:
+        raise ValueError(
+            f"parallel stage sub-stages must all be the same type (warm or fusion), "
+            f"got {sub_types}"
+        )
+
+
+# ── Parallel warm ─────────────────────────────────────────────────────────────
+
+def _parallel_warm(
+    sub_cfgs, cfg, towers, stage_models, data, split,
+    label_filter, num_classes, device, fold, cfg_stages, _make_loader, _balanced_sampler,
+):
+    bs = cfg["training"]["batch_size"]
+
+    contexts = []
+    for sc in sub_cfgs:
+        tower_name   = sc["tower"]
+        level        = sc["level"]
+        shell_filter = sc.get("shell_filter", {})
+        n_epochs     = sc.get("epochs", 0)
+
+        s_train = data.build_shells(split.train, level=level, label_filter=label_filter,
+                                    **shell_filter)
+        loader  = _make_loader(s_train, {tower_name: towers[tower_name]},
+                               batch_size=bs, shuffle=False,
+                               sampler=_balanced_sampler(s_train))
+        head_name = sc.get("head_name")
+        if head_name:
+            head_cfg = next((s for s in cfg_stages if s.get("name") == head_name), None)
+            if head_cfg is None:
+                raise ValueError(f"warm stage requested head_name={head_name!r}, but no such head exists")
+            h_mod = importlib.import_module(head_cfg.get("module", "v4.classes.heads.classifier"))
+            h_cls = getattr(h_mod, head_cfg.get("class", "ClassificationHead"))
+            probe = stage_models.get(head_name)
+            if probe is None:
+                probe = h_cls(towers[tower_name].out_dim, num_classes, **head_cfg.get("args", {}))
+            probe = probe.to(device)
+        else:
+            probe = torch.nn.Linear(towers[tower_name].out_dim, num_classes).to(device)
+        opt   = torch.optim.Adam(
+            list(towers[tower_name].parameters()) + list(probe.parameters()),
+            lr=cfg["training"]["lr"],
+        )
+        cw = class_weights_from_shell(
+            s_train, num_classes, device,
+            enabled=cfg["training"].get("class_weighted", False),
+        )
+        contexts.append({
+            "name": tower_name, "n_epochs": n_epochs,
+            "loader": loader, "probe": probe, "opt": opt,
+            "head_name": head_name, "class_weights": cw,
+        })
+
+    active = {c["name"] for c in contexts if c["n_epochs"] > 0}
+    for name, t in towers.items():
+        for p in t.parameters():
+            p.requires_grad_(name in active)
+
+    max_epochs = max((c["n_epochs"] for c in contexts), default=0)
+    for epoch in range(max_epochs):
+        for ctx in contexts:
+            if epoch >= ctx["n_epochs"]:
+                continue
+            tower_name = ctx["name"]
+            towers[tower_name].train()
+            total_loss = total_correct = total_n = 0
+            for batch in ctx["loader"]:
+                y = batch.get("label")
+                x = batch.get(tower_name)
+                if not torch.is_tensor(y) or not torch.is_tensor(x):
+                    continue
+                y_t    = to_label_tensor(y, device)
+                logits = ctx["probe"](towers[tower_name](x.to(device)))
+                loss   = F.cross_entropy(logits, y_t, weight=ctx["class_weights"])
+                ctx["opt"].zero_grad(); loss.backward(); ctx["opt"].step()
+                total_loss    += loss.item() * len(y_t)
+                total_correct += int((logits.argmax(1) == y_t).sum())
+                total_n       += len(y_t)
+            if total_n:
+                print(
+                    f"  fold{fold+1} [warm/{tower_name}]"
+                    f" ep{epoch+1:03d}/{ctx['n_epochs']}"
+                    f"  loss={total_loss/total_n:.4f}  acc={total_correct/total_n:.3f}",
+                    flush=True,
+                )
+
+    for t in towers.values():
+        for p in t.parameters():
+            p.requires_grad_(True)
+
+    updated = dict(stage_models)
+    for ctx in contexts:
+        if ctx.get("head_name"):
+            updated[ctx["head_name"]] = ctx["probe"]
+
+    return updated, {}, {}
+
+
+# ── Parallel fusion ───────────────────────────────────────────────────────────
+
+def _parallel_fusion(
+    sub_cfgs, cfg, towers, stage_models, data, split,
+    label_filter, num_classes, device, fold, cfg_stages, _make_loader,
+):
+    nan      = float("nan")
+    bs       = cfg["training"]["batch_size"]
+    bcd_prob = cfg["training"].get("bcd_prob", 0.5)
+
+    # Freeze all prior stage models once, before building any bridges.
+    for m in stage_models.values():
+        for p in m.parameters():
+            p.requires_grad_(False)
+        m.eval()
+
+    # ── Per-sub-stage setup ───────────────────────────────────────────────────
+    contexts = []
+    for sc in sub_cfgs:
+        name       = sc["name"]
+        level      = sc["level"]
+        inputs     = sc["inputs"]
+        epochs     = sc["epochs"]
+        eye_filter = sc.get("eye_filter", None)
+
+        s_train = data.build_shells(split.train, level=level, label_filter=label_filter,
+                                    eye_filter=eye_filter)
+        s_val   = data.build_shells(split.val,   level=level, label_filter=label_filter,
+                                    eye_filter=eye_filter)
+        s_test  = (data.build_shells(split.test, level=level, label_filter=label_filter,
+                                     eye_filter=eye_filter)
+                   if split.test is not None else LoaderShell(entries=[]))
+
+        if not s_val.entries:
+            print(f"  fold{fold+1}: no val samples for stage {name!r}, skipping.", flush=True)
+            continue
+
+        input_dims = resolve_input_dims(inputs, towers, stage_models)
+        bmod       = importlib.import_module(sc["module"])
+        bridge     = getattr(bmod, sc["class"])(input_dims, **sc.get("args", {})).to(device)
+
+        # Head stages for this sub-stage.
+        head_stage_cfgs = [s for s in cfg_stages if s["type"] == "head"
+                           and s.get("train_with") == name]
+        head_models: dict[str, torch.nn.Module] = {}
+        for hs in head_stage_cfgs:
+            h_dim = get_out_dim(hs["input"], towers, {**stage_models, name: bridge})
+            h_mod = importlib.import_module(hs.get("module", "v4.classes.heads.classifier"))
+            h_cls = getattr(h_mod, hs.get("class", "ClassificationHead"))
+            existing = stage_models.get(hs["name"])
+            head_models[hs["name"]] = (
+                existing.to(device) if existing is not None
+                else h_cls(h_dim, num_classes, **hs.get("args", {})).to(device)
+            )
+        for h in head_models.values():
+            for p in h.parameters():
+                p.requires_grad_(True)
+
+        primary_hs_cfg = next((hs for hs in head_stage_cfgs if not hs.get("bcd", False)), None)
+        bcd_head_cfgs  = [hs for hs in head_stage_cfgs if hs.get("bcd", False)]
+
+        if primary_hs_cfg is None:
+            print(f"  WARNING: no primary head for stage {name!r}; skipping.", flush=True)
+            continue
+
+        primary_head = head_models[primary_hs_cfg["name"]]
+
+        train_towers  = sc.get("train_towers", False)
+        direct_inputs = inputs if isinstance(inputs, list) else list(inputs.values())
+        tower_params  = ([p for n in direct_inputs if n in towers
+                          for p in towers[n].parameters()]
+                         if train_towers else [])
+        opt_params    = (tower_params + list(bridge.parameters()) +
+                         [p for h in head_models.values() for p in h.parameters()])
+        opt = torch.optim.Adam(opt_params, lr=cfg["training"]["lr"])
+
+        warmup_cfg = sc.get("warmup", {})
+        wt         = warmup_cfg.get("tower_epochs", 0)
+        wf         = warmup_cfg.get("fused_epochs", 0)
+
+        train_loader = _make_loader(s_train, towers, batch_size=bs, shuffle=True)
+        val_loader   = _make_loader(s_val,   towers, batch_size=bs, shuffle=False)
+        test_loader  = (_make_loader(s_test, towers, batch_size=bs, shuffle=False)
+                        if s_test.entries else None)
+
+        cw = class_weights_from_shell(
+            s_train, num_classes, device,
+            enabled=cfg["training"].get("class_weighted", False),
+        )
+        contexts.append({
+            "name": name, "inputs": inputs, "epochs": epochs,
+            "bridge": bridge, "head_models": head_models,
+            "primary_head": primary_head, "primary_hs_cfg": primary_hs_cfg,
+            "bcd_head_cfgs": bcd_head_cfgs,
+            "opt": opt, "wt": wt, "wf": wf,
+            "train_towers": train_towers, "direct_inputs": direct_inputs,
+            "train_loader": train_loader, "val_loader": val_loader,
+            "test_loader": test_loader, "sc": sc,
+            "class_weights": cw,
+        })
+
+    if not contexts:
+        return stage_models, {}, {}
+
+    max_epochs = max(c["epochs"] for c in contexts)
+
+    # ── Shared epoch loop ─────────────────────────────────────────────────────
+    for epoch in range(max_epochs):
+        for ctx in contexts:
+            if epoch >= ctx["epochs"]:
+                continue
+
+            name   = ctx["name"]
+            bridge = ctx["bridge"]
+            inputs = ctx["inputs"]
+            wt, wf = ctx["wt"], ctx["wf"]
+            phase  = phase_for_epoch(epoch, wt, wf)
+
+            bridge.train()
+            for h in ctx["head_models"].values():
+                h.train()
+            if ctx["train_towers"]:
+                for n in ctx["direct_inputs"]:
+                    if n in towers:
+                        towers[n].train()
+            else:
+                for n in ctx["direct_inputs"]:
+                    if n in towers:
+                        towers[n].eval()
+
+            if hasattr(bridge, "set_phase"):
+                bridge.set_phase(phase)
+
+            total_loss = total_correct = total_n = 0
+
+            for batch in ctx["train_loader"]:
+                y = batch.get("label")
+                if not torch.is_tensor(y):
+                    continue
+                y_t = to_label_tensor(y, device)
+                if y_t.numel() == 0:
+                    continue
+
+                local_embs = {n: towers[n](batch[n].to(device)) for n in inputs
+                              if n in batch and torch.is_tensor(batch[n])}
+                if len(local_embs) != len(inputs):
+                    continue
+                local_embs[name] = bridge(list(local_embs[n] for n in inputs))
+
+                head_logits = {
+                    hs["name"]: ctx["head_models"][hs["name"]](local_embs[hs["input"]])
+                    for hs in ([ctx["primary_hs_cfg"]] + ctx["bcd_head_cfgs"])
+                    if hs["input"] in local_embs
+                }
+
+                if phase == "fused_warmup":
+                    logits = head_logits.get(ctx["primary_hs_cfg"]["name"])
+                elif phase == "tower_warmup" and ctx["bcd_head_cfgs"]:
+                    losses = [F.cross_entropy(head_logits[hs["name"]], y_t,
+                                              weight=ctx["class_weights"])
+                              for hs in ctx["bcd_head_cfgs"] if hs["name"] in head_logits]
+                    if not losses:
+                        continue
+                    loss = sum(losses) / len(losses)
+                    ctx["opt"].zero_grad(); loss.backward(); ctx["opt"].step()
+                    total_loss += loss.item() * len(y_t)
+                    total_n    += len(y_t)
+                    continue
+                else:
+                    if ctx["bcd_head_cfgs"] and _random() < bcd_prob:
+                        logits = head_logits.get(choice(ctx["bcd_head_cfgs"])["name"])
+                    else:
+                        logits = head_logits.get(ctx["primary_hs_cfg"]["name"])
+
+                if logits is None:
+                    continue
+                loss = F.cross_entropy(logits, y_t, weight=ctx["class_weights"])
+                ctx["opt"].zero_grad(); loss.backward(); ctx["opt"].step()
+                total_correct += int((logits.argmax(1) == y_t).sum())
+                total_loss    += loss.item() * len(y_t)
+                total_n       += len(y_t)
+
+            tr_loss = total_loss / total_n if total_n else nan
+            tr_acc  = total_correct / total_n if total_n else nan
+
+            y_v, p_v, _, _ = collect_probs(bridge, ctx["primary_head"], ctx["sc"], towers,
+                                           stage_models, cfg_stages, ctx["val_loader"],
+                                           device, num_classes)
+            _, val_auc, _ = score_arrays(y_v, p_v, num_classes) if y_v.size else (nan, nan, nan)
+            print(
+                f"  fold{fold+1} [{name}] ep{epoch+1:03d}/{ctx['epochs']} [{phase:14s}]"
+                f"  loss={tr_loss:.4f}  acc={tr_acc:.3f}  val_auc={val_auc:.4f}",
+                flush=True,
+            )
+
+    # ── Final eval + collect results ──────────────────────────────────────────
+    updated    = dict(stage_models)
+    all_metrics: dict = {}
+    all_preds:   dict = {}
+
+    for ctx in contexts:
+        name         = ctx["name"]
+        bridge       = ctx["bridge"]
+        primary_head = ctx["primary_head"]
+
+        updated[name] = bridge
+        updated.update(ctx["head_models"])
+
+        y_val, p_val, ids_val, z_val = collect_probs(bridge, primary_head, ctx["sc"], towers,
+                                                      stage_models, cfg_stages, ctx["val_loader"],
+                                                      device, num_classes)
+        val_acc, val_auc, val_n = (score_arrays(y_val, p_val, num_classes)
+                                   if y_val.size else (nan, nan, nan))
+        ext = compute_extended_metrics(y_val, p_val, num_classes) if y_val.size else {}
+
+        val_threshold = 0.5
+        if (cfg["training"].get("tune_binary_threshold")
+                and num_classes == 2 and y_val.size >= 2):
+            val_threshold = tune_binary_threshold(y_val, p_val[:, 1])
+
+        y_te = p_te = ids_te = z_te = None
+        test_auc = test_acc = test_n = nan
+        if ctx["test_loader"] is not None:
+            y_te, p_te, ids_te, z_te = collect_probs(bridge, primary_head, ctx["sc"], towers,
+                                                      stage_models, cfg_stages, ctx["test_loader"],
+                                                      device, num_classes)
+            test_acc, test_auc, test_n = (score_arrays(y_te, p_te, num_classes)
+                                          if y_te.size else (nan, nan, nan))
+
+        all_metrics.update({
+            f"{name}_val_auc":       val_auc,
+            f"{name}_val_acc":       val_acc,
+            f"{name}_val_n":         val_n,
+            f"{name}_val_kappa":     ext.get("kappa", nan),
+            f"{name}_val_mcc":       ext.get("mcc", nan),
+            f"{name}_val_f1":        ext.get("macro_f1", nan),
+            f"{name}_val_threshold": val_threshold,
+            f"{name}_test_auc":      test_auc,
+            f"{name}_test_acc":      test_acc,
+            f"{name}_test_n":        test_n,
+        })
+        all_preds[name] = {
+            "val_y": y_val, "val_p": p_val, "val_ids": ids_val, "val_z": z_val,
+            "test_y": y_te, "test_p": p_te, "test_ids": ids_te, "test_z": z_te,
+        }
+
+    return updated, all_metrics, all_preds

v4/classes/stages/warm.py

@@ -1,11 +1,13 @@
-"""stages/warm — warm stage runner: pre-trains a single tower with a temporary probe."""
+"""stages/warm — warm stage runner: pre-trains a single tower and optional real head."""
 from __future__ import annotations
 
+import importlib
+
 import torch
 import torch.nn.functional as F
 
 from v4.classes.dataset import to_label_tensor
-from v4.classes.stages.helpers import phase_for_epoch
+from v4.classes.stages.helpers import class_weights_from_shell
 
 
 def run(
@@ -20,16 +22,27 @@ def run(
     fold:        int,
     _make_loader,
     _balanced_sampler,
-) -> None:
-    """Pre-train one tower using a temporary linear probe (probe discarded after)."""
-    tower_name = stage_cfg["tower"]
-    n_epochs   = stage_cfg.get("epochs", 0)
-    level      = stage_cfg["level"]
+    stage_models: dict | None = None,
+    cfg_stages:   list[dict] | None = None,
+) -> dict:
+    """Pre-train one tower.
+
+    If ``head_name`` is set on the stage config, train that real downstream head
+    and return it in ``stage_models``. Otherwise, fall back to a temporary linear
+    probe for backward-compatible representation warmup.
+    """
+    tower_name   = stage_cfg["tower"]
+    n_epochs     = stage_cfg.get("epochs", 0)
+    level        = stage_cfg["level"]
+    shell_filter = stage_cfg.get("shell_filter", {})
+    stage_models = dict(stage_models or {})
+    cfg_stages   = list(cfg_stages or [])
 
     if n_epochs == 0:
-        return
+        return stage_models
 
-    s_train = data.build_shells(split.train, level=level, label_filter=label_filter)
+    s_train = data.build_shells(split.train, level=level, label_filter=label_filter,
+                                **shell_filter)
     bs      = cfg["training"]["batch_size"]
     loader  = _make_loader(
         s_train, {tower_name: towers[tower_name]},
@@ -41,13 +54,32 @@ def run(
         for p in t.parameters():
             p.requires_grad_(n == tower_name)
 
-    probe = torch.nn.Linear(towers[tower_name].out_dim, num_classes).to(device)
+    head_name = stage_cfg.get("head_name")
+    if head_name:
+        head_cfg = next((s for s in cfg_stages if s.get("name") == head_name), None)
+        if head_cfg is None:
+            raise ValueError(f"warm stage requested head_name={head_name!r}, but no such head exists")
+        h_mod = importlib.import_module(head_cfg.get("module", "v4.classes.heads.classifier"))
+        h_cls = getattr(h_mod, head_cfg.get("class", "ClassificationHead"))
+        probe = stage_models.get(head_name)
+        if probe is None:
+            probe = h_cls(towers[tower_name].out_dim, num_classes, **head_cfg.get("args", {}))
+        probe = probe.to(device)
+    else:
+        probe = torch.nn.Linear(towers[tower_name].out_dim, num_classes).to(device)
+
     opt   = torch.optim.Adam(
         list(towers[tower_name].parameters()) + list(probe.parameters()),
         lr=cfg["training"]["lr"],
     )
 
+    cw = class_weights_from_shell(
+        s_train, num_classes, device,
+        enabled=cfg["training"].get("class_weighted", False),
+    )
+
     towers[tower_name].train()
+    probe.train()
     for epoch in range(n_epochs):
         total_loss = total_correct = total_n = 0
         for batch in loader:
@@ -57,7 +89,7 @@ def run(
                 continue
             y_t    = to_label_tensor(y, device)
             logits = probe(towers[tower_name](x.to(device)))
-            loss   = F.cross_entropy(logits, y_t)
+            loss   = F.cross_entropy(logits, y_t, weight=cw)
             opt.zero_grad(); loss.backward(); opt.step()
             total_loss    += loss.item() * len(y_t)
             total_correct += int((logits.argmax(1) == y_t).sum())
@@ -71,3 +103,7 @@ def run(
     for t in towers.values():
         for p in t.parameters():
             p.requires_grad_(True)
+
+    if head_name:
+        stage_models[head_name] = probe
+    return stage_models

v4/classes/towers/clinical_tower.py

@@ -2,6 +2,25 @@
 
 Self-contained: no v3 dependencies.
 Inherits get_sample dispatch from TowerBase.
+
+Geometry injection (EPC consumption)
+--------------------------------------
+When geom_dim > 0, ClinicalEncoder requests the "geometry_vectors" key from EPC
+during early_pass and appends the geometry features to every clinical vector.
+The input layer is sized to clinical_data.feature_dim + geom_dim automatically.
+
+Config example (cd tower consuming geometry):
+{
+  "name":         "cd",
+  "module":       "v4.classes.towers.clinical_tower",
+  "class":        "ClinicalEncoder",
+  "data_source":  "matrix",
+  "epc_requests": ["geometry_vectors"],
+  "args": {
+    "hidden_dim": 128,
+    "geom_dim":   5
+  }
+}
 """
 from __future__ import annotations
 
@@ -14,7 +33,7 @@ from v4.classes.accessory.se_block import SEBlock
 
 
 class ClinicalEncoder(TowerBase):
-    """MLP over tabular clinical features.
+    """MLP over tabular clinical features, with optional geometry vector injection.
 
     clinical_data : ClinicalDataView — provides feature_dim, vectorize_entity, side_map
     hidden_dim    : output embedding dimensionality
@@ -22,21 +41,28 @@ class ClinicalEncoder(TowerBase):
     use_se        : wrap output with SEBlock channel gating
     se_reduction  : SEBlock bottleneck factor
     se_pre_norm   : apply LayerNorm before SEBlock
+    geom_dim      : number of geometry features to append from EPC (0 = disabled)
+                    requires epc_requests: ["geometry_vectors"] in tower config
     """
 
+    EPC_GEOMETRY_KEY = "geometry_vectors"
+
     def __init__(
         self,
         clinical_data,
-        hidden_dim: int = 128,
-        dropout: float = 0.1,
-        use_se: bool = False,
-        se_reduction: int = 16,
-        se_pre_norm: bool = True,
+        hidden_dim:  int   = 128,
+        dropout:     float = 0.1,
+        use_se:      bool  = False,
+        se_reduction: int  = 16,
+        se_pre_norm: bool  = True,
+        geom_dim:    int   = 0,
     ):
         super().__init__()
-        self.clinical_data = clinical_data
-        self._out_dim      = hidden_dim
-        feature_dim        = clinical_data.feature_dim
+        self.clinical_data  = clinical_data
+        self._out_dim       = hidden_dim
+        self._geom_dim      = geom_dim
+        self._geom_vectors: dict | None = None  # filled by early_pass when geom_dim > 0
+        feature_dim = clinical_data.feature_dim + geom_dim
 
         self.block0 = nn.Sequential(
             nn.Linear(feature_dim, hidden_dim),
@@ -53,6 +79,12 @@ class ClinicalEncoder(TowerBase):
         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
 
+    # ── EPC early_pass ───────────────────────────────────────────────────────
+
+    def early_pass(self, context) -> None:
+        if self._geom_dim > 0:
+            self._geom_vectors = context.require(self.EPC_GEOMETRY_KEY)
+
     # ── TowerBase interface ──────────────────────────────────────────────────
 
     @property
@@ -65,6 +97,14 @@ class ClinicalEncoder(TowerBase):
 
     def _get(self, *ids) -> torch.Tensor:
         arr = self.clinical_data.vectorize_entity(*ids)
+        if self._geom_dim > 0 and self._geom_vectors is not None:
+            pid  = int(ids[0])
+            eye  = str(ids[1]) if len(ids) > 1 else "OD"
+            geom = self._geom_vectors.get(
+                (pid, eye),
+                np.zeros(self._geom_dim, dtype=np.float32),
+            )
+            arr = np.concatenate([arr, geom[: self._geom_dim]])
         return torch.from_numpy(arr.astype(np.float32, copy=False))
 
     # ── nn.Module forward ────────────────────────────────────────────────────

v4/classes/towers/geometry_tower.py

@@ -0,0 +1,237 @@
+"""geometry_tower — GeometrySegEncoder for v4.
+
+A CNN tower that takes a per-eye disc/cup *segmentation map* as input (rather
+than the raw fundus image) and contributes its pooled embedding to fusion.
+
+Seg maps are produced by an underlying loader (GT contour rasterisation or
+UNet inference) during early_pass, then cached per fold.
+
+UNet fine-tuning lives in early_pass too — the loader's `finetune(train_samples)`
+call uses only the training split, then precompute() runs inference on all
+fold samples (train + val + test).
+
+Config example:
+{
+  "name":         "geom",
+  "module":       "v4.classes.towers.geometry_tower",
+  "class":        "GeometrySegEncoder",
+  "data_source":  "image",
+  "args": {
+    "backbone":          "resnet18",
+    "channels":          3,
+    "target_size":       224,
+    "augment":           true,
+    "seg_source":        "gt",
+    "contour_dir":       "Papila/ExpertsSegmentations/Contours"
+  }
+}
+"""
+from __future__ import annotations
+
+import math
+import sys
+from pathlib import Path
+from typing import Any
+
+import numpy as np
+import torch
+from torch import nn
+from torchvision import models
+
+_REPO_ROOT = Path(__file__).resolve().parents[3]
+if str(_REPO_ROOT) not in sys.path:
+    sys.path.insert(0, str(_REPO_ROOT))
+
+from v4.classes.accessory.backbones import build_backbone
+from v4.classes.towerbase import TowerBase
+
+
+class GeometrySegEncoder(TowerBase):
+    """CNN tower over disc/cup segmentation maps.
+
+    image_data    : ImageDataView — provides get_image_path(*ids) and side_map.
+                    Must implement build_seg_map_loader(source, **kwargs).
+    backbone      : backbone key (see accessory/backbones.py)
+    channels      : 1 (label map in [0,1]) or 3 (one-hot bg/rim/cup)
+    target_size   : CNN input spatial size (cached arrays already at this size)
+    augment       : random flip + 90° rotation at training time
+    freeze_ratio  : fraction of early backbone blocks to freeze in [0, 1]
+    seg_source    : passed to image_data.build_seg_map_loader (e.g. "gt", "unet")
+    **seg_kwargs  : forwarded to build_seg_map_loader
+    """
+
+    def __init__(
+        self,
+        image_data,
+        backbone:     str   = "resnet18",
+        channels:     int   = 3,
+        target_size:  int   = 224,
+        augment:      bool  = True,
+        freeze_ratio: float = 0.0,
+        seg_source:   str   = "gt",
+        **seg_kwargs: Any,
+    ):
+        super().__init__()
+        self.image_data   = image_data
+        self._channels    = channels
+        self._target_size = target_size
+        self._augment     = augment
+
+        if not hasattr(image_data, "build_seg_map_loader"):
+            raise TypeError(
+                f"GeometrySegEncoder requires image_data to implement "
+                f"build_seg_map_loader(), but {type(image_data).__name__} does not."
+            )
+        loader_kwargs = {
+            "channels":     channels,
+            "target_size":  target_size,
+            **seg_kwargs,
+        }
+        self._loader = image_data.build_seg_map_loader(seg_source, **loader_kwargs)
+        self._seg_cache: dict = {}
+        self._seg_source = seg_source
+
+        self.backbone, self._base_dim, self._blocks = build_backbone(backbone, freeze_ratio)
+        if channels != 3:
+            self._adapt_first_conv(channels)
+
+        print(
+            f"[GeometrySegEncoder] backbone={backbone}  channels={channels}  "
+            f"target_size={target_size}  seg_source={seg_source}",
+            flush=True,
+        )
+
+    # ── TowerBase interface ──────────────────────────────────────────────────
+
+    @property
+    def out_dim(self) -> int:
+        return self._base_dim
+
+    @property
+    def _side_map(self) -> dict[str, str]:
+        return self.image_data.side_map
+
+    def _get(self, *ids) -> torch.Tensor:
+        key = tuple(ids)
+        arr = self._seg_cache.get(key)
+        if arr is None:
+            arr = np.zeros(
+                (self._channels, self._target_size, self._target_size),
+                dtype=np.float32,
+            )
+        if self.training and self._augment:
+            arr = self._augment_array(arr)
+        return torch.from_numpy(np.ascontiguousarray(arr))
+
+    # ── EPC early_pass ───────────────────────────────────────────────────────
+
+    def early_pass(self, context) -> None:
+        data   = context.require("data")
+        split  = context.require("split")
+
+        train_samples = self._collect_samples(split.train, data)
+        all_samples   = self._collect_samples(split.train, data)
+        all_samples  += self._collect_samples(split.val,   data)
+        if split.test is not None:
+            all_samples += self._collect_samples(split.test, data)
+
+        # Reset per-fold state if loader supports it (UNet only).
+        if hasattr(self._loader, "reset_cache"):
+            self._loader.reset_cache()
+        if hasattr(self._loader, "reset_weights"):
+            self._loader.reset_weights()
+        if hasattr(self._loader, "finetune"):
+            self._loader.finetune(train_samples)
+
+        self._loader.precompute(all_samples)
+        self._seg_cache = self._loader.all_seg_maps()
+        print(
+            f"[GeometrySegEncoder] cached {len(self._seg_cache)} seg maps for fold",
+            flush=True,
+        )
+
+    # ── nn.Module forward ────────────────────────────────────────────────────
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        y = self.backbone(x)
+        if y.dim() > 2:
+            y = y.flatten(1)
+        return y
+
+    # ── Utilities ────────────────────────────────────────────────────────────
+
+    def set_freeze_ratio(self, ratio: float) -> None:
+        r        = max(0.0, min(1.0, float(ratio)))
+        n_freeze = int(math.floor(len(self._blocks) * r))
+        for b in self._blocks:
+            for p in b.parameters():
+                p.requires_grad = True
+        for b in self._blocks[:n_freeze]:
+            for p in b.parameters():
+                p.requires_grad = False
+
+    # ── Internals ────────────────────────────────────────────────────────────
+
+    def _collect_samples(self, df, data) -> list:
+        """Build (pid, eye, image_path) tuples from a split DataFrame."""
+        if df is None or len(df) == 0:
+            return []
+        pc = data.patient_col
+        out = []
+        for _, row in df.iterrows():
+            pid = int(row[pc])
+            eye = str(row.get("eyeID", "OD"))
+            out.append((pid, eye, data.image.get_image_path(pid, eye)))
+        return out
+
+    @staticmethod
+    def _augment_array(arr: np.ndarray) -> np.ndarray:
+        """Random flip + 90° rotation on a (C, H, W) seg-map array."""
+        if np.random.rand() < 0.5:
+            arr = arr[:, :, ::-1]
+        if np.random.rand() < 0.5:
+            arr = arr[:, ::-1, :]
+        k = int(np.random.randint(0, 4))
+        if k:
+            arr = np.rot90(arr, k=k, axes=(1, 2))
+        return arr
+
+    def _adapt_first_conv(self, in_channels: int) -> None:
+        """Replace the first Conv2d to accept a non-3-channel input.
+
+        Pretrained weights are averaged across the original input channels and
+        broadcast across the new ones.
+        """
+        first = self._find_first_conv(self.backbone)
+        new   = nn.Conv2d(
+            in_channels,
+            first.out_channels,
+            kernel_size=first.kernel_size,
+            stride=first.stride,
+            padding=first.padding,
+            bias=first.bias is not None,
+        )
+        with torch.no_grad():
+            new.weight.copy_(
+                first.weight.mean(dim=1, keepdim=True).expand_as(new.weight)
+            )
+            if first.bias is not None:
+                new.bias.copy_(first.bias)
+        self._replace_first_conv(self.backbone, new)
+
+    @staticmethod
+    def _find_first_conv(module: nn.Module) -> nn.Conv2d:
+        for m in module.modules():
+            if isinstance(m, nn.Conv2d):
+                return m
+        raise RuntimeError("No Conv2d found in backbone")
+
+    @classmethod
+    def _replace_first_conv(cls, module: nn.Module, new_conv: nn.Conv2d) -> bool:
+        for name, child in module.named_children():
+            if isinstance(child, nn.Conv2d):
+                setattr(module, name, new_conv)
+                return True
+            if cls._replace_first_conv(child, new_conv):
+                return True
+        return False

v4/classes/towers/image_tower.py

@@ -2,50 +2,128 @@
 
 Self-contained: no v3 dependencies.
 Inherits get_sample dispatch from TowerBase.
+
+Geometry injection (EPC supply)
+--------------------------------
+When geometry_source is set, ImageEncoder asks the image_data view for a loader
+via image_data.build_geometry_loader(source, **kwargs).  The view is responsible
+for understanding what that source means for its specific domain (fundus contours,
+U-Net segmentations, cat ear landmarks, etc.).
+
+During early_pass the loader pre-computes all per-entity geometry vectors and
+publishes them to the EarlyPassContext under the key "geometry_vectors"
+({(entity_id...): np.ndarray of length geom_dim}).  ClinicalEncoder (or any
+other tower with epc_requests: ["geometry_vectors"]) can then consume them.
+
+The tower reads feature_dim and feature_names from the loader instance, so it
+can log geometry info without knowing anything about CDR, disc masks, or other
+domain-specific concepts.
+
+Config example:
+{
+  "name":         "img",
+  "module":       "v4.classes.towers.image_tower",
+  "class":        "ImageEncoder",
+  "data_source":  "image",
+  "epc_supplies": ["geometry_vectors"],
+  "args": {
+    "backbone":        "refugelike",
+    "augment":         true,
+    "geometry_source": "gt",
+    "contour_dir":     "Papila/ExpertsSegmentations/Contours"
+  }
+}
 """
 from __future__ import annotations
 
 import math
+import sys
+from pathlib import Path
+from typing import Any
 
 import torch
 from torch import nn
 
+_REPO_ROOT = Path(__file__).resolve().parents[3]
+if str(_REPO_ROOT) not in sys.path:
+    sys.path.insert(0, str(_REPO_ROOT))
+
 from v4.classes.towerbase import TowerBase
 from v4.classes.accessory.backbones import build_backbone
 from v4.classes.accessory.se_block import SEBlock
-from v4.classes.accessory.transforms import build_backbone_transform, build_eval_transform
+from v4.classes.accessory.transforms import (
+    build_backbone_transform, build_eval_transform, build_split_transforms,
+)
 
 
 class ImageEncoder(TowerBase):
     """Vision backbone → pooled feature vector.
 
-    image_data   : ImageDataView — provides load_image(*ids) and side_map
-    backbone     : backbone key (see accessory/backbones.py)
-    freeze_ratio : fraction of early blocks to freeze in [0, 1]
-    use_se       : apply SE attention over the pooled feature vector
-    augment      : include random flip/rotation/jitter in the train transform
+    image_data        : ImageDataView — provides load_image(*ids) and side_map.
+                        Must implement build_geometry_loader(source, **kwargs)
+                        if geometry_source is set.
+    backbone          : backbone key (see accessory/backbones.py)
+    freeze_ratio      : fraction of early blocks to freeze in [0, 1]
+    use_se            : apply SE attention over the pooled feature vector
+    augment           : include random flip/rotation/jitter in the train transform
+    cache_transformed : if True, cache resized + ToTensor'd float32 [0, 1] CHW
+                        tensors per fold.  Per-batch cost drops to augment +
+                        Normalize on tensors only (no PIL, no Resize, no decode).
+                        Memory: ~3 × crop_size² × 4B per cached image.
+                        Cache is rebuilt at the start of every fold via early_pass.
+    geometry_source   : source key passed to image_data.build_geometry_loader()
+                        (e.g. "gt", "unet").  None = geometry disabled.
+    **geom_kwargs     : forwarded verbatim to build_geometry_loader() — e.g.
+                        contour_dir="Papila/ExpertsSegmentations/Contours"
     """
 
+    EPC_GEOMETRY_KEY = "geometry_vectors"
+
     def __init__(
         self,
         image_data,
-        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,
+        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,
+        cache_transformed: bool       = False,
+        geometry_source:   str | None = None,
+        **geom_kwargs: Any,
     ):
         super().__init__()
         self.image_data     = image_data
         self._name          = backbone
         self.backbone, self._base_dim, self._blocks = build_backbone(backbone, freeze_ratio)
-        self.transform      = build_backbone_transform(backbone, augment=augment)
-        self.eval_transform = build_eval_transform(backbone)
+
+        self._cache_transformed = cache_transformed
+        if cache_transformed:
+            self._precache_tf, self._post_train_tf = build_split_transforms(backbone, augment=augment)
+            _,                 self._post_eval_tf  = build_split_transforms(backbone, augment=False)
+            self._tensor_cache: dict[tuple, torch.Tensor] = {}
+        else:
+            self.transform      = build_backbone_transform(backbone, augment=augment)
+            self.eval_transform = build_eval_transform(backbone)
 
         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
 
+        self._geom_loader = None
+        if geometry_source is not None:
+            if not hasattr(image_data, "build_geometry_loader"):
+                raise TypeError(
+                    f"ImageEncoder geometry_source={geometry_source!r} requires "
+                    f"image_data to implement build_geometry_loader(), "
+                    f"but {type(image_data).__name__} does not."
+                )
+            self._geom_loader = image_data.build_geometry_loader(geometry_source, **geom_kwargs)
+            print(
+                f"[ImageEncoder] geometry_source={geometry_source!r}  "
+                f"features={self._geom_loader.feature_names}",
+                flush=True,
+            )
+
     # ── TowerBase interface ──────────────────────────────────────────────────
 
     @property
@@ -57,10 +135,61 @@ class ImageEncoder(TowerBase):
         return self.image_data.side_map
 
     def _get(self, *ids) -> torch.Tensor:
+        if self._cache_transformed:
+            key = tuple(ids)
+            cached = self._tensor_cache.get(key)
+            if cached is None:
+                cached = self._precache_tf(self.image_data.load_image(*ids))
+                self._tensor_cache[key] = cached
+            tail = self._post_train_tf if self.training else self._post_eval_tf
+            return tail(cached)
         img = self.image_data.load_image(*ids)
         t   = self.transform if self.training else self.eval_transform
         return t(img)
 
+    # ── EPC early_pass ───────────────────────────────────────────────────────
+
+    def early_pass(self, context) -> None:
+        """Per-fold setup: warm tensor cache (if enabled), publish geometry vectors."""
+        data = context.require("data")
+
+        if self._cache_transformed:
+            self._tensor_cache.clear()
+            n = self._warm_tensor_cache(data, context.require("split"))
+            print(
+                f"[ImageEncoder] warmed transformed-tensor cache for {n} entries "
+                f"({self._name})",
+                flush=True,
+            )
+
+        if self._geom_loader is None:
+            return
+        self._geom_loader.precompute(data.df, patient_col=data.patient_col)
+        vecs = self._geom_loader.all_vectors()
+        context.put(self.EPC_GEOMETRY_KEY, vecs)
+        print(
+            f"[ImageEncoder] published {len(vecs)} geometry vectors "
+            f"(dim={self._geom_loader.feature_dim}) to EPC key '{self.EPC_GEOMETRY_KEY}'",
+            flush=True,
+        )
+
+    def _warm_tensor_cache(self, data, split) -> int:
+        """Pre-fill the per-tower tensor cache for all entries in this fold's splits."""
+        seen: set[tuple] = set()
+        for df in (split.train, split.val, split.test):
+            if df is None or len(df) == 0:
+                continue
+            pc = data.patient_col
+            for _, row in df.iterrows():
+                pid = int(row[pc])
+                eye = str(row.get("eyeID", "OD"))
+                key = (pid, eye)
+                if key in self._tensor_cache or key in seen:
+                    continue
+                self._tensor_cache[key] = self._precache_tf(self.image_data.load_image(pid, eye))
+                seen.add(key)
+        return len(self._tensor_cache)
+
     # ── nn.Module forward ────────────────────────────────────────────────────
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:

v4/classes/v4_hypertower.py

@@ -29,7 +29,8 @@ sys.path.insert(0, str(REPO_ROOT))
 from v4.classes.dataset import LoaderShell, HTDataset, ht_collate
 from v4.classes.utils import seed_everything, choose_device
 from v4.classes.split_manager import SplitManager
-from v4.classes.stages import warm, fusion
+from v4.classes.stages import warm, fusion, parallel
+from v4.classes.logging.prediction_store import PredictionStore, FeatureStore
 
 
 # ---------------------------------------------------------------------------
@@ -167,28 +168,49 @@ def run_fold(fold: int, splits, cfg: dict, data, num_classes: int, device) -> di
         if hasattr(tower, "early_pass"):
             tower.early_pass(context)
 
+    # Flatten parallel wrappers so sub-stage configs are addressable by name.
+    flat_stages: list[dict] = []
+    for s in cfg_stages:
+        if s.get("type") == "parallel":
+            flat_stages.extend(s["stages"])
+        else:
+            flat_stages.append(s)
+
     stage_models: dict = {}
     fold_result         = {"fold": fold}
+    fold_preds:   dict  = {}   # stage_name → pred_data
 
     for stage_cfg in cfg_stages:
         stype = stage_cfg["type"]
 
         if stype == "warm":
-            warm.run(stage_cfg, towers, data, split, label_filter,
-                     cfg, num_classes, device, fold,
-                     _make_loader, _balanced_sampler)
+            stage_models = warm.run(
+                stage_cfg, towers, data, split, label_filter,
+                cfg, num_classes, device, fold,
+                _make_loader, _balanced_sampler, stage_models, flat_stages,
+            )
 
         elif stype == "fusion":
-            stage_models, metrics = fusion.run(
+            stage_models, metrics, preds = fusion.run(
                 stage_cfg, cfg, towers, stage_models, data, split,
-                label_filter, num_classes, device, fold, cfg_stages,
+                label_filter, num_classes, device, fold, flat_stages,
                 _make_loader,
             )
             fold_result.update(metrics)
+            fold_preds.update(preds)
+
+        elif stype == "parallel":
+            stage_models, metrics, preds = parallel.run(
+                stage_cfg, cfg, towers, stage_models, data, split,
+                label_filter, num_classes, device, fold, flat_stages,
+                _make_loader, _balanced_sampler,
+            )
+            fold_result.update(metrics)
+            fold_preds.update(preds)
 
         # head stages are handled inside fusion.run
 
-    return fold_result
+    return fold_result, fold_preds
 
 
 # ---------------------------------------------------------------------------
@@ -239,8 +261,11 @@ def main():
         out_dir = out_dir / tag
     out_dir.mkdir(parents=True, exist_ok=True)
 
-    eval_stage   = cfg.get("eval_stage", "hb")
-    fold_results = []
+    eval_stage        = cfg.get("eval_stage", "hb")
+    save_predictions  = cfg.get("save_predictions", False)
+    save_features     = cfg.get("save_features", False)
+    fold_results      = []
+    eval_stage_preds  = []   # list[dict] — one per fold, only for eval_stage
     t0 = time.time()
 
     for fold in range(cfg.get("folds", 5)):
@@ -248,8 +273,10 @@ def main():
         n_train = split.train[group_col].nunique() if group_col else len(split.train)
         print(f"\n── fold {fold+1}/{cfg.get('folds', 5)}  train_groups={n_train} ──",
               flush=True)
-        result = run_fold(fold, splits, cfg, data, num_classes, device)
+        result, fold_preds = run_fold(fold, splits, cfg, data, num_classes, device)
         fold_results.append(result)
+        if save_predictions and eval_stage in fold_preds:
+            eval_stage_preds.append(fold_preds[eval_stage])
         print(
             f"  fold{fold+1} DONE"
             f"  val_auc={result.get(f'{eval_stage}_val_auc', float('nan')):.4f}"
@@ -273,6 +300,7 @@ def main():
             "elapsed_s":     round(time.time() - t0, 1),
             "fold_results":  fold_results,
         }
+        out_dir.mkdir(parents=True, exist_ok=True)
         summary_path = out_dir / "summary.json"
         summary_path.write_text(json.dumps(summary, indent=2))
         print(f"\n{'='*60}", flush=True)
@@ -280,6 +308,81 @@ def main():
         print(f"Test AUC: {summary['mean_test_auc']:.4f} ± {summary['std_test_auc']:.4f}", flush=True)
         print(f"Saved: {summary_path}", flush=True)
 
+        if save_predictions and eval_stage_preds:
+            # Collect all unique entity_ids across val+test sets of all folds.
+            seen, all_ids, id_to_y = set(), [], {}
+            for fp in eval_stage_preds:
+                for eid, y in zip(fp["val_ids"], fp["val_y"]):
+                    k = str(eid)
+                    if k not in seen:
+                        seen.add(k); all_ids.append(eid)
+                    id_to_y[k] = int(y)
+                if fp.get("test_ids"):
+                    for eid, y in zip(fp["test_ids"], fp["test_y"]):
+                        k = str(eid)
+                        if k not in seen:
+                            seen.add(k); all_ids.append(eid)
+                        id_to_y[k] = int(y)
+
+            y_true = np.array([id_to_y.get(str(e), -1) for e in all_ids], dtype=np.int64)
+            store  = PredictionStore(n_folds=len(eval_stage_preds), n_classes=num_classes)
+            store.register_phase(
+                phase=eval_stage,
+                entity_ids=all_ids,
+                y_true=y_true,
+                head_names=[f"{eval_stage}_head"],
+                n_epochs=1,
+            )
+            for fold_idx, fp in enumerate(eval_stage_preds):
+                store.record(eval_stage, fold_idx, 0, fp["val_ids"],
+                             f"{eval_stage}_head", fp["val_p"])
+                store.set_split(eval_stage, fold_idx, fp["val_ids"], "val")
+                if fp.get("test_ids"):
+                    store.record(eval_stage, fold_idx, 0, fp["test_ids"],
+                                 f"{eval_stage}_head", fp["test_p"])
+                    store.set_split(eval_stage, fold_idx, fp["test_ids"], "test")
+
+            pred_path = out_dir / "predictions.h5"
+            store.save(pred_path)
+            print(f"Predictions saved: {pred_path}", flush=True)
+
+        if save_features and eval_stage_preds:
+            emb_dim = eval_stage_preds[0]["val_z"].shape[-1]
+            fstore  = FeatureStore(n_folds=len(eval_stage_preds))
+
+            # Build entity_id / y_true universe (same as predictions).
+            seen, all_ids, id_to_y = set(), [], {}
+            for fp in eval_stage_preds:
+                for eid, y in zip(fp["val_ids"], fp["val_y"]):
+                    k = str(eid)
+                    if k not in seen:
+                        seen.add(k); all_ids.append(eid)
+                    id_to_y[k] = int(y)
+                if fp.get("test_ids"):
+                    for eid, y in zip(fp["test_ids"], fp["test_y"]):
+                        k = str(eid)
+                        if k not in seen:
+                            seen.add(k); all_ids.append(eid)
+                        id_to_y[k] = int(y)
+
+            y_true = np.array([id_to_y.get(str(e), -1) for e in all_ids], dtype=np.int64)
+            fstore.register_phase(phase=eval_stage, entity_ids=all_ids, y_true=y_true)
+            fstore.register_head(phase=eval_stage, head=f"{eval_stage}_embedding",
+                                 n_epochs=1, embedding_dim=emb_dim)
+
+            for fold_idx, fp in enumerate(eval_stage_preds):
+                fstore.record(eval_stage, fold_idx, 0, fp["val_ids"],
+                              f"{eval_stage}_embedding", fp["val_z"])
+                fstore.set_split(eval_stage, fold_idx, fp["val_ids"], "val")
+                if fp.get("test_ids") and fp.get("test_z") is not None:
+                    fstore.record(eval_stage, fold_idx, 0, fp["test_ids"],
+                                  f"{eval_stage}_embedding", fp["test_z"])
+                    fstore.set_split(eval_stage, fold_idx, fp["test_ids"], "test")
+
+            feat_path = out_dir / "features.h5"
+            fstore.save(feat_path)
+            print(f"Features saved: {feat_path}", flush=True)
+
 
 if __name__ == "__main__":
     main()