Refactor geometry feature loaders and enhance distributed client status reporting

- Updated GTGeometryLoader to streamline geometry vector computation and caching.
- Introduced UNetGeometryLoader for UNet-derived geometry vectors.
- Added sample collection method in PapilaBundle for better data handling.
- Refined GeometrySegEncoder and ImageEncoder to utilize new sample collection.
- Enhanced distributed client with heartbeat mechanism for improved job tracking.
- Added new configuration files for UNet-derived geometry integration.

v4/classes/profiles/fundus_images.py

@@ -124,71 +124,47 @@ def compute_geometry_features(disc_mask: np.ndarray, cup_mask: np.ndarray) -> np
 # ---------------------------------------------------------------------------
 
 class GTGeometryLoader:
-    """Pre-computes per-eye geometry vectors from PAPILA GT contour annotations.
+    """Per-eye 5-feature CDR vectors from PAPILA GT contour annotations.
 
-    File naming: RET{pid:03d}{eye}_{disc|cup}_exp{n}.txt
+    Loads contour coordinates from per-expert text files, rasterises them at
-    Averages exp1 and exp2 when both are present; zero vector for missing entries.
+    the original image's pixel space (so polygons aren't clipped), then
-
+    computes the CDR feature vector.  Expert masks are merged before feature
-    Usage:
+    computation to match the seg-map path.
-        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:
+    def __init__(self, contour_dir: str | Path, *, mask_size: int = _MASK_SIZE[0]) -> None:
-        self._dir   = Path(contour_dir)
+        self._dir       = Path(contour_dir)
+        self._mask_size = mask_size
         self._cache: dict[tuple, np.ndarray] = {}
 
-    def precompute(self, df, patient_col: str = "Patient ID") -> None:
+    def reset_cache(self) -> None:
+        self._cache.clear()
+
+    def precompute(self, samples: Iterable[Tuple[int, str, Path]]) -> None:
         n_ok = 0
-        for _, row in df.iterrows():
+        for pid, eye, image_path in samples:
-            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)
+            with Image.open(image_path) as img:
-            self._cache[key] = vec if vec is not None else np.zeros(self.feature_dim, dtype=np.float32)
+                image_size = img.size  # (W, H)
-            if vec is not None:
+            res = _papila_disc_cup_masks(
+                pid, eye, self._dir, image_size, self._mask_size,
+            )
+            if res is None:
+                self._cache[key] = np.zeros(self.feature_dim, dtype=np.float32)
+            else:
+                disc, cup = res
+                self._cache[key] = compute_geometry_features(disc, cup)
                 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)
@@ -411,12 +387,103 @@ class _UNetFTDataset(Dataset):
         return self._imgs[idx], self._masks[idx]
 
 
+class _PapilaUNetMaskPipeline:
+    """Per-fold UNet mask producer.
+
+    Owns a UNetSegmenter, handles the per-fold lifecycle:
+      - reset_weights() restores REFUGE base state (call at start of each fold)
+      - finetune(train_samples) fine-tunes on the train split's GT contours
+      - predict(samples) returns {(pid, eye): (disc_mask, cup_mask)} raw masks
+
+    Downstream loaders interpret the raw masks differently — seg-map loader
+    crops/resizes/encodes for a CNN, geometry-vector loader computes 5 CDR
+    features.  Sharing this pipeline avoids duplicating UNet load + finetune
+    when both seg-map and feature-vector outputs are needed in one config.
+    """
+
+    def __init__(
+        self,
+        weights_path:        str | Path,
+        *,
+        contour_dir:         str | Path,
+        unet_size:           int   = 512,
+        normalize:           str   = "per_image",
+        threshold:           float = 0.5,
+        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._contour_dir   = Path(contour_dir)
+        self._threshold     = threshold
+        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(Path(weights_path))
+        self._base_state = copy.deepcopy(self._segmenter.model.state_dict())
+
+    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."""
+        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"[UNetMaskPipeline] 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 predict(
+        self, samples: Iterable[Tuple[int, str, Path]],
+    ) -> dict[tuple, Tuple[np.ndarray, np.ndarray]]:
+        """Run inference. Returns {(pid, eye): (disc_mask, cup_mask)} raw uint8."""
+        import time
+        samples = list(samples)
+        if not samples:
+            return {}
+        print(
+            f"[UNetMaskPipeline] running UNet inference on {len(samples)} images...",
+            flush=True,
+        )
+        result: dict[tuple, Tuple[np.ndarray, np.ndarray]] = {}
+        t0     = time.time()
+        report = max(1, len(samples) // 4)
+        for i, (pid, eye, image_path) in enumerate(samples, 1):
+            with Image.open(image_path) as raw:
+                disc, cup = self._segmenter.predict(raw, threshold=self._threshold)
+            result[(pid, eye)] = (disc, cup)
+            if i % report == 0 or i == len(samples):
+                print(
+                    f"  [UNet inf] {i}/{len(samples)}  ({time.time() - t0:.1f}s)",
+                    flush=True,
+                )
+        return result
+
+
 class UNetSegMapLoader:
-    """Pre-computes per-eye seg maps via a REFUGE-pretrained UNet.
+    """Per-eye CNN-ready seg maps via a REFUGE-pretrained UNet.
 
-    Optionally fine-tunes the UNet per fold on the training split's GT contours.
+    Wraps a `_PapilaUNetMaskPipeline` and post-processes raw masks into
+    (C, H, W) float32 arrays sized for a downstream CNN.
 
-    Output: dict {(pid, eye): np.ndarray (C, H, W) float32} cached for the fold.
+    Output: dict {(pid, eye): np.ndarray (C, H, W) float32}
     """
 
     def __init__(
@@ -435,22 +502,20 @@ class UNetSegMapLoader:
         finetune_batch_size: int   = 4,
         device:              str | None = None,
     ) -> None:
-        from v4.classes.accessory.unet import UNetSegmenter
+        self._pipeline = _PapilaUNetMaskPipeline(
-
+            weights_path,
-        self._weights_path        = Path(weights_path)
+            contour_dir=contour_dir,
-        self._contour_dir         = Path(contour_dir)
+            unet_size=unet_size,
-        self._channels            = channels
+            normalize=normalize,
-        self._target_size         = target_size
+            threshold=threshold,
-        self._threshold           = threshold
+            finetune_epochs=finetune_epochs,
-        self._crop                = crop_to_disc
+            finetune_lr=finetune_lr,
-        self._ft_epochs           = finetune_epochs
+            finetune_batch_size=finetune_batch_size,
-        self._ft_lr               = finetune_lr
+            device=device,
-        self._ft_batch_size       = finetune_batch_size
+        )
-
+        self._channels    = channels
-        self._segmenter = UNetSegmenter(
+        self._target_size = target_size
-            target_size=unet_size, normalize=normalize, device=device,
+        self._crop        = crop_to_disc
-        ).load_weights(self._weights_path)
-        self._base_state = copy.deepcopy(self._segmenter.model.state_dict())
         self._cache: dict[tuple, np.ndarray] = {}
 
     @property
@@ -458,89 +523,130 @@ class UNetSegMapLoader:
         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._pipeline.reset_weights()
-        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.
+        self._pipeline.finetune(train_samples)
-
-        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."""
+        todo  = [s for s in samples if (s[0], s[1]) not in self._cache]
-        import time
+        masks = self._pipeline.predict(todo)
-        samples = list(samples)
+        for (pid, eye), (disc, cup) in masks.items():
-        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):
+        if masks:
-                print(
+            print(
-                    f"  [UNet inf] {i}/{len(todo)}  ({time.time() - t0:.1f}s)",
+                f"[UNetSegMapLoader] {len(masks)} seg maps cached "
-                    flush=True,
+                f"(channels={self._channels}, target={self._target_size})",
-                )
+                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
 
 
+class UNetGeometryLoader:
+    """Per-eye 5-feature CDR vectors derived from UNet-predicted masks.
+
+    Same UNet lifecycle as `UNetSegMapLoader` but the output is a 5-vector
+    (compute_geometry_features over the predicted disc/cup masks) rather
+    than a CNN-ready seg map.  Designed to slot into ImageEncoder's
+    geometry_source mechanism for vector-style geometry injection.
+    """
+
+    feature_dim   = _FEATURE_DIM
+    feature_names = _FEATURE_NAMES
+
+    def __init__(
+        self,
+        weights_path: str | Path,
+        *,
+        contour_dir:         str | Path,
+        unet_size:           int   = 512,
+        normalize:           str   = "per_image",
+        threshold:           float = 0.5,
+        finetune_epochs:     int   = 0,
+        finetune_lr:         float = 1e-5,
+        finetune_batch_size: int   = 4,
+        device:              str | None = None,
+    ) -> None:
+        self._pipeline = _PapilaUNetMaskPipeline(
+            weights_path,
+            contour_dir=contour_dir,
+            unet_size=unet_size,
+            normalize=normalize,
+            threshold=threshold,
+            finetune_epochs=finetune_epochs,
+            finetune_lr=finetune_lr,
+            finetune_batch_size=finetune_batch_size,
+            device=device,
+        )
+        self._cache: dict[tuple, np.ndarray] = {}
+
+    def reset_cache(self) -> None:
+        self._cache.clear()
+
+    def reset_weights(self) -> None:
+        self._pipeline.reset_weights()
+
+    def finetune(self, train_samples: list) -> None:
+        self._pipeline.finetune(train_samples)
+
+    def precompute(self, samples: Iterable[Tuple[int, str, Path]]) -> None:
+        todo  = [s for s in samples if (s[0], s[1]) not in self._cache]
+        masks = self._pipeline.predict(todo)
+        for (pid, eye), (disc, cup) in masks.items():
+            self._cache[(pid, eye)] = compute_geometry_features(disc, cup)
+        if masks:
+            print(
+                f"[UNetGeometryLoader] {len(masks)} geometry vectors cached "
+                f"(dim={self.feature_dim})",
+                flush=True,
+            )
+
+    def all_vectors(self) -> dict:
+        return dict(self._cache)
+
+
 # ---------------------------------------------------------------------------
 # Factories
 # ---------------------------------------------------------------------------
 
 def build_geometry_loader(source: str, **kwargs):
-    """Return the appropriate geometry-vector loader for the given source string.
+    """Return the appropriate geometry-vector loader for the given source.
 
     Parameters
     ----------
-    source      : "gt" | "unet"
+    source : "gt" | "unet"
-    contour_dir : (gt) path to contour annotation directory
+
+    GT kwargs:
+        contour_dir
+    UNet kwargs:
+        weights_path, contour_dir, unet_size=512, normalize="per_image",
+        threshold=0.5, finetune_epochs=0, finetune_lr=1e-5,
+        finetune_batch_size=4, device=None
     """
     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)
+    if source == "unet":
+        if "weights_path" not in kwargs:
+            raise ValueError("build_geometry_loader source='unet' requires weights_path")
+        if "contour_dir" not in kwargs:
+            raise ValueError(
+                "build_geometry_loader source='unet' requires contour_dir "
+                "(needed for per-fold fine-tuning, even if finetune_epochs=0)"
+            )
+        return UNetGeometryLoader(**kwargs)
     raise NotImplementedError(f"build_geometry_loader: source={source!r} not implemented")
 
 

v4/classes/profiles/v4papila.py

@@ -507,6 +507,25 @@ class PapilaBundle:
         """
         return [self._bundle.patient_col, "eyeID"]
 
+    # ── Sample collection (for tower early_pass) ─────────────────────────────
+
+    def collect_samples(self, df: pd.DataFrame | None) -> list[tuple]:
+        """Build (pid, eye, image_path) tuples from a split DataFrame.
+
+        Used by tower early_pass implementations that need per-eye image paths
+        (UNet inference, contour rasterisation, etc.).  Returns [] for an
+        empty/None df.
+        """
+        if df is None or len(df) == 0:
+            return []
+        pc = self._bundle.patient_col
+        out: list[tuple] = []
+        for _, row in df.iterrows():
+            pid = int(row[pc])
+            eye = str(row.get("eyeID", "OD"))
+            out.append((pid, eye, self.image.get_image_path(pid, eye)))
+        return out
+
     # ── Backward-compat delegates ────────────────────────────────────────────
 
     @property

v4/classes/towers/geometry_tower.py

@@ -126,16 +126,14 @@ class GeometrySegEncoder(TowerBase):
     # ── EPC early_pass ───────────────────────────────────────────────────────
 
     def early_pass(self, context) -> None:
-        data   = context.require("data")
+        data  = context.require("data")
-        split  = context.require("split")
+        split = context.require("split")
 
-        train_samples = self._collect_samples(split.train, data)
+        train_samples = data.collect_samples(split.train)
-        all_samples   = self._collect_samples(split.train, data)
+        all_samples   = train_samples + data.collect_samples(split.val)
-        all_samples  += self._collect_samples(split.val,   data)
         if split.test is not None:
-            all_samples += self._collect_samples(split.test, data)
+            all_samples += data.collect_samples(split.test)
 
-        # 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"):
@@ -172,18 +170,6 @@ class GeometrySegEncoder(TowerBase):
 
     # ── 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."""

v4/classes/towers/image_tower.py

@@ -151,11 +151,12 @@ class ImageEncoder(TowerBase):
 
     def early_pass(self, context) -> None:
         """Per-fold setup: warm tensor cache (if enabled), publish geometry vectors."""
-        data = context.require("data")
+        data  = context.require("data")
+        split = context.require("split")
 
         if self._cache_transformed:
             self._tensor_cache.clear()
-            n = self._warm_tensor_cache(data, context.require("split"))
+            n = self._warm_tensor_cache(data, split)
             print(
                 f"[ImageEncoder] warmed transformed-tensor cache for {n} entries "
                 f"({self._name})",
@@ -164,7 +165,20 @@ class ImageEncoder(TowerBase):
 
         if self._geom_loader is None:
             return
-        self._geom_loader.precompute(data.df, patient_col=data.patient_col)
+
+        train_samples = data.collect_samples(split.train)
+        all_samples   = train_samples + data.collect_samples(split.val)
+        if split.test is not None:
+            all_samples += data.collect_samples(split.test)
+
+        if hasattr(self._geom_loader, "reset_cache"):
+            self._geom_loader.reset_cache()
+        if hasattr(self._geom_loader, "reset_weights"):
+            self._geom_loader.reset_weights()
+        if hasattr(self._geom_loader, "finetune"):
+            self._geom_loader.finetune(train_samples)
+
+        self._geom_loader.precompute(all_samples)
         vecs = self._geom_loader.all_vectors()
         context.put(self.EPC_GEOMETRY_KEY, vecs)
         print(

v4/configs/ensemble_fused_geom_unet.json

@@ -0,0 +1,142 @@
+{
+  "_notes": [
+    "Geometry vector injection via UNet (no GT, no seg-CNN tower).",
+    "img tower runs UNet per-fold to predict disc/cup masks, computes 5 CDR",
+    "features per eye, publishes to EPC; cd tower consumes them.",
+    "Same shape as ensemble_fused.json — just adds the UNet-derived geometry hook."
+  ],
+  "run_name": "v4/ensemble_fused_geom_unet",
+  "num_classes": 2,
+  "label_filter": [0, 1],
+  "split_identity_level": 1,
+  "eval_stage": "hb",
+  "save_predictions": true,
+  "seed": 1234,
+  "folds": 5,
+  "fold_seed": 100,
+  "output_root": "v4/results",
+  "out_dir_tags": ["binary"],
+
+  "data": {
+    "module": "v4.classes.profiles.v4papila",
+    "args": {
+      "image_dir": "Papila/FundusImages",
+      "clinical_dir": "Papila/ClinicalData",
+      "label_col": "Diagnosis",
+      "iop_corr_method": "ratio",
+      "iop_drop_raw": true,
+      "exclude_cols": ["Axial_Length"],
+      "in_memory_cache": true
+    }
+  },
+
+  "towers": [
+    {
+      "name": "img",
+      "module": "v4.classes.towers.image_tower",
+      "class": "ImageEncoder",
+      "data_source": "image",
+      "epc_supplies": ["geometry_vectors"],
+      "args": {
+        "backbone":            "refugelike",
+        "freeze_ratio":        0.0,
+        "augment":             true,
+        "geometry_source":     "unet",
+        "weights_path":        "models/v2/refuge/segmentation/per_image/best.pt",
+        "contour_dir":         "Papila/ExpertsSegmentations/Contours",
+        "unet_size":           512,
+        "normalize":           "per_image",
+        "threshold":           0.5,
+        "finetune_epochs":     10,
+        "finetune_lr":         1e-5,
+        "finetune_batch_size": 4
+      }
+    },
+    {
+      "name": "cd",
+      "module": "v4.classes.towers.clinical_tower",
+      "class": "ClinicalEncoder",
+      "data_source": "matrix",
+      "epc_requests": ["geometry_vectors"],
+      "args": {
+        "hidden_dim": 128,
+        "geom_dim":   5
+      }
+    }
+  ],
+
+  "stages": [
+    {
+      "name": "cd_warm",
+      "type": "warm",
+      "tower": "cd",
+      "head_name": "cd_aux",
+      "level": "eye",
+      "epochs": 40
+    },
+    {
+      "name": "img_aux",
+      "type": "head",
+      "input": "img",
+      "train_with": "nt",
+      "bcd": true
+    },
+    {
+      "name": "cd_aux",
+      "type": "head",
+      "input": "cd",
+      "train_with": "nt",
+      "bcd": true
+    },
+    {
+      "name": "nt",
+      "type": "fusion",
+      "module": "v4.classes.bridges.fusion_bridge",
+      "class": "FusionBridge",
+      "inputs": ["img", "cd"],
+      "level": "eye",
+      "epochs": 36,
+      "train_towers": true,
+      "warmup": {
+        "tower_epochs": 3,
+        "fused_epochs": 3
+      },
+      "args": {
+        "fusion_dim": 256
+      }
+    },
+    {
+      "name": "nt_head",
+      "type": "head",
+      "input": "nt",
+      "train_with": "nt"
+    },
+    {
+      "name": "hb",
+      "type": "fusion",
+      "module": "v4.classes.bridges.hyperbridge",
+      "class": "HyperBridge",
+      "inputs": { "a": "nt", "b": "nt" },
+      "level": "patient",
+      "epochs": 10,
+      "args": {
+        "hidden_dim": 256,
+        "mode": "embedding_mlp"
+      }
+    },
+    {
+      "name": "hb_head",
+      "type": "head",
+      "input": "hb",
+      "train_with": "hb",
+      "args": { "dropout": 0.3 }
+    }
+  ],
+
+  "training": {
+    "lr": 1e-4,
+    "batch_size": 8,
+    "bcd_prob": 0.5,
+    "tune_binary_threshold": true
+  }
+}

v4/distributed/cli.py

@@ -126,11 +126,12 @@ def _clients_table(api: _API) -> str:
             c["hostname"],
             c["gpu_info"][:30],
             s["state"],
+            (s.get("job_id") or "-")[:12],
             s.get("run_name") or "-",
             prog,
             _ago(c["last_seen"]),
         ])
-    headers = ["ID", "HOST", "GPU", "STATE", "RUN", "PROGRESS", "SEEN"]
+    headers = ["ID", "HOST", "GPU", "STATE", "JOB_ID", "RUN", "PROGRESS", "SEEN"]
     widths = [max(len(str(r[i])) for r in ([headers] + rows)) for i in range(len(headers))]
     sep = "  "
     lines = []

v4/distributed/client.py

@@ -232,6 +232,20 @@ def _run_job(job: JobSpec, server: _Server,
     log_dir.mkdir(parents=True, exist_ok=True)
     log_file = log_dir / f"job_{job.job_id}.log"
     ctx: dict = {}
+    last_push = [time.time()]  # mutable holder so _tail and _heartbeat share it
+
+    def _push():
+        server.push_status(StatusPush(
+            state="running",
+            job_id=job.job_id,
+            run_name=job.run_name,
+            fold=ctx.get("fold"),
+            stage=ctx.get("stage"),
+            epoch=ctx.get("epoch"),
+            total_epochs=ctx.get("total_epochs"),
+            last_val_auc=ctx.get("last_val_auc"),
+        ))
+        last_push[0] = time.time()
 
     def _tail(path: Path):
         with open(path, "r") as f:
@@ -242,16 +256,7 @@ def _run_job(job: JobSpec, server: _Server,
                     info = _parse_line(raw)
                     ctx.update(info)
                     if "epoch" in info:
-                        server.push_status(StatusPush(
+                        _push()
-                            state="running",
-                            job_id=job.job_id,
-                            run_name=job.run_name,
-                            fold=ctx.get("fold"),
-                            stage=ctx.get("stage"),
-                            epoch=ctx.get("epoch"),
-                            total_epochs=ctx.get("total_epochs"),
-                            last_val_auc=ctx.get("last_val_auc"),
-                        ))
                 elif proc.poll() is not None:
                     for raw in f:
                         print(raw, end="", flush=True)
@@ -259,6 +264,18 @@ def _run_job(job: JobSpec, server: _Server,
                 else:
                     time.sleep(0.05)
 
+    def _heartbeat():
+        # Push a status update every ~30s even when no log line is parsed.
+        # Prevents the server's reaper from declaring this client stale during
+        # long deterministic blocks (UNet fine-tune, data load, etc.).
+        while proc.poll() is None:
+            time.sleep(5)
+            if time.time() - last_push[0] >= 30:
+                try:
+                    _push()
+                except Exception:
+                    pass
+
     with open(log_file, "w") as logf:
         proc = subprocess.Popen(
             cmd,
@@ -268,10 +285,13 @@ def _run_job(job: JobSpec, server: _Server,
             start_new_session=True,
         )
 
-    tailer = threading.Thread(target=_tail, args=(log_file,), daemon=True)
+    tailer    = threading.Thread(target=_tail,      args=(log_file,), daemon=True)
+    heartbeat = threading.Thread(target=_heartbeat,                     daemon=True)
     tailer.start()
+    heartbeat.start()
     proc.wait()
     tailer.join(timeout=5)
+    heartbeat.join(timeout=5)
     log_file.unlink(missing_ok=True)
 
     success = proc.returncode == 0

v4/distributed/server.py

@@ -48,13 +48,20 @@ _DB_PATH: Path = Path("v4/distributed/jobs.db")
 _REPO_ROOT: Path = Path.cwd()
 _CLIENT_TTL: int = 120  # seconds before a client is considered gone
 _MAX_ATTEMPTS: int = 3  # max times a job is retried before being left as failed
+_SERVER_START_TS: float = 0.0  # set in main(); used as a reaper grace window
 
 _clients: dict[str, ClientInfo] = {}
 _clients_lock = threading.Lock()
 
 
 def _reap_stale_clients():
-    """Background thread: remove silent clients and re-queue their running jobs."""
+    """Background thread: remove silent clients and re-queue their running jobs.
+
+    On startup, the in-memory `_clients` dict is empty until clients re-register
+    via the `please_reregister` mechanism.  We skip the running-job re-queue pass
+    for the first `_CLIENT_TTL` seconds after startup so still-alive clients have
+    time to come back; otherwise the reaper would orphan their jobs.
+    """
     while True:
         time.sleep(30)
         cutoff = datetime.now(timezone.utc).timestamp() - _CLIENT_TTL
@@ -73,6 +80,10 @@ def _reap_stale_clients():
                 del _clients[cid]
             known_ids = set(_clients.keys())
 
+        in_grace = (time.time() - _SERVER_START_TS) < _CLIENT_TTL
+        if in_grace:
+            continue
+
         with _db() as conn:
             rows = conn.execute(
                 "SELECT job_id, assigned_to FROM jobs WHERE state='running'"
@@ -454,12 +465,13 @@ def main():
     if not args.token:
         ap.error("--token is required (or set HT_TOKEN)")
 
-    global _TOKEN, _DB_PATH, _REPO_ROOT, _CLIENT_TTL, _MAX_ATTEMPTS
+    global _TOKEN, _DB_PATH, _REPO_ROOT, _CLIENT_TTL, _MAX_ATTEMPTS, _SERVER_START_TS
     _TOKEN = args.token
     _DB_PATH = Path(args.db)
     _REPO_ROOT = Path(args.root).resolve() if args.root else Path.cwd()
     _CLIENT_TTL = args.client_ttl
     _MAX_ATTEMPTS = args.max_attempts
+    _SERVER_START_TS = time.time()
     _init_db()
 
     reaper = threading.Thread(target=_reap_stale_clients, daemon=True)

v4/scripts/experiments/tri_v1/vec_geom_gt.json

@@ -0,0 +1,7 @@
+[
+  {
+    "_note": "GT-derived 5-feature CDR vector injected into cd via EPC. Re-run with bug fix (rasterisation now in original image space, not 512x512 canvas). 3 reps.",
+    "run_name": "experiments/tri_v1/geom_vec_gt",
+    "reps": 3
+  }
+]

v4/scripts/experiments/tri_v1/vec_geom_unet.json

@@ -0,0 +1,7 @@
+[
+  {
+    "_note": "UNet-derived 5-feature CDR vector injected into cd via EPC. 3 reps.",
+    "run_name": "experiments/tri_v1/geom_vec_unet",
+    "reps": 3
+  }
+]