diff --git a/TODO.md b/TODO.md new file mode 100644 index 0000000..27cf8f8 --- /dev/null +++ b/TODO.md @@ -0,0 +1,23 @@ +# Project TODO + +## Paper + +- [ ] **Learning curve analysis** — train on 25/50/75/100% of training data, plot AUC vs n. + Motivation: empirical evidence that the model is data-starved, which justifies the decision + not to pursue attention-gating (transformer) extensions to the NTowerHT bridge. + If the curve is still ascending at full data → supports the argument that a more expressive + architecture would overfit at this sample size. Generates a figure for the paper. + +- [ ] **GradCAM nasal-side analysis** — re-run GradCAM separately for OD and OS eyes rather + than aggregated. The current aggregation mirrors the two eyes against each other, washing out + any directional bias. Clinically, we would expect GradCAM attention offset from the disc center + to trend toward the nasal side (where RNFL loss presents earliest in glaucoma). If the model + has learned this, it would only be visible in per-side heatmaps — OD and OS are mirror images + so the nasal direction is opposite for each. This could be a strong interpretability result + for the paper if the bias is present. + +- [ ] **Quantify attention-gating as future work** — use the learning curve result + parameter + count ratio (Q/K/V projections over fusion_dim vs training n) to formally justify the choice. + Frame in paper as: "we identify cross-attention inside the NTowerHT bridge as a promising + extension, but our sample size (N≈400 training patients) is insufficient to avoid overfitting + a more expressive interaction layer" — cite the learning curve figure as evidence. diff --git a/v3/PLAN_v4_hypertower.md b/v3/PLAN_v4_hypertower.md new file mode 100644 index 0000000..c94a5e6 --- /dev/null +++ b/v3/PLAN_v4_hypertower.md @@ -0,0 +1,179 @@ +# v4 HyperTower Planning Document + +## Goals +Rebuild the orchestrator using `run_ntower_cv` as the architectural foundation, with four key improvements: JSON config, decoupled data sources, declarative tower lists, and a cross-tower communication protocol. + +--- + +## 1. JSON Config (replace argparse) + +The orchestrator receives a single JSON config file. It has no hardcoded knowledge of what args individual towers or data modules need — it just forwards the relevant subtrees. + +```json +{ + "run_name": "v4/ensemble_fused", + "eval_mode": "binary", + "epochs": 30, + "fusion_epochs": 10, + "fold_seed": 100, + "seed": 1234, + "data": { + "module": "v4.papila.v4papila", + "args": { + "iop_corr_method": "ratio", + "iop_drop_raw": true, + "exclude_cols": ["Axial_Length"] + } + }, + "towers": [ + { + "name": "img", + "module": "v3.classes.image_towers", + "class": "ImageEncoder", + "args": { "backbone": "refugelike", "freeze_ratio": 0.5, "augment": true }, + "warmup_epochs": 0 + }, + { + "name": "cd", + "module": "v3.classes.clinical_towers", + "class": "ClinicalEncoder", + "args": { "hidden_dim": 128 }, + "warmup_epochs": 40 + } + ], + "bridge": { + "mode": "embedding_mlp", + "fusion_dim": 256, + "hidden_dim": 256 + }, + "training": { + "lr": 1e-4, + "batch_size": 16, + "bcd_prob": 0.5, + "warmup_tower_epochs": 3, + "warmup_fused_epochs": 3 + } +} +``` + +The orchestrator loads this with `json.load`, then calls `importlib.import_module(cfg["data"]["module"]).build_data(cfg["data"]["args"])` and similarly instantiates towers. No argparse anywhere in the orchestrator. + +--- + +## 2. Decoupled Data Sources + +`v3/classes/papila_builders.py` → clone to `v4/papila/v4papila.py`. + +Merge in the relevant logic from `papila_data.py` (preprocessing, feature typing, IOP correction, etc.) so `v4papila.py` is self-contained. + +Contract: every data module must expose: +```python +def build_data(args: dict) -> DataBundle: + ... +``` +The orchestrator calls `build_data` and gets back a `DataBundle`. It knows nothing else about the data source. Future modules (e.g. `v4/eyepacs/eyepacs_data.py`) just implement the same function. + +--- + +## 3. Declarative Tower List + +Towers are loaded from the `"towers"` list in the JSON and stored as an ordered dict keyed by `name`. The orchestrator never imports a tower class directly. + +```python +towers = {} +for t_cfg in cfg["towers"]: + mod = importlib.import_module(t_cfg["module"]) + cls = getattr(mod, t_cfg["class"]) + # some tower constructors need data (e.g. ClinicalEncoder needs feature_dim) + # pass data as an optional kwarg; tower ignores it if not needed + towers[t_cfg["name"]] = cls(data=data, **t_cfg["args"]) +``` + +Tower-specific training metadata (warmup epochs, batch key) lives entirely in the JSON, not in the orchestrator. + +--- + +## 4. Cross-Tower Communication: `early_pass` Protocol + +**Problem:** GeometryTower needs to precompute segmentation maps from images, then inject them into other towers' sample dicts before loaders are built. This is currently done imperatively in the orchestrator. + +**Proposed solution: `early_pass` connector interface** + +Each tower optionally implements: +```python +class TowerBase: + def early_pass(self, context: EarlyPassContext) -> None: + """Called once per fold before loaders are built. + Tower can read from / write to shared context.""" + pass +``` + +`EarlyPassContext` is a shared mutable object passed to all towers in order: +```python +@dataclass +class EarlyPassContext: + eye_train: list[dict] + bilat_train: list[dict] + bilat_val: list[dict] + bilat_test: list[dict] + image_preprocessor: object + image_cache: object + device: torch.device + store: dict = field(default_factory=dict) # cross-tower key-value store +``` + +Example: GeometryTower's `early_pass` computes seg maps and injects them into the sample dicts directly (modifying `eye_train` etc. in place), exactly as it does today — but now the orchestrator just calls: +```python +for tower in towers.values(): + tower.early_pass(context) +``` + +The cross-talk case the user described (img_tower outputs geometry → cd_tower reads it) uses `context.store`: +```python +# ImageTower.early_pass: +context.store["geometry_maps"] = self._compute_geometry(context) + +# ClinicalTower.early_pass: +geo = context.store.get("geometry_maps") +if geo is not None: + self._inject_geometry(context, geo) +``` + +Tower ordering in the JSON list determines execution order, so dependencies are declared implicitly. If a tower has no `early_pass`, the default no-op in `TowerBase` is used. + +**Alternative considered:** explicit dependency graph / DAG execution. Rejected for now — JSON ordering is simpler and sufficient for current needs. Can revisit if cross-tower dependencies become non-linear. + +--- + +## 5. File Layout + +``` +v4/ + hypertower/ + v4_hypertower.py # orchestrator (no argparse, no tower imports) + split_manager.py # copy/adapt from v3 (or just import) + papila/ + v4papila.py # merged papila_builders + papila_data + configs/ + ensemble_fused.json # example config +``` + +Existing `v3/classes/` tower implementations are reused directly — no duplication needed since they're importable by the JSON `"module"` field. + +--- + +## 6. Open Questions / Decisions Needed + +- **DataBundle API**: Does `build_data` need to return anything beyond the current `DataBundle`? Or should `DataBundle` grow a `profile` factory method? +- **Per-tower batch_key convention**: Currently `EYE_KEY_MAP = {"img": "image_1", "cd": "matrix_1"}` is hardcoded. Should this be declared in the tower JSON config or inferred from tower type? +- **cd_warmup loader**: Slot stripping (`if k != "image_1"`) is currently img-tower-aware. Under the new design, each tower should declare which slots it needs for warmup vs full training, so the orchestrator can build the right loader without knowing about `image_1`. +- **Geometry injection today vs `early_pass`**: Geometry currently mutates sample dicts; `early_pass` formalizes this. Needs a migration plan for existing GeometryTower. + +--- + +## Implementation Order (once ntower_cv is validated) + +1. Write `v4papila.py` (merge papila_builders + papila_data, expose `build_data(args)`) +2. Add `early_pass(context)` no-op to `TowerBase`; implement in `GeometryTower` +3. Write `v4_hypertower.py` orchestrator using JSON config + importlib tower loading +4. Port one config (ensemble_fused) end-to-end and compare outputs against ntower_cv diff --git a/v3/classes/__init__.py b/v3/classes/__init__.py index 01a94bd..2b8a8aa 100644 --- a/v3/classes/__init__.py +++ b/v3/classes/__init__.py @@ -21,14 +21,6 @@ from .slot_dataset import SlotDataset, slot_collate from .papila_data import PapilaData from .papila_builders import build_papila_data from .data_bundle import DataBundle -from .dataset import ClinicalDataset -from .config_builder import ( - ConfigAssembly, - assemble_config, - load_config, - resolve_imports, -) -from .filters import RegexFilter, ColumnFilter, apply_regex_filters, apply_column_filters from .transforms import ( ImageTransformConfig, backbone_transform_config, @@ -43,11 +35,35 @@ from .transforms import ( TRANSFORM_REGISTRY, build_transform_chain, ) -from .model_builder import V2ModelBundle, build_model_bundle -from .towers import ImageTower, ClinicalTower, SiameseImageTower, build_backbone -from .bridges import Bridge, VoteBridge -from .models import SingleEyeHT, BilateralHT -from .v2_hypertower import V2HyperTower, V2ModeComparisonOps, V2ModeComparator +from .towerbase import TowerBase, build_backbone, train_towers_epoch, collect_probs_towers +from .image_towers import ImageEncoder, SiameseImageTower, ImageTower +from .clinical_towers import ClinicalEncoder, ClinicalDataTower +from .geometry_towers import GeometryTower +from .hypertower_models import ( + SingleEyeHT, + BilateralHT, + SiameseHT, + FusedEnsembleHT, + LogitMLPEnsembleHT, + EmbeddingMLPEnsembleHT, + NTowerHT, + NLateralHT, + MonoTowerHT, + train_single_epoch, + train_bilateral_epoch, + train_siamese_epoch, + train_fusion_epoch, + train_ntower_epoch, + train_mono_epoch, + collect_probs_classic, + collect_probs_ensemble, + collect_probs_bilateral, + collect_probs_siamese, + collect_probs_ntower, + collect_probs_mono, + V2ModeComparisonOps, +) +from .bridges import Bridge, HTClassifier, HyperBridge, VoteBridge from .hypertower_logger import HypertowerLogger __all__ = [ @@ -66,18 +82,9 @@ __all__ = [ "PapilaData", "build_papila_data", "DataBundle", - "ClinicalDataset", "SlotLoaderFactory", "SlotDataset", "slot_collate", - "ConfigAssembly", - "assemble_config", - "load_config", - "resolve_imports", - "RegexFilter", - "ColumnFilter", - "apply_regex_filters", - "apply_column_filters", "ImageTransformConfig", "backbone_transform_config", "build_backbone_transform", @@ -90,18 +97,41 @@ __all__ = [ "UnetMaskProvider", "TRANSFORM_REGISTRY", "build_transform_chain", - "V2ModelBundle", - "build_model_bundle", - "ImageTower", - "ClinicalTower", - "SiameseImageTower", + "TowerBase", "build_backbone", - "Bridge", - "VoteBridge", + "train_towers_epoch", + "collect_probs_towers", + "ImageEncoder", + "SiameseImageTower", + "ImageTower", + "ClinicalEncoder", + "ClinicalDataTower", + "GeometryTower", "SingleEyeHT", "BilateralHT", - "V2HyperTower", + "SiameseHT", + "FusedEnsembleHT", + "LogitMLPEnsembleHT", + "EmbeddingMLPEnsembleHT", + "NTowerHT", + "NLateralHT", + "MonoTowerHT", + "train_single_epoch", + "train_bilateral_epoch", + "train_siamese_epoch", + "train_fusion_epoch", + "train_ntower_epoch", + "train_mono_epoch", + "collect_probs_classic", + "collect_probs_ensemble", + "collect_probs_bilateral", + "collect_probs_siamese", + "collect_probs_ntower", + "collect_probs_mono", + "Bridge", + "HTClassifier", + "HyperBridge", + "VoteBridge", "V2ModeComparisonOps", - "V2ModeComparator", "HypertowerLogger", ] diff --git a/v3/classes/bridges.py b/v3/classes/bridges.py index dd0a8de..c80f47a 100644 --- a/v3/classes/bridges.py +++ b/v3/classes/bridges.py @@ -1,19 +1,64 @@ from __future__ import annotations + import torch import torch.nn as nn from v3.classes.SE_attention import SEBlock, SEGateLogger +# --------------------------------------------------------------------------- +# HTClassifier — standalone classification head +# --------------------------------------------------------------------------- + +class HTClassifier(nn.Module): + """Minimal classification head: ReLU → Dropout → Linear(in_dim → num_classes). + + Used as the output stage of Bridge, HyperBridge, and any vehicle that needs + a reusable, identifiable classifier type. + """ + + def __init__(self, in_dim: int, num_classes: int, dropout: float = 0.5): + super().__init__() + self.head = nn.Sequential( + nn.ReLU(), + nn.Dropout(dropout), + nn.Linear(in_dim, num_classes), + ) + + def forward(self, z: torch.Tensor) -> torch.Tensor: + return self.head(z) + + +# --------------------------------------------------------------------------- +# Bridge — N-tower fusion +# --------------------------------------------------------------------------- + class Bridge(nn.Module): + """ + N-tower fusion bridge. + + Takes a list of tower embeddings, projects each to a common ``fusion_dim``, + element-wise multiplies all projections, optionally applies an SE gate, then + classifies the fused representation via an HTClassifier. + + Each tower also gets an auxiliary classification head (used for BCD training). + + Construction + ------------ + ``tower_dims`` is an ordered list of embedding dimensionalities — one entry per + embedding slot that will be passed to ``fuse()`` or ``forward()``. + + Tower slots are accessed by index: ``W[i]``, ``ln[i]``, ``aux_heads[i]``. + The bridge has no knowledge of what modality each slot carries. + """ + def __init__( self, - img_dim, - meta_dim, - num_classes, - fusion_dim=256, - mode="fused", + tower_dims: list[int], + num_classes: int, + fusion_dim: int = 256, + mode: str = "fused", dropout: float = 0.5, use_se: bool = True, se_reduction: int = 16, @@ -22,29 +67,30 @@ class Bridge(nn.Module): super().__init__() self.mode = mode self.use_se = use_se + self.tower_dims = list(tower_dims) - # project towers to equal width - self.W_img = nn.Linear(img_dim, fusion_dim) - self.W_md = nn.Linear(meta_dim, fusion_dim) + # Per-tower projection heads: each projects dim_i → fusion_dim + self.W = nn.ModuleList([nn.Linear(d, fusion_dim) for d in tower_dims]) + self.ln = nn.ModuleList( + [nn.LayerNorm(fusion_dim) if se_pre_norm else nn.Identity() + for _ in tower_dims] + ) - # optional: layernorm before SE - self.ln_img = nn.LayerNorm(fusion_dim) if se_pre_norm else nn.Identity() - self.ln_md = nn.LayerNorm(fusion_dim) if se_pre_norm else nn.Identity() + # Per-tower auxiliary classifiers (for BCD training) + self.aux_heads = nn.ModuleList([nn.Linear(d, num_classes) for d in tower_dims]) # SE gate on the fused vector self.se = SEBlock(fusion_dim, reduction=se_reduction, residual=True) if use_se else None self.se_log = SEGateLogger(enabled=use_se, track_channels=False, dim=fusion_dim) - # heads - self.classifier_fused = nn.Sequential( - nn.ReLU(), - nn.Dropout(dropout), - nn.Linear(fusion_dim, num_classes), - ) - self.classifier_img = nn.Linear(img_dim, num_classes) - self.classifier_cd = nn.Linear(meta_dim, num_classes) + # Fused classifier head + self.classifier_fused = HTClassifier(fusion_dim, num_classes, dropout) - def reset_se_stats(self): + # ------------------------------------------------------------------ + # SE helpers + # ------------------------------------------------------------------ + + def reset_se_stats(self) -> None: """Call at epoch start.""" if getattr(self, "se_log", None): self.se_log.reset() @@ -55,41 +101,205 @@ class Bridge(nn.Module): return self.se_log.get(reset=reset) return None - def _compute_fused(self, img_feats, md_feats): - """Return z_fused embedding (before classifier_fused). Used by encode() and forward().""" - hi = self.ln_img(self.W_img(img_feats)) - hm = self.ln_md(self.W_md(md_feats)) - fused = hi * hm + # ------------------------------------------------------------------ + # Core fusion + # ------------------------------------------------------------------ + + def _compute_fused(self, embeddings: list[torch.Tensor]) -> torch.Tensor: + """Return z_fused embedding (before classifier_fused).""" + assert len(embeddings) == len(self.W), ( + f"Bridge expects {len(self.W)} embeddings, got {len(embeddings)}" + ) + h = self.ln[0](self.W[0](embeddings[0])) + for i in range(1, len(embeddings)): + h = h * self.ln[i](self.W[i](embeddings[i])) if self.se is not None: - fused, gates = self.se(fused) + h, gates = self.se(h) if self.se_log.enabled: self.se_log.accumulate(gates) - return fused + return h - def encode(self, img_feats, md_feats) -> torch.Tensor: - """Return z_fused embedding without applying the classifier head.""" - assert self.mode == "fused", "encode() only valid in fused mode" - return self._compute_fused(img_feats, md_feats) + # ------------------------------------------------------------------ + # N-tower API + # ------------------------------------------------------------------ - def forward(self, img_feats, md_feats): - out_img = None if self.mode == "clinical_only" else self.classifier_img(img_feats) - out_md = None if self.mode == "image_only" else self.classifier_cd(md_feats) + def fuse( + self, embeddings: list[torch.Tensor] + ) -> tuple[torch.Tensor, list[torch.Tensor]]: + """ + N-tower forward pass. - if self.mode == "fused": - fused = self._compute_fused(img_feats, md_feats) - out_f = self.classifier_fused(fused) - return out_f, out_img, out_md - # if ablation modes: - if self.mode == "image_only": - return out_img, out_img, None - if self.mode == "clinical_only": - return out_md, None, out_md + Parameters + ---------- + embeddings : list of Tensor — one per tower slot (same order as tower_dims). + Returns + ------- + logits_fused : Tensor [B, num_classes] + aux_logits : list of Tensor — one per tower slot, each [B, num_classes] + """ + z_fused = self._compute_fused(embeddings) + logits_fused = self.classifier_fused(z_fused) + aux = [head(e) for head, e in zip(self.aux_heads, embeddings)] + return logits_fused, aux + + def encode(self, embeddings: list[torch.Tensor]) -> torch.Tensor: + """Return z_fused without applying the classifier head.""" + return self._compute_fused(embeddings) + + def set_phase(self, phase: str) -> None: + """ + Set requires_grad on bridge sub-modules according to training phase. + + - ``cd_warmup`` — freeze everything in the bridge + - ``tower_warmup`` — aux_heads trainable, projections + fused head frozen + - ``fused_warmup`` — projections + fused head trainable, aux_heads frozen + - ``main`` / other — everything trainable + """ + def _rg(module, enabled): + for p in module.parameters(): + p.requires_grad = enabled + + if phase == "cd_warmup": + _rg(self, False) + return + if phase == "tower_warmup": + for head in self.aux_heads: + _rg(head, True) + for W_i in self.W: + _rg(W_i, False) + for ln_i in self.ln: + _rg(ln_i, False) + _rg(self.classifier_fused, False) + if self.se is not None: + _rg(self.se, False) + return + if phase == "fused_warmup": + for head in self.aux_heads: + _rg(head, False) + for W_i in self.W: + _rg(W_i, True) + for ln_i in self.ln: + _rg(ln_i, True) + _rg(self.classifier_fused, True) + if self.se is not None: + _rg(self.se, True) + return + _rg(self, True) + + def forward( + self, embeddings: list[torch.Tensor] + ) -> tuple[torch.Tensor, list[torch.Tensor]]: + """N-tower forward. Delegates to fuse().""" + return self.fuse(embeddings) + + +# --------------------------------------------------------------------------- +# HyperBridge — higher-order bridge over HT module outputs +# --------------------------------------------------------------------------- + +class HyperBridge(nn.Module): + """Higher-order bridge that fuses z_fused embeddings from multiple HT modules. + + Operates at the HT output level (z_fused from each HT's Bridge.encode()) + rather than raw tower embedding level. + + Modes + ----- + embedding_mlp (default) + Concatenate all z_fused inputs → MLP → logits. + Analogous to EmbeddingMLPEnsembleHT, generalised to N inputs. + ``Linear(N*fusion_dim → hidden_dim) → ReLU → Dropout → Linear(hidden_dim → num_classes)`` + + classic_bridge + Project each input to ``hidden_dim``, Hadamard product, HTClassifier. + Analogous to Bridge operating at the HT level — handles inputs of + differing dims via per-input projection layers. + ``W[i](z_i) → LayerNorm → Hadamard → ReLU → Dropout → Linear(hidden_dim → num_classes)`` + + Both modes expose per-input auxiliary HTClassifier heads for BCD-style training. + + Parameters + ---------- + input_dims : ordered dict {name: dim} for each HT input. + In embedding_mlp mode, dims may differ. + In classic_bridge mode, all dims must be equal (shared space). + num_classes : output classes + hidden_dim : hidden dim for the embedding_mlp MLP head + mode : "embedding_mlp" | "classic_bridge" + dropout : dropout throughout + """ + + def __init__( + self, + input_dims: dict[str, int], + num_classes: int, + hidden_dim: int = 256, + mode: str = "embedding_mlp", + dropout: float = 0.3, + ): + super().__init__() + self.input_names = list(input_dims.keys()) + self.mode = mode + dims = list(input_dims.values()) + + if mode == "embedding_mlp": + total_dim = sum(dims) + self.head = nn.Sequential( + nn.Linear(total_dim, hidden_dim), nn.ReLU(), + nn.Dropout(dropout), nn.Linear(hidden_dim, num_classes), + ) + elif mode == "classic_bridge": + # Project each input to shared fusion_dim space, then Hadamard + self.W = nn.ModuleList([nn.Linear(d, hidden_dim) for d in dims]) + self.ln = nn.ModuleList([nn.LayerNorm(hidden_dim) for _ in dims]) + self.head = HTClassifier(hidden_dim, num_classes, dropout) + else: + raise ValueError(f"Unknown HyperBridge mode: {mode!r}") + + # Per-input auxiliary classifiers (both modes) + self.aux_heads = nn.ModuleList([ + HTClassifier(d, num_classes, dropout) for d in dims + ]) + + def forward( + self, + inputs: dict[str, torch.Tensor], + ) -> tuple[torch.Tensor, dict[str, torch.Tensor]]: + """Fuse HT-level embeddings. + + Parameters + ---------- + inputs : {name: z_fused [B, dim]} — z_fused from each HT's encode() + + Returns + ------- + logits : [B, num_classes] + aux_dict : {name: [B, num_classes]} — per-input aux head logits + """ + ordered = [inputs[name] for name in self.input_names] + + if self.mode == "embedding_mlp": + logits = self.head(torch.cat(ordered, dim=1)) + else: # classic_bridge: project → Hadamard → classify + h = self.ln[0](self.W[0](ordered[0])) + for i in range(1, len(ordered)): + h = h * self.ln[i](self.W[i](ordered[i])) + logits = self.head(h) + + aux = {name: head(z) for name, head, z + in zip(self.input_names, self.aux_heads, ordered)} + return logits, aux + + +# --------------------------------------------------------------------------- +# VoteBridge +# --------------------------------------------------------------------------- class VoteBridge(nn.Module): def __init__(self, num_classes): super().__init__() - self.vote_combiner = nn.Linear(num_classes * 2, num_classes) # two sets of logits + self.vote_combiner = nn.Linear(num_classes * 2, num_classes) def forward(self, out_img, out_md): votes = torch.cat([out_img, out_md], dim=1) diff --git a/v3/classes/clinical_towers.py b/v3/classes/clinical_towers.py new file mode 100644 index 0000000..67d484f --- /dev/null +++ b/v3/classes/clinical_towers.py @@ -0,0 +1,162 @@ +"""clinical_towers — ClinicalEncoder and ClinicalDataTower. + +Self-contained: defines ClinicalEncoder directly (does not import it from +towers.py). Imports only TowerBase from towerbase plus infrastructure +(SEBlock, DataBundle). +""" +from __future__ import annotations + +from typing import Optional + +import torch +from torch import nn + +from v3.classes.towerbase import TowerBase +from v3.classes.SE_attention import SEBlock +from v3.classes.data_bundle import DataBundle + + +# --------------------------------------------------------------------------- +# ClinicalEncoder — MLP over tabular features +# --------------------------------------------------------------------------- + +class ClinicalEncoder(nn.Module): + """MLP over DataBundle.vectorize_row outputs (converts to torch inside tower).""" + + def __init__( + self, + clinical_data: DataBundle, + hidden_dim: int = 128, + dropout: float = 0.1, + use_se: bool = False, + se_reduction: int = 16, + se_pre_norm: bool = True, + ): + super().__init__() + self.feature_dim = clinical_data.feature_dim + self.out_dim = hidden_dim + # Two-block MLP so we can optionally freeze/thaw per block. + self.block0 = nn.Sequential( + nn.Linear(self.feature_dim, hidden_dim), + nn.LayerNorm(hidden_dim), + nn.ReLU(inplace=True), + nn.Dropout(dropout), + ) + self.block1 = nn.Sequential( + nn.Linear(hidden_dim, hidden_dim), + nn.ReLU(inplace=True), + ) + self.net = nn.Sequential(self.block0, self.block1) + self.tower_ln = nn.LayerNorm(hidden_dim) if se_pre_norm else nn.Identity() + self.tower_se = ( + SEBlock(hidden_dim, reduction=se_reduction, residual=True) + if use_se + else None + ) + + def forward(self, meta_np_or_torch) -> torch.Tensor: + if isinstance(meta_np_or_torch, torch.Tensor): + x = meta_np_or_torch + else: + x = torch.as_tensor(meta_np_or_torch, dtype=torch.float32) + h = self.net(x) + if self.tower_se is not None: + h, _ = self.tower_se(self.tower_ln(h)) + return h + + def set_freeze_ratio(self, ratio: float): + """Optionally freeze earliest blocks of the MLP.""" + r = max(0.0, min(1.0, float(ratio))) + for p in self.block0.parameters(): + p.requires_grad = True + for p in self.block1.parameters(): + p.requires_grad = True + if r >= 0.5: + for p in self.block0.parameters(): + p.requires_grad = False + if r >= 1.0: + for p in self.block1.parameters(): + p.requires_grad = False + + +# --------------------------------------------------------------------------- +# ClinicalDataTower — TowerBase implementation +# --------------------------------------------------------------------------- + +class ClinicalDataTower(TowerBase, nn.Module): + """ + TowerBase implementation for the clinical metadata modality. + + Wraps ClinicalEncoder (MLP over tabular features). + Contributes one embedding per eye slot: [z_cd]. + + Implements ``cd_warmup_embedding`` so train_towers_epoch can identify + this tower for cd_warmup phase via duck typing rather than isinstance checks. + """ + + def __init__( + self, + *, + clinical_data, + cd_hidden_dim: int = 128, + cd_dropout: float = 0.1, + use_se: bool = False, + ): + nn.Module.__init__(self) + self._encoder = ClinicalEncoder( + clinical_data=clinical_data, + hidden_dim=cd_hidden_dim, + dropout=cd_dropout, + use_se=use_se, + ) + + @property + def out_dim(self) -> int: + return self._encoder.out_dim + + @property + def embed_dims(self) -> list[int]: + return [self._encoder.out_dim] + + def set_phase(self, phase: str) -> None: + enabled = phase not in ("fused_warmup",) + for p in self._encoder.parameters(): + p.requires_grad = enabled + + def cd_warmup_embedding( + self, + batch: dict, + *, + device: torch.device, + ) -> Optional[torch.Tensor]: + """Return clinical embedding for slot 1, or None if matrix_1 is absent.""" + m = batch.get("matrix_1") + if not torch.is_tensor(m): + return None + return self._encoder(m.to(device)) + + def embed_batch( + self, + batch: dict, + *, + device: torch.device, + slot: int = 1, + ) -> list[torch.Tensor]: + m = batch.get(f"matrix_{slot}") + if not torch.is_tensor(m): + raise ValueError(f"ClinicalDataTower.embed_batch: matrix_{slot} is missing or not a tensor") + return [self._encoder(m.to(device))] + + def prepare_fold( + self, + *, + eye_train, + bilat_train, + bilat_val, + bilat_test, + image_preprocessor, + image_cache, + device, + args, + ) -> None: + pass # shares loader with ImageTower; no per-fold setup needed diff --git a/v3/classes/config_builder.py b/v3/classes/config_builder.py deleted file mode 100644 index 9ec4df6..0000000 --- a/v3/classes/config_builder.py +++ /dev/null @@ -1,276 +0,0 @@ -from __future__ import annotations - -from dataclasses import dataclass -from pathlib import Path -from typing import Any, Dict, Iterable, List, Optional - -import json - -from v3.classes.papila_data import PapilaData - - -@dataclass -class ImportSpec: - id: str - class_name: str - params: Dict[str, Any] - - -@dataclass -class DataSourceSpec: - node_id: str - label: str - output_type: str - source: Optional[Dict[str, Any]] - source_ref: Optional[Dict[str, Any]] - - -@dataclass -class TransformSpec: - node_id: str - label: str - transform_type: str - params: Dict[str, Any] - - -@dataclass -class LoaderSpec: - node_id: str - label: str - input_type: str - input_index: str - input_key: str - output_key: str - transforms: List[TransformSpec] - data_source: Optional[DataSourceSpec] - - -@dataclass -class TowerSpec: - node_id: str - label: str - tower_type: str - params: Dict[str, Any] - - -@dataclass -class BridgeSpec: - node_id: str - label: str - method: str - params: Dict[str, Any] - - -@dataclass -class ClassifierSpec: - node_id: str - label: str - - -@dataclass -class ConfigAssembly: - raw: Dict[str, Any] - imports: Dict[str, ImportSpec] - data_sources: Dict[str, DataSourceSpec] - transforms: Dict[str, TransformSpec] - loaders: Dict[str, LoaderSpec] - towers: Dict[str, TowerSpec] - bridges: Dict[str, BridgeSpec] - classifiers: Dict[str, ClassifierSpec] - - -def load_config(path: Path) -> Dict[str, Any]: - payload = json.loads(Path(path).read_text()) - if not isinstance(payload, dict): - raise ValueError("Config JSON must be an object.") - return payload - - -def assemble_config(path: Path) -> ConfigAssembly: - config = load_config(path) - meta = config.get("meta", {}) - imports = _build_imports(meta.get("imports", [])) - nodes = {node["id"]: node for node in config.get("nodes", [])} - edges = config.get("edges", []) - - data_sources: Dict[str, DataSourceSpec] = {} - transforms: Dict[str, TransformSpec] = {} - loaders: Dict[str, LoaderSpec] = {} - towers: Dict[str, TowerSpec] = {} - bridges: Dict[str, BridgeSpec] = {} - classifiers: Dict[str, ClassifierSpec] = {} - - for node in nodes.values(): - ntype = node.get("type") - if ntype == "data": - data_sources[node["id"]] = DataSourceSpec( - node_id=node["id"], - label=node.get("label", ""), - output_type=node.get("outputType", ""), - source=node.get("source"), - source_ref=node.get("sourceRef"), - ) - elif ntype == "transform": - transforms[node["id"]] = TransformSpec( - node_id=node["id"], - label=node.get("label", ""), - transform_type=node.get("transformType", ""), - params=_extract_transform_params(node), - ) - elif ntype == "loader": - loaders[node["id"]] = LoaderSpec( - node_id=node["id"], - label=node.get("label", ""), - input_type=node.get("inputType", ""), - input_index=node.get("inputIndex", ""), - input_key=node.get("inputKey", ""), - output_key=node.get("outputKey", ""), - transforms=[], - data_source=None, - ) - elif ntype in ("image_tower", "metadata_tower"): - towers[node["id"]] = TowerSpec( - node_id=node["id"], - label=node.get("label", ""), - tower_type=node.get("towerType", "image" if ntype == "image_tower" else "clinical data"), - params=_extract_tower_params(node), - ) - elif ntype == "bridge": - bridges[node["id"]] = BridgeSpec( - node_id=node["id"], - label=node.get("label", ""), - method=node.get("bridgeMethod", "fusion"), - params=_extract_bridge_params(node), - ) - elif ntype == "classifier": - classifiers[node["id"]] = ClassifierSpec( - node_id=node["id"], - label=node.get("label", ""), - ) - - # attach transforms + data sources to loaders by walking upstream - for loader_id, loader in loaders.items(): - chain = _upstream_chain(loader_id, nodes, edges) - for node_id in reversed(chain): - if node_id in transforms: - loader.transforms.append(transforms[node_id]) - if node_id in data_sources: - loader.data_source = data_sources[node_id] - - return ConfigAssembly( - raw=config, - imports=imports, - data_sources=data_sources, - transforms=transforms, - loaders=loaders, - towers=towers, - bridges=bridges, - classifiers=classifiers, - ) - - -def resolve_imports(assembly: ConfigAssembly) -> Dict[str, Any]: - resolved: Dict[str, Any] = {} - for import_id, spec in assembly.imports.items(): - if spec.class_name == "PapilaData": - params = spec.params - resolved[import_id] = PapilaData.from_dirs( - image_dir=params.get("image_dir", "Papila/FundusImages"), - clinical_dir=params.get("clinical_dir", "Papila/ClinicalData"), - label_col=params.get("label_col", "Diagnosis"), - cat_cols=params.get("cat_cols", ["Gender", "Phakic/Pseudophakic"]), - ) - else: - raise ValueError(f"Unsupported import class {spec.class_name!r}") - return resolved - - -def _build_imports(entries: Iterable[Dict[str, Any]]) -> Dict[str, ImportSpec]: - specs: Dict[str, ImportSpec] = {} - for entry in entries or []: - import_id = entry.get("id") - if not import_id: - continue - specs[import_id] = ImportSpec( - id=import_id, - class_name=entry.get("className", ""), - params=entry.get("params", {}) or {}, - ) - return specs - - -def _extract_transform_params(node: Dict[str, Any]) -> Dict[str, Any]: - return { - "transformType": node.get("transformType"), - "roiMaskSource": node.get("roiMaskSource"), - "roiScale": node.get("roiScale"), - "roiTargetSize": node.get("roiTargetSize"), - "roiFallback": node.get("roiFallback"), - "centerCropSize": node.get("centerCropSize"), - "jitterHFlip": node.get("jitterHFlip"), - "jitterVFlip": node.get("jitterVFlip"), - "jitterRotation": node.get("jitterRotation"), - "jitterColorEnabled": node.get("jitterColorEnabled"), - "jitterColor": node.get("jitterColor"), - "resizeSize": node.get("resizeSize"), - } - - -def _extract_tower_params(node: Dict[str, Any]) -> Dict[str, Any]: - if node.get("towerType") == "clinical data": - return { - "hidden_dim": node.get("mdHiddenDim"), - "dropout": node.get("mdDropout"), - "use_se": node.get("mdUseSe"), - "se_reduction": node.get("mdSeReduction"), - "se_pre_norm": node.get("mdSePreNorm"), - "freeze_ratio": node.get("mdFreezeRatio"), - } - return { - "backbone": node.get("imageBackbone"), - "freeze_ratio": node.get("imageFreezeRatio"), - "augment": node.get("imageAugment"), - "geometry_dim": node.get("imageGeometryDim"), - "use_se": node.get("imageUseSe"), - "se_reduction": node.get("imageSeReduction"), - "se_pre_norm": node.get("imageSePreNorm"), - } - - -def _extract_bridge_params(node: Dict[str, Any]) -> Dict[str, Any]: - return { - "fusion_dim": node.get("bridgeFusionDim"), - "use_se": node.get("bridgeUseSe"), - "se_reduction": node.get("bridgeSeReduction"), - "se_pre_norm": node.get("bridgeSePreNorm"), - } - - -def _edge_from(edge: Dict[str, Any]) -> Optional[str]: - return edge.get("from") or edge.get("source") - - -def _edge_to(edge: Dict[str, Any]) -> Optional[str]: - return edge.get("to") or edge.get("target") - - -def _upstream_chain(start_id: str, nodes: Dict[str, Dict[str, Any]], edges: List[Dict[str, Any]]) -> List[str]: - chain: List[str] = [] - visited = set() - current = start_id - while True: - if current in visited: - break - visited.add(current) - incoming = [edge for edge in edges if _edge_to(edge) == current] - if not incoming: - break - # prefer first incoming edge for now - current = _edge_from(incoming[0]) - if not current: - break - chain.append(current) - node = nodes.get(current) - if node and node.get("type") == "data": - break - return chain diff --git a/v3/classes/dataset.py b/v3/classes/dataset.py deleted file mode 100644 index c8c2f1e..0000000 --- a/v3/classes/dataset.py +++ /dev/null @@ -1,115 +0,0 @@ -from torch.utils.data import Dataset -from PIL import Image -import numpy as np -import torch - - -class ClinicalDataset(Dataset): - """Generic dataset wrapping a DataBundle-like instance. - Returns (img_tensor, meta_tensor, label).""" - - def __init__( - self, - clinical_data, - img_transform, - meta_transform=None, - image_preprocessor=None, - geometry_provider=None, - geometry_dim: int = 0, - image_cache: "dict | None" = None, - ): - self.clinical = clinical_data - self.transform_image = img_transform - self.meta_transform = meta_transform or (lambda x: x) - self.image_preprocessor = image_preprocessor - self.geometry_provider = geometry_provider - self.geometry_dim = geometry_dim if geometry_provider is not None else 0 - self.image_cache = image_cache - - def __len__(self): - return len(self.clinical.df) - - def __getitem__(self, idx: int): - row = self.clinical.df.iloc[idx] - # load & transform image - img_path = self.clinical.get_image_path(row) - cache_key = str(img_path) - if self.image_cache is not None and cache_key in self.image_cache: - orig_img = Image.fromarray(self.image_cache[cache_key]) - else: - orig_img = Image.open(img_path).convert("RGB") - if self.image_cache is not None: - self.image_cache[cache_key] = np.asarray(orig_img, dtype=np.uint8) - img = orig_img - if self.image_preprocessor is not None: - img = self.image_preprocessor(img, img_path) - img_t = self.transform_image(img) - # encode & transform metadata - meta = self.clinical.encode_metadata(row) - meta_t = self.meta_transform(meta) - # label - label = self.clinical.get_label(row) - if self.geometry_dim > 0: - features = None - if self.geometry_provider is not None and hasattr(self.geometry_provider, "geometry_features"): - features = self.geometry_provider.geometry_features(orig_img, img_path) - if features is None: - geom_vec = torch.zeros(self.geometry_dim, dtype=torch.float32) - else: - features = np.asarray(features, dtype=np.float32) - if features.shape[0] != self.geometry_dim: - geom_vec = torch.zeros(self.geometry_dim, dtype=torch.float32) - else: - geom_vec = torch.from_numpy(features) - return img_t, meta_t, geom_vec, label - return img_t, meta_t, label - - -# --------------------------------------------------------------------------- -# _ClinicalView — shim used by V2HyperTower._run_fold -# --------------------------------------------------------------------------- - -from .data_bundle import DataBundle # noqa: E402 - - -class _ClinicalView: - """Minimal shim so ClinicalDataset can iterate an epoch-specific DataFrame - while still delegating encoding/paths/labels to the DataBundle object.""" - - def __init__(self, base: DataBundle, df): - self.base = base - self.df = df - - @property - def image_dir(self): - return self.base.image_dir - - @property - def clinical_dir(self): - return self.base.clinical_dir - - @property - def id_cols(self): - return ("Patient ID", "eyeID") - - @property - def label_col(self): - return self.base.label_col - - @property - def filename_template(self): - return getattr(self.base, "filename_template", "RET{pid:03d}{eye}.jpg") - - @property - def dim(self): - return self.base.feature_dim - - def encode_metadata(self, row): - vec = self.base.vectorize_row(row) - return torch.as_tensor(vec, dtype=torch.float32) - - def get_image_path(self, row): - return self.base.get_image_path(row) - - def get_label(self, row): - return int(row[self.base.label_col]) diff --git a/v3/classes/filters.py b/v3/classes/filters.py deleted file mode 100644 index 96b6922..0000000 --- a/v3/classes/filters.py +++ /dev/null @@ -1,119 +0,0 @@ -from __future__ import annotations - -from dataclasses import dataclass -from typing import Iterable, List, Sequence, Tuple, Union -import re - -import pandas as pd - - -@dataclass -class RegexFilter: - pattern: str - flags: int = 0 - - def apply_paths(self, paths: Sequence[str]) -> Tuple[List[str], List[str]]: - if not self.pattern: - return list(paths), [] - try: - regex = re.compile(self.pattern, self.flags) - except re.error as err: - return list(paths), [f'Invalid regex "{self.pattern}": {err}'] - filtered = [p for p in paths if regex.search(p)] - return filtered, [] - - -@dataclass -class ColumnFilter: - column: str - operator: str - value: str - case_insensitive: bool = True - - def apply_df(self, df: pd.DataFrame) -> Tuple[pd.DataFrame, List[str]]: - warnings: List[str] = [] - if not self.column: - return df, ["Column filter missing column name."] - columns = list(df.columns) - col_index = _resolve_column_index(columns, self.column, warnings) - if col_index is None: - return df, warnings - col_name = columns[col_index] - if self.value is None or self.value == "": - return df, [f'Column filter "{self.column}" missing value.'] - series = df[col_name] - mask = series.apply( - lambda cell: compare_cell( - cell, self.value, self.operator, case_insensitive=self.case_insensitive - ) - ) - return df[mask], warnings - - -FilterSpec = Union[RegexFilter, ColumnFilter] - - -def apply_regex_filters(paths: Sequence[str], filters: Iterable[RegexFilter]) -> Tuple[List[str], List[str]]: - filtered = list(paths) - warnings: List[str] = [] - for filt in filters: - filtered, warn = filt.apply_paths(filtered) - warnings.extend(warn) - return filtered, warnings - - -def apply_column_filters(df: pd.DataFrame, filters: Iterable[ColumnFilter]) -> Tuple[pd.DataFrame, List[str]]: - filtered = df - warnings: List[str] = [] - for filt in filters: - filtered, warn = filt.apply_df(filtered) - warnings.extend(warn) - return filtered, warnings - - -def compare_cell(cell, raw_value: str, operator: str, case_insensitive: bool = True) -> bool: - cell_str = "" if cell is None else str(cell).strip() - value_str = "" if raw_value is None else str(raw_value).strip() - if case_insensitive: - cell_str = cell_str.lower() - value_str = value_str.lower() - if operator == "=": - return cell_str == value_str - if operator == "!=": - return cell_str != value_str - cell_num = _to_float(cell_str) - value_num = _to_float(value_str) - if cell_num is None or value_num is None: - return False - if operator == ">": - return cell_num > value_num - if operator == ">=": - return cell_num >= value_num - if operator == "<": - return cell_num < value_num - if operator == "<=": - return cell_num <= value_num - return False - - -def _resolve_column_index(columns: Sequence[str], column: str, warnings: List[str]) -> int | None: - try: - return columns.index(column) - except ValueError: - lower = column.lower() - matches = [idx for idx, col in enumerate(columns) if str(col).lower() == lower] - if matches: - if len(matches) > 1: - warnings.append( - f'Column "{column}" matched multiple headers; using "{columns[matches[0]]}".' - ) - return matches[0] - warnings.append(f'Column "{column}" not found.') - return None - - -def _to_float(value: str) -> float | None: - try: - return float(value) - except (TypeError, ValueError): - return None diff --git a/v3/classes/seg_cnn.py b/v3/classes/geometry_towers.py similarity index 51% rename from v3/classes/seg_cnn.py rename to v3/classes/geometry_towers.py index 6ea3a02..2a6f960 100644 --- a/v3/classes/seg_cnn.py +++ b/v3/classes/geometry_towers.py @@ -1,19 +1,22 @@ -"""Segmentation-map CNN for glaucoma grading. +"""geometry_towers — GeometryTower and all segmentation-map infrastructure. -Trains a CNN on combined disc/cup segmentation maps — pixel values - 0 = background, 1 = disc (rim only), 2 = cup -— instead of raw RGB fundus images, forcing the model to learn purely -from optic nerve head geometry (CDR, rim width, cup location, etc.). +Self-contained: absorbs everything that was in seg_cnn.py so that file can +eventually be removed. Does not import from seg_cnn.py or any other tower file. +Imports only TowerBase from towerbase plus standard infrastructure. -Two segmentation sources are supported: - gt – rasterise expert contour/mask annotations directly (pure NumPy/PIL, - no CUDA — safe in DataLoader worker processes) - unet – run a trained UNetSegmenter on the raw fundus image - -Usage (import from training script): - from v3.classes.seg_cnn import SegMapRecord, SegMapDataset, SegCNN, seg_map_to_tensor +Contents +-------- +SegMapRecord — labelled-eye data record +_combine_masks — merge disc/cup binary masks → 3-class label map +crop_to_disc — tight bounding-box crop +seg_map_to_tensor — (H,W) uint8 → (C,H,W) float32 tensor +load_gt_masks — load GT disc+cup masks from contour/mask files +UNetFineTuneDataset — Dataset for fine-tuning the UNet on GT annotations +precompute_unet_seg_maps — batch UNet inference helper +SegMapDataset — Dataset yielding (seg_tensor, label) pairs +SegCNN — pretrained CNN adapted for segmentation-map input +GeometryTower — TowerBase implementation (the main class to use) """ - from __future__ import annotations from dataclasses import dataclass @@ -29,6 +32,10 @@ from torch.utils.data import Dataset from torchvision import models, transforms from tqdm import tqdm +import pandas as pd + +from v3.classes.towerbase import TowerBase + # --------------------------------------------------------------------------- # Data record @@ -53,8 +60,7 @@ class SegMapRecord: # --------------------------------------------------------------------------- def _combine_masks(disc_mask: np.ndarray, cup_mask: np.ndarray) -> np.ndarray: - """ - Combine binary disc and cup masks into a 3-class label map. + """Combine binary disc and cup masks into a 3-class label map. Returns a uint8 array with values: 0 — background @@ -69,8 +75,7 @@ def _combine_masks(disc_mask: np.ndarray, cup_mask: np.ndarray) -> np.ndarray: def crop_to_disc(seg_map: np.ndarray) -> np.ndarray: - """ - Crop a seg map tightly to the disc bounding box. + """Crop a seg map tightly to the disc bounding box. The disc is anywhere seg_map > 0 (i.e. rim or cup). Returns the original array unchanged if no disc is found. @@ -89,8 +94,7 @@ def seg_map_to_tensor( channels: int, target_size: int, ) -> torch.Tensor: - """ - Convert an (H, W) seg map with values {0, 1, 2} to a float tensor. + """Convert an (H, W) seg map with values {0, 1, 2} to a float tensor. channels=1 → (1, H, W) float in [0, 1] (values 0/0.5/1.0) channels=3 → (3, H, W) one-hot binary channels [bg, disc_rim, cup] @@ -118,13 +122,15 @@ def seg_map_to_tensor( def _load_contour(path: Path) -> np.ndarray: """Load x,y contour pairs from a whitespace- or comma-delimited text file.""" + arr = np.zeros((0, 2), dtype=np.float32) for delimiter in (",", None): try: - arr = np.loadtxt(str(path), delimiter=delimiter, comments="#", dtype=np.float32) - if arr.size > 0: + candidate = np.loadtxt(str(path), delimiter=delimiter, comments="#", dtype=np.float32) + if candidate.size > 0: + arr = candidate break except Exception: - arr = np.zeros((0, 2), dtype=np.float32) + pass if arr.size == 0 or arr.ndim == 1: return np.zeros((0, 2), dtype=np.float32) if arr.shape[1] < 2: @@ -135,13 +141,10 @@ def _load_contour(path: Path) -> np.ndarray: def _contour_to_mask( coords: np.ndarray, image_size: Tuple[int, int], target_size: int ) -> np.ndarray: - """ - Rasterise a polygon defined by (x, y) coords into a binary mask. + """Rasterise a polygon defined by (x, y) coords into a binary mask. image_size is the (width, height) of the original fundus image — the - coordinate space the contour was annotated in. The mask is drawn at - that resolution then resized to target_size, matching UNetSegmenter's - behaviour and avoiding off-canvas clipping. + coordinate space the contour was annotated in. """ if coords is None or len(coords) < 3: return np.zeros((target_size, target_size), dtype=np.uint8) @@ -155,8 +158,7 @@ def _contour_to_mask( def _extract_masks_from_image( mask_path: Path, target_size: int ) -> Tuple[np.ndarray, np.ndarray]: - """ - Extract disc and cup binary masks from a segmentation image file. + """Extract disc and cup binary masks from a segmentation image file. Handles both grayscale label images (e.g. REFUGE .bmp) and RGB colour-coded masks. Returns (disc_mask, cup_mask) both at @@ -166,7 +168,6 @@ def _extract_masks_from_image( arr = np.array(raw) if arr.ndim == 2: - # Grayscale: identify background from edge statistics edges = np.concatenate([arr[0], arr[-1], arr[:, 0], arr[:, -1]]) bg_val = int(np.argmax(np.bincount(edges.astype(np.int64).clip(0, 255), minlength=256))) disc_arr = (arr != bg_val).astype(np.uint8) @@ -182,9 +183,7 @@ def _extract_masks_from_image( img_rgb = raw.convert("RGB") arr = np.array(img_rgb) h, w, c = arr.shape - edges_rgb = np.concatenate( - [arr[0], arr[-1], arr[:, 0], arr[:, -1]], axis=0 - ) + edges_rgb = np.concatenate([arr[0], arr[-1], arr[:, 0], arr[:, -1]], axis=0) edge_colors, edge_counts = np.unique(edges_rgb.reshape(-1, c), axis=0, return_counts=True) bg_color = edge_colors[int(np.argmax(edge_counts))] colors, counts = np.unique(arr.reshape(-1, c), axis=0, return_counts=True) @@ -200,7 +199,6 @@ def _extract_masks_from_image( cup_color = colors[order[1]] cup_arr[np.all(arr == cup_color, axis=-1)] = 1 - # Resize to target_size with nearest-neighbour to preserve binary values def _resize(m: np.ndarray) -> np.ndarray: pil = Image.fromarray((m > 0).astype(np.uint8) * 255) pil = pil.resize((target_size, target_size), Resampling.NEAREST) @@ -209,9 +207,8 @@ def _extract_masks_from_image( return _resize(disc_arr), _resize(cup_arr) -def load_gt_masks(rec: "SegMapRecord", target_size: int) -> Tuple[np.ndarray, np.ndarray]: - """ - Load GT disc + cup masks for one record. +def load_gt_masks(rec: SegMapRecord, target_size: int) -> Tuple[np.ndarray, np.ndarray]: + """Load GT disc + cup masks for one record. Handles annotation_type "contour" (x,y text file) and "mask" (image file). Returns (disc_mask, cup_mask) as uint8 arrays of shape (target_size, target_size). @@ -219,7 +216,6 @@ def load_gt_masks(rec: "SegMapRecord", target_size: int) -> Tuple[np.ndarray, np disc_mask: Optional[np.ndarray] = None cup_mask: Optional[np.ndarray] = None - # Get original image size so contour coordinates are drawn in the right space with Image.open(rec.image_path) as _img: image_size = _img.size # (width, height) @@ -246,7 +242,6 @@ def load_gt_masks(rec: "SegMapRecord", target_size: int) -> Tuple[np.ndarray, np if cup_mask is None: cup_mask = np.zeros((target_size, target_size), dtype=np.uint8) - # Structural prior: cup must lie within disc cup_mask = (cup_mask > 0) & (disc_mask > 0) return disc_mask.astype(np.uint8), cup_mask.astype(np.uint8) @@ -256,11 +251,7 @@ def load_gt_masks(rec: "SegMapRecord", target_size: int) -> Tuple[np.ndarray, np # --------------------------------------------------------------------------- class UNetFineTuneDataset(Dataset): - """ - Loads (image_tensor, mask_tensor) pairs for fine-tuning the U-Net on - PAPILA GT annotations. Uses the same preprocessing as UNetSegmenter - so the fine-tuned weights are compatible with inference. - """ + """Loads (image_tensor, mask_tensor) pairs for fine-tuning the U-Net.""" def __init__( self, @@ -292,7 +283,6 @@ class UNetFineTuneDataset(Dataset): image = Image.open(rec.image_path).convert("RGB") image = image.resize((self.target_size, self.target_size), Resampling.BILINEAR) img_tensor = self._normalize(self.to_tensor(image)) - disc_mask, cup_mask = load_gt_masks(rec, self.target_size) mask_tensor = torch.from_numpy( np.stack([disc_mask, cup_mask], axis=0).astype(np.float32) @@ -301,20 +291,15 @@ class UNetFineTuneDataset(Dataset): # --------------------------------------------------------------------------- -# U-Net precomputation (run once per full record list, not per fold) +# U-Net precomputation # --------------------------------------------------------------------------- def precompute_unet_seg_maps( - records: List["SegMapRecord"], + records: List[SegMapRecord], segmenter, threshold: float = 0.5, ) -> List[np.ndarray]: - """ - Run the U-Net on every record and return a list of combined seg maps. - - Call this once before the CV loop and pass the results to each fold's - SegMapDataset via precomputed_seg_maps, so the U-Net isn't re-run per fold. - """ + """Run the U-Net on every record and return a list of combined seg maps.""" to_tensor = transforms.ToTensor() seg_maps = [] for rec in tqdm(records, desc="U-Net inference", unit="img", leave=False): @@ -334,28 +319,11 @@ def precompute_unet_seg_maps( # --------------------------------------------------------------------------- -# Dataset +# SegMapDataset # --------------------------------------------------------------------------- class SegMapDataset(Dataset): - """ - PyTorch Dataset that yields (seg_tensor, label) pairs. - - Parameters - ---------- - records : list of SegMapRecord - target_size : CNN input spatial size (images are resized to this) - channels : 1 = single-channel label map; 3 = one-hot three channels - augment : apply random flips + rotation (for training set) - unet_segmenter : if provided, use U-Net predictions instead of GT masks; - must be a loaded UNetSegmenter with model weights set - unet_threshold : threshold for U-Net logit → binary mask - seg_target_size: resolution at which GT masks are rasterised (or U-Net - output size). Default 512 matches UNetSegmenter default. - crop_to_disc : crop the seg map tightly to the disc bounding box before - resizing to target_size (default True — eliminates the - background zeros that make up most of the full image) - """ + """PyTorch Dataset that yields (seg_tensor, label) pairs.""" def __init__( self, @@ -366,7 +334,7 @@ class SegMapDataset(Dataset): unet_segmenter=None, unet_threshold: float = 0.5, seg_target_size: int = 512, - crop_to_disc: bool = True, + crop_to_disc_flag: bool = True, precomputed_seg_maps: Optional[List[np.ndarray]] = None, ) -> None: self.records = records @@ -374,7 +342,7 @@ class SegMapDataset(Dataset): self.channels = channels self.augment = augment self.seg_target_size = seg_target_size - self.crop_to_disc = crop_to_disc + self.crop_to_disc_flag = crop_to_disc_flag if precomputed_seg_maps is not None: self._seg_maps = precomputed_seg_maps @@ -385,13 +353,10 @@ class SegMapDataset(Dataset): else: self._seg_maps = None - # ------------------------------------------------------------------ def __len__(self) -> int: return len(self.records) - # ------------------------------------------------------------------ def _augment(self, seg_map: np.ndarray) -> np.ndarray: - """Random flips + 90° rotations (label-safe since NEAREST resize).""" if np.random.rand() < 0.5: seg_map = np.fliplr(seg_map) if np.random.rand() < 0.5: @@ -401,19 +366,16 @@ class SegMapDataset(Dataset): seg_map = np.rot90(seg_map, k=k) return np.ascontiguousarray(seg_map) - # ------------------------------------------------------------------ def __getitem__(self, idx: int): rec = self.records[idx] - if self._seg_maps is not None: seg_map = self._seg_maps[idx] else: disc_mask, cup_mask = load_gt_masks(rec, self.seg_target_size) seg_map = _combine_masks(disc_mask, cup_mask) - if self.crop_to_disc: + if self.crop_to_disc_flag: seg_map = crop_to_disc(seg_map) - if self.augment: seg_map = self._augment(seg_map) @@ -422,19 +384,24 @@ class SegMapDataset(Dataset): # --------------------------------------------------------------------------- -# Model +# SegCNN # --------------------------------------------------------------------------- +_SEGCNN_FEAT_DIM = { + "resnet18": 512, + "resnet50": 2048, + "efficientnet_b0": 1280, +} + + class SegCNN(nn.Module): - """ - Pretrained CNN backbone adapted for segmentation-map input. + """Pretrained CNN backbone adapted for segmentation-map input. Parameters ---------- num_classes : output classes (2 for binary glaucoma grading) backbone : "resnet18" | "resnet50" | "efficientnet_b0" - pretrained : initialise with ImageNet weights (recommended even for - non-RGB input — transfer generalises across domains) + pretrained : initialise with ImageNet weights in_channels : 1 (single label map) or 3 (one-hot channels) dropout : dropout rate before the final classifier head """ @@ -448,7 +415,6 @@ class SegCNN(nn.Module): dropout: float = 0.3, ) -> None: super().__init__() - weights_arg = "DEFAULT" if pretrained else None if backbone == "resnet18": @@ -466,7 +432,6 @@ class SegCNN(nn.Module): else: raise ValueError(f"Unknown backbone: {backbone!r}") - # Adapt first conv layer if in_channels ≠ 3 if in_channels != 3: first_conv = self._find_first_conv(base) new_conv = nn.Conv2d( @@ -478,7 +443,6 @@ class SegCNN(nn.Module): bias=first_conv.bias is not None, ) if pretrained: - # Average pretrained RGB weights across channel dim with torch.no_grad(): new_conv.weight.copy_( first_conv.weight.mean(dim=1, keepdim=True).expand_as(new_conv.weight) @@ -491,7 +455,6 @@ class SegCNN(nn.Module): nn.Linear(feat_dim, num_classes), ) - # ------------------------------------------------------------------ @staticmethod def _find_first_conv(module: nn.Module) -> nn.Conv2d: for m in module.modules(): @@ -501,7 +464,6 @@ class SegCNN(nn.Module): @staticmethod def _replace_first_conv(module: nn.Module, new_conv: nn.Conv2d) -> None: - """Replace the first Conv2d in-place (handles resnet and efficientnet).""" for name, child in module.named_children(): if isinstance(child, nn.Conv2d): setattr(module, name, new_conv) @@ -513,9 +475,304 @@ class SegCNN(nn.Module): pass raise RuntimeError("Could not replace first Conv2d") - # ------------------------------------------------------------------ def forward(self, x: torch.Tensor) -> torch.Tensor: feats = self.backbone(x) if feats.dim() > 2: feats = feats.flatten(1) return self.head(feats) + + +# --------------------------------------------------------------------------- +# GeometryTower — TowerBase implementation +# --------------------------------------------------------------------------- + +class GeometryTower(TowerBase, nn.Module): + """TowerBase implementation for the optic-disc/cup segmentation modality. + + Encodes a 3-class disc/cup segmentation map (bg=0, rim=1, cup=2) through a + CNN backbone, contributing one spatial embedding to the bridge. + + The seg map is produced from GT annotations (manifest-based) or from a + trained U-Net, depending on ``geometry_source``. + + ``prepare_fold`` builds a seg-map generator, pre-computes all maps for the + fold, and caches them keyed by image path. ``augment_samples`` then + injects ``seg_map_1`` / ``seg_map_2`` float32 numpy arrays (shape C×H×W) + into each sample dict so the DataLoader delivers them as tensors to + ``embed_batch``. + + Parameters + ---------- + backbone : CNN backbone — "resnet18" | "resnet50" | "efficientnet_b0" + in_channels : 1 (label map) or 3 (one-hot disc/rim/cup channels) + pretrained : initialise backbone with ImageNet weights + frozen : if True, backbone is always frozen + target_size : spatial size the seg map tensor is resized to + seg_target_size : resolution at which GT masks are rasterised / U-Net runs + crop_to_disc : crop seg map tightly to disc bounding box before resizing + geometry_source : "gt" (manifest annotations) or "unet" (U-Net predictions) + manifest_path : path to the geometry manifest CSV (required) + weights_path : path to UNet checkpoint (required when source="unet") + unet_normalize : UNet normalisation mode (default "per_image") + unet_threshold : UNet mask threshold (default 0.5) + finetune_unet_epochs : epochs to fine-tune U-Net per fold (0 = disabled) + finetune_unet_lr : learning rate for U-Net fine-tuning + """ + + def __init__( + self, + *, + backbone: str = "resnet18", + in_channels: int = 3, + pretrained: bool = True, + frozen: bool = False, + target_size: int = 224, + seg_target_size: int = 512, + crop_to_disc: bool = True, + geometry_source: str = "gt", + manifest_path=None, + weights_path=None, + unet_normalize: str = "per_image", + unet_threshold: float = 0.5, + finetune_unet_epochs: int = 0, + finetune_unet_lr: float = 1e-5, + ): + nn.Module.__init__(self) + self._backbone_name = backbone + self._in_channels = in_channels + self._frozen = frozen + self._target_size = target_size + self._seg_target_size = seg_target_size + self._crop_to_disc = crop_to_disc + self._geometry_source = geometry_source + self._manifest_path = Path(manifest_path) if manifest_path is not None else None + self._weights_path = Path(weights_path) if weights_path is not None else None + self._unet_normalize = unet_normalize + self._unet_threshold = unet_threshold + self._finetune_unet_epochs = finetune_unet_epochs + self._finetune_unet_lr = finetune_unet_lr + + self._out_dim = _SEGCNN_FEAT_DIM.get(backbone, 512) + self._seg_cnn = SegCNN( + num_classes=2, + backbone=backbone, + pretrained=pretrained, + in_channels=in_channels, + ) + self._seg_cache: dict = {} # image_path_str → float32 (C, H, W) numpy array + + # ------------------------------------------------------------------ + # TowerBase interface + # ------------------------------------------------------------------ + + @property + def embed_dims(self) -> list[int]: + return [self._out_dim] + + @property + def total_epochs(self) -> int: + return 0 if self._frozen else 1 + + def set_phase(self, phase: str) -> None: + trainable = not self._frozen and phase in ("tower_warmup", "main") + for p in self._seg_cnn.parameters(): + p.requires_grad = trainable + + def prepare_fold( + self, + *, + eye_train, + bilat_train, + bilat_val, + bilat_test, + image_preprocessor, + image_cache, + device, + args, + ) -> None: + """Build seg-map generator and pre-compute maps for all fold images.""" + if self._manifest_path is None: + raise ValueError("GeometryTower requires manifest_path") + + all_paths: dict = {} + for split in (eye_train, bilat_train, bilat_val, bilat_test): + for s in split: + for slot in ("image_1", "image_2"): + p = s.get(slot) + if p is not None: + all_paths[str(Path(p).resolve())] = None + + if self._geometry_source == "unet": + self._prepare_fold_unet(list(all_paths.keys()), eye_train, device) + else: + self._prepare_fold_gt(list(all_paths.keys())) + + def augment_samples(self, samples: list) -> list: + """Inject ``seg_map_1`` / ``seg_map_2`` float32 arrays into each sample dict. + + Arrays have shape (C, H, W) and are collated by the DataLoader into + (B, C, H, W) tensors delivered to ``embed_batch``. + """ + blank = np.zeros( + (self._in_channels, self._target_size, self._target_size), dtype=np.float32 + ) + for s in samples: + for img_slot, seg_slot in (("image_1", "seg_map_1"), ("image_2", "seg_map_2")): + img_path = s.get(img_slot) + if img_path is None: + continue + key = str(Path(img_path).resolve()) + s[seg_slot] = self._seg_cache.get(key, blank) + return samples + + def embed_batch( + self, + batch: dict, + *, + device: torch.device, + slot: int = 1, + ) -> list[torch.Tensor]: + seg = batch.get(f"seg_map_{slot}") + if seg is None or not torch.is_tensor(seg): + ref = batch.get(f"image_{slot}") + bs = ref.shape[0] if torch.is_tensor(ref) else 1 + return [torch.zeros(bs, self._out_dim, device=device)] + return [self._seg_cnn.backbone(seg.float().to(device))] + + # ------------------------------------------------------------------ + # Internal helpers + # ------------------------------------------------------------------ + + def _seg_map_to_array(self, seg_map: np.ndarray) -> np.ndarray: + """Apply crop + resize and return a (C, H, W) float32 numpy array.""" + if self._crop_to_disc: + seg_map = crop_to_disc(seg_map) + return seg_map_to_tensor(seg_map, self._in_channels, self._target_size).numpy() + + def _prepare_fold_gt(self, image_paths: list) -> None: + """Pre-compute GT seg maps from manifest annotations.""" + manifest_df = pd.read_csv(self._manifest_path) + manifest_df["_img_key"] = manifest_df["image_path"].apply( + lambda p: str(Path(p).resolve()) + ) + manifest_index = manifest_df.set_index("_img_key").to_dict("index") + + print( + f"[GeometryTower] pre-computing GT seg maps for {len(image_paths)} images...", + flush=True, + ) + n_ok = 0 + blank = np.zeros((self._seg_target_size, self._seg_target_size), dtype=np.uint8) + for img_path in image_paths: + entry = manifest_index.get(img_path) + if entry is None: + self._seg_cache[img_path] = self._seg_map_to_array(blank) + continue + rec = SegMapRecord( + sample_id="", + image_path=Path(img_path), + annotation_disc=Path(entry["annotation_disc"]), + annotation_cup=Path(entry["annotation_cup"]), + annotation_type_disc=entry["annotation_type_disc"], + annotation_type_cup=entry["annotation_type_cup"], + patient_id=0, + eye="", + label=0, + ) + try: + disc_mask, cup_mask = load_gt_masks(rec, self._seg_target_size) + self._seg_cache[img_path] = self._seg_map_to_array( + _combine_masks(disc_mask, cup_mask) + ) + n_ok += 1 + except Exception: + self._seg_cache[img_path] = self._seg_map_to_array(blank) + print(f"[GeometryTower] {n_ok}/{len(image_paths)} GT seg maps computed", flush=True) + + def _prepare_fold_unet(self, image_paths: list, eye_train: list, device) -> None: + """Pre-compute U-Net seg maps, with optional per-fold fine-tuning.""" + from v3.classes.unet_segmenter import UNetSegmenter + from torch.utils.data import DataLoader as _DL + + if self._weights_path is None: + raise ValueError("GeometryTower(source='unet') requires weights_path") + + segmenter = UNetSegmenter( + manifest_path=self._manifest_path, + normalize=self._unet_normalize, + ) + state = torch.load(self._weights_path, map_location=segmenter.device) + segmenter.model.load_state_dict(state.get("model", state)) + segmenter.model.to(segmenter.device).eval() + + if self._finetune_unet_epochs > 0: + print( + f"[GeometryTower] fine-tuning U-Net for {self._finetune_unet_epochs} epochs...", + flush=True, + ) + ft_loader = _DL( + UNetFineTuneDataset( + self._build_records_from_samples(eye_train), + target_size=segmenter.target_size, + normalize=self._unet_normalize, + ), + batch_size=4, shuffle=True, num_workers=0, + ) + optimizer = torch.optim.Adam(segmenter.model.parameters(), lr=self._finetune_unet_lr) + criterion = torch.nn.BCEWithLogitsLoss() + segmenter.model.train() + for _ in range(self._finetune_unet_epochs): + for images, masks in ft_loader: + images, masks = images.to(segmenter.device), masks.to(segmenter.device) + optimizer.zero_grad() + criterion(segmenter.model(images), masks).backward() + optimizer.step() + segmenter.model.eval() + + print( + f"[GeometryTower] running U-Net inference on {len(image_paths)} images...", + flush=True, + ) + records = [ + SegMapRecord( + sample_id="", image_path=Path(p), + annotation_disc=Path(p), annotation_cup=Path(p), + annotation_type_disc="", annotation_type_cup="", + patient_id=0, eye="", label=0, + ) + for p in image_paths + ] + seg_maps = precompute_unet_seg_maps(records, segmenter, self._unet_threshold) + for img_path, seg_map in zip(image_paths, seg_maps): + self._seg_cache[img_path] = self._seg_map_to_array(seg_map) + print(f"[GeometryTower] {len(seg_maps)} U-Net seg maps cached", flush=True) + + def _build_records_from_samples(self, samples: list) -> list: + """Build SegMapRecord list from HyperTower sample dicts (for U-Net fine-tuning).""" + manifest_df = pd.read_csv(self._manifest_path) + manifest_df["_img_key"] = manifest_df["image_path"].apply( + lambda p: str(Path(p).resolve()) + ) + manifest_index = manifest_df.set_index("_img_key").to_dict("index") + records = [] + for s in samples: + for slot in ("image_1", "image_2"): + p = s.get(slot) + if p is None: + continue + key = str(Path(p).resolve()) + entry = manifest_index.get(key) + if entry is None: + continue + records.append(SegMapRecord( + sample_id="", + image_path=Path(p), + annotation_disc=Path(entry["annotation_disc"]), + annotation_cup=Path(entry["annotation_cup"]), + annotation_type_disc=entry["annotation_type_disc"], + annotation_type_cup=entry["annotation_type_cup"], + patient_id=int(s.get("patient_id", 0)), + eye=str(s.get("eye", "")), + label=int(s.get("label", 0)), + )) + return records diff --git a/v3/classes/hypertower_models.py b/v3/classes/hypertower_models.py new file mode 100644 index 0000000..96e01d0 --- /dev/null +++ b/v3/classes/hypertower_models.py @@ -0,0 +1,1187 @@ +"""hypertower_models — HyperTower vehicle classes and their training/eval helpers. + +A "vehicle" wires one or more tower encoders together with a Bridge to form +a complete trainable model. Vehicles can in principle be run standalone; +the v3_hypertower orchestrator drives them through the full fold/epoch loop. + +Contents +-------- +SingleEyeHT — ImageEncoder + ClinicalEncoder + Bridge (eye-level) +BilateralHT — shared eye towers + joint fusion layers + Bridge +SiameseHT — SiameseImageTower + Bridge (image-only bilateral) +FusedEnsembleHT — SingleEyeHT base + per-eye attention scorer +LogitMLPEnsembleHT — MLP head over concatenated per-eye logits +EmbeddingMLPEnsembleHT — MLP head over concatenated per-eye z_fused embeddings +NTowerHT — N named encoders + Bridge (general, key-mapped) +NLateralHT — N same-type inputs through shared encoder + joint MLP +MonoTowerHT — single tower + direct classifier (no bridge) + +Training helpers : train_single_epoch, train_bilateral_epoch, + train_siamese_epoch, train_fusion_epoch, + train_ntower_epoch, train_mono_epoch +Inference helpers : collect_probs_classic, collect_probs_ensemble, + collect_probs_ensemble_pereye, collect_probs_bilateral, + collect_probs_bilateral_components, collect_probs_siamese, + collect_probs_fused, collect_probs_single_components, + collect_probs_eye_level, collect_probs_ntower, + collect_probs_mono +""" +from __future__ import annotations + +from random import random +from typing import Optional + +import numpy as np +import torch +import torch.nn.functional as F +from torch import nn +from torch.utils.data import DataLoader + +from v3.classes.towerbase import _set_requires_grad, _to_label_tensor +from v3.classes.image_towers import ImageEncoder, SiameseImageTower +from v3.classes.clinical_towers import ClinicalEncoder +from v3.classes.bridges import Bridge, HTClassifier, HyperBridge + + +# --------------------------------------------------------------------------- +# Vehicle classes +# --------------------------------------------------------------------------- + +class SingleEyeHT(nn.Module): + """ImageEncoder + ClinicalEncoder + Bridge, trained on eye-level samples.""" + + def __init__( + self, + *, + backbone: str, + freeze_ratio: float, + augment: bool, + clinical_data, + num_classes: int, + cd_hidden_dim: int = 128, + fusion_dim: int = 256, + bridge_mode: str = "fused", + bridge_dropout: float = 0.5, + cd_dropout: float = 0.1, + se_img_tower: bool = False, + se_cd_tower: bool = False, + se_bridge: bool = False, + ): + super().__init__() + self.img_tower = ImageEncoder( + backbone=backbone, freeze_ratio=freeze_ratio, + augment=augment, use_se=se_img_tower, + ) + self.cd_tower = ClinicalEncoder( + clinical_data=clinical_data, hidden_dim=cd_hidden_dim, + dropout=cd_dropout, use_se=se_cd_tower, + ) + self.bridge = Bridge( + tower_dims=[self.img_tower.out_dim, self.cd_tower.out_dim], + num_classes=num_classes, fusion_dim=fusion_dim, + mode=bridge_mode, dropout=bridge_dropout, use_se=se_bridge, + ) + + @property + def transform(self): + return self.img_tower.transform + + def encode(self, x: torch.Tensor, meta: torch.Tensor) -> torch.Tensor: + """Return z_fused (fusion_dim) without the classifier head.""" + return self.bridge.encode([self.img_tower(x), self.cd_tower(meta)]) + + def forward(self, x: torch.Tensor, meta: torch.Tensor) -> torch.Tensor: + img_feats = None if self.bridge.mode == "clinical_only" else self.img_tower(x) + md_feats = None if self.bridge.mode == "image_only" else self.cd_tower(meta) + out_f, _ = self.bridge.fuse([img_feats, md_feats]) + return out_f + + +class NTowerHT(nn.Module): + """General N-tower vehicle: any named encoder modules fused through a Bridge. + + Parameters + ---------- + towers : ordered dict of {name: encoder_module}. Each module must + expose ``.out_dim``. Bridge slot order follows dict order. + num_classes : number of output classes + fusion_dim : projection dimensionality inside the bridge + dropout : dropout before the fused output head + use_se : SE gate on the fused vector + + Forward contract + ---------------- + ``forward(embeddings)`` takes a ``dict[str, Tensor]`` of pre-computed + per-tower embeddings (keyed by tower name) and returns + ``(logits_fused, aux_dict)`` where ``aux_dict`` maps each tower name to + its auxiliary head logits. + + The vehicle does not know how to extract embeddings from raw data — + that is the driver's responsibility. Use ``embed_batch`` (TowerBase API) + or custom extraction logic in the training loop, then pass the result here. + + Type-aware helpers + ------------------ + ``transform`` — returns the ``.transform`` of the first tower that has one + (typically the image tower), for use by data loaders. + """ + + def __init__( + self, + towers: dict[str, nn.Module], + num_classes: int, + fusion_dim: int = 256, + dropout: float = 0.5, + use_se: bool = False, + ): + super().__init__() + self.towers = nn.ModuleDict(towers) + self.bridge = Bridge( + tower_dims=[t.out_dim for t in self.towers.values()], + num_classes=num_classes, + fusion_dim=fusion_dim, + dropout=dropout, + use_se=use_se, + ) + + @property + def transform(self): + """Image transform from the first tower that exposes one, or None.""" + for t in self.towers.values(): + if hasattr(t, "transform"): + return t.transform + return None + + def encode(self, embeddings: dict[str, torch.Tensor]) -> torch.Tensor: + """Return z_fused (pre-classifier) from a dict of per-tower embeddings.""" + return self.bridge.encode([embeddings[name] for name in self.towers]) + + def forward( + self, + embeddings: dict[str, torch.Tensor], + ) -> tuple[torch.Tensor, dict[str, torch.Tensor]]: + """Fuse pre-computed embeddings through the bridge. + + Returns + ------- + logits_fused : Tensor [B, num_classes] + aux_logits : dict mapping tower name → Tensor [B, num_classes] + """ + ordered = [embeddings[name] for name in self.towers] + logits_fused, aux = self.bridge.fuse(ordered) + return logits_fused, {name: aux[i] for i, name in enumerate(self.towers)} + + +class NLateralHT(nn.Module): + """N same-type inputs through a shared encoder, jointly compressed, then classified. + + Intended for encoding multiple instances of the same modality together — + e.g. left + right fundus images, or OD + OS clinical vectors. + All inputs share the same encoder weights (one forward pass per input). + + Output matches NTowerHT's contract: ``(logits, {name: aux_logits})``. + Aux logits are per-input classifications taken before the joint MLP, + useful for BCD-style training. + + Parameters + ---------- + encoder : shared encoder module with ``.out_dim`` + input_names : ordered names for each input slot (e.g. ``["od", "os"]``) + num_classes : output classes + fusion_dim : hidden dim of the joint MLP + dropout : dropout in joint MLP and classifier head + """ + + def __init__( + self, + encoder: nn.Module, + input_names: list[str], + num_classes: int, + fusion_dim: int = 256, + dropout: float = 0.5, + ): + super().__init__() + self.encoder = encoder + self.input_names = list(input_names) + n = len(input_names) + in_dim: int = encoder.out_dim # type: ignore[assignment] + + # Joint MLP: [z0 ‖ z1 ‖ ... ‖ z_{n-1}] → fusion_dim → in_dim + self.joint = nn.Sequential( + nn.Linear(n * in_dim, fusion_dim), nn.LayerNorm(fusion_dim), + nn.ReLU(), nn.Dropout(dropout), nn.Linear(fusion_dim, in_dim), + ) + # Per-input aux heads — classify each input before joining + self.aux_heads = nn.ModuleList([ + nn.Linear(in_dim, num_classes) for _ in range(n) + ]) + # Main classifier on the joint representation + self.head = nn.Sequential( + nn.ReLU(), nn.Dropout(dropout), nn.Linear(in_dim, num_classes), + ) + + @property + def transform(self): + return getattr(self.encoder, "transform", None) + + def encode(self, inputs: dict[str, torch.Tensor]) -> torch.Tensor: + """Return joint embedding (post-MLP, pre-classifier).""" + zs = [self.encoder(inputs[name]) for name in self.input_names] + return self.joint(torch.cat(zs, dim=1)) + + def forward( + self, + inputs: dict[str, torch.Tensor], + ) -> tuple[torch.Tensor, dict[str, torch.Tensor]]: + """Encode inputs jointly and classify. + + Parameters + ---------- + inputs : ``{name: raw_tensor}`` for each input slot + + Returns + ------- + logits : [B, num_classes] + aux_dict : ``{name: [B, num_classes]}`` — per-input pre-join logits + """ + zs = [self.encoder(inputs[name]) for name in self.input_names] + z_joint = self.joint(torch.cat(zs, dim=1)) + logits = self.head(z_joint) + aux = {name: head(z) + for name, head, z in zip(self.input_names, self.aux_heads, zs)} + return logits, aux + + +class BilateralHT(nn.Module): + """Bilateral vehicle: shared eye-level towers + joint fusion layers + Bridge.""" + + def __init__( + self, + *, + backbone: str, + freeze_ratio: float, + augment: bool, + clinical_data, + num_classes: int, + cd_hidden_dim: int = 128, + fusion_dim: int = 256, + ): + super().__init__() + self.eye_img_tower = ImageEncoder( + backbone=backbone, freeze_ratio=freeze_ratio, augment=augment, use_se=False, + ) + self.eye_cd_tower = ClinicalEncoder( + clinical_data=clinical_data, hidden_dim=cd_hidden_dim, use_se=False, + ) + img_dim = self.eye_img_tower.out_dim + md_dim = self.eye_cd_tower.out_dim + self.joint_img = nn.Sequential( + nn.Linear(2 * img_dim, fusion_dim), nn.LayerNorm(fusion_dim), + nn.ReLU(), nn.Dropout(0.3), nn.Linear(fusion_dim, img_dim), + ) + self.joint_md = nn.Sequential( + nn.Linear(2 * md_dim, fusion_dim), nn.LayerNorm(fusion_dim), + nn.ReLU(), nn.Dropout(0.3), nn.Linear(fusion_dim, md_dim), + ) + self.bridge = Bridge(tower_dims=[img_dim, md_dim], num_classes=num_classes, + fusion_dim=fusion_dim, mode="fused", use_se=False) + self.aux_img = nn.Linear(img_dim, num_classes) + self.aux_md = nn.Linear(md_dim, num_classes) + + @property + def transform(self): + return self.eye_img_tower.transform + + def encode_joint( + self, + x_od: torch.Tensor, meta_od: torch.Tensor, + x_os: torch.Tensor, meta_os: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor]: + img_od = self.eye_img_tower(x_od); md_od = self.eye_cd_tower(meta_od) + img_os = self.eye_img_tower(x_os); md_os = self.eye_cd_tower(meta_os) + return (self.joint_img(torch.cat([img_od, img_os], dim=1)), + self.joint_md(torch.cat([md_od, md_os], dim=1))) + + def forward( + self, + x_od: torch.Tensor, meta_od: torch.Tensor, + x_os: torch.Tensor, meta_os: torch.Tensor, + ) -> torch.Tensor: + joint_img, joint_md = self.encode_joint(x_od, meta_od, x_os, meta_os) + out_f, _ = self.bridge.fuse([joint_img, joint_md]) + return out_f + + +class SiameseHT(nn.Module): + """Bilateral vehicle using a shared-weight SiameseImageTower (mean+delta).""" + + def __init__( + self, + *, + backbone: str, + freeze_ratio: float, + augment: bool, + num_classes: int, + fusion_dim: int = 256, + ): + super().__init__() + self.img_tower = SiameseImageTower( + backbone=backbone, freeze_ratio=freeze_ratio, augment=augment, use_se=False, + ) + img_dim = self.img_tower.out_dim + self.bridge = Bridge(tower_dims=[img_dim], num_classes=num_classes, + fusion_dim=fusion_dim, use_se=False) + self.aux_img = nn.Linear(img_dim, num_classes) + + @property + def transform(self): + return self.img_tower.transform + + def encode(self, x_od: torch.Tensor, x_os: torch.Tensor) -> torch.Tensor: + return self.img_tower(x_od, x_os) + + def forward(self, x_od: torch.Tensor, x_os: torch.Tensor) -> torch.Tensor: + out_f, _ = self.bridge.fuse([self.encode(x_od, x_os)]) + return out_f + + +class FusedEnsembleHT(nn.Module): + """SingleEyeHT base with a per-eye attention scorer for bilateral fusion.""" + + def __init__(self, base: SingleEyeHT, num_classes: int): + super().__init__() + self.base = base + self.eye_scorer = nn.Linear(num_classes, 1, bias=True) + + @property + def head(self) -> nn.Module: + return self.eye_scorer + + def forward( + self, + x_od: torch.Tensor, meta_od: torch.Tensor, + x_os: torch.Tensor, meta_os: torch.Tensor, + ) -> torch.Tensor: + logit_od = self.base(x_od, meta_od) + logit_os = self.base(x_os, meta_os) + scores = torch.cat([self.eye_scorer(logit_od), self.eye_scorer(logit_os)], dim=1) + alpha = torch.softmax(scores, dim=1) + return alpha[:, 0:1] * logit_od + alpha[:, 1:2] * logit_os + + +class LogitMLPEnsembleHT(nn.Module): + """MLP head trained on concatenated per-eye logits.""" + + def __init__(self, base: SingleEyeHT, num_classes: int, hidden: int = 64): + super().__init__() + self.base = base + self.head = nn.Sequential( + nn.Linear(2 * num_classes, hidden), nn.ReLU(), + nn.Dropout(0.3), nn.Linear(hidden, num_classes), + ) + + def forward( + self, + x_od: torch.Tensor, meta_od: torch.Tensor, + x_os: torch.Tensor, meta_os: torch.Tensor, + ) -> torch.Tensor: + return self.head(torch.cat([self.base(x_od, meta_od), self.base(x_os, meta_os)], dim=1)) + + +class EmbeddingMLPEnsembleHT(nn.Module): + """MLP head trained on concatenated per-eye z_fused embeddings.""" + + def __init__(self, base: SingleEyeHT, num_classes: int, hidden: int = 256): + super().__init__() + self.base = base + fusion_dim = base.bridge.W[0].out_features + self.head = nn.Sequential( + nn.Linear(2 * fusion_dim, hidden), nn.ReLU(), + nn.Dropout(0.3), nn.Linear(hidden, num_classes), + ) + + def forward( + self, + x_od: torch.Tensor, meta_od: torch.Tensor, + x_os: torch.Tensor, meta_os: torch.Tensor, + ) -> torch.Tensor: + return self.head(torch.cat([self.base.encode(x_od, meta_od), + self.base.encode(x_os, meta_os)], dim=1)) + + +# --------------------------------------------------------------------------- +# MonoTowerHT — single tower + Bridge(N=1) +# --------------------------------------------------------------------------- + +class MonoTowerHT(nn.Module): + """Single-tower vehicle: tower output fed directly into a classifier head. + + No bridge projection — the tower's embedding goes straight to + ``ReLU → Dropout → Linear(out_dim → num_classes)``. This is the + minimal architecture: just the tower's learned representation with a + classification head attached. + + Contrast with ``NTowerHT(N=1)``, which still projects through the bridge's + shared ``fusion_dim`` space. These are distinct architectures and may + yield different results. + + Parameters + ---------- + tower : any nn.Module with ``.out_dim`` (ImageEncoder, ClinicalEncoder, …) + num_classes : number of output classes + dropout : dropout before the output linear layer + """ + + def __init__( + self, + tower: nn.Module, + num_classes: int, + dropout: float = 0.5, + ): + super().__init__() + self.tower = tower + out_dim: int = tower.out_dim # type: ignore[assignment] + self.head = nn.Sequential( + nn.ReLU(), + nn.Dropout(dropout), + nn.Linear(out_dim, num_classes), + ) + + def classify(self, z: torch.Tensor) -> torch.Tensor: + """Classify a pre-computed embedding.""" + return self.head(z) + + def forward(self, *args, **kwargs) -> torch.Tensor: + """Pass inputs through tower then classifier head.""" + return self.head(self.tower(*args, **kwargs)) + + +# --------------------------------------------------------------------------- +# Phase control +# --------------------------------------------------------------------------- + +def _set_single_phase(model: SingleEyeHT, phase: str) -> None: + bridge_mode = model.bridge.mode + if bridge_mode in ("image_only", "clinical_only") and phase == "fused_warmup": + phase = "tower_warmup" + if phase == "cd_warmup": + _set_requires_grad(model.img_tower, False) + _set_requires_grad(model.cd_tower, True) + _set_requires_grad(model.bridge.aux_heads[0], False) + _set_requires_grad(model.bridge.aux_heads[1], True) + _set_requires_grad(model.bridge.W[0], False) + _set_requires_grad(model.bridge.W[1], False) + _set_requires_grad(model.bridge.classifier_fused, False) + return + if phase == "tower_warmup": + _set_requires_grad(model.img_tower, bridge_mode != "clinical_only") + _set_requires_grad(model.cd_tower, bridge_mode != "image_only") + _set_requires_grad(model.bridge.aux_heads[0], bridge_mode != "clinical_only") + _set_requires_grad(model.bridge.aux_heads[1], bridge_mode != "image_only") + _set_requires_grad(model.bridge.W[0], False) + _set_requires_grad(model.bridge.W[1], False) + _set_requires_grad(model.bridge.classifier_fused, False) + return + if phase == "fused_warmup": + _set_requires_grad(model.img_tower, False) + _set_requires_grad(model.cd_tower, False) + _set_requires_grad(model.bridge.aux_heads[0], False) + _set_requires_grad(model.bridge.aux_heads[1], False) + _set_requires_grad(model.bridge.W[0], True) + _set_requires_grad(model.bridge.W[1], True) + _set_requires_grad(model.bridge.classifier_fused, True) + return + _set_requires_grad(model, True) + + +def _set_bilateral_phase(model: BilateralHT, phase: str) -> None: + if phase == "tower_warmup": + for m in (model.eye_img_tower, model.eye_cd_tower, + model.joint_img, model.joint_md, model.aux_img, model.aux_md): + _set_requires_grad(m, True) + _set_requires_grad(model.bridge, False) + return + if phase == "fused_warmup": + for m in (model.eye_img_tower, model.eye_cd_tower, + model.joint_img, model.joint_md, model.aux_img, model.aux_md): + _set_requires_grad(m, False) + _set_requires_grad(model.bridge, True) + return + _set_requires_grad(model, True) + + +# --------------------------------------------------------------------------- +# Training helpers +# --------------------------------------------------------------------------- + +def train_single_epoch( + model: SingleEyeHT, + loader: DataLoader, + opt, + device: torch.device, + *, + phase: str, + bcd_prob: float = 0.5, + tower_loss_mode: str = "bcd", +) -> tuple[float, float]: + model.train() + _set_single_phase(model, phase) + total_loss = total_correct = total_n = 0 + for batch in loader: + x = batch.get("image_1"); m = batch.get("matrix_1"); y = batch.get("label_1") + if phase == "cd_warmup": + if not torch.is_tensor(m): + continue + y = _to_label_tensor(y, device) + logits = model.bridge.aux_heads[1](model.cd_tower(m.to(device))) + loss = F.cross_entropy(logits, y) + opt.zero_grad(); loss.backward(); opt.step() + bs = y.shape[0] + total_loss += float(loss.item()) * bs + total_correct += int((logits.argmax(1) == y).sum()) + total_n += bs + continue + if not (torch.is_tensor(x) and torch.is_tensor(m)): + continue + x = x.to(device); m = m.to(device); y = _to_label_tensor(y, device) + bridge_mode = model.bridge.mode + img_feats = None if bridge_mode == "clinical_only" else model.img_tower(x) + md_feats = None if bridge_mode == "image_only" else model.cd_tower(m) + + if phase == "tower_warmup": + if bridge_mode == "clinical_only": + logits = model.bridge.aux_heads[1](md_feats); loss = F.cross_entropy(logits, y) + elif bridge_mode == "image_only": + logits = model.bridge.aux_heads[0](img_feats); loss = F.cross_entropy(logits, y) + else: + li = model.bridge.aux_heads[0](img_feats); lm = model.bridge.aux_heads[1](md_feats) + loss = 0.5 * (F.cross_entropy(li, y) + F.cross_entropy(lm, y)) + logits = 0.5 * (F.softmax(li, dim=1) + F.softmax(lm, dim=1)) + elif phase == "fused_warmup": + logits, _ = model.bridge.fuse([img_feats, md_feats]); loss = F.cross_entropy(logits, y) + else: + if bridge_mode == "clinical_only": + logits = model.bridge.aux_heads[1](md_feats); loss = F.cross_entropy(logits, y) + elif bridge_mode == "image_only": + logits = model.bridge.aux_heads[0](img_feats); loss = F.cross_entropy(logits, y) + elif tower_loss_mode == "all": + li = model.bridge.aux_heads[0](img_feats); lm = model.bridge.aux_heads[1](md_feats) + logits, _ = model.bridge.fuse([img_feats, md_feats]) + loss = F.cross_entropy(logits, y) + F.cross_entropy(li, y) + F.cross_entropy(lm, y) + elif random() < bcd_prob: + logits = (model.bridge.aux_heads[0](img_feats) if random() < 0.5 + else model.bridge.aux_heads[1](md_feats)) + loss = F.cross_entropy(logits, y) + else: + logits, _ = model.bridge.fuse([img_feats, md_feats]); loss = F.cross_entropy(logits, y) + + opt.zero_grad(); loss.backward(); opt.step() + bs = y.shape[0] + total_loss += float(loss.item()) * bs + total_correct += int((logits.argmax(1) == y).sum()) + total_n += bs + return (total_loss / total_n if total_n else float("nan"), + total_correct / total_n if total_n else float("nan")) + + +def train_bilateral_epoch( + model: BilateralHT, + loader: DataLoader, + opt, + device: torch.device, + *, + phase: str, + bcd_prob: float = 0.5, + tower_loss_mode: str = "bcd", +) -> tuple[float, float]: + model.train() + _set_bilateral_phase(model, phase) + total_loss = total_correct = total_n = 0 + for batch in loader: + x1 = batch.get("image_1"); m1 = batch.get("matrix_1") + x2 = batch.get("image_2"); m2 = batch.get("matrix_2"); y = batch.get("label_1") + if not (torch.is_tensor(x1) and torch.is_tensor(m1) and + torch.is_tensor(x2) and torch.is_tensor(m2)): + continue + x1 = x1.to(device); m1 = m1.to(device) + x2 = x2.to(device); m2 = m2.to(device); y = _to_label_tensor(y, device) + joint_img, joint_md = model.encode_joint(x1, m1, x2, m2) + + if phase == "tower_warmup": + li = model.aux_img(joint_img); lm = model.aux_md(joint_md) + loss = 0.5 * (F.cross_entropy(li, y) + F.cross_entropy(lm, y)) + logits = 0.5 * (F.softmax(li, dim=1) + F.softmax(lm, dim=1)) + elif phase == "fused_warmup": + logits, _ = model.bridge.fuse([joint_img, joint_md]); loss = F.cross_entropy(logits, y) + else: + if tower_loss_mode == "all": + li = model.aux_img(joint_img); lm = model.aux_md(joint_md) + logits, _ = model.bridge.fuse([joint_img, joint_md]) + loss = F.cross_entropy(logits, y) + F.cross_entropy(li, y) + F.cross_entropy(lm, y) + elif random() < bcd_prob: + logits = model.aux_img(joint_img) if random() < 0.5 else model.aux_md(joint_md) + loss = F.cross_entropy(logits, y) + else: + logits, _ = model.bridge.fuse([joint_img, joint_md]); loss = F.cross_entropy(logits, y) + + opt.zero_grad(); loss.backward(); opt.step() + bs = y.shape[0] + total_loss += float(loss.item()) * bs + total_correct += int((logits.argmax(1) == y).sum()) + total_n += bs + return (total_loss / total_n if total_n else float("nan"), + total_correct / total_n if total_n else float("nan")) + + +def train_siamese_epoch( + model: SiameseHT, + loader: DataLoader, + opt, + device: torch.device, + *, + bcd_prob: float = 0.5, + tower_loss_mode: str = "bcd", +) -> tuple[float, float]: + model.train() + total_loss = total_correct = total_n = 0 + for batch in loader: + x1 = batch.get("image_1"); x2 = batch.get("image_2"); y = batch.get("label_1") + if not (torch.is_tensor(x1) and torch.is_tensor(x2)): + continue + x1 = x1.to(device); x2 = x2.to(device); y = _to_label_tensor(y, device) + feats = model.encode(x1, x2) + + if tower_loss_mode == "all": + out_f, _ = model.bridge.fuse([feats, None]) + loss = F.cross_entropy(out_f, y) + F.cross_entropy(model.aux_img(feats), y) + logits = out_f + elif random() < bcd_prob: + logits = model.aux_img(feats); loss = F.cross_entropy(logits, y) + else: + logits, _ = model.bridge.fuse([feats, None]); loss = F.cross_entropy(logits, y) + + opt.zero_grad(); loss.backward(); opt.step() + bs = y.shape[0] + total_loss += float(loss.item()) * bs + total_correct += int((logits.argmax(1) == y).sum()) + total_n += bs + return (total_loss / total_n if total_n else float("nan"), + total_correct / total_n if total_n else float("nan")) + + +def train_fusion_epoch( + model, # FusedEnsembleHT | LogitMLPEnsembleHT | EmbeddingMLPEnsembleHT + loader: DataLoader, + opt, + device: torch.device, +) -> tuple[float, float]: + """Train only the fusion head; base SingleEyeHT is frozen in eval mode.""" + model.base.eval(); model.head.train() + total_loss = total_correct = total_n = 0 + for batch in loader: + x1 = batch.get("image_1"); m1 = batch.get("matrix_1") + x2 = batch.get("image_2"); m2 = batch.get("matrix_2"); y = batch.get("label_1") + if not (torch.is_tensor(x1) and torch.is_tensor(m1) and + torch.is_tensor(x2) and torch.is_tensor(m2)): + continue + y_t = _to_label_tensor(y, device) + out = model(x1.to(device), m1.to(device), x2.to(device), m2.to(device)) + loss = F.cross_entropy(out, y_t) + opt.zero_grad(); loss.backward(); opt.step() + bs = y_t.shape[0] + total_loss += float(loss.item()) * bs + total_correct += int((out.argmax(1) == y_t).sum()) + total_n += bs + return (total_loss / total_n if total_n else float("nan"), + total_correct / total_n if total_n else float("nan")) + + +# --------------------------------------------------------------------------- +# Inference helpers +# --------------------------------------------------------------------------- + +def collect_probs_classic( + model: SingleEyeHT, loader: DataLoader, device: torch.device, +) -> tuple[np.ndarray, np.ndarray]: + model.eval() + y_c, p_c = [], [] + with torch.no_grad(): + for batch in loader: + x1 = batch.get("image_1"); m1 = batch.get("matrix_1") + x2 = batch.get("image_2"); m2 = batch.get("matrix_2"); y = batch.get("label_1") + if not (torch.is_tensor(x1) and torch.is_tensor(m1) and + torch.is_tensor(x2) and torch.is_tensor(m2)): + continue + y_t = _to_label_tensor(y, device) + p_od = F.softmax(model(x1.to(device), m1.to(device)), dim=1) + p_os = F.softmax(model(x2.to(device), m2.to(device)), dim=1) + y_np = y_t.cpu().numpy() + y_c += [y_np, y_np]; p_c += [p_od.cpu().numpy(), p_os.cpu().numpy()] + if not y_c: + return np.array([], dtype=np.int64), np.zeros((0, 0), dtype=np.float32) + return np.concatenate(y_c), np.concatenate(p_c, axis=0) + + +def collect_probs_ensemble( + model: SingleEyeHT, loader: DataLoader, device: torch.device, +) -> tuple[np.ndarray, np.ndarray]: + model.eval() + y_c, p_c = [], [] + with torch.no_grad(): + for batch in loader: + x1 = batch.get("image_1"); m1 = batch.get("matrix_1") + x2 = batch.get("image_2"); m2 = batch.get("matrix_2"); y = batch.get("label_1") + if not (torch.is_tensor(x1) and torch.is_tensor(m1) and + torch.is_tensor(x2) and torch.is_tensor(m2)): + continue + y_t = _to_label_tensor(y, device) + p_od = F.softmax(model(x1.to(device), m1.to(device)), dim=1) + p_os = F.softmax(model(x2.to(device), m2.to(device)), dim=1) + y_c.append(y_t.cpu().numpy()); p_c.append((0.5 * (p_od + p_os)).cpu().numpy()) + if not y_c: + return np.array([], dtype=np.int64), np.zeros((0, 0), dtype=np.float32) + return np.concatenate(y_c), np.concatenate(p_c, axis=0) + + +def collect_probs_ensemble_pereye( + model: SingleEyeHT, loader: DataLoader, device: torch.device, *, return_ids: bool = False, +): + model.eval() + y_c = []; pf_od_c, pi_od_c, pm_od_c = [], [], []; pf_os_c, pi_os_c, pm_os_c = [], [], [] + id_c: list = [] + with torch.no_grad(): + for batch in loader: + x1 = batch.get("image_1"); m1 = batch.get("matrix_1") + x2 = batch.get("image_2"); m2 = batch.get("matrix_2"); y = batch.get("label_1") + if not (torch.is_tensor(x1) and torch.is_tensor(m1) and + torch.is_tensor(x2) and torch.is_tensor(m2)): + continue + y_t = _to_label_tensor(y, device) + + def _fwd(x, m): + img_f = None if model.bridge.mode == "clinical_only" else model.img_tower(x.to(device)) + md_f = None if model.bridge.mode == "image_only" else model.cd_tower(m.to(device)) + out_f, aux = model.bridge.fuse([img_f, md_f]) + out_i, out_m = aux[0], aux[1] + pf = F.softmax(out_f, dim=1) + pi = F.softmax(out_i, dim=1) if out_i is not None else pf + pm = F.softmax(out_m, dim=1) if out_m is not None else pf + return pf, pi, pm + + pf_od, pi_od, pm_od = _fwd(x1, m1); pf_os, pi_os, pm_os = _fwd(x2, m2) + y_c.append(y_t.cpu().numpy()) + for lst, t in ((pf_od_c, pf_od), (pi_od_c, pi_od), (pm_od_c, pm_od), + (pf_os_c, pf_os), (pi_os_c, pi_os), (pm_os_c, pm_os)): + lst.append(t.cpu().numpy()) + if return_ids: + ids = batch.get("id_1", [""] * len(y_t)) + id_c.extend([str(i) for i in (ids.tolist() if torch.is_tensor(ids) else ids)]) + + if not y_c: + z = np.zeros((0, 0), dtype=np.float32); ei = np.array([], dtype=np.int64) + base = (ei, z, z, z, z, z, z) + return base + (np.array([], dtype=object),) if return_ids else base + + y = np.concatenate(y_c) + pf_od, pi_od, pm_od = (np.concatenate(c, axis=0) for c in (pf_od_c, pi_od_c, pm_od_c)) + pf_os, pi_os, pm_os = (np.concatenate(c, axis=0) for c in (pf_os_c, pi_os_c, pm_os_c)) + if return_ids: + return y, pf_od, pi_od, pm_od, pf_os, pi_os, pm_os, np.array(id_c, dtype=object) + return y, pf_od, pi_od, pm_od, pf_os, pi_os, pm_os + + +def collect_probs_bilateral( + model: BilateralHT, loader: DataLoader, device: torch.device, +) -> tuple[np.ndarray, np.ndarray]: + model.eval() + y_c, p_c = [], [] + with torch.no_grad(): + for batch in loader: + x1 = batch.get("image_1"); m1 = batch.get("matrix_1") + x2 = batch.get("image_2"); m2 = batch.get("matrix_2"); y = batch.get("label_1") + if not (torch.is_tensor(x1) and torch.is_tensor(m1) and + torch.is_tensor(x2) and torch.is_tensor(m2)): + continue + y_t = _to_label_tensor(y, device) + p = F.softmax(model(x1.to(device), m1.to(device), + x2.to(device), m2.to(device)), dim=1) + y_c.append(y_t.cpu().numpy()); p_c.append(p.cpu().numpy()) + if not y_c: + return np.array([], dtype=np.int64), np.zeros((0, 0), dtype=np.float32) + return np.concatenate(y_c), np.concatenate(p_c, axis=0) + + +def collect_probs_bilateral_components( + model: BilateralHT, loader: DataLoader, device: torch.device, +) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]: + model.eval() + y_c = []; pf_c, pi_c, pm_c = [], [], [] + with torch.no_grad(): + for batch in loader: + x1 = batch.get("image_1"); m1 = batch.get("matrix_1") + x2 = batch.get("image_2"); m2 = batch.get("matrix_2"); y = batch.get("label_1") + if not (torch.is_tensor(x1) and torch.is_tensor(m1) and + torch.is_tensor(x2) and torch.is_tensor(m2)): + continue + y_t = _to_label_tensor(y, device) + ji, jm = model.encode_joint(x1.to(device), m1.to(device), + x2.to(device), m2.to(device)) + out_f, _ = model.bridge.fuse([ji, jm]) + y_c.append(y_t.cpu().numpy()) + pf_c.append(F.softmax(out_f, dim=1).cpu().numpy()) + pi_c.append(F.softmax(model.aux_img(ji), dim=1).cpu().numpy()) + pm_c.append(F.softmax(model.aux_md(jm), dim=1).cpu().numpy()) + if not y_c: + z = np.zeros((0, 0), dtype=np.float32) + return np.array([], dtype=np.int64), z, z, z + return (np.concatenate(y_c), np.concatenate(pf_c, axis=0), + np.concatenate(pi_c, axis=0), np.concatenate(pm_c, axis=0)) + + +def collect_probs_siamese( + model: SiameseHT, loader: DataLoader, device: torch.device, +) -> tuple[np.ndarray, np.ndarray]: + model.eval() + y_c, p_c = [], [] + with torch.no_grad(): + for batch in loader: + x1 = batch.get("image_1"); x2 = batch.get("image_2"); y = batch.get("label_1") + if not (torch.is_tensor(x1) and torch.is_tensor(x2)): + continue + feats = model.encode(x1.to(device), x2.to(device)) + logits, _ = model.bridge.fuse([feats]) + y_c.append(np.array(y) if not torch.is_tensor(y) else y.cpu().numpy()) + p_c.append(torch.softmax(logits, dim=1).cpu().numpy()) + return np.concatenate(y_c, axis=0), np.concatenate(p_c, axis=0) + + +def collect_probs_fused( + model: FusedEnsembleHT, loader: DataLoader, device: torch.device, +) -> tuple[np.ndarray, np.ndarray]: + model.eval() + y_c, p_c = [], [] + with torch.no_grad(): + for batch in loader: + x1 = batch.get("image_1"); m1 = batch.get("matrix_1") + x2 = batch.get("image_2"); m2 = batch.get("matrix_2"); y = batch.get("label_1") + if not (torch.is_tensor(x1) and torch.is_tensor(m1) and + torch.is_tensor(x2) and torch.is_tensor(m2)): + continue + y_t = _to_label_tensor(y, device) + p = F.softmax(model(x1.to(device), m1.to(device), + x2.to(device), m2.to(device)), dim=1) + y_c.append(y_t.cpu().numpy()); p_c.append(p.cpu().numpy()) + if not y_c: + return np.array([], dtype=np.int64), np.zeros((0, 0), dtype=np.float32) + return np.concatenate(y_c), np.concatenate(p_c, axis=0) + + +def collect_probs_single_components( + model: SingleEyeHT, loader: DataLoader, device: torch.device, + *, aggregate_patient: bool, return_logits: bool = False, +): + model.eval() + y_c = []; pf_c, pi_c, pm_c = [], [], []; lf_c, li_c, lm_c = [], [], [] + with torch.no_grad(): + for batch in loader: + x1 = batch.get("image_1"); m1 = batch.get("matrix_1") + x2 = batch.get("image_2"); m2 = batch.get("matrix_2"); y = batch.get("label_1") + if not (torch.is_tensor(x1) and torch.is_tensor(m1) and + torch.is_tensor(x2) and torch.is_tensor(m2)): + continue + y_t = _to_label_tensor(y, device) + + def _per_eye(x, m): + img_f = None if model.bridge.mode == "clinical_only" else model.img_tower(x.to(device)) + md_f = None if model.bridge.mode == "image_only" else model.cd_tower(m.to(device)) + out_f, aux = model.bridge.fuse([img_f, md_f]) + out_i, out_m = aux[0], aux[1] + pf = F.softmax(out_f, dim=1) + pi = F.softmax(out_i, dim=1) if out_i is not None else pf + pm = F.softmax(out_m, dim=1) if out_m is not None else pf + return pf, pi, pm, out_f, (out_i if out_i is not None else out_f), (out_m if out_m is not None else out_f) + + pf_od, pi_od, pm_od, lf_od, li_od, lm_od = _per_eye(x1, m1) + pf_os, pi_os, pm_os, lf_os, li_os, lm_os = _per_eye(x2, m2) + + if aggregate_patient: + y_c.append(y_t.cpu().numpy()) + pf_c.append((0.5*(pf_od+pf_os)).cpu().numpy()) + pi_c.append((0.5*(pi_od+pi_os)).cpu().numpy()) + pm_c.append((0.5*(pm_od+pm_os)).cpu().numpy()) + lf_c.append((0.5*(lf_od+lf_os)).cpu().numpy()) + li_c.append((0.5*(li_od+li_os)).cpu().numpy()) + lm_c.append((0.5*(lm_od+lm_os)).cpu().numpy()) + else: + y_np = y_t.cpu().numpy(); y_c += [y_np, y_np] + pf_c += [pf_od.cpu().numpy(), pf_os.cpu().numpy()] + pi_c += [pi_od.cpu().numpy(), pi_os.cpu().numpy()] + pm_c += [pm_od.cpu().numpy(), pm_os.cpu().numpy()] + lf_c += [lf_od.cpu().numpy(), lf_os.cpu().numpy()] + li_c += [li_od.cpu().numpy(), li_os.cpu().numpy()] + lm_c += [lm_od.cpu().numpy(), lm_os.cpu().numpy()] + + if not y_c: + z = np.zeros((0, 0), dtype=np.float32) + if return_logits: + return np.array([], dtype=np.int64), z, z, z, z, z, z + return np.array([], dtype=np.int64), z, z, z + + y = np.concatenate(y_c) + pf, pi, pm = (np.concatenate(c, axis=0) for c in (pf_c, pi_c, pm_c)) + if return_logits: + lf, li, lm = (np.concatenate(c, axis=0) for c in (lf_c, li_c, lm_c)) + return y, pf, pi, pm, lf, li, lm + return y, pf, pi, pm + + +def collect_probs_eye_level( + model: SingleEyeHT, loader: DataLoader, device: torch.device, *, return_ids: bool = False, +): + model.eval() + y_c, pf_c, pi_c, pm_c, id_c = [], [], [], [], [] + with torch.no_grad(): + for batch in loader: + x = batch.get("image_1"); m = batch.get("matrix_1"); y = batch.get("label_1") + if not (torch.is_tensor(x) and torch.is_tensor(m)): + continue + y_t = _to_label_tensor(y, device) + img_f = None if model.bridge.mode == "clinical_only" else model.img_tower(x.to(device)) + md_f = None if model.bridge.mode == "image_only" else model.cd_tower(m.to(device)) + out_f, aux = model.bridge.fuse([img_f, md_f]) + out_i, out_m = aux[0], aux[1] + pf = F.softmax(out_f, dim=1) + pi = F.softmax(out_i, dim=1) if out_i is not None else pf + pm = F.softmax(out_m, dim=1) if out_m is not None else pf + y_c.append(y_t.cpu().numpy()) + pf_c.append(pf.cpu().numpy()); pi_c.append(pi.cpu().numpy()); pm_c.append(pm.cpu().numpy()) + if return_ids: + ids = batch.get("id_1", [""] * len(y_t)) + eyes = batch.get("eye_id_1", [""] * len(y_t)) + if torch.is_tensor(ids): ids = ids.tolist() + if torch.is_tensor(eyes): eyes = eyes.tolist() + id_c.extend([f"{p}{e}" for p, e in zip(ids, eyes)]) + + if not y_c: + z = np.zeros((0, 0), dtype=np.float32) + if return_ids: + return np.array([], dtype=np.int64), z, z, z, np.array([], dtype=object) + return np.array([], dtype=np.int64), z, z, z + + y = np.concatenate(y_c) + pf, pi, pm = (np.concatenate(c, axis=0) for c in (pf_c, pi_c, pm_c)) + if return_ids: + return y, pf, pi, pm, np.array(id_c, dtype=object) + return y, pf, pi, pm + + +# --------------------------------------------------------------------------- +# MonoTowerHT training / inference helpers +# --------------------------------------------------------------------------- + +def train_mono_epoch( + model: MonoTowerHT, + loader: DataLoader, + opt, + device: torch.device, +) -> tuple[float, float]: + """One training epoch for MonoTowerHT. + + Uses the tower's ``embed_batch`` (TowerBase API) to extract the embedding, + then classifies directly with the head. Loss is cross-entropy on the + head output only — no aux head. + + Returns + ------- + (mean_loss, accuracy) + """ + model.train() + total_loss = total_correct = total_n = 0 + for batch in loader: + y = _to_label_tensor(batch.get("label_1"), device) + if y.numel() == 0: + continue + [z] = model.tower.embed_batch(batch, device=device) + logits = model.classify(z) + loss = F.cross_entropy(logits, y) + opt.zero_grad() + loss.backward() + opt.step() + total_loss += loss.item() * len(y) + total_correct += int((logits.argmax(1) == y).sum()) + total_n += len(y) + mean_loss = total_loss / total_n if total_n else float("nan") + accuracy = total_correct / total_n if total_n else float("nan") + return mean_loss, accuracy + + +def collect_probs_mono( + model: MonoTowerHT, + loader: DataLoader, + device: torch.device, +) -> tuple[np.ndarray, np.ndarray]: + """Collect predictions from a MonoTowerHT. + + Returns + ------- + y_true : int64 array [N] + probs : float32 array [N, num_classes] — softmax of fused head + """ + model.eval() + y_all, p_all = [], [] + with torch.no_grad(): + for batch in loader: + y = _to_label_tensor(batch.get("label_1"), device) + if y.numel() == 0: + continue + [z] = model.tower.embed_batch(batch, device=device) + logits = model.classify(z) + y_all.append(y.cpu().numpy()) + p_all.append(F.softmax(logits, dim=1).cpu().numpy()) + if not y_all: + return np.zeros(0, dtype=np.int64), np.zeros((0, 0), dtype=np.float32) + return np.concatenate(y_all), np.concatenate(p_all, axis=0) + + +# --------------------------------------------------------------------------- +# NTowerHT training / inference helpers +# --------------------------------------------------------------------------- + +def train_ntower_epoch( + model: NTowerHT, + loader: DataLoader, + opt, + device: torch.device, + *, + batch_key_map: dict[str, str], + phase: str, + bcd_prob: float = 0.5, + tower_loss_mode: str = "bcd", +) -> tuple[float, float]: + """One training epoch for NTowerHT. + + Parameters + ---------- + batch_key_map : {tower_name: batch_key} — maps each tower slot to the key + in the DataLoader batch dict that carries its input tensor. + Example: {'od': 'image_1', 'os': 'image_2', 'cd': 'matrix_1'} + phase : 'tower_warmup' | 'fused_warmup' | 'main' + bcd_prob : probability of using a random aux head instead of the fused + head in 'main' phase (BCD training). + tower_loss_mode : 'bcd' | 'all' — 'all' adds all aux + fused losses each step. + """ + from random import choice + model.train() + model.bridge.set_phase(phase) + # Mirror _set_single_phase tower freezing: + # tower_warmup → towers trainable; fused_warmup → towers frozen; main → trainable + towers_trainable = phase != "fused_warmup" + for t in model.towers.values(): + for p in t.parameters(): + p.requires_grad_(towers_trainable) + tower_names = list(model.towers.keys()) + total_loss = total_correct = total_n = 0 + + for batch in loader: + tensors = {name: batch.get(key) for name, key in batch_key_map.items()} + if not all(torch.is_tensor(t) for t in tensors.values()): + continue + y = _to_label_tensor(batch.get("label_1"), device) + if y.numel() == 0: + continue + embeddings = {name: model.towers[name](t.to(device)) for name, t in tensors.items()} + + if phase == "tower_warmup": + aux_logits = [model.bridge.aux_heads[i](embeddings[name]) + for i, name in enumerate(tower_names)] + loss = sum(F.cross_entropy(l, y) for l in aux_logits) / len(aux_logits) + avg_probs = sum(F.softmax(l, dim=1) for l in aux_logits) / len(aux_logits) + logits = avg_probs # for accuracy tracking + elif phase == "fused_warmup": + logits, _ = model(embeddings) + loss = F.cross_entropy(logits, y) + else: # main + if tower_loss_mode == "all": + logits, aux_dict = model(embeddings) + loss = F.cross_entropy(logits, y) + sum( + F.cross_entropy(a, y) for a in aux_dict.values() + ) + elif random() < bcd_prob: + # BCD: pick one random aux head + name = choice(tower_names) + idx = tower_names.index(name) + logits = model.bridge.aux_heads[idx](embeddings[name]) + loss = F.cross_entropy(logits, y) + else: + logits, _ = model(embeddings) + loss = F.cross_entropy(logits, y) + + opt.zero_grad() + loss.backward() + opt.step() + if logits.ndim == 2: + total_correct += int((logits.argmax(1) == y).sum()) + total_loss += loss.item() * len(y) + total_n += len(y) + + return ( + total_loss / total_n if total_n else float("nan"), + total_correct / total_n if total_n else float("nan"), + ) + + +def collect_probs_ntower( + model: NTowerHT, + loader: DataLoader, + device: torch.device, + *, + batch_key_map: dict[str, str], +) -> tuple[np.ndarray, np.ndarray]: + """Collect fused-head predictions from NTowerHT. + + Parameters + ---------- + batch_key_map : same mapping used during training. + + Returns + ------- + y_true : int64 array [N] + probs : float32 array [N, num_classes] — softmax of fused head + """ + model.eval() + y_all, p_all = [], [] + with torch.no_grad(): + for batch in loader: + tensors = {name: batch.get(key) for name, key in batch_key_map.items()} + if not all(torch.is_tensor(t) for t in tensors.values()): + continue + y = _to_label_tensor(batch.get("label_1"), device) + if y.numel() == 0: + continue + embeddings = {name: model.towers[name](t.to(device)) for name, t in tensors.items()} + logits, _ = model(embeddings) + y_all.append(y.cpu().numpy()) + p_all.append(F.softmax(logits, dim=1).cpu().numpy()) + if not y_all: + return np.zeros(0, dtype=np.int64), np.zeros((0, 0), dtype=np.float32) + return np.concatenate(y_all), np.concatenate(p_all, axis=0) + + +# --------------------------------------------------------------------------- +# V2ModeComparisonOps — backward-compat namespace +# --------------------------------------------------------------------------- + +class V2ModeComparisonOps: + """Namespace kept for backward-compatibility imports.""" + _set_requires_grad = staticmethod(_set_requires_grad) + _set_single_phase = staticmethod(_set_single_phase) + _set_bilateral_phase = staticmethod(_set_bilateral_phase) + train_single_epoch = staticmethod(train_single_epoch) + train_bilateral_epoch = staticmethod(train_bilateral_epoch) + collect_probs_classic = staticmethod(collect_probs_classic) + collect_probs_ensemble = staticmethod(collect_probs_ensemble) + collect_probs_bilateral = staticmethod(collect_probs_bilateral) + + @staticmethod + def _to_label_tensor(labels, device: torch.device) -> torch.Tensor: + return _to_label_tensor(labels, device) diff --git a/v3/classes/image_towers.py b/v3/classes/image_towers.py new file mode 100644 index 0000000..e275e8d --- /dev/null +++ b/v3/classes/image_towers.py @@ -0,0 +1,252 @@ +"""image_towers — ImageEncoder, SiameseImageTower, and ImageTower (TowerBase). + +Self-contained image-modality tower layer. No dependencies on other tower +files, bridge, or model classes. Clear contract: accepts an image batch, +returns a fixed-size embedding vector. +""" +from __future__ import annotations + +import math +from typing import Optional + +import torch +from torch import nn +from torch.utils.data import DataLoader + +from v3.classes.towerbase import TowerBase, build_backbone +from v3.classes.backbones import BACKBONES +from v3.classes.SE_attention import SEBlock + + +# --------------------------------------------------------------------------- +# ImageEncoder — vision backbone → pooled feature vector +# --------------------------------------------------------------------------- + +class ImageEncoder(nn.Module): + """Vision backbone → pooled feature vector. + + Wraps a torchvision backbone (default weights), strips the classifier, + and optionally appends an SE attention block and/or a geometry vector. + + Parameters + ---------- + backbone : backbone key (see 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 transform + geometry_dim : if > 0, concatenate a geometry vector of this length + """ + + def __init__( + self, + backbone: str = "efficientnet_b0", + freeze_ratio: float = 0.0, + use_se: bool = False, + se_reduction: int = 16, + se_pre_norm: bool = True, + augment: bool = True, + geometry_dim: int = 0, + ): + super().__init__() + self.backbone, base_dim, self.transform = build_backbone( + backbone, freeze_ratio, augment=augment + ) + self._name = backbone + key = (self._name or "").lower() + self._spec = BACKBONES[key] + self._blocks = self._spec.blocks(self.backbone) + self.base_dim = base_dim + self.geometry_dim = max(0, int(geometry_dim)) + self.out_dim = self.base_dim + self.geometry_dim + self.tower_ln = nn.LayerNorm(self.base_dim) if se_pre_norm else nn.Identity() + self.tower_se = ( + SEBlock(self.base_dim, reduction=se_reduction, residual=True) + if use_se else None + ) + + def forward( + self, x: torch.Tensor, geometry: Optional[torch.Tensor] = None + ) -> torch.Tensor: + y = self.backbone(x) + assert y.dim() == 2 and y.size(1) == self.base_dim, ( + f"Expected features [N,{self.base_dim}], got {tuple(y.shape)}" + ) + if self.tower_se is not None: + y, _ = self.tower_se(self.tower_ln(y)) + if self.geometry_dim > 0: + if geometry is None or geometry.numel() == 0: + geom = torch.zeros(y.size(0), self.geometry_dim, device=y.device, dtype=y.dtype) + else: + geom = geometry.unsqueeze(0) if geometry.dim() == 1 else geometry + geom = geom.to(device=y.device, dtype=y.dtype) + if geom.size(0) != y.size(0): + raise ValueError(f"Geometry batch size mismatch: {geom.size(0)} vs {y.size(0)}") + if geom.size(1) != self.geometry_dim: + raise ValueError(f"Expected geometry dim {self.geometry_dim}, got {geom.size(1)}") + y = torch.cat([y, geom], dim=1) + return y + + def set_freeze_ratio(self, ratio: float) -> None: + """Dynamically freeze earliest floor(N*ratio) backbone blocks.""" + r = max(0.0, min(1.0, float(ratio))) + freeze_n = 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[:freeze_n]: + for p in b.parameters(): + p.requires_grad = False + + +# --------------------------------------------------------------------------- +# SiameseImageTower — shared-weight bilateral image encoder +# --------------------------------------------------------------------------- + +class SiameseImageTower(nn.Module): + """Shared-weight bilateral image tower. + + Runs OD and OS images through a single shared backbone and returns + cat([f_mean, f_delta]) where: + f_mean = (f_od + f_os) / 2 — shared bilateral representation + f_delta = f_od - f_os — signed asymmetry (OD-relative) + + out_dim = 2 × backbone_out_dim. When x_os is None the tower degrades + gracefully: f_mean = f_od, f_delta = zeros. + """ + + def __init__( + self, + backbone: str = "efficientnet_b0", + freeze_ratio: float = 0.0, + use_se: bool = False, + se_reduction: int = 16, + se_pre_norm: bool = True, + augment: bool = True, + ): + super().__init__() + self._tower = ImageEncoder( + backbone=backbone, + freeze_ratio=freeze_ratio, + use_se=use_se, + se_reduction=se_reduction, + se_pre_norm=se_pre_norm, + augment=augment, + ) + self.out_dim = self._tower.out_dim * 2 + self.transform = self._tower.transform + + def forward( + self, + x_od: torch.Tensor, + x_os: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + f_od = self._tower(x_od) + if x_os is None: + return torch.cat([f_od, torch.zeros_like(f_od)], dim=1) + f_os = self._tower(x_os) + return torch.cat([(f_od + f_os) * 0.5, f_od - f_os], dim=1) + + def set_freeze_ratio(self, ratio: float) -> None: + self._tower.set_freeze_ratio(ratio) + + +# --------------------------------------------------------------------------- +# ImageTower — TowerBase implementation +# --------------------------------------------------------------------------- + +class ImageTower(TowerBase, nn.Module): + """TowerBase implementation for the fundus image modality. + + Wraps ImageEncoder (backbone → pooled features). + Contributes one embedding per eye slot: [z_img]. + """ + + def __init__( + self, + *, + backbone: str, + freeze_ratio: float = 0.0, + augment: bool = True, + use_se: bool = False, + ): + nn.Module.__init__(self) + self._encoder = ImageEncoder(backbone=backbone, freeze_ratio=freeze_ratio, + augment=augment, use_se=use_se) + self._train_loader = None + + @property + def transform(self): + return self._encoder.transform + + @property + def out_dim(self) -> int: + return self._encoder.out_dim + + @property + def embed_dims(self) -> list[int]: + return [self._encoder.out_dim] + + @property + def total_epochs(self) -> int: + return 0 + + @property + def train_loader(self) -> Optional[DataLoader]: + return self._train_loader + + def set_phase(self, phase: str) -> None: + enabled = phase not in ("cd_warmup", "fused_warmup") + for p in self._encoder.parameters(): + p.requires_grad = enabled + + def embed_batch( + self, + batch: dict, + *, + device: torch.device, + slot: int = 1, + ) -> list[torch.Tensor]: + x = batch.get(f"image_{slot}") + if not torch.is_tensor(x): + raise ValueError(f"ImageTower.embed_batch: image_{slot} missing or not a tensor") + return [self._encoder(x.to(device))] + + def prepare_fold( + self, + *, + eye_train, + bilat_train, + bilat_val, + bilat_test, + image_preprocessor, + image_cache, + device, + args, + ) -> None: + from v3.classes.loader_factory import ( + build_balanced_sampler, + filter_eye_samples, + make_loader, + ) + from v3.classes.profiles import build_papila_profile + + profile_eye = build_papila_profile( + patient_col="Patient ID", label_col=args.label_col, sample_mode="eye" + ) + slots_eye = profile_eye.slot_descriptors() + use_balanced = bool(getattr(args, "balanced_sampling", False)) + _persistent = args.num_workers > 0 + loader_kw = dict( + batch_size=args.batch_size, + num_workers=args.num_workers, + image_cache=image_cache, + persistent_workers=_persistent, + ) + eye_samples = filter_eye_samples(eye_train) + sampler = build_balanced_sampler(eye_samples) if use_balanced else None + self._train_loader = make_loader( + eye_samples, slots_eye, + image_transform=self.transform, + image_preprocessor=image_preprocessor, + shuffle=True, sampler=sampler, **loader_kw, + ) diff --git a/v3/classes/loader_factory.py b/v3/classes/loader_factory.py index 6174a3e..981a440 100644 --- a/v3/classes/loader_factory.py +++ b/v3/classes/loader_factory.py @@ -4,7 +4,7 @@ from dataclasses import dataclass from typing import Any, Callable, Optional import torch -from torch.utils.data import DataLoader, WeightedRandomSampler +from torch.utils.data import DataLoader, Sampler, WeightedRandomSampler from .network_manager import LoaderBundle, PatientSplit from .slot_dataset import SlotDataset, slot_collate @@ -210,7 +210,7 @@ def make_loader( batch_size: int, shuffle: bool, num_workers: int, - sampler: Optional[WeightedRandomSampler] = None, + sampler: Optional[Sampler] = None, persistent_workers: bool = False, ) -> DataLoader: ds = SlotDataset( diff --git a/v3/classes/model_builder.py b/v3/classes/model_builder.py deleted file mode 100644 index 8bac831..0000000 --- a/v3/classes/model_builder.py +++ /dev/null @@ -1,148 +0,0 @@ -from __future__ import annotations - -from dataclasses import dataclass -from typing import Any, Callable, Optional - -import torch -from torch import nn - -from v3.classes.bridges import Bridge, VoteBridge -from v3.classes.towers import ImageTower, ClinicalTower - -from .config_builder import ConfigAssembly -from .transforms import build_transform_chain - - -@dataclass -class V2ModelBundle: - image_tower: Optional[ImageTower] - metadata_tower: Optional[ClinicalTower] - bridge: Optional[nn.Module] - classifier: Optional[nn.Module] - image_transform: Optional[Callable] - matrix_transform: Optional[Callable] - - -def build_model_bundle( - assembly: ConfigAssembly, - clinical: Any, - *, - device: Optional[torch.device] = None, - strict: bool = True, -) -> V2ModelBundle: - """ - Build torch modules and input transforms from a V2 config assembly. - """ - image_tower_spec = _pick_tower(assembly, "image") - cd_tower_spec = _pick_tower(assembly, "clinical data") - bridge_spec = _pick_bridge(assembly) - image_loader = _pick_loader(assembly, input_type="image") - - clinical_core = getattr(clinical, "clinical", clinical) - num_classes = _infer_num_classes(clinical) - - img_tower = None - if image_tower_spec is not None: - img_tower = ImageTower( - backbone=image_tower_spec.params.get("backbone", "efficientnet_b0"), - freeze_ratio=float(image_tower_spec.params.get("freeze_ratio", 0.0) or 0.0), - use_se=bool(image_tower_spec.params.get("use_se", False)), - se_reduction=int(image_tower_spec.params.get("se_reduction", 16) or 16), - se_pre_norm=bool(image_tower_spec.params.get("se_pre_norm", True)), - augment=bool(image_tower_spec.params.get("augment", True)), - geometry_dim=int(image_tower_spec.params.get("geometry_dim", 0) or 0), - ) - if device is not None: - img_tower = img_tower.to(device) - - cd_tower = None - if cd_tower_spec is not None: - cd_tower = ClinicalTower( - clinical_core, - hidden_dim=int(cd_tower_spec.params.get("hidden_dim", 128) or 128), - dropout=float(cd_tower_spec.params.get("dropout", 0.1) or 0.1), - use_se=bool(cd_tower_spec.params.get("use_se", False)), - se_reduction=int(cd_tower_spec.params.get("se_reduction", 16) or 16), - se_pre_norm=bool(cd_tower_spec.params.get("se_pre_norm", True)), - ) - if device is not None: - cd_tower = cd_tower.to(device) - - bridge = None - if bridge_spec is not None and img_tower is not None and cd_tower is not None: - if bridge_spec.method == "consensus": - bridge = VoteBridge(num_classes=num_classes) - else: - bridge = Bridge( - img_dim=img_tower.out_dim, - meta_dim=cd_tower.out_dim, - num_classes=num_classes, - fusion_dim=int(bridge_spec.params.get("fusion_dim", 256) or 256), - mode="fused", - use_se=bool(bridge_spec.params.get("use_se", True)), - se_reduction=int(bridge_spec.params.get("se_reduction", 16) or 16), - se_pre_norm=bool(bridge_spec.params.get("se_pre_norm", True)), - ) - if device is not None: - bridge = bridge.to(device) - - classifier = None - if assembly.classifiers: - classifier = nn.Identity() - if device is not None: - classifier = classifier.to(device) - - image_transform = None - if image_loader is not None and image_tower_spec is not None: - image_transform = build_transform_chain( - image_loader.transforms, - backbone_name=image_tower_spec.params.get("backbone", "efficientnet_b0"), - augment=bool(image_tower_spec.params.get("augment", True)), - strict=strict, - ) - - return V2ModelBundle( - image_tower=img_tower, - metadata_tower=cd_tower, - bridge=bridge, - classifier=classifier, - image_transform=image_transform, - matrix_transform=None, - ) - - -def _pick_tower(assembly: ConfigAssembly, tower_type: str): - matches = [tower for tower in assembly.towers.values() if tower.tower_type == tower_type] - if not matches: - return None - if len(matches) > 1: - raise ValueError(f"Multiple {tower_type} towers found; only one is supported for now.") - return matches[0] - - -def _pick_bridge(assembly: ConfigAssembly): - if not assembly.bridges: - return None - if len(assembly.bridges) > 1: - raise ValueError("Multiple bridges found; only one is supported for now.") - return next(iter(assembly.bridges.values())) - - -def _pick_loader(assembly: ConfigAssembly, input_type: str): - matches = [loader for loader in assembly.loaders.values() if loader.input_type == input_type] - if not matches: - return None - if len(matches) > 1: - raise ValueError(f"Multiple loaders with input_type={input_type!r} found.") - return matches[0] - - -def _infer_num_classes(clinical: Any) -> int: - df = getattr(clinical, "df", None) - label_col = getattr(clinical, "label_col", None) - if df is None and hasattr(clinical, "clinical"): - df = clinical.clinical.df - label_col = clinical.clinical.label_col - if df is None or label_col is None or label_col not in df.columns: - return 2 - return int(df[label_col].dropna().nunique()) diff --git a/v3/classes/models.py b/v3/classes/models.py deleted file mode 100644 index 5fd5070..0000000 --- a/v3/classes/models.py +++ /dev/null @@ -1,1074 +0,0 @@ -"""V2 model classes and training/inference helpers.""" -from __future__ import annotations - -from random import random -from typing import Optional - -import numpy as np -import torch -import torch.nn.functional as F -from torch import nn -from torch.utils.data import DataLoader - -from v3.classes.bridges import Bridge -from v3.classes.towers import ImageTower, ClinicalTower, SiameseImageTower - - -# --------------------------------------------------------------------------- -# Model classes -# --------------------------------------------------------------------------- - -class SingleEyeHT(nn.Module): - """ - ImageTower + ClinicalTower + Bridge, trained on eye-level samples. - Supports both Classic (eye-level) and Ensemble (patient-level averaging) eval. - """ - - def __init__( - self, - *, - backbone: str, - freeze_ratio: float, - augment: bool, - clinical_data, - num_classes: int, - cd_hidden_dim: int = 128, - fusion_dim: int = 256, - bridge_mode: str = "fused", - bridge_dropout: float = 0.5, - cd_dropout: float = 0.1, - se_img_tower: bool = False, - se_cd_tower: bool = False, - se_bridge: bool = False, - ): - super().__init__() - self.img_tower = ImageTower( - backbone=backbone, - freeze_ratio=freeze_ratio, - augment=augment, - use_se=se_img_tower, - ) - self.cd_tower = ClinicalTower( - clinical_data=clinical_data, - hidden_dim=cd_hidden_dim, - dropout=cd_dropout, - use_se=se_cd_tower, - ) - self.bridge = Bridge( - img_dim=self.img_tower.out_dim, - meta_dim=self.cd_tower.out_dim, - num_classes=num_classes, - fusion_dim=fusion_dim, - mode=bridge_mode, - dropout=bridge_dropout, - use_se=se_bridge, - ) - - @property - def transform(self): - return self.img_tower.transform - - def encode(self, x: torch.Tensor, meta: torch.Tensor) -> torch.Tensor: - """Return z_fused embedding (fusion_dim) without applying the classifier head.""" - img_feats = self.img_tower(x) - md_feats = self.cd_tower(meta) - return self.bridge.encode(img_feats, md_feats) - - def forward(self, x: torch.Tensor, meta: torch.Tensor) -> torch.Tensor: - img_feats = None if self.bridge.mode == "clinical_only" else self.img_tower(x) - md_feats = None if self.bridge.mode == "image_only" else self.cd_tower(meta) - out_f, _, _ = self.bridge(img_feats, md_feats) - return out_f - - -class BilateralHT(nn.Module): - """ - Bilateral mode with joint towers: - - shared eye-level towers encode OD/OS independently - - joint image and clinical data towers combine OD/OS embeddings - - standard Bridge fuses joint image + joint image + joint clinical data embeddings - """ - - def __init__( - self, - *, - backbone: str, - freeze_ratio: float, - augment: bool, - clinical_data, - num_classes: int, - cd_hidden_dim: int = 128, - fusion_dim: int = 256, - ): - super().__init__() - self.eye_img_tower = ImageTower( - backbone=backbone, - freeze_ratio=freeze_ratio, - augment=augment, - use_se=False, - ) - self.eye_cd_tower = ClinicalTower( - clinical_data=clinical_data, - hidden_dim=cd_hidden_dim, - use_se=False, - ) - img_dim = self.eye_img_tower.out_dim - md_dim = self.eye_cd_tower.out_dim - self.joint_img = nn.Sequential( - nn.Linear(2 * img_dim, fusion_dim), - nn.LayerNorm(fusion_dim), - nn.ReLU(), - nn.Dropout(0.3), - nn.Linear(fusion_dim, img_dim), - ) - self.joint_md = nn.Sequential( - nn.Linear(2 * md_dim, fusion_dim), - nn.LayerNorm(fusion_dim), - nn.ReLU(), - nn.Dropout(0.3), - nn.Linear(fusion_dim, md_dim), - ) - self.bridge = Bridge( - img_dim=img_dim, - meta_dim=md_dim, - num_classes=num_classes, - fusion_dim=fusion_dim, - mode="fused", - use_se=False, - ) - # Auxiliary heads for tower warmup / BCD tower steps. - self.aux_img = nn.Linear(img_dim, num_classes) - self.aux_md = nn.Linear(md_dim, num_classes) - - @property - def transform(self): - return self.eye_img_tower.transform - - def encode_joint( - self, - x_od: torch.Tensor, - meta_od: torch.Tensor, - x_os: torch.Tensor, - meta_os: torch.Tensor, - ) -> tuple[torch.Tensor, torch.Tensor]: - img_od = self.eye_img_tower(x_od) - md_od = self.eye_cd_tower(meta_od) - img_os = self.eye_img_tower(x_os) - md_os = self.eye_cd_tower(meta_os) - joint_img = self.joint_img(torch.cat([img_od, img_os], dim=1)) - joint_md = self.joint_md(torch.cat([md_od, md_os], dim=1)) - return joint_img, joint_md - - def forward( - self, - x_od: torch.Tensor, - meta_od: torch.Tensor, - x_os: torch.Tensor, - meta_os: torch.Tensor, - ) -> torch.Tensor: - joint_img, joint_md = self.encode_joint(x_od, meta_od, x_os, meta_os) - out_f, _, _ = self.bridge(joint_img, joint_md) - return out_f - - -class SiameseHT(nn.Module): - """ - Bilateral model using a shared-weight SiameseImageTower (mean+delta). - - Both eyes pass through the same backbone; features are combined as - cat([mean(f_od, f_os), f_od - f_os]) giving the model both a shared - bilateral representation and an asymmetry signal. - - Uses the same bilateral loader and training loop as BilateralHT. - The bridge operates in image_only mode (no clinical data in phase 4). - """ - - def __init__( - self, - *, - backbone: str, - freeze_ratio: float, - augment: bool, - num_classes: int, - fusion_dim: int = 256, - bridge_mode: str = "image_only", - ): - super().__init__() - self.img_tower = SiameseImageTower( - backbone=backbone, - freeze_ratio=freeze_ratio, - augment=augment, - use_se=False, - ) - img_dim = self.img_tower.out_dim # 2 * backbone_out_dim - self.bridge = Bridge( - img_dim=img_dim, - meta_dim=1, # dummy — not used in image_only mode - num_classes=num_classes, - fusion_dim=fusion_dim, - mode="image_only", - use_se=False, - ) - self.aux_img = nn.Linear(img_dim, num_classes) - - @property - def transform(self): - return self.img_tower.transform - - def encode(self, x_od: torch.Tensor, x_os: torch.Tensor) -> torch.Tensor: - return self.img_tower(x_od, x_os) - - def forward(self, x_od: torch.Tensor, x_os: torch.Tensor) -> torch.Tensor: - feats = self.encode(x_od, x_os) - out_f, _, _ = self.bridge(feats, None) - return out_f - - -class FusedEnsembleHT(nn.Module): - """ - SingleEyeHT base with a per-eye attention scorer for bilateral fusion. - - The base model is trained eye-level (identical to ensemble mode). - After base training completes, the base is frozen and only the - eye_scorer is trained on bilateral (patient-level) samples. - - At inference, eye_scorer is applied independently to each eye's logit - vector to produce a scalar attention score. Softmax over the two scores - gives attention weights; the final logit is a weighted sum: - - score_od = eye_scorer(logit_od) # [B, 1] - score_os = eye_scorer(logit_os) # [B, 1] - alpha = softmax([score_od, score_os]) # [B, 2], sums to 1 - out = alpha[:,0:1]*logit_od + alpha[:,1:2]*logit_os - - Because eye_scorer is applied to each eye with the same weights, the - mechanism is permutation-equivariant — there is no left/right positional - bias. Through training on bilateral labels the scorer learns to give high - scores to logits that point strongly toward the GC class, creating the - desired asymmetry: a confidently GC eye dominates the patient prediction - more than a comparably confident healthy eye would. - """ - - def __init__(self, base: SingleEyeHT, num_classes: int): - super().__init__() - self.base = base - # Applied independently to each eye's logit → scalar attention score. - # Learns the GC-direction in logit space from bilateral labels. - self.eye_scorer = nn.Linear(num_classes, 1, bias=True) - - @property - def head(self) -> nn.Module: - return self.eye_scorer - - def forward( - self, - x_od: torch.Tensor, - meta_od: torch.Tensor, - x_os: torch.Tensor, - meta_os: torch.Tensor, - ) -> torch.Tensor: - logit_od = self.base(x_od, meta_od) # [B, C] - logit_os = self.base(x_os, meta_os) # [B, C] - scores = torch.cat([self.eye_scorer(logit_od), - self.eye_scorer(logit_os)], dim=1) # [B, 2] - alpha = torch.softmax(scores, dim=1) # [B, 2] - return alpha[:, 0:1] * logit_od + alpha[:, 1:2] * logit_os # [B, C] - - -class LogitMLPEnsembleHT(nn.Module): - """ - MLP head trained on concatenated per-eye logits. - - Both eyes pass through the frozen base independently, producing per-eye - logit vectors. These are concatenated and fed through a small MLP: - - cat([logit_od, logit_os]) [B, 2C] - → Linear(2C, hidden) → ReLU → Dropout → Linear(hidden, C) - - Permutation-variant by design: the model can learn left/right asymmetries - directly from the concatenated pair, at the cost of needing a consistent - OD-first input ordering. - """ - - def __init__(self, base: SingleEyeHT, num_classes: int, hidden: int = 64): - super().__init__() - self.base = base - self.head = nn.Sequential( - nn.Linear(2 * num_classes, hidden), - nn.ReLU(), - nn.Dropout(0.3), - nn.Linear(hidden, num_classes), - ) - - def forward( - self, - x_od: torch.Tensor, meta_od: torch.Tensor, - x_os: torch.Tensor, meta_os: torch.Tensor, - ) -> torch.Tensor: - logit_od = self.base(x_od, meta_od) - logit_os = self.base(x_os, meta_os) - return self.head(torch.cat([logit_od, logit_os], dim=1)) - - -class EmbeddingMLPEnsembleHT(nn.Module): - """ - MLP head trained on concatenated per-eye z_fused embeddings. - - Both eyes pass through the frozen base independently, and their bridge - embeddings (pre-classifier, shape [B, fusion_dim]) are concatenated and - fed through an MLP: - - cat([z_od, z_os]) [B, 2 * fusion_dim] - → Linear(2*fusion_dim, hidden) → ReLU → Dropout → Linear(hidden, C) - - Richer than logit-level: the head sees pre-softmax feature vectors rather - than the compressed C-dimensional output, giving it more signal to work - with when fusion_dim >> C. - """ - - def __init__(self, base: SingleEyeHT, num_classes: int, hidden: int = 256): - super().__init__() - self.base = base - fusion_dim = base.bridge.W_img.out_features - self.head = nn.Sequential( - nn.Linear(2 * fusion_dim, hidden), - nn.ReLU(), - nn.Dropout(0.3), - nn.Linear(hidden, num_classes), - ) - - def forward( - self, - x_od: torch.Tensor, meta_od: torch.Tensor, - x_os: torch.Tensor, meta_os: torch.Tensor, - ) -> torch.Tensor: - z_od = self.base.encode(x_od, meta_od) - z_os = self.base.encode(x_os, meta_os) - return self.head(torch.cat([z_od, z_os], dim=1)) - - -# --------------------------------------------------------------------------- -# Phase control -# --------------------------------------------------------------------------- - -def _set_requires_grad(module: nn.Module, enabled: bool) -> None: - for p in module.parameters(): - p.requires_grad = enabled - - -def _set_single_phase(model: SingleEyeHT, phase: str) -> None: - bridge_mode = model.bridge.mode - # Ablation modes have no fusion bridge; fused_warmup is meaningless — treat as tower_warmup - if bridge_mode in ("image_only", "clinical_only") and phase == "fused_warmup": - phase = "tower_warmup" - if phase == "cd_warmup": - _set_requires_grad(model.img_tower, False) - _set_requires_grad(model.cd_tower, True) - _set_requires_grad(model.bridge.classifier_img, False) - _set_requires_grad(model.bridge.classifier_cd, True) - _set_requires_grad(model.bridge.W_img, False) - _set_requires_grad(model.bridge.W_md, False) - _set_requires_grad(model.bridge.classifier_fused, False) - return - if phase == "tower_warmup": - _set_requires_grad(model.img_tower, bridge_mode != "clinical_only") - _set_requires_grad(model.cd_tower, bridge_mode != "image_only") - _set_requires_grad(model.bridge.classifier_img, bridge_mode != "clinical_only") - _set_requires_grad(model.bridge.classifier_cd, bridge_mode != "image_only") - _set_requires_grad(model.bridge.W_img, False) - _set_requires_grad(model.bridge.W_md, False) - _set_requires_grad(model.bridge.classifier_fused, False) - return - if phase == "fused_warmup": - _set_requires_grad(model.img_tower, False) - _set_requires_grad(model.cd_tower, False) - _set_requires_grad(model.bridge.classifier_img, False) - _set_requires_grad(model.bridge.classifier_cd, False) - _set_requires_grad(model.bridge.W_img, True) - _set_requires_grad(model.bridge.W_md, True) - _set_requires_grad(model.bridge.classifier_fused, True) - return - _set_requires_grad(model, True) - - -def _set_bilateral_phase(model: BilateralHT, phase: str) -> None: - if phase == "tower_warmup": - _set_requires_grad(model.eye_img_tower, True) - _set_requires_grad(model.eye_cd_tower, True) - _set_requires_grad(model.joint_img, True) - _set_requires_grad(model.joint_md, True) - _set_requires_grad(model.aux_img, True) - _set_requires_grad(model.aux_md, True) - _set_requires_grad(model.bridge, False) - return - if phase == "fused_warmup": - _set_requires_grad(model.eye_img_tower, False) - _set_requires_grad(model.eye_cd_tower, False) - _set_requires_grad(model.joint_img, False) - _set_requires_grad(model.joint_md, False) - _set_requires_grad(model.aux_img, False) - _set_requires_grad(model.aux_md, False) - _set_requires_grad(model.bridge, True) - return - _set_requires_grad(model, True) - - -# --------------------------------------------------------------------------- -# Training helpers -# --------------------------------------------------------------------------- - -def train_single_epoch( - model: SingleEyeHT, - loader: DataLoader, - opt, - device: torch.device, - *, - phase: str, - bcd_prob: float = 0.5, - tower_loss_mode: str = "bcd", -) -> tuple[float, float]: - model.train() - _set_single_phase(model, phase) - total_loss = total_correct = total_n = 0 - for batch in loader: - x = batch.get("image_1") - m = batch.get("matrix_1") - y = batch.get("label_1") - if phase == "cd_warmup": - if not torch.is_tensor(m): - continue - m = m.to(device) - y = _to_label_tensor(y, device) - md_feats = model.cd_tower(m) - logits = model.bridge.classifier_cd(md_feats) - loss = F.cross_entropy(logits, y) - opt.zero_grad(); loss.backward(); opt.step() - bs = y.shape[0] - total_loss += float(loss.item()) * bs - total_correct += int((logits.argmax(1) == y).sum()) - total_n += bs - continue - if not torch.is_tensor(x) or not torch.is_tensor(m): - continue - x = x.to(device) - m = m.to(device) - y = _to_label_tensor(y, device) - bridge_mode = model.bridge.mode - - img_feats = None if bridge_mode == "clinical_only" else model.img_tower(x) - md_feats = None if bridge_mode == "image_only" else model.cd_tower(m) - - if phase == "tower_warmup": - if bridge_mode == "clinical_only": - logits = model.bridge.classifier_cd(md_feats) - loss = F.cross_entropy(logits, y) - elif bridge_mode == "image_only": - logits = model.bridge.classifier_img(img_feats) - loss = F.cross_entropy(logits, y) - else: - logits_i = model.bridge.classifier_img(img_feats) - logits_m = model.bridge.classifier_cd(md_feats) - loss = 0.5 * (F.cross_entropy(logits_i, y) + F.cross_entropy(logits_m, y)) - logits = 0.5 * (F.softmax(logits_i, dim=1) + F.softmax(logits_m, dim=1)) - elif phase == "fused_warmup": - logits, _, _ = model.bridge(img_feats, md_feats) - loss = F.cross_entropy(logits, y) - else: - if bridge_mode == "clinical_only": - logits = model.bridge.classifier_cd(md_feats) - loss = F.cross_entropy(logits, y) - elif bridge_mode == "image_only": - logits = model.bridge.classifier_img(img_feats) - loss = F.cross_entropy(logits, y) - elif tower_loss_mode == "all": - loss_i = F.cross_entropy(model.bridge.classifier_img(img_feats), y) - loss_m = F.cross_entropy(model.bridge.classifier_cd(md_feats), y) - logits, _, _ = model.bridge(img_feats, md_feats) - loss = F.cross_entropy(logits, y) + loss_i + loss_m - elif random() < bcd_prob: - if random() < 0.5: - logits = model.bridge.classifier_img(img_feats) - else: - logits = model.bridge.classifier_cd(md_feats) - loss = F.cross_entropy(logits, y) - else: - logits, _, _ = model.bridge(img_feats, md_feats) - loss = F.cross_entropy(logits, y) - - opt.zero_grad() - loss.backward() - opt.step() - bs = y.shape[0] - total_loss += float(loss.item()) * bs - total_correct += int((logits.argmax(1) == y).sum()) - total_n += bs - return ( - total_loss / total_n if total_n else float("nan"), - total_correct / total_n if total_n else float("nan"), - ) - - -def train_bilateral_epoch( - model: BilateralHT, - loader: DataLoader, - opt, - device: torch.device, - *, - phase: str, - bcd_prob: float = 0.5, - tower_loss_mode: str = "bcd", -) -> tuple[float, float]: - model.train() - _set_bilateral_phase(model, phase) - total_loss = total_correct = total_n = 0 - for batch in loader: - x1 = batch.get("image_1") - m1 = batch.get("matrix_1") - x2 = batch.get("image_2") - m2 = batch.get("matrix_2") - y = batch.get("label_1") - if not (torch.is_tensor(x1) and torch.is_tensor(m1) and torch.is_tensor(x2) and torch.is_tensor(m2)): - continue - x1 = x1.to(device); m1 = m1.to(device) - x2 = x2.to(device); m2 = m2.to(device) - y = _to_label_tensor(y, device) - joint_img, joint_md = model.encode_joint(x1, m1, x2, m2) - - if phase == "tower_warmup": - logits_i = model.aux_img(joint_img) - logits_m = model.aux_md(joint_md) - loss = 0.5 * (F.cross_entropy(logits_i, y) + F.cross_entropy(logits_m, y)) - logits = 0.5 * (F.softmax(logits_i, dim=1) + F.softmax(logits_m, dim=1)) - elif phase == "fused_warmup": - logits, _, _ = model.bridge(joint_img, joint_md) - loss = F.cross_entropy(logits, y) - else: - if tower_loss_mode == "all": - loss_i = F.cross_entropy(model.aux_img(joint_img), y) - loss_m = F.cross_entropy(model.aux_md(joint_md), y) - logits, _, _ = model.bridge(joint_img, joint_md) - loss = F.cross_entropy(logits, y) + loss_i + loss_m - elif random() < bcd_prob: - if random() < 0.5: - logits = model.aux_img(joint_img) - else: - logits = model.aux_md(joint_md) - loss = F.cross_entropy(logits, y) - else: - logits, _, _ = model.bridge(joint_img, joint_md) - loss = F.cross_entropy(logits, y) - - opt.zero_grad() - loss.backward() - opt.step() - bs = y.shape[0] - total_loss += float(loss.item()) * bs - total_correct += int((logits.argmax(1) == y).sum()) - total_n += bs - return ( - total_loss / total_n if total_n else float("nan"), - total_correct / total_n if total_n else float("nan"), - ) - - -def train_siamese_epoch( - model: SiameseHT, - loader: DataLoader, - opt, - device: torch.device, - *, - bcd_prob: float = 0.5, - tower_loss_mode: str = "bcd", -) -> tuple[float, float]: - """Train one epoch of SiameseHT on bilateral (patient-level) samples.""" - model.train() - total_loss = total_correct = total_n = 0 - for batch in loader: - x1 = batch.get("image_1") - x2 = batch.get("image_2") - y = batch.get("label_1") - if not (torch.is_tensor(x1) and torch.is_tensor(x2)): - continue - x1 = x1.to(device); x2 = x2.to(device) - y = _to_label_tensor(y, device) - feats = model.encode(x1, x2) - - if tower_loss_mode == "all": - logits_i = model.aux_img(feats) - out_f, _, _ = model.bridge(feats, None) - loss = F.cross_entropy(out_f, y) + F.cross_entropy(logits_i, y) - logits = out_f - elif random() < bcd_prob: - logits = model.aux_img(feats) - loss = F.cross_entropy(logits, y) - else: - logits, _, _ = model.bridge(feats, None) - loss = F.cross_entropy(logits, y) - - opt.zero_grad() - loss.backward() - opt.step() - bs = y.shape[0] - total_loss += float(loss.item()) * bs - total_correct += int((logits.argmax(1) == y).sum()) - total_n += bs - return ( - total_loss / total_n if total_n else float("nan"), - total_correct / total_n if total_n else float("nan"), - ) - - -def collect_probs_siamese( - model: SiameseHT, - loader: DataLoader, - device: torch.device, -) -> tuple[np.ndarray, np.ndarray]: - """Patient-level probs from a bilateral loader using SiameseHT.""" - model.eval() - y_c, p_c = [], [] - with torch.no_grad(): - for batch in loader: - x1 = batch.get("image_1") - x2 = batch.get("image_2") - y = batch.get("label_1") - if not (torch.is_tensor(x1) and torch.is_tensor(x2)): - continue - x1 = x1.to(device); x2 = x2.to(device) - feats = model.encode(x1, x2) - logits, _, _ = model.bridge(feats, None) - probs = torch.softmax(logits, dim=1) - y_c.append(np.array(y) if not torch.is_tensor(y) else y.cpu().numpy()) - p_c.append(probs.cpu().numpy()) - return ( - np.concatenate(y_c, axis=0), - np.concatenate(p_c, axis=0), - ) - - -def train_fusion_epoch( - model, # FusedEnsembleHT | LogitMLPEnsembleHT | EmbeddingMLPEnsembleHT - loader: DataLoader, - opt, - device: torch.device, -) -> tuple[float, float]: - """Train only the fusion head; the base SingleEyeHT is frozen in eval mode.""" - model.base.eval() - model.head.train() - total_loss = total_correct = total_n = 0 - for batch in loader: - x1 = batch.get("image_1"); m1 = batch.get("matrix_1") - x2 = batch.get("image_2"); m2 = batch.get("matrix_2") - y = batch.get("label_1") - if not (torch.is_tensor(x1) and torch.is_tensor(m1) and - torch.is_tensor(x2) and torch.is_tensor(m2)): - continue - y_t = _to_label_tensor(y, device) - out = model(x1.to(device), m1.to(device), x2.to(device), m2.to(device)) - loss = F.cross_entropy(out, y_t) - opt.zero_grad() - loss.backward() - opt.step() - bs = y_t.shape[0] - total_loss += float(loss.item()) * bs - total_correct += int((out.argmax(1) == y_t).sum()) - total_n += bs - return ( - total_loss / total_n if total_n else float("nan"), - total_correct / total_n if total_n else float("nan"), - ) - - -# --------------------------------------------------------------------------- -# Inference helpers -# --------------------------------------------------------------------------- - -def _to_label_tensor(labels, device: torch.device) -> torch.Tensor: - if torch.is_tensor(labels): - return labels.to(device=device, dtype=torch.long) - return torch.as_tensor(labels, dtype=torch.long, device=device) - - -def collect_probs_classic( - model: SingleEyeHT, - loader: DataLoader, - device: torch.device, -) -> tuple[np.ndarray, np.ndarray]: - """ - Classic eye-level eval using the bilateral val loader. - OD and OS are treated as independent samples (both contribute to the - arrays with the same patient label). Returns (y_true [2N], probs [2N, C]). - """ - model.eval() - y_chunks, p_chunks = [], [] - with torch.no_grad(): - for batch in loader: - x1 = batch.get("image_1"); m1 = batch.get("matrix_1") - x2 = batch.get("image_2"); m2 = batch.get("matrix_2") - y = batch.get("label_1") - if not (torch.is_tensor(x1) and torch.is_tensor(m1) and torch.is_tensor(x2) and torch.is_tensor(m2)): - continue - y_t = _to_label_tensor(y, device) - p_od = F.softmax(model(x1.to(device), m1.to(device)), dim=1) - p_os = F.softmax(model(x2.to(device), m2.to(device)), dim=1) - y_np = y_t.cpu().numpy() - y_chunks += [y_np, y_np] - p_chunks += [p_od.cpu().numpy(), p_os.cpu().numpy()] - if not y_chunks: - return np.array([], dtype=np.int64), np.zeros((0, 0), dtype=np.float32) - return np.concatenate(y_chunks), np.concatenate(p_chunks, axis=0) - - -def collect_probs_ensemble_pereye( - model: "SingleEyeHT", - loader: DataLoader, - device: torch.device, - *, - return_ids: bool = False, -): - """ - Per-patient, per-eye probs for all 3 heads from a bilateral loader (ensemble mode). - - OD corresponds to image_1/matrix_1; OS to image_2/matrix_2. - Arrays are in patient order (not interleaved at sample level). - - Returns: - (y, pf_od, pi_od, pm_od, pf_os, pi_os, pm_os) - or, when return_ids=True: - (y, pf_od, pi_od, pm_od, pf_os, pi_os, pm_os, patient_ids) - - Patient-level averaged ensemble probs can be recovered as: - p_en = 0.5 * (pf_od + pf_os) - """ - model.eval() - y_chunks: list = [] - pf_od_c, pi_od_c, pm_od_c = [], [], [] - pf_os_c, pi_os_c, pm_os_c = [], [], [] - id_chunks: list[str] = [] - - with torch.no_grad(): - for batch in loader: - x1 = batch.get("image_1"); m1 = batch.get("matrix_1") - x2 = batch.get("image_2"); m2 = batch.get("matrix_2") - y = batch.get("label_1") - if not (torch.is_tensor(x1) and torch.is_tensor(m1) and - torch.is_tensor(x2) and torch.is_tensor(m2)): - continue - y_t = _to_label_tensor(y, device) - - def _fwd(x, m): - img_feats = None if model.bridge.mode == "clinical_only" else model.img_tower(x.to(device)) - md_feats = None if model.bridge.mode == "image_only" else model.cd_tower(m.to(device)) - out_f, out_i, out_m = model.bridge(img_feats, md_feats) - pf = F.softmax(out_f, dim=1) - pi = F.softmax(out_i, dim=1) if out_i is not None else pf - pm = F.softmax(out_m, dim=1) if out_m is not None else pf - return pf, pi, pm - - pf_od, pi_od, pm_od = _fwd(x1, m1) - pf_os, pi_os, pm_os = _fwd(x2, m2) - - y_chunks.append(y_t.cpu().numpy()) - pf_od_c.append(pf_od.cpu().numpy()); pi_od_c.append(pi_od.cpu().numpy()); pm_od_c.append(pm_od.cpu().numpy()) - pf_os_c.append(pf_os.cpu().numpy()); pi_os_c.append(pi_os.cpu().numpy()); pm_os_c.append(pm_os.cpu().numpy()) - - if return_ids: - ids = batch.get("id_1", [""] * len(y_t)) - if torch.is_tensor(ids): - ids = ids.tolist() - id_chunks.extend([str(i) for i in ids]) - - if not y_chunks: - z = np.zeros((0, 0), dtype=np.float32) - empty_i = np.array([], dtype=np.int64) - base = (empty_i, z, z, z, z, z, z) - return base + (np.array([], dtype=object),) if return_ids else base - - y = np.concatenate(y_chunks) - pf_od = np.concatenate(pf_od_c, axis=0); pi_od = np.concatenate(pi_od_c, axis=0); pm_od = np.concatenate(pm_od_c, axis=0) - pf_os = np.concatenate(pf_os_c, axis=0); pi_os = np.concatenate(pi_os_c, axis=0); pm_os = np.concatenate(pm_os_c, axis=0) - if return_ids: - return y, pf_od, pi_od, pm_od, pf_os, pi_os, pm_os, np.array(id_chunks, dtype=object) - return y, pf_od, pi_od, pm_od, pf_os, pi_os, pm_os - - -def collect_probs_ensemble( - model: SingleEyeHT, - loader: DataLoader, - device: torch.device, -) -> tuple[np.ndarray, np.ndarray]: - """ - Patient-level ensemble eval: average OD and OS softmax probabilities. - Returns (y_true [N], probs [N, C]). - """ - model.eval() - y_chunks, p_chunks = [], [] - with torch.no_grad(): - for batch in loader: - x1 = batch.get("image_1"); m1 = batch.get("matrix_1") - x2 = batch.get("image_2"); m2 = batch.get("matrix_2") - y = batch.get("label_1") - if not (torch.is_tensor(x1) and torch.is_tensor(m1) and torch.is_tensor(x2) and torch.is_tensor(m2)): - continue - y_t = _to_label_tensor(y, device) - p_od = F.softmax(model(x1.to(device), m1.to(device)), dim=1) - p_os = F.softmax(model(x2.to(device), m2.to(device)), dim=1) - p = 0.5 * (p_od + p_os) - y_chunks.append(y_t.cpu().numpy()) - p_chunks.append(p.cpu().numpy()) - if not y_chunks: - return np.array([], dtype=np.int64), np.zeros((0, 0), dtype=np.float32) - return np.concatenate(y_chunks), np.concatenate(p_chunks, axis=0) - - -def collect_probs_bilateral( - model: BilateralHT, - loader: DataLoader, - device: torch.device, -) -> tuple[np.ndarray, np.ndarray]: - """Patient-level bilateral eval. Returns (y_true [N], probs [N, C]).""" - model.eval() - y_chunks, p_chunks = [], [] - with torch.no_grad(): - for batch in loader: - x1 = batch.get("image_1"); m1 = batch.get("matrix_1") - x2 = batch.get("image_2"); m2 = batch.get("matrix_2") - y = batch.get("label_1") - if not (torch.is_tensor(x1) and torch.is_tensor(m1) and torch.is_tensor(x2) and torch.is_tensor(m2)): - continue - y_t = _to_label_tensor(y, device) - p = F.softmax(model(x1.to(device), m1.to(device), x2.to(device), m2.to(device)), dim=1) - y_chunks.append(y_t.cpu().numpy()) - p_chunks.append(p.cpu().numpy()) - if not y_chunks: - return np.array([], dtype=np.int64), np.zeros((0, 0), dtype=np.float32) - return np.concatenate(y_chunks), np.concatenate(p_chunks, axis=0) - - -def collect_probs_fused( - model: FusedEnsembleHT, - loader: DataLoader, - device: torch.device, -) -> tuple[np.ndarray, np.ndarray]: - """Patient-level fused-head eval. Returns (y_true [N], probs [N, C]).""" - model.eval() - y_chunks, p_chunks = [], [] - with torch.no_grad(): - for batch in loader: - x1 = batch.get("image_1"); m1 = batch.get("matrix_1") - x2 = batch.get("image_2"); m2 = batch.get("matrix_2") - y = batch.get("label_1") - if not (torch.is_tensor(x1) and torch.is_tensor(m1) and - torch.is_tensor(x2) and torch.is_tensor(m2)): - continue - y_t = _to_label_tensor(y, device) - p = F.softmax(model(x1.to(device), m1.to(device), - x2.to(device), m2.to(device)), dim=1) - y_chunks.append(y_t.cpu().numpy()) - p_chunks.append(p.cpu().numpy()) - if not y_chunks: - return np.array([], dtype=np.int64), np.zeros((0, 0), dtype=np.float32) - return np.concatenate(y_chunks), np.concatenate(p_chunks, axis=0) - - -def collect_probs_single_components( - model: SingleEyeHT, - loader: DataLoader, - device: torch.device, - *, - aggregate_patient: bool, - return_logits: bool = False, -): - """ - Collect fused/img/md probabilities (and optionally raw logits) for SingleEyeHT. - - aggregate_patient=False: eye-level (OD/OS as independent samples) - - aggregate_patient=True : patient-level (average OD/OS per head) - - return_logits=False: returns (y, probs_f, probs_i, probs_m) - - return_logits=True: returns (y, probs_f, probs_i, probs_m, - logits_f, logits_i, logits_m) - Note: logits are averaged across eyes when aggregate_patient=True, - which is equivalent to averaging in logit space (before softmax). - """ - model.eval() - y_chunks = [] - pf_chunks, pi_chunks, pm_chunks = [], [], [] - lf_chunks, li_chunks, lm_chunks = [], [], [] - with torch.no_grad(): - for batch in loader: - x1 = batch.get("image_1"); m1 = batch.get("matrix_1") - x2 = batch.get("image_2"); m2 = batch.get("matrix_2") - y = batch.get("label_1") - if not (torch.is_tensor(x1) and torch.is_tensor(m1) and torch.is_tensor(x2) and torch.is_tensor(m2)): - continue - y_t = _to_label_tensor(y, device) - - def _per_eye(x, m): - img_feats = None if model.bridge.mode == "clinical_only" else model.img_tower(x.to(device)) - md_feats = None if model.bridge.mode == "image_only" else model.cd_tower(m.to(device)) - out_f, out_i, out_m = model.bridge(img_feats, md_feats) - pf = F.softmax(out_f, dim=1) - pi = F.softmax(out_i, dim=1) if out_i is not None else pf - pm = F.softmax(out_m, dim=1) if out_m is not None else pf - lf = out_f - li = out_i if out_i is not None else out_f - lm = out_m if out_m is not None else out_f - return pf, pi, pm, lf, li, lm - - pf_od, pi_od, pm_od, lf_od, li_od, lm_od = _per_eye(x1, m1) - pf_os, pi_os, pm_os, lf_os, li_os, lm_os = _per_eye(x2, m2) - - if aggregate_patient: - y_chunks.append(y_t.cpu().numpy()) - pf_chunks.append((0.5 * (pf_od + pf_os)).cpu().numpy()) - pi_chunks.append((0.5 * (pi_od + pi_os)).cpu().numpy()) - pm_chunks.append((0.5 * (pm_od + pm_os)).cpu().numpy()) - lf_chunks.append((0.5 * (lf_od + lf_os)).cpu().numpy()) - li_chunks.append((0.5 * (li_od + li_os)).cpu().numpy()) - lm_chunks.append((0.5 * (lm_od + lm_os)).cpu().numpy()) - else: - y_np = y_t.cpu().numpy() - y_chunks += [y_np, y_np] - pf_chunks += [pf_od.cpu().numpy(), pf_os.cpu().numpy()] - pi_chunks += [pi_od.cpu().numpy(), pi_os.cpu().numpy()] - pm_chunks += [pm_od.cpu().numpy(), pm_os.cpu().numpy()] - lf_chunks += [lf_od.cpu().numpy(), lf_os.cpu().numpy()] - li_chunks += [li_od.cpu().numpy(), li_os.cpu().numpy()] - lm_chunks += [lm_od.cpu().numpy(), lm_os.cpu().numpy()] - - if not y_chunks: - z = np.zeros((0, 0), dtype=np.float32) - if return_logits: - return np.array([], dtype=np.int64), z, z, z, z, z, z - return np.array([], dtype=np.int64), z, z, z - - y = np.concatenate(y_chunks) - pf = np.concatenate(pf_chunks, axis=0) - pi = np.concatenate(pi_chunks, axis=0) - pm = np.concatenate(pm_chunks, axis=0) - if return_logits: - lf = np.concatenate(lf_chunks, axis=0) - li = np.concatenate(li_chunks, axis=0) - lm = np.concatenate(lm_chunks, axis=0) - return y, pf, pi, pm, lf, li, lm - return y, pf, pi, pm - - -def collect_probs_eye_level( - model: "SingleEyeHT", - loader: DataLoader, - device: torch.device, - *, - return_ids: bool = False, -): - """ - Collect fused/img/md probabilities from a single-eye loader (image_1/matrix_1 only). - Used for eval-mode passes over the training set. - - Returns (y, probs_f, probs_i, probs_m) or, when return_ids=True, - (y, probs_f, probs_i, probs_m, sample_ids) where sample_ids is an - array of strings like "2OD", "4OS". - """ - model.eval() - y_chunks, pf_chunks, pi_chunks, pm_chunks, id_chunks = [], [], [], [], [] - with torch.no_grad(): - for batch in loader: - x = batch.get("image_1") - m = batch.get("matrix_1") - y = batch.get("label_1") - if not (torch.is_tensor(x) and torch.is_tensor(m)): - continue - y_t = _to_label_tensor(y, device) - img_feats = None if model.bridge.mode == "clinical_only" else model.img_tower(x.to(device)) - md_feats = None if model.bridge.mode == "image_only" else model.cd_tower(m.to(device)) - out_f, out_i, out_m = model.bridge(img_feats, md_feats) - pf = F.softmax(out_f, dim=1) - pi = F.softmax(out_i, dim=1) if out_i is not None else pf - pm = F.softmax(out_m, dim=1) if out_m is not None else pf - y_chunks.append(y_t.cpu().numpy()) - pf_chunks.append(pf.cpu().numpy()) - pi_chunks.append(pi.cpu().numpy()) - pm_chunks.append(pm.cpu().numpy()) - if return_ids: - ids = batch.get("id_1", [""] * len(y_t)) - eyes = batch.get("eye_id_1", [""] * len(y_t)) - # ids/eyes may be tensors (int) or lists of strings - if torch.is_tensor(ids): - ids = ids.tolist() - if torch.is_tensor(eyes): - eyes = eyes.tolist() - id_chunks.extend( - [f"{pid}{eye}" for pid, eye in zip(ids, eyes)] - ) - - if not y_chunks: - z = np.zeros((0, 0), dtype=np.float32) - empty_ids = np.array([], dtype=object) - if return_ids: - return np.array([], dtype=np.int64), z, z, z, empty_ids - return np.array([], dtype=np.int64), z, z, z - - y = np.concatenate(y_chunks) - pf = np.concatenate(pf_chunks, axis=0) - pi = np.concatenate(pi_chunks, axis=0) - pm = np.concatenate(pm_chunks, axis=0) - if return_ids: - return y, pf, pi, pm, np.array(id_chunks, dtype=object) - return y, pf, pi, pm - - -def collect_probs_bilateral_components( - model: BilateralHT, - loader: DataLoader, - device: torch.device, -) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]: - """Collect fused/img/md probabilities for bilateral joint-tower model.""" - model.eval() - y_chunks = [] - pf_chunks, pi_chunks, pm_chunks = [], [], [] - with torch.no_grad(): - for batch in loader: - x1 = batch.get("image_1"); m1 = batch.get("matrix_1") - x2 = batch.get("image_2"); m2 = batch.get("matrix_2") - y = batch.get("label_1") - if not (torch.is_tensor(x1) and torch.is_tensor(m1) and torch.is_tensor(x2) and torch.is_tensor(m2)): - continue - y_t = _to_label_tensor(y, device) - joint_img, joint_md = model.encode_joint( - x1.to(device), m1.to(device), x2.to(device), m2.to(device) - ) - out_f, _, _ = model.bridge(joint_img, joint_md) - out_i = model.aux_img(joint_img) - out_m = model.aux_md(joint_md) - y_chunks.append(y_t.cpu().numpy()) - pf_chunks.append(F.softmax(out_f, dim=1).cpu().numpy()) - pi_chunks.append(F.softmax(out_i, dim=1).cpu().numpy()) - pm_chunks.append(F.softmax(out_m, dim=1).cpu().numpy()) - if not y_chunks: - z = np.zeros((0, 0), dtype=np.float32) - return np.array([], dtype=np.int64), z, z, z - return ( - np.concatenate(y_chunks), - np.concatenate(pf_chunks, axis=0), - np.concatenate(pi_chunks, axis=0), - np.concatenate(pm_chunks, axis=0), - ) - - -# --------------------------------------------------------------------------- -# V2ModeComparisonOps — thin class wrapper kept for external import compat -# --------------------------------------------------------------------------- - -class V2ModeComparisonOps: - """Namespace wrapper kept for backward-compatibility imports.""" - - _set_requires_grad = staticmethod(_set_requires_grad) - _set_single_phase = staticmethod(_set_single_phase) - _set_bilateral_phase = staticmethod(_set_bilateral_phase) - train_single_epoch = staticmethod(train_single_epoch) - train_bilateral_epoch = staticmethod(train_bilateral_epoch) - collect_probs_classic = staticmethod(collect_probs_classic) - collect_probs_ensemble = staticmethod(collect_probs_ensemble) - collect_probs_bilateral = staticmethod(collect_probs_bilateral) - - @staticmethod - def _to_label_tensor(labels, device: torch.device) -> torch.Tensor: - return _to_label_tensor(labels, device) diff --git a/v3/classes/towerbase.py b/v3/classes/towerbase.py new file mode 100644 index 0000000..5e4d8a8 --- /dev/null +++ b/v3/classes/towerbase.py @@ -0,0 +1,455 @@ +"""towerbase — unified TowerBase ABC, backbone factory, and modular training utilities. + +This module is the structural backbone of the HyperTower v3 architecture. +It provides the abstract tower interface plus the training/eval helpers that +operate on any list of TowerBase instances. + +Design principles +----------------- +* No concrete tower classes are defined here (ImageEncoder, ClinicalEncoder, etc. + live in their respective tower files). +* No imports from any tower file — this module is self-contained with respect to + the tower layer. Tower files import from here; this file does not import from them. +* Removing any tower file leaves this module fully intact. +* Duck-typing via optional TowerBase methods (e.g. cd_warmup_embedding) replaces + isinstance checks so new tower types never require changes here. +""" + +from __future__ import annotations + +import math +from abc import ABC, abstractmethod +from random import random +from dataclasses import dataclass, field +from typing import Optional + +import numpy as np +import torch +import torch.nn.functional as F +from torch import nn +from torch.utils.data import DataLoader +from torchvision import transforms + +from v3.classes.backbones import BACKBONES, list_names, load_backbone_weights +from v3.classes.bridges import Bridge +from random import random + + +# --------------------------------------------------------------------------- +# Cross-tower communication context +# --------------------------------------------------------------------------- + + +@dataclass +class EarlyPassContext: + eye_train: list[dict] + bilat_train: list[dict] + bilat_val: list[dict] + bilat_test: list[dict] + image_preprocessor: object + image_cache: object + device: torch.device + store: dict = field(default_factory=dict) # cross-tower key-value store + + +# --------------------------------------------------------------------------- +# Backbone factory +# --------------------------------------------------------------------------- + + +def build_backbone(name: str, freeze_ratio: float = 0.0, augment: bool = True): + """ + Operational builder: + - instantiate with DEFAULT weights + - strip classifier → features + - apply ratio-based freezing over coarse blocks + - return (model, out_dim, transform) + """ + key = (name or "").lower() + if key not in BACKBONES: + raise ValueError( + f"Unsupported backbone '{name}'. Valid options: {list_names()}" + ) + + spec = BACKBONES[key] + m = spec.ctor(weights=spec.weights_default) + out_dim, m = spec.strip(m) + load_backbone_weights(key, m) + + mean = getattr(spec.weights_default, "meta", {}).get("mean", (0.485, 0.456, 0.406)) + std = getattr(spec.weights_default, "meta", {}).get("std", (0.229, 0.224, 0.225)) + crop = 299 if key == "inception_v3" else 224 + + if augment: + transform = transforms.Compose( + [ + transforms.Resize(256), + transforms.CenterCrop(crop), + transforms.RandomHorizontalFlip(), + transforms.RandomVerticalFlip(), + transforms.RandomRotation(15), + transforms.ColorJitter(0.1, 0.1, 0.1, 0.05), + transforms.ToTensor(), + transforms.Normalize(mean=mean, std=std), + ] + ) + else: + transform = transforms.Compose( + [ + transforms.Resize(256), + transforms.CenterCrop(crop), + transforms.ToTensor(), + transforms.Normalize(mean=mean, std=std), + ] + ) + + fr = max(0.0, min(1.0, float(freeze_ratio))) + blocks = spec.blocks(m) + n = len(blocks) + freeze_n = int(math.floor(n * fr)) + for b in blocks[:freeze_n]: + for p in b.parameters(): + p.requires_grad = False + + return m, out_dim, transform + + +# --------------------------------------------------------------------------- +# TowerBase ABC +# --------------------------------------------------------------------------- + + +class TowerBase(ABC): + """Abstract base class for a HyperTower tower. + + Concrete sub-classes must implement ``embed_dims``, ``embed_batch``, and + ``prepare_fold``. Everything else has a sensible default no-op. + """ + + # ------------------------------------------------------------------ + # Required interface + # ------------------------------------------------------------------ + + @property + @abstractmethod + def embed_dims(self) -> list[int]: + """Ordered list of embedding dimensionalities contributed to the bridge. + + Most towers contribute one embedding (e.g. GeometryTower → [geom_hidden]). + ImageClinicalTower contributes two (image + clinical → [img_dim, cd_dim]). + """ + + @abstractmethod + def embed_batch( + self, + batch: dict, + *, + device: torch.device, + slot: int = 1, + ) -> list[torch.Tensor]: + """Return a list of embeddings for one eye slot in *batch*. + + Parameters + ---------- + batch : dict — batch produced by a SlotDataset loader + device : torch.device + slot : 1 (OD / image_1 / matrix_1) or 2 (OS / image_2 / matrix_2) + + Returns + ------- + list of Tensor — same length and order as ``embed_dims`` + """ + + @abstractmethod + def prepare_fold( + self, + *, + eye_train: list, + bilat_train: list, + bilat_val: list, + bilat_test: list, + image_preprocessor, + image_cache, + device: torch.device, + args, + ) -> None: + """Called once per fold before the main epoch loop.""" + + def early_pass(self, context: EarlyPassContext) -> None: + """Optional: called once per fold before loaders are built.""" + pass + + def get_sample(self, entry) -> "torch.Tensor | dict": # noqa: ARG002 + """Optional: called by HTDataset to retrieve one sample for this tower. + + entry : ShellEntry — carries entity_id, label, side, paired. + + Single mode (entry.paired=False): return one Tensor. + Paired mode (entry.paired=True): return {"a": Tensor, "b": Tensor}. + + Default raises NotImplementedError. Towers that participate in the + v4 HTDataset pipeline must implement this. + """ + raise NotImplementedError( + f"{type(self).__name__}.get_sample() is not implemented. " + "Implement it to use this tower with HTDataset." + ) + + # ------------------------------------------------------------------ + # Optional interface + # ------------------------------------------------------------------ + + def augment_samples(self, samples: list) -> list: + """Optional: add modality-specific keys to sample dicts before loaders are built. + + Called by the orchestrator on each sample list (eye_train, bilat_train, + bilat_val, bilat_test) *after* ``prepare_fold`` completes. + + The default implementation is a no-op. GeometryTower overrides this to + inject ``seg_map_1`` / ``seg_map_2`` numpy arrays so the loader can deliver + them as tensors alongside the image and clinical slots. + """ + return samples + + def cd_warmup_embedding( + self, + batch: dict, + *, + device: torch.device, + ) -> Optional[torch.Tensor]: + """Return the clinical embedding for cd_warmup phase, or None if not applicable. + + ClinicalDataTower overrides this to return its encoder output. + All other towers return None (the default). + + This replaces isinstance(t, ClinicalDataTower) checks in train_towers_epoch, + so new tower types never require changes to towerbase.py. + """ + return None + + def set_phase(self, phase: str) -> None: + """Control requires_grad on this tower's parameters for *phase*. + + Phases: ``cd_warmup``, ``tower_warmup``, ``fused_warmup``, ``main``. + Default: no-op (tower parameters always trainable unless overridden). + """ + + @property + def total_epochs(self) -> int: + """How many epochs this tower participates in the main loop.""" + return 0 + + @property + def train_loader(self) -> Optional[DataLoader]: + """Single-eye training loader, or None if not applicable.""" + return None + + @property + def cd_only_loader(self) -> Optional[DataLoader]: + """Clinical-data-only loader for cd_warmup phase, or None.""" + return None + + def finalize_fold( + self, + *, + bridge, + bilat_train_loader: Optional[DataLoader] = None, + val_loader: Optional[DataLoader] = None, + device: torch.device, + args, + ) -> None: + """Optional post-epoch-loop operations (e.g. fused head training).""" + + +# --------------------------------------------------------------------------- +# Shared training utilities +# --------------------------------------------------------------------------- + + +def _set_requires_grad(module: nn.Module, enabled: bool) -> None: + for p in module.parameters(): + p.requires_grad = enabled + + +def _to_label_tensor(labels, device: torch.device) -> torch.Tensor: + if torch.is_tensor(labels): + return labels.to(device=device, dtype=torch.long) + return torch.as_tensor(labels, dtype=torch.long, device=device) + + +# --------------------------------------------------------------------------- +# Modular tower training / evaluation (multi-tower interface) +# --------------------------------------------------------------------------- + + +def train_towers_epoch( + towers: "list[TowerBase]", + bridge: Bridge, + loader: DataLoader, + optimizer, + device: torch.device, + *, + phase: str, + bcd_prob: float = 0.5, + tower_loss_mode: str = "bcd", +) -> tuple[float, float]: + """Train one epoch using the modular tower interface. + + All towers and the bridge are set to the given phase via their + ``set_phase`` methods. BCD / fused loss semantics mirror the + existing ``train_single_epoch`` logic: + + cd_warmup — clinical encoder aux head only (slot index 1) + tower_warmup — img + cd aux heads equally + fused_warmup — fused bridge output only + main — BCD (randomly img or cd aux) vs fused, per tower_loss_mode + + Duck typing: towers that implement ``cd_warmup_embedding`` participate in + cd_warmup; all others are skipped for that phase. No isinstance checks. + """ + for t in towers: + t.set_phase(phase) + bridge.set_phase(phase) + + total_loss = total_correct = total_n = 0 + + for batch in loader: + y = batch.get("label_1") + if y is None: + continue + + # ---- cd_warmup: train clinical encoder via its aux head ---- + if phase == "cd_warmup": + z_cd = None + for t in towers: + z = t.cd_warmup_embedding(batch, device=device) + if z is not None: + z_cd = z + break + if z_cd is None: + continue + logits = bridge.aux_heads[1](z_cd) + y_t = _to_label_tensor(y, device) + loss = F.cross_entropy(logits, y_t) + optimizer.zero_grad() + loss.backward() + optimizer.step() + bs = y_t.shape[0] + total_loss += float(loss.item()) * bs + total_correct += int((logits.argmax(1) == y_t).sum()) + total_n += bs + continue + + # ---- collect all slot-1 embeddings from every tower ---- + all_embs = [] + for t in towers: + all_embs.extend(t.embed_batch(batch, device=device, slot=1)) + + if any(e is None for e in all_embs): + continue + + y_t = _to_label_tensor(y, device) + + if phase == "tower_warmup": + # Average loss across all available aux heads + aux_logits = [bridge.aux_heads[i](emb) for i, emb in enumerate(all_embs)] + loss = sum(F.cross_entropy(l, y_t) for l in aux_logits) / len(aux_logits) + # Softmax average for metrics + logits = sum(F.softmax(l, dim=1) for l in aux_logits) / len(aux_logits) + + elif phase == "fused_warmup": + logits_fused, _ = bridge.fuse(all_embs) + loss = F.cross_entropy(logits_fused, y_t) + logits = logits_fused + + else: # main + if tower_loss_mode == "all": + logits_fused, aux = bridge.fuse(all_embs) + loss = F.cross_entropy(logits_fused, y_t) + for aux_l in aux: + loss = loss + F.cross_entropy(aux_l, y_t) + logits = logits_fused + elif random() < bcd_prob: + # Randomly pick ONE tower to train (Generalized BCD) + idx = int(random() * len(all_embs)) + logits = bridge.aux_heads[idx](all_embs[idx]) + loss = F.cross_entropy(logits, y_t) + else: + logits_fused, _ = bridge.fuse(all_embs) + loss = F.cross_entropy(logits_fused, y_t) + logits = logits_fused + + optimizer.zero_grad() + loss.backward() + optimizer.step() + bs = y_t.shape[0] + total_loss += float(loss.item()) * bs + total_correct += int((logits.argmax(1) == y_t).sum()) + total_n += bs + + return ( + total_loss / total_n if total_n else float("nan"), + total_correct / total_n if total_n else float("nan"), + ) + + +def collect_probs_towers( + towers: "list[TowerBase]", + bridge: Bridge, + loader: DataLoader, + device: torch.device, + *, + tower_mode: str = "ensemble", +) -> tuple[np.ndarray, np.ndarray]: + """Evaluate using the modular tower interface. + + Returns ``(y_true, probs)`` with patient-level probabilities + (OD + OS averaged for ensemble mode, OD-only for single mode). + """ + for t in towers: + if isinstance(t, nn.Module): + t.eval() + bridge.eval() + + y_chunks, p_chunks = [], [] + with torch.no_grad(): + for batch in loader: + y = batch.get("label_1") + if y is None: + continue + if not ( + torch.is_tensor(batch.get("image_1")) + and torch.is_tensor(batch.get("image_2")) + ): + continue + + all_embs_od = [] + all_embs_os = [] + for t in towers: + all_embs_od.extend(t.embed_batch(batch, device=device, slot=1)) + all_embs_os.extend(t.embed_batch(batch, device=device, slot=2)) + + if any(e is None for e in all_embs_od + all_embs_os): + continue + + logits_od, _ = bridge.fuse(all_embs_od) + logits_os, _ = bridge.fuse(all_embs_os) + + if tower_mode == "ensemble": + probs = 0.5 * ( + F.softmax(logits_od, dim=1) + F.softmax(logits_os, dim=1) + ) + else: + probs = F.softmax(logits_od, dim=1) + + y_chunks.append(_to_label_tensor(y, device).cpu().numpy()) + p_chunks.append(probs.cpu().numpy()) + + for t in towers: + if isinstance(t, nn.Module): + t.train() + bridge.train() + + if not y_chunks: + return np.array([], dtype=np.int64), np.zeros((0, 0), dtype=np.float32) + return np.concatenate(y_chunks), np.concatenate(p_chunks, axis=0) diff --git a/v3/classes/towers.py b/v3/classes/towers.py deleted file mode 100644 index 4b0868a..0000000 --- a/v3/classes/towers.py +++ /dev/null @@ -1,279 +0,0 @@ -from __future__ import annotations - -import math -from typing import Optional - -import torch -from torch import nn -from torchvision import transforms - -from v3.classes.backbones import BACKBONES, list_names, load_backbone_weights -from v3.classes.SE_attention import SEBlock -from v3.classes.data_bundle import DataBundle - - -def build_backbone(name: str, freeze_ratio: float = 0.0, augment: bool = True): - """ - Operational builder: - - instantiate with DEFAULT weights - - strip classifier → features - - apply ratio-based freezing over coarse blocks - - return (model, out_dim, transform) - """ - key = (name or "").lower() - if key not in BACKBONES: - raise ValueError(f"Unsupported backbone '{name}'. Valid options: {list_names()}") - - spec = BACKBONES[key] - m = spec.ctor(weights=spec.weights_default) - out_dim, m = spec.strip(m) - load_backbone_weights(key, m) - - # transforms: use the weights’ mean/std, but keep your augmentation pipeline - mean = getattr(spec.weights_default, "meta", {}).get("mean", (0.485, 0.456, 0.406)) - std = getattr(spec.weights_default, "meta", {}).get("std", (0.229, 0.224, 0.225)) - crop = 299 if key == "inception_v3" else 224 - - if augment: - transform = transforms.Compose( - [ - transforms.Resize(256), - transforms.CenterCrop(crop), - transforms.RandomHorizontalFlip(), - transforms.RandomVerticalFlip(), - transforms.RandomRotation(15), - transforms.ColorJitter(0.1, 0.1, 0.1, 0.05), - transforms.ToTensor(), - transforms.Normalize(mean=mean, std=std), - ] - ) - else: - transform = transforms.Compose( - [ - transforms.Resize(256), - transforms.CenterCrop(crop), - transforms.ToTensor(), - transforms.Normalize(mean=mean, std=std), - ] - ) - - # ratio-based freezing: freeze earliest floor(N * freeze_ratio) blocks - fr = max(0.0, min(1.0, float(freeze_ratio))) - blocks = spec.blocks(m) - n = len(blocks) - freeze_n = int(math.floor(n * fr)) - for b in blocks[:freeze_n]: - for p in b.parameters(): - p.requires_grad = False - - return m, out_dim, transform - - -class ImageTower(nn.Module): - """ - Vision backbone → pooled features. - - backbone: one of list_names() (default 'efficientnet_b0') - - always DEFAULT torchvision weights - - freeze_ratio ∈ [0,1] freezes earliest floor(N*freeze_ratio) blocks - - returns [N, out_dim] features from backbone forward - """ - - def __init__( - self, - backbone: str = "efficientnet_b0", - freeze_ratio: float = 0.0, - use_se: bool = False, - se_reduction: int = 16, - se_pre_norm: bool = True, - augment: bool = True, - geometry_dim: int = 0, - ): - super().__init__() - self.backbone, base_dim, self.transform = build_backbone( - backbone, freeze_ratio, augment=augment - ) - self._name = backbone - # Keep ordered blocks for dynamic freezing/thawing - key = (self._name or "").lower() - self._spec = BACKBONES[key] - self._blocks = self._spec.blocks(self.backbone) - # Optional tower-level SE over the final feature vector - self.base_dim = base_dim - self.geometry_dim = max(0, int(geometry_dim)) - self.out_dim = self.base_dim + self.geometry_dim - self.tower_ln = nn.LayerNorm(self.base_dim) if se_pre_norm else nn.Identity() - self.tower_se = ( - SEBlock(self.base_dim, reduction=se_reduction, residual=True) - if use_se - else None - ) - - def forward( - self, x: torch.Tensor, geometry: Optional[torch.Tensor] = None - ) -> torch.Tensor: - y = self.backbone(x) - # sanity: pooled features, not logits - assert y.dim() == 2 and y.size(1) == self.base_dim, ( - f"Expected features [N,{self.base_dim}], got {tuple(y.shape)}" - ) - if self.tower_se is not None: - y, _ = self.tower_se(self.tower_ln(y)) - if self.geometry_dim > 0: - if geometry is None or geometry.numel() == 0: - geom = torch.zeros( - y.size(0), self.geometry_dim, device=y.device, dtype=y.dtype - ) - else: - if geometry.dim() == 1: - geom = geometry.unsqueeze(0) - else: - geom = geometry - geom = geom.to(device=y.device, dtype=y.dtype) - if geom.size(0) != y.size(0): - raise ValueError( - f"Geometry batch size mismatch: {geom.size(0)} vs {y.size(0)}" - ) - if geom.size(1) != self.geometry_dim: - raise ValueError( - f"Expected geometry dim {self.geometry_dim}, got {geom.size(1)}" - ) - y = torch.cat([y, geom], dim=1) - return y - - def set_freeze_ratio(self, ratio: float): - """Dynamically freeze earliest floor(N*ratio) backbone blocks.""" - r = max(0.0, min(1.0, float(ratio))) - n = len(self._blocks) - freeze_n = int(math.floor(n * r)) - # Unfreeze all first - for b in self._blocks: - for p in b.parameters(): - p.requires_grad = True - # Freeze earliest blocks - for b in self._blocks[:freeze_n]: - for p in b.parameters(): - p.requires_grad = False - - -class SiameseImageTower(nn.Module): - """ - Shared-weight bilateral image tower. - - Runs OD and OS images through a single shared backbone, then returns - cat([f_mean, f_delta]) where: - f_mean = (f_od + f_os) / 2 -- shared bilateral representation - f_delta = f_od - f_os -- asymmetry, signed OD-relative - - out_dim = 2 * backbone_out_dim - - When x_os is None (single-eye fallback): - f_mean = f_od - f_delta = zeros - so the module degrades gracefully when only one eye is available. - - The shared backbone means both eyes contribute to every gradient update, - effectively doubling the training signal for the visual pathway without - doubling parameters. - """ - - def __init__( - self, - backbone: str = "efficientnet_b0", - freeze_ratio: float = 0.0, - use_se: bool = False, - se_reduction: int = 16, - se_pre_norm: bool = True, - augment: bool = True, - ): - super().__init__() - self._tower = ImageTower( - backbone=backbone, - freeze_ratio=freeze_ratio, - use_se=use_se, - se_reduction=se_reduction, - se_pre_norm=se_pre_norm, - augment=augment, - geometry_dim=0, - ) - self.out_dim = self._tower.out_dim * 2 - self.transform = self._tower.transform - - def forward( - self, - x_od: torch.Tensor, - x_os: Optional[torch.Tensor] = None, - ) -> torch.Tensor: - f_od = self._tower(x_od) - if x_os is None: - f_mean = f_od - f_delta = torch.zeros_like(f_od) - else: - f_os = self._tower(x_os) - f_mean = (f_od + f_os) * 0.5 - f_delta = f_od - f_os - return torch.cat([f_mean, f_delta], dim=1) - - def set_freeze_ratio(self, ratio: float) -> None: - """Delegates to the shared inner tower.""" - self._tower.set_freeze_ratio(ratio) - - -class ClinicalTower(nn.Module): - """MLP over DataBundle.vectorize_row outputs (convert to torch inside tower).""" - - def __init__( - self, - clinical_data: DataBundle, - hidden_dim: int = 128, - dropout: float = 0.1, - use_se: bool = False, - se_reduction: int = 16, - se_pre_norm: bool = True, - ): - super().__init__() - self.feature_dim = clinical_data.feature_dim - self.out_dim = hidden_dim - # two-block MLP so we can optionally freeze/thaw per block - self.block0 = nn.Sequential( - nn.Linear(self.feature_dim, hidden_dim), - nn.LayerNorm(hidden_dim), - nn.ReLU(inplace=True), - nn.Dropout(dropout), - ) - self.block1 = nn.Sequential( - nn.Linear(hidden_dim, hidden_dim), - nn.ReLU(inplace=True), - ) - self.net = nn.Sequential(self.block0, self.block1) - self.tower_ln = nn.LayerNorm(hidden_dim) if se_pre_norm else nn.Identity() - self.tower_se = ( - SEBlock(hidden_dim, reduction=se_reduction, residual=True) - if use_se - else None - ) - - def forward(self, meta_np_or_torch) -> torch.Tensor: - if isinstance(meta_np_or_torch, torch.Tensor): - x = meta_np_or_torch - else: - x = torch.as_tensor(meta_np_or_torch, dtype=torch.float32) - h = self.net(x) - if self.tower_se is not None: - h, _ = self.tower_se(self.tower_ln(h)) - return h - - def set_freeze_ratio(self, ratio: float): - """Optionally freeze earliest blocks of the MLP.""" - r = max(0.0, min(1.0, float(ratio))) - # Unfreeze all - for p in self.block0.parameters(): - p.requires_grad = True - for p in self.block1.parameters(): - p.requires_grad = True - # Freeze earliest blocks based on ratio threshold - if r >= 0.5: - for p in self.block0.parameters(): - p.requires_grad = False - if r >= 1.0: - for p in self.block1.parameters(): - p.requires_grad = False diff --git a/v3/classes/v2_hypertower.py b/v3/classes/v2_hypertower.py deleted file mode 100644 index 8733b11..0000000 --- a/v3/classes/v2_hypertower.py +++ /dev/null @@ -1,1853 +0,0 @@ -"""V2HyperTower — central orchestrator for the V2 mode-comparison pipeline. - -All model classes, metric helpers, croppers and data utilities live in their -respective category-specific modules. This file owns only: - - V2HyperTower (the big-picture orchestrator) - - V2ModeComparator (thin backward-compat shim) - - module-level ``build_parser`` alias (used by old scripts that import it - directly; will be removed once all callers are updated) -""" -from __future__ import annotations - -import argparse -import copy -import csv -import json -import time -from pathlib import Path -from types import SimpleNamespace -from typing import Optional - -import numpy as np -import torch - -from v3.classes.croppers import build_image_preprocessor_from_args -from v3.classes.dataset import _ClinicalView # noqa: F401 (re-exported for compat) -from v3.classes.loader_factory import ( - build_balanced_sampler, - filter_bilateral_samples, - filter_eye_samples, - make_loader, -) -from v3.classes.metrics import _score_arrays, _svf, _tune_and_snap -from v3.classes.models import ( - BilateralHT, - FusedEnsembleHT, - SingleEyeHT, - V2ModeComparisonOps, - collect_probs_bilateral, - collect_probs_bilateral_components, - collect_probs_classic, - collect_probs_ensemble, - collect_probs_ensemble_pereye, - collect_probs_eye_level, - collect_probs_fused, - collect_probs_single_components, - train_bilateral_epoch, - train_fusion_epoch, - train_single_epoch, -) -from v3.classes.papila_builders import build_papila_data -from v3.classes.predictions import PredictionStore, head_names_for_mode -from v3.classes.profiles import build_papila_profile -from v3.classes.results import FoldArtifacts, FoldResult, _f, _nan, _sv -from v3.classes.split_manager import PatientFirstSplitManager -from v3.classes.transforms import build_eval_transform -from v3.classes.utils import ( - _drop_mixed_label_patients, - _relabel_mixed_patients_to_max, - choose_device, - seed_everything, -) -from v3.classes.hypertower_logger import HypertowerLogger - - -# --------------------------------------------------------------------------- -# Fusion-event helper -# --------------------------------------------------------------------------- - -def _fusion_events( - y: np.ndarray, - pf: np.ndarray, - pi: np.ndarray, - pm: np.ndarray, -) -> tuple[int, int]: - """ - Count fusion corrections and errors. - - correction: fused correct, both img and md wrong - - error: fused wrong, both img and md correct - Returns (n_corrections, n_errors). - """ - pred_f = pf.argmax(1); pred_i = pi.argmax(1); pred_m = pm.argmax(1) - corr = int(((pred_f == y) & (pred_i != y) & (pred_m != y)).sum()) - err = int(((pred_f != y) & (pred_i == y) & (pred_m == y)).sum()) - return corr, err - - -def _cm_cells(y: np.ndarray, p: np.ndarray, num_classes: int) -> dict[str, int]: - """ - Return confusion matrix cells as a flat dict. - Binary: keys tn/fp/fn/tp - Multiclass: keys cm_{i}_{j} for true class i, predicted class j - Returns empty dict if arrays are empty or wrong shape. - """ - if not y.size or p.ndim < 2 or p.shape[1] != num_classes: - return {} - pred = p.argmax(1) - if num_classes == 2: - tn = int(((pred == 0) & (y == 0)).sum()) - fp = int(((pred == 1) & (y == 0)).sum()) - fn = int(((pred == 0) & (y == 1)).sum()) - tp = int(((pred == 1) & (y == 1)).sum()) - return {"tn": tn, "fp": fp, "fn": fn, "tp": tp} - # multiclass: full NxN matrix - out: dict[str, int] = {} - for i in range(num_classes): - for j in range(num_classes): - out[f"cm_{i}_{j}"] = int(((y == i) & (pred == j)).sum()) - return out - - -# --------------------------------------------------------------------------- -# Per-sample prediction logging -# --------------------------------------------------------------------------- - -def _save_predictions_csv( - fold_dir: Path, - eval_mode: str, - y_true: np.ndarray, - heads: dict, # {"fused": probs_array, "img": probs_array, "md": probs_array, ...} - suffix: str = "", # e.g. "_pereye" -) -> None: - """ - Save a per-sample CSV with predicted class, per-class probabilities, - and TP/FP/TN/FN (binary) or correct flag (multiclass) for every head. - """ - N = len(y_true) - num_classes = next(p.shape[1] for p in heads.values() if p is not None) - rows = [] - for i in range(N): - true = int(y_true[i]) - row: dict = {"idx": i, "y_true": true} - for head_name, probs in heads.items(): - if probs is None: - continue - pred = int(probs[i].argmax()) - row[f"pred_{head_name}"] = pred - for c in range(num_classes): - row[f"prob_{head_name}_c{c}"] = float(probs[i, c]) - if eval_mode == "binary": - row[f"tp_{head_name}"] = int(pred == 1 and true == 1) - row[f"fp_{head_name}"] = int(pred == 1 and true == 0) - row[f"tn_{head_name}"] = int(pred == 0 and true == 0) - row[f"fn_{head_name}"] = int(pred == 0 and true == 1) - else: - row[f"correct_{head_name}"] = int(pred == true) - rows.append(row) - - if not rows: - return - csv_path = fold_dir / f"predictions{suffix}.csv" - with csv_path.open("w", newline="") as f: - writer = csv.DictWriter(f, fieldnames=list(rows[0].keys())) - writer.writeheader() - writer.writerows(rows) - - -# --------------------------------------------------------------------------- -# V2HyperTower -# --------------------------------------------------------------------------- - -class V2HyperTower: - """Central orchestrator. Construct with ``V2HyperTower(args)``, call ``.run()``.""" - - # ------------------------------------------------------------------ - # CLI - # ------------------------------------------------------------------ - - @staticmethod - def build_parser() -> argparse.ArgumentParser: - ap = argparse.ArgumentParser( - description=( - "Three HyperTower modes: Classic (eye-level), Ensemble (patient-level avg), " - "Bilateral (BilateralBridge with shared towers). Pure k-fold CV." - ) - ) - ap.add_argument("--image-dir", default="Papila/FundusImages") - ap.add_argument("--clinical-dir", default="Papila/ClinicalData") - ap.add_argument("--label-col", default="Diagnosis") - ap.add_argument("--cat-cols", nargs="*", default=["Gender", "Phakic/Pseudophakic"]) - ap.add_argument("--exclude-cols", nargs="*", default=[], - help="Feature columns to exclude entirely from the clinical feature matrix.") - ap.add_argument("--eval-mode", choices=["binary", "multiclass"], default="multiclass") - ap.add_argument( - "--tower-mode", choices=["single", "ensemble", "bilateral", "classic"], - default="single", - help="Train/evaluate a single tower mode.", - ) - ap.add_argument("--n-splits", type=int, default=5) - ap.add_argument("--fold-seed", type=int, default=42) - ap.add_argument("--holdout-per-class", type=int, default=5, - help="Patients per class reserved for holdout before train/test split (0 disables)") - ap.add_argument("--holdout-seed", type=int, default=123, - help="Random seed for holdout sampling") - ap.add_argument( - "--folds", type=int, default=None, - help="Optional cap on how many folds to run (default: all --n-splits).", - ) - ap.add_argument("--epochs", type=int, default=40) - ap.add_argument( - "--warmup-tower-epochs", type=int, default=None, - help="Extra tower warmup epochs (added before main epochs). Default: auto by mode.", - ) - ap.add_argument( - "--warmup-fused-epochs", type=int, default=None, - help="Extra fused warmup epochs (added before main epochs). Default: auto by mode.", - ) - ap.add_argument("--single-warmup-tower-epochs", type=int, default=None, - help="Single-eye model tower warmup (overrides --warmup-tower-epochs).") - ap.add_argument("--single-warmup-fused-epochs", type=int, default=None, - help="Single-eye model fused warmup (overrides --warmup-fused-epochs).") - ap.add_argument("--warmup-cd-epochs", type=int, default=0, - help="MD-only warmup epochs before tower warmup. Trains only cd_tower + " - "classifier_cd (no CNN forward pass, so 50-100 epochs is cheap).") - ap.add_argument("--bilat-warmup-tower-epochs", type=int, default=None, - help="Bilateral model tower warmup (overrides --warmup-tower-epochs).") - ap.add_argument("--bilat-warmup-fused-epochs", type=int, default=None, - help="Bilateral model fused warmup (overrides --warmup-fused-epochs).") - ap.add_argument("--batch-size", type=int, default=8) - ap.add_argument("--lr", type=float, default=1e-4) - ap.add_argument("--bcd-prob", type=float, default=0.5, - help="Tower-only step probability during main phase (per model).") - ap.add_argument("--tower-loss-mode", choices=["bcd", "all"], default="bcd", - help="Main-phase tower loss strategy: " - "'bcd' (Block Coordinate Descent — randomly train one tower or fused per step) " - "or 'all' (sum all three losses — fused + img + md — every step).") - ap.add_argument("--backbone", default="refugelike") - ap.add_argument("--freeze-ratio", type=float, default=0.0) - ap.add_argument("--augment", action="store_true") - ap.add_argument("--balanced-sampling", action="store_true", - help="Use WeightedRandomSampler during training to equalise class frequency (default: off).") - ap.add_argument("--num-workers", type=int, default=4) - ap.add_argument("--in-memory-cache", action="store_true", default=True, - help="Cache preprocessed images in RAM (default: on).") - ap.add_argument("--no-in-memory-cache", action="store_false", dest="in_memory_cache", - help="Disable in-memory image cache.") - ap.add_argument("--cache-workers", type=int, default=4, - help="Threads for prebuilding in-memory image cache (default: 4).") - ap.add_argument("--device", choices=["auto", "cpu", "cuda"], default="auto") - ap.add_argument("--seed", type=int, default=1234) - ap.add_argument("--run-name", default=None) - ap.add_argument("--output-root", default="analysis_data") - # Optional ROI cropping - ap.add_argument("--img-crop-manifest", type=str, default=None, - help="Path to crop manifest CSV for ROI cropping.") - ap.add_argument("--img-crop-gt", action="store_true", - help="Use ground-truth masks/contours from manifest for ROI crop.") - ap.add_argument("--img-crop-weights", type=str, default=None, - help="UNet weights path for ROI cropping from predicted masks.") - ap.add_argument("--img-crop-normalize", type=str, default="per_image", - choices=["per_image", "imagenet"], - help="UNet input normalization mode.") - ap.add_argument("--img-crop-threshold", type=float, default=0.5, - help="UNet mask threshold for ROI extraction.") - ap.add_argument("--img-crop-tta", action="store_true", - help="Enable flip-TTA during UNet mask inference.") - ap.add_argument("--img-crop-scale", type=float, default=2.5, - help="Disc-radius multiplier for square crop.") - ap.add_argument("--img-crop-size", type=int, default=224, - help="Output ROI size before tower transforms.") - ap.add_argument("--img-crop-cache", type=str, default="cache_data/hypertower_crops", - help="Cache directory for cropped images and geometry sidecars.") - ap.add_argument("--persist-img-crop-cache", action="store_true", - help="Keep existing cached crop .npz files instead of clearing at run start.") - # Architecture - ap.add_argument("--cd-hidden-dim", type=int, default=128, - help="ClinicalTower hidden dimension.") - ap.add_argument("--fusion-dim", type=int, default=256, - help="Bridge/BilateralBridge fusion dimension.") - ap.add_argument("--bridge-mode", default="fused", - choices=["fused", "image_only", "clinical_only"], - help="Bridge fusion mode: fused (default), image_only, or clinical_only.") - # Mixed patients - ap.add_argument( - "--exclude-mixed-patients", - dest="exclude_mixed_patients", action="store_true", - help="Drop patients whose two eyes have different labels before splitting.", - ) - ap.add_argument( - "--include-mixed-patients", - dest="exclude_mixed_patients", action="store_false", - ) - ap.add_argument( - "--relabel-mixed-patients-to-max", - dest="relabel_mixed_patients_to_max", - action="store_true", - help="When mixed patients are included, relabel both eyes to patient max severity.", - ) - ap.add_argument( - "--keep-mixed-raw-labels", - dest="relabel_mixed_patients_to_max", - action="store_false", - help=argparse.SUPPRESS, - ) - ap.set_defaults(exclude_mixed_patients=False, relabel_mixed_patients_to_max=False) - # Tuning - ap.add_argument( - "--tune-binary-threshold", action="store_true", - help="Tune per-model binary threshold on validation each epoch.", - ) - ap.add_argument( - "--tune-multiclass-bias", action="store_true", - help="Tune per-model multiclass log-prob bias on validation each epoch.", - ) - ap.add_argument("--ece-bins", type=int, default=10) - ap.add_argument("--log-every", type=int, default=1) - ap.add_argument("--save-checkpoints", action=argparse.BooleanOptionalAction, default=True, - help="Save best_single.pt / best_holdout_single.pt per fold (use --no-save-checkpoints to disable)") - ap.add_argument("--use-last-epoch", action="store_true", default=False, - help="Score using the final epoch's model state rather than the best-AUC checkpoint.") - # IOP feature options - ap.add_argument( - "--iop-corr-method", - choices=["ratio", "ols", "lad", "multi"], - default="ratio", - help="Perkins→Pneumatic conversion method: ratio (default), ols, lad, or multi (+CCT).", - ) - ap.add_argument( - "--iop-drop-raw", - action="store_true", - default=False, - help="Exclude IOP_raw from the feature matrix (keep only IOP_corr).", - ) - ap.add_argument( - "--fused-head", action="store_true", - help="(ensemble mode only) After base SingleEyeHT training, freeze it and train a " - "small logit-level MLP fusion head on bilateral samples instead of averaging " - "OD/OS softmax probabilities.", - ) - ap.add_argument( - "--fusion-epochs", type=int, default=10, - help="Number of epochs to train the fusion head (--fused-head, ensemble mode only).", - ) - return ap - - # ------------------------------------------------------------------ - # Construction - # ------------------------------------------------------------------ - - def __init__(self, args) -> None: - self.args = args - self.device = choose_device(args.device) - seed_everything(args.seed) - - print(f"Device: {self.device}", flush=True) - print("Loading PAPILA data...", flush=True) - self.data = build_papila_data( - image_dir=args.image_dir, - clinical_dir=args.clinical_dir, - label_col=args.label_col, - cat_cols=list(args.cat_cols), - n_splits=args.n_splits, - random_seed=args.fold_seed, - iop_corr_method=getattr(args, "iop_corr_method", "ratio"), - iop_drop_raw=getattr(args, "iop_drop_raw", False), - exclude_cols=list(getattr(args, "exclude_cols", []) or []), - ) - print(f"Loaded: {len(self.data.df)} rows feature_dim={self.data.feature_dim}", flush=True) - self.image_preprocessor = build_image_preprocessor_from_args(args) - self.profile_eye = build_papila_profile( - patient_col="Patient ID", label_col=args.label_col, sample_mode="eye" - ) - self.profile_patient = build_papila_profile( - patient_col="Patient ID", label_col=args.label_col, sample_mode="patient" - ) - - # ------------------------------------------------------------------ - # Orchestration - # ------------------------------------------------------------------ - - def run(self) -> Path: - """Execute the full fold loop for one eval_mode × tower_mode combination.""" - args = self.args - ts = time.strftime("%Y%m%d_%H%M%S") - run_name = args.run_name or f"hypertower_modes_{ts}" - out_dir = Path(args.output_root) / run_name - out_dir.mkdir(parents=True, exist_ok=True) - - mode = args.eval_mode - tower_mode = "single" if args.tower_mode == "classic" else args.tower_mode - df_mode = self.data.df.copy() - - if args.exclude_mixed_patients: - before = df_mode["Patient ID"].nunique() - df_mode, mixed = _drop_mixed_label_patients( - df_mode, patient_col="Patient ID", label_col=args.label_col - ) - print( - f"[{mode}] dropped {len(mixed)} mixed-label patients " - f"({before} → {df_mode['Patient ID'].nunique()})", - flush=True, - ) - else: - if args.relabel_mixed_patients_to_max: - before_rows = len(df_mode) - df_mode, changed_rows, still_mixed = _relabel_mixed_patients_to_max( - df_mode, patient_col="Patient ID", label_col=args.label_col - ) - print( - f"[{mode}] relabeled mixed patients to max severity " - f"(changed={changed_rows}, rows={before_rows}→{len(df_mode)}, " - f"remaining_mixed={len(still_mixed)}).", - flush=True, - ) - else: - print(f"[{mode}] keeping mixed-label patients with raw per-eye labels.", flush=True) - - if mode == "binary": - df_mode = df_mode[df_mode[args.label_col].isin([0, 1])].reset_index(drop=True) - - num_classes = 2 if mode == "binary" else int(df_mode[args.label_col].nunique()) - print( - f"\n[{mode}] num_classes={num_classes} rows={len(df_mode)} " - f"patients={df_mode['Patient ID'].nunique()}", - flush=True, - ) - - split_manager = PatientFirstSplitManager( - patient_col="Patient ID", label_col=args.label_col - ) - split_args = SimpleNamespace( - eval_mode=mode, - holdout_per_class=args.holdout_per_class, - holdout_seed=args.holdout_seed, - n_splits=args.n_splits, - fold_seed=args.fold_seed, - ) - clinical_ns = SimpleNamespace(df=df_mode, label_col=args.label_col) - plans = split_manager.build_plans(clinical=clinical_ns, args=split_args, profile=None) - requested_folds = args.n_splits if args.folds is None else int(args.folds) - n_folds = min(requested_folds, len(plans)) - - tm_dir = out_dir / mode / tower_mode - tm_dir.mkdir(parents=True, exist_ok=True) - fold_results: list[FoldResult] = [] - - # Override profiles with df_mode slice. - profile_eye = build_papila_profile( - patient_col="Patient ID", label_col=args.label_col, sample_mode="eye" - ) - profile_patient = build_papila_profile( - patient_col="Patient ID", label_col=args.label_col, sample_mode="patient" - ) - - # ---- PredictionStore — build once before fold loop --------------- - fused_head = getattr(args, "fused_head", False) - _head_names = head_names_for_mode(tower_mode, fused_head=fused_head) - fusion_epochs = int(getattr(args, "fusion_epochs", 10)) if fused_head else 0 - _global_warmup_tower = getattr(args, "warmup_tower_epochs", None) - _global_warmup_fused = getattr(args, "warmup_fused_epochs", None) - _warmup_tower = ( - int(args.single_warmup_tower_epochs) - if getattr(args, "single_warmup_tower_epochs", None) is not None - else int(_global_warmup_tower) if _global_warmup_tower is not None else 2 - ) - _warmup_fused = ( - int(args.single_warmup_fused_epochs) - if getattr(args, "single_warmup_fused_epochs", None) is not None - else int(_global_warmup_fused) if _global_warmup_fused is not None else 2 - ) - _warmup_cd = int(getattr(args, "warmup_cd_epochs", 0)) - _total_epochs = _warmup_cd + _warmup_tower + _warmup_fused + int(args.epochs) + fusion_epochs - - # sample IDs depend on mode: single uses eye IDs, others use patient IDs - if tower_mode in ("single", "classic"): - _sample_ids = [ - f"{row['Patient ID']}{row['eyeID']}" - for _, row in df_mode.iterrows() - ] - _y_true = df_mode[args.label_col].tolist() - else: - # one row per patient (deduplicate — take first occurrence per patient) - _pat_df = df_mode.drop_duplicates(subset="Patient ID") - _sample_ids = _pat_df["Patient ID"].astype(str).tolist() - _y_true = _pat_df[args.label_col].tolist() - - pred_store = PredictionStore( - sample_ids=_sample_ids, - y_true=_y_true, - head_names=_head_names, - n_folds=n_folds, - n_epochs=_total_epochs, - n_classes=num_classes, - ) - - image_cache: dict | None = {} if getattr(args, "in_memory_cache", False) else None - - for fold in range(n_folds): - seed_everything(args.seed + fold * 100) - fold_dir = tm_dir / f"fold{fold}" - fold_dir.mkdir(exist_ok=True) - - print(f"\n[{mode}:{tower_mode}] fold {fold+1}/{n_folds}", flush=True) - result, artifacts = self._run_fold( - fold=fold, - split=plans[fold], - mode=mode, - data=self.data, - num_classes=num_classes, - profile_eye=profile_eye, - profile_patient=profile_patient, - fold_dir=fold_dir, - tower_mode=tower_mode, - pred_store=pred_store, - image_cache=image_cache, - ) - fold_results.append(result) - if artifacts.y_true_ensemble is not None: - np.save(fold_dir / "y_true.npy", artifacts.y_true_ensemble) - if artifacts.probs_ensemble is not None: - np.save(fold_dir / "probs_fused.npy", artifacts.probs_ensemble) - if artifacts.probs_ensemble_img is not None: - np.save(fold_dir / "probs_img.npy", artifacts.probs_ensemble_img) - if artifacts.probs_ensemble_md is not None: - np.save(fold_dir / "probs_cd.npy", artifacts.probs_ensemble_md) - if artifacts.y_true_classic is not None: - np.save(fold_dir / "y_true.npy", artifacts.y_true_classic) - if artifacts.probs_classic is not None: - np.save(fold_dir / "probs_classic.npy", artifacts.probs_classic) - if artifacts.probs_classic_img is not None: - np.save(fold_dir / "probs_classic_img.npy", artifacts.probs_classic_img) - if artifacts.probs_classic_md is not None: - np.save(fold_dir / "probs_classic_md.npy", artifacts.probs_classic_md) - if artifacts.y_true_ensemble_pereye is not None: - np.save(fold_dir / "y_true_pereye.npy", artifacts.y_true_ensemble_pereye) - if artifacts.probs_ensemble_pereye is not None: - np.save(fold_dir / "probs_fused_pereye.npy", artifacts.probs_ensemble_pereye) - if artifacts.probs_ensemble_img_pereye is not None: - np.save(fold_dir / "probs_img_pereye.npy", artifacts.probs_ensemble_img_pereye) - if artifacts.probs_ensemble_md_pereye is not None: - np.save(fold_dir / "probs_cd_pereye.npy", artifacts.probs_ensemble_md_pereye) - if artifacts.logits_ensemble is not None: - np.save(fold_dir / "logits_fused.npy", artifacts.logits_ensemble) - if artifacts.logits_ensemble_img is not None: - np.save(fold_dir / "logits_img.npy", artifacts.logits_ensemble_img) - if artifacts.logits_ensemble_md is not None: - np.save(fold_dir / "logits_cd.npy", artifacts.logits_ensemble_md) - if artifacts.logits_classic is not None: - np.save(fold_dir / "logits_classic.npy", artifacts.logits_classic) - if artifacts.logits_classic_img is not None: - np.save(fold_dir / "logits_classic_img.npy", artifacts.logits_classic_img) - if artifacts.logits_classic_md is not None: - np.save(fold_dir / "logits_classic_md.npy", artifacts.logits_classic_md) - if artifacts.logits_ensemble_pereye is not None: - np.save(fold_dir / "logits_fused_pereye.npy", artifacts.logits_ensemble_pereye) - if artifacts.logits_ensemble_img_pereye is not None: - np.save(fold_dir / "logits_img_pereye.npy", artifacts.logits_ensemble_img_pereye) - if artifacts.logits_ensemble_md_pereye is not None: - np.save(fold_dir / "logits_cd_pereye.npy", artifacts.logits_ensemble_md_pereye) - if artifacts.probs_bilat is not None: - np.save(fold_dir / "probs_bilat.npy", artifacts.probs_bilat) - if artifacts.probs_fused is not None: - np.save(fold_dir / "probs_fused_head.npy", artifacts.probs_fused) - # y_true is shared across all heads for the same fold - if artifacts.y_true_bilat is not None and artifacts.y_true_ensemble is None: - np.save(fold_dir / "y_true.npy", artifacts.y_true_bilat) - # per-sample prediction CSVs - if artifacts.y_true_ensemble is not None: - _save_predictions_csv( - fold_dir, mode, artifacts.y_true_ensemble, - {"fused": artifacts.probs_ensemble, - "img": artifacts.probs_ensemble_img, - "md": artifacts.probs_ensemble_md}, - ) - if artifacts.y_true_ensemble_pereye is not None: - _save_predictions_csv( - fold_dir, mode, artifacts.y_true_ensemble_pereye, - {"fused": artifacts.probs_ensemble_pereye, - "img": artifacts.probs_ensemble_img_pereye, - "md": artifacts.probs_ensemble_md_pereye}, - suffix="_pereye", - ) - if artifacts.y_true_classic is not None: - _save_predictions_csv( - fold_dir, mode, artifacts.y_true_classic, - {"fused": artifacts.probs_classic, - "img": artifacts.probs_classic_img, - "md": artifacts.probs_classic_md}, - suffix="_classic", - ) - - fold_csv = tm_dir / "fold_results.csv" - csv_fields = list(FoldResult.__dataclass_fields__.keys()) - with fold_csv.open("w", newline="", encoding="utf-8") as fh: - w = csv.DictWriter(fh, fieldnames=csv_fields) - w.writeheader() - for r in fold_results: - w.writerow({k: getattr(r, k) for k in csv_fields}) - - summary = self._summary(fold_results) - self._print_summary(f"{mode}:{tower_mode}", summary, tower_mode=tower_mode) - metric_key = { - "single": "classic_val_auc", - "ensemble": "ensemble_val_auc", - "bilateral": "bilat_val_auc", - }[tower_mode] - fold_metrics = [] - best_vals = [] - for r in fold_results: - best_val = getattr(r, metric_key) - fold_metrics.append({ - "fold": r.fold, - "best_metric_value": _f(best_val), - "best_epoch": (r.best_epoch_bilat if tower_mode == "bilateral" else r.best_epoch_single), - "monitor": metric_key, - }) - if not np.isnan(float(best_val)): - best_vals.append(float(best_val)) - - ts_now = time.strftime("%Y%m%d_%H%M%S") - mode_summary = { - "run_id": run_name, - "backbone": args.backbone, - "epochs": args.epochs, - "warmup_cd_epochs": getattr(args, "warmup_cd_epochs", 0), - "warmup_tower_epochs": args.warmup_tower_epochs, - "warmup_fused_epochs": args.warmup_fused_epochs, - "single_warmup_tower_epochs": args.single_warmup_tower_epochs, - "single_warmup_fused_epochs": args.single_warmup_fused_epochs, - "bilat_warmup_tower_epochs": args.bilat_warmup_tower_epochs, - "bilat_warmup_fused_epochs": args.bilat_warmup_fused_epochs, - "batch_size": args.batch_size, - "lr": args.lr, - "eval_mode": mode, - "tower_mode": tower_mode, - "n_splits": n_folds, - "best_metric": metric_key, - "best_metric_mode": "max", - "best_metric_mean": (float(np.mean(best_vals)) if best_vals else None), - "best_metric_std": (float(np.std(best_vals)) if best_vals else None), - "fold_metrics": fold_metrics, - "mode_summary": summary, - } - (tm_dir / "summary.json").write_text(json.dumps(mode_summary, indent=2), encoding="utf-8") - pred_store.save(tm_dir / "predictions.npz") - - root_summary_path = out_dir / "summary.json" - if root_summary_path.exists(): - try: - payload = json.loads(root_summary_path.read_text(encoding="utf-8")) - except Exception: - payload = {} - else: - payload = {} - payload.setdefault("run_name", run_name) - payload.setdefault("timestamp", ts_now) - payload["config"] = vars(args) - payload.setdefault("summaries", {}) - payload["summaries"][f"{mode}:{tower_mode}"] = summary - root_summary_path.write_text(json.dumps(payload, indent=2), encoding="utf-8") - - print(f"\nOutputs written to: {out_dir}") - return out_dir - - # ------------------------------------------------------------------ - # Fold runner - # ------------------------------------------------------------------ - - def _run_fold( - self, - fold: int, - split, - mode: str, - data, - num_classes: int, - profile_eye, - profile_patient, - fold_dir: Path, - tower_mode: str, - pred_store: "PredictionStore | None" = None, - image_cache: "dict | None" = None, - ) -> tuple[FoldResult, FoldArtifacts]: - args = self.args - device = self.device - image_preprocessor = self.image_preprocessor - nan = _nan() - tower_mode = "single" if tower_mode == "classic" else tower_mode - run_single = tower_mode in ("single", "ensemble") - run_bilat = tower_mode == "bilateral" - run_fused = (tower_mode == "ensemble") and bool(getattr(args, "fused_head", False)) - - global_warmup_tower = getattr(args, "warmup_tower_epochs", None) - global_warmup_fused = getattr(args, "warmup_fused_epochs", None) - single_warmup_tower = ( - int(args.single_warmup_tower_epochs) - if getattr(args, "single_warmup_tower_epochs", None) is not None - else int(global_warmup_tower) if global_warmup_tower is not None else 2 - ) - single_warmup_fused = ( - int(args.single_warmup_fused_epochs) - if getattr(args, "single_warmup_fused_epochs", None) is not None - else int(global_warmup_fused) if global_warmup_fused is not None else 2 - ) - bilat_warmup_tower = ( - int(args.bilat_warmup_tower_epochs) - if getattr(args, "bilat_warmup_tower_epochs", None) is not None - else int(global_warmup_tower) if global_warmup_tower is not None else 4 - ) - bilat_warmup_fused = ( - int(args.bilat_warmup_fused_epochs) - if getattr(args, "bilat_warmup_fused_epochs", None) is not None - else int(global_warmup_fused) if global_warmup_fused is not None else 3 - ) - single_warmup_cd = int(getattr(args, "warmup_cd_epochs", 0)) if run_single else 0 - if not run_single: - single_warmup_tower = 0 - single_warmup_fused = 0 - if not run_bilat: - bilat_warmup_tower = 0 - bilat_warmup_fused = 0 - main_epochs = int(args.epochs) - total_single_epochs = (single_warmup_cd + single_warmup_tower + single_warmup_fused + main_epochs) if run_single else 0 - total_bilat_epochs = (bilat_warmup_tower + bilat_warmup_fused + main_epochs) if run_bilat else 0 - total_epochs = max(total_single_epochs, total_bilat_epochs) - - # ---- samples --------------------------------------------------- - eye_train = filter_eye_samples(profile_eye.build_samples(df=split.train, clinical=data)) - bilat_train = filter_bilateral_samples(profile_patient.build_samples(df=split.train, clinical=data)) - bilat_val = filter_bilateral_samples(profile_patient.build_samples(df=split.val, clinical=data)) - - # Register split labels in the prediction store - if pred_store is not None: - if tower_mode in ("single", "classic"): - # eye-level IDs: "{patient_id}{eyeID}" - train_sids = [f"{s['id_1']}{s.get('eye_id_1','')}" for s in eye_train] - val_sids = [f"{s['id_1']}{s.get('eye_id_1','')}" for s in bilat_val] - else: - train_sids = [str(s["id_1"]) for s in bilat_train] - val_sids = [str(s["id_1"]) for s in bilat_val] - pred_store.set_split(fold, train_sids, "train") - pred_store.set_split(fold, val_sids, "val") - - if len(bilat_val) == 0: - print(f" [fold {fold+1}] WARNING: no bilateral val samples; skipping fold.", flush=True) - empty = FoldResult( - mode=mode, fold=fold, - best_epoch_single=0, best_epoch_bilat=0, - classic_val_auc=nan, classic_val_acc=nan, classic_val_kappa=nan, - classic_val_mcc=nan, classic_val_f1=nan, classic_val_recall=None, - classic_val_ece=nan, classic_val_threshold=nan, classic_val_bias=None, - classic_val_n=0, - ensemble_val_auc=nan, ensemble_val_acc=nan, ensemble_val_kappa=nan, - ensemble_val_mcc=nan, ensemble_val_f1=nan, ensemble_val_recall=None, - ensemble_val_ece=nan, ensemble_val_threshold=nan, ensemble_val_bias=None, - ensemble_val_n=0, - bilat_val_auc=nan, bilat_val_acc=nan, bilat_val_kappa=nan, - bilat_val_mcc=nan, bilat_val_f1=nan, bilat_val_recall=None, - bilat_val_ece=nan, bilat_val_threshold=nan, bilat_val_bias=None, - bilat_val_n=0, - classic_holdout_auc=nan, classic_holdout_acc=nan, - ensemble_holdout_auc=nan, ensemble_holdout_acc=nan, - bilat_holdout_auc=nan, bilat_holdout_acc=nan, - holdout_n=0, - single_train_n=len(eye_train), bilat_train_n=len(bilat_train), - ) - return empty, FoldArtifacts( - y_true_classic=None, probs_classic=None, - y_true_ensemble=None, probs_ensemble=None, - y_true_bilat=None, probs_bilat=None, - ) - - # ---- models ---------------------------------------------------- - single = None - bilateral = None - if run_single: - single = SingleEyeHT( - backbone=args.backbone, freeze_ratio=args.freeze_ratio, - augment=args.augment, clinical_data=data, - num_classes=num_classes, - cd_hidden_dim=args.cd_hidden_dim, fusion_dim=args.fusion_dim, - bridge_mode=getattr(args, "bridge_mode", "fused"), - ).to(device) - if run_bilat: - bilateral = BilateralHT( - backbone=args.backbone, freeze_ratio=args.freeze_ratio, - augment=args.augment, clinical_data=data, - num_classes=num_classes, - cd_hidden_dim=args.cd_hidden_dim, fusion_dim=args.fusion_dim, - ).to(device) - - slots_eye = profile_eye.slot_descriptors() - slots_patient = profile_patient.slot_descriptors() - loader_kw = dict(batch_size=args.batch_size, num_workers=args.num_workers, - image_cache=image_cache) - - # ---- loaders --------------------------------------------------- - use_balanced = bool(getattr(args, "balanced_sampling", False)) - train_single_loader = None - train_eval_loader = None # non-shuffled, no sampler — for per-epoch train logging - train_bilat_loader = None - cd_only_loader = None # image-free loader for md_warmup phase - if run_single: - single_sampler = build_balanced_sampler(eye_train) if use_balanced else None - train_single_loader = make_loader( - eye_train, slots_eye, - image_transform=single.transform, - image_preprocessor=image_preprocessor, - shuffle=True, - sampler=single_sampler, - **loader_kw, - ) - train_eval_loader = make_loader( - eye_train, slots_eye, - image_transform=build_eval_transform(args.backbone), - image_preprocessor=image_preprocessor, - shuffle=False, - **loader_kw, - ) - if single_warmup_cd > 0: - # MD-only loader: drop image_1 so PIL never opens files during md_warmup. - # Always use balanced sampling for md_warmup — MD features alone are weaker - # than images and collapse to majority class without class balancing. - slots_cd_only = {k: v for k, v in slots_eye.items() if k != "image_1"} - md_warmup_sampler = single_sampler if single_sampler is not None else build_balanced_sampler(eye_train) - cd_only_loader = make_loader( - eye_train, slots_cd_only, - image_transform=None, - image_preprocessor=None, - shuffle=True, - sampler=md_warmup_sampler, - **loader_kw, - ) - if run_bilat: - bilat_sampler = build_balanced_sampler(bilat_train) if use_balanced else None - train_bilat_loader = make_loader( - bilat_train, slots_patient, - image_transform=bilateral.transform, - image_preprocessor=image_preprocessor, - shuffle=True, - sampler=bilat_sampler, - **loader_kw, - ) - elif run_fused: - # Fused head trains on bilateral samples using the single model's transform. - fused_sampler = build_balanced_sampler(bilat_train) if use_balanced else None - train_bilat_loader = make_loader( - bilat_train, slots_patient, - image_transform=single.transform, - image_preprocessor=image_preprocessor, - shuffle=True, - sampler=fused_sampler, - **loader_kw, - ) - eval_transform = build_eval_transform(args.backbone) - val_loader = make_loader( - bilat_val, slots_patient, - image_transform=eval_transform, - image_preprocessor=image_preprocessor, - shuffle=False, - **loader_kw, - ) - - # ---- holdout loader (optional) --------------------------------- - holdout_bilat: list = [] - holdout_loader = None - if split.holdout is not None and not split.holdout.empty: - holdout_bilat = filter_bilateral_samples( - profile_patient.build_samples(df=split.holdout, clinical=data) - ) - if holdout_bilat: - holdout_loader = make_loader( - holdout_bilat, slots_patient, - image_transform=eval_transform, - image_preprocessor=image_preprocessor, - shuffle=False, - **loader_kw, - ) - print(f" [fold {fold+1}] holdout_n={len(holdout_bilat)} (bilateral patients)", flush=True) - if pred_store is not None: - pred_store.set_split(fold, [str(s["id_1"]) for s in holdout_bilat], "holdout") - - # ---- prebuild in-memory image cache (fold 0 only; shared dict fills for later folds) ---- - if image_cache is not None: - cache_workers = int(getattr(args, "cache_workers", 4)) - _loaders_to_warm = [ - train_single_loader, train_bilat_loader, val_loader, holdout_loader, - ] - for _ldr in _loaders_to_warm: - if _ldr is not None: - _ldr.dataset.prebuild_image_cache(cache_workers=cache_workers) - - opt_single = torch.optim.Adam(single.parameters(), lr=args.lr) if run_single else None - opt_bilateral = torch.optim.Adam(bilateral.parameters(), lr=args.lr) if run_bilat else None - - # ---- epoch log ------------------------------------------------- - epoch_fields = [ - "fold", "epoch", - "phase_single", "phase_bilat", - "main_epoch_single", "main_epoch_bilat", - "single_active", "bilat_active", - "single_train_loss", "single_train_acc", - # val — fused head (existing) - "classic_val_auc", "classic_val_acc", "classic_val_n", - "ensemble_val_auc", "ensemble_val_acc", "ensemble_val_n", - "bilat_train_loss", "bilat_train_acc", - "bilat_val_auc", "bilat_val_acc", "bilat_val_n", - # val — img/md heads + fusion events - "classic_val_auc_img", "classic_val_acc_img", - "classic_val_auc_cd", "classic_val_acc_cd", - "classic_val_fe_corr", "classic_val_fe_err", - "ensemble_val_auc_img", "ensemble_val_acc_img", - "ensemble_val_auc_cd", "ensemble_val_acc_cd", - "ensemble_val_fe_corr", "ensemble_val_fe_err", - "bilat_val_auc_img", "bilat_val_acc_img", - "bilat_val_auc_cd", "bilat_val_acc_cd", - "bilat_val_fe_corr", "bilat_val_fe_err", - # holdout — fused head (existing) - "classic_holdout_auc", "classic_holdout_acc", - "ensemble_holdout_auc", "ensemble_holdout_acc", - "bilat_holdout_auc", "bilat_holdout_acc", - # holdout — img/md heads + fusion events - "classic_holdout_auc_img", "classic_holdout_acc_img", - "classic_holdout_auc_cd", "classic_holdout_acc_cd", - "classic_holdout_fe_corr", "classic_holdout_fe_err", - "ensemble_holdout_auc_img", "ensemble_holdout_acc_img", - "ensemble_holdout_auc_cd", "ensemble_holdout_acc_cd", - "ensemble_holdout_fe_corr", "ensemble_holdout_fe_err", - # train-set eval pass (eval mode, all 3 heads) - "train_auc_fused", "train_acc_fused", - "train_auc_img", "train_acc_img", - "train_auc_cd", "train_acc_cd", - "train_fe_corr", "train_fe_err", - "train_n", - "is_best_single", "is_best_bilat", - "is_best_holdout_single", "is_best_holdout_bilat", - ] - # CM columns — named by num_classes so binary and multiclass both work - if num_classes == 2: - _cm_keys = ["tn", "fp", "fn", "tp"] - else: - _cm_keys = [f"cm_{i}_{j}" for i in range(num_classes) for j in range(num_classes)] - for _split in ("classic_val", "ensemble_val", "classic_holdout", "ensemble_holdout", "train"): - for _head in ("fused", "img", "md"): - for _k in _cm_keys: - epoch_fields.append(f"{_split}_{_head}_{_k}") - fold_logger = HypertowerLogger(run_dir=fold_dir) - - # per-epoch accumulation for npy tensors - _epoch_train_pf: list[np.ndarray] = [] - _epoch_train_pi: list[np.ndarray] = [] - _epoch_train_pm: list[np.ndarray] = [] - _epoch_train_ids: list[np.ndarray] = [] - _epoch_train_y: list[np.ndarray] = [] - # per-eye val accumulators (ensemble mode: OD and OS separate) - _epoch_val_pf_od: list[np.ndarray] = [] - _epoch_val_pi_od: list[np.ndarray] = [] - _epoch_val_pm_od: list[np.ndarray] = [] - _epoch_val_pf_os: list[np.ndarray] = [] - _epoch_val_pi_os: list[np.ndarray] = [] - _epoch_val_pm_os: list[np.ndarray] = [] - _epoch_val_y: list[np.ndarray] = [] - _epoch_val_ids: list[np.ndarray] = [] - - # ---- best-epoch trackers --------------------------------------- - best_single_auc = -1.0 - best_bilat_auc = -1.0 - best_epoch_single = 0 - best_epoch_bilat = 0 - best_single_state: Optional[dict] = None - best_bilat_state: Optional[dict] = None - snap_classic: dict = {} - snap_ensemble: dict = {} - snap_bilat: dict = {} - snap_holdout_single: dict = {} - snap_holdout_bilat: dict = {} - snap_fused: dict = {} - snap_holdout_fused: dict = {} - best_holdout_single_auc = -1.0 - best_holdout_bilat_auc = -1.0 - best_epoch_holdout_single = 0 - best_epoch_holdout_bilat = 0 - best_holdout_single_state: Optional[dict] = None - best_holdout_bilat_state: Optional[dict] = None - - if run_single: - print( - f" [fold {fold+1}] single_train_n={len(eye_train)} (eye-level) " - f"val_n={len(bilat_val)} " - f"single_warmup=md{single_warmup_cd}+twr{single_warmup_tower}+fus{single_warmup_fused} total={total_single_epochs}", - flush=True, - ) - else: - print( - f" [fold {fold+1}] bilat_train_n={len(bilat_train)} (bilateral) " - f"val_n={len(bilat_val)} " - f"bilat_warmup={bilat_warmup_tower}+{bilat_warmup_fused} total={total_bilat_epochs}", - flush=True, - ) - - # ---- epoch loop ------------------------------------------------ - _prev_phase_single = "inactive" # used to detect md_warmup → next phase transition - for epoch in range(total_epochs): - _epoch_t0 = time.time() - if not run_single: - phase_single, main_epoch_single, single_active = "inactive", 0, False - elif epoch < single_warmup_cd: - phase_single, main_epoch_single, single_active = "cd_warmup", 0, True - elif epoch < (single_warmup_cd + single_warmup_tower): - phase_single, main_epoch_single, single_active = "tower_warmup", 0, True - elif epoch < (single_warmup_cd + single_warmup_tower + single_warmup_fused): - phase_single, main_epoch_single, single_active = "fused_warmup", 0, True - elif epoch < total_single_epochs: - phase_single, main_epoch_single, single_active = ( - "main", - epoch - single_warmup_cd - single_warmup_tower - single_warmup_fused + 1, - True, - ) - else: - phase_single, main_epoch_single, single_active = "done", main_epochs, False - - if not run_bilat: - phase_bilat, main_epoch_bilat, bilat_active = "inactive", 0, False - elif epoch < bilat_warmup_tower: - phase_bilat, main_epoch_bilat, bilat_active = "tower_warmup", 0, True - elif epoch < (bilat_warmup_tower + bilat_warmup_fused): - phase_bilat, main_epoch_bilat, bilat_active = "fused_warmup", 0, True - elif epoch < total_bilat_epochs: - phase_bilat, main_epoch_bilat, bilat_active = ( - "main", - epoch - bilat_warmup_tower - bilat_warmup_fused + 1, - True, - ) - else: - phase_bilat, main_epoch_bilat, bilat_active = "done", main_epochs, False - - if run_single and single_active: - _active_loader = cd_only_loader if phase_single == "cd_warmup" else train_single_loader - sl_loss, sl_acc = train_single_epoch( - single, _active_loader, opt_single, device, - phase=phase_single, bcd_prob=float(args.bcd_prob), - tower_loss_mode=args.tower_loss_mode, - ) - else: - sl_loss, sl_acc = nan, nan - - if run_bilat and bilat_active: - bl_loss, bl_acc = train_bilateral_epoch( - bilateral, train_bilat_loader, opt_bilateral, device, - phase=phase_bilat, bcd_prob=float(args.bcd_prob), - tower_loss_mode=args.tower_loss_mode, - ) - else: - bl_loss, bl_acc = nan, nan - - _skip_val_eval = (phase_single == "cd_warmup") - - if run_single and tower_mode == "single" and not _skip_val_eval: - y_cl, p_cl, p_cl_img, p_cl_cd = collect_probs_single_components( - single, val_loader, device, aggregate_patient=False - ) - cl_acc, cl_auc, cl_n = _score_arrays(y_cl, p_cl, num_classes) - cl_acc_img = float((p_cl_img.argmax(1) == y_cl).mean()) if y_cl.size else nan - cl_acc_cd = float((p_cl_cd.argmax(1) == y_cl).mean()) if y_cl.size else nan - _, cl_auc_img, _ = _score_arrays(y_cl, p_cl_img, num_classes) - _, cl_auc_cd, _ = _score_arrays(y_cl, p_cl_cd, num_classes) - y_en = np.array([], dtype=np.int64) - p_en = p_en_img = p_en_cd = np.zeros((0, num_classes), dtype=np.float32) - en_acc = en_auc = nan - en_n = 0 - en_acc_img = en_acc_cd = en_auc_img = en_auc_cd = nan - elif run_single and tower_mode == "ensemble" and not _skip_val_eval: - (y_en, - _p_en_f_od, _p_en_i_od, _p_en_m_od, - _p_en_f_os, _p_en_i_os, _p_en_m_os, - _en_pat_ids) = collect_probs_ensemble_pereye( - single, val_loader, device, return_ids=True - ) - # patient-level averages (used for metrics, same as before) - p_en = 0.5 * (_p_en_f_od + _p_en_f_os) - p_en_img = 0.5 * (_p_en_i_od + _p_en_i_os) - p_en_cd = 0.5 * (_p_en_m_od + _p_en_m_os) - en_acc, en_auc, en_n = _score_arrays(y_en, p_en, num_classes) - en_acc_img = float((p_en_img.argmax(1) == y_en).mean()) if y_en.size else nan - en_acc_cd = float((p_en_cd.argmax(1) == y_en).mean()) if y_en.size else nan - _, en_auc_img, _ = _score_arrays(y_en, p_en_img, num_classes) - _, en_auc_cd, _ = _score_arrays(y_en, p_en_cd, num_classes) - y_cl = np.array([], dtype=np.int64) - p_cl = p_cl_img = p_cl_cd = np.zeros((0, num_classes), dtype=np.float32) - cl_acc = cl_auc = nan - cl_n = 0 - cl_acc_img = cl_acc_cd = cl_auc_img = cl_auc_cd = nan - else: - y_cl = y_en = np.array([], dtype=np.int64) - p_cl = p_cl_img = p_cl_cd = np.zeros((0, num_classes), dtype=np.float32) - p_en = p_en_img = p_en_cd = np.zeros((0, num_classes), dtype=np.float32) - cl_acc = cl_auc = en_acc = en_auc = nan - cl_n = en_n = 0 - cl_acc_img = cl_acc_cd = en_acc_img = en_acc_cd = nan - cl_auc_img = cl_auc_cd = en_auc_img = en_auc_cd = nan - - if run_bilat and not _skip_val_eval: - y_bi, p_bi, p_bi_img, p_bi_cd = collect_probs_bilateral_components( - bilateral, val_loader, device - ) - bi_acc, bi_auc, bi_n = _score_arrays(y_bi, p_bi, num_classes) - bi_acc_img = float((p_bi_img.argmax(1) == y_bi).mean()) if y_bi.size else nan - bi_acc_cd = float((p_bi_cd.argmax(1) == y_bi).mean()) if y_bi.size else nan - _, bi_auc_img, _ = _score_arrays(y_bi, p_bi_img, num_classes) - _, bi_auc_cd, _ = _score_arrays(y_bi, p_bi_cd, num_classes) - else: - y_bi = np.array([], dtype=np.int64) - p_bi = np.zeros((0, 0), dtype=np.float32) - bi_acc = bi_auc = nan - bi_n = 0 - bi_acc_img = bi_acc_cd = bi_auc_img = bi_auc_cd = nan - - # --- holdout evaluation ------------------------------------ - # defaults (overwritten below when holdout_loader is not None) - _z2 = np.zeros((0, num_classes), dtype=np.float32) - _e2 = np.array([], dtype=np.int64) - y_cl_h = y_en_h = _e2 - p_cl_h = p_cl_h_img = p_cl_h_cd = _z2 - p_en_h = p_en_h_img = p_en_h_cd = _z2 - - if holdout_loader is not None and not _skip_val_eval: - if run_single and tower_mode == "single": - y_cl_h, p_cl_h, p_cl_h_img, p_cl_h_cd = collect_probs_single_components( - single, holdout_loader, device, aggregate_patient=False - ) - _, cl_auc_h, _ = _score_arrays(y_cl_h, p_cl_h, num_classes) - cl_acc_h = float((p_cl_h.argmax(1) == y_cl_h).mean()) if y_cl_h.size else nan - _, cl_auc_h_img, _ = _score_arrays(y_cl_h, p_cl_h_img, num_classes) - cl_acc_h_img = float((p_cl_h_img.argmax(1) == y_cl_h).mean()) if y_cl_h.size else nan - _, cl_auc_h_md, _ = _score_arrays(y_cl_h, p_cl_h_cd, num_classes) - cl_acc_h_md = float((p_cl_h_cd.argmax(1) == y_cl_h).mean()) if y_cl_h.size else nan - cl_fe_h_corr, cl_fe_h_err = _fusion_events(y_cl_h, p_cl_h, p_cl_h_img, p_cl_h_cd) - en_auc_h = en_acc_h = nan - en_auc_h_img = en_acc_h_img = en_auc_h_md = en_acc_h_md = nan - en_fe_h_corr = en_fe_h_err = 0 - elif run_single and tower_mode == "ensemble": - y_en_h, p_en_h, p_en_h_img, p_en_h_cd = collect_probs_single_components( - single, holdout_loader, device, aggregate_patient=True - ) - _, en_auc_h, _ = _score_arrays(y_en_h, p_en_h, num_classes) - en_acc_h = float((p_en_h.argmax(1) == y_en_h).mean()) if y_en_h.size else nan - _, en_auc_h_img, _ = _score_arrays(y_en_h, p_en_h_img, num_classes) - en_acc_h_img = float((p_en_h_img.argmax(1) == y_en_h).mean()) if y_en_h.size else nan - _, en_auc_h_md, _ = _score_arrays(y_en_h, p_en_h_cd, num_classes) - en_acc_h_md = float((p_en_h_cd.argmax(1) == y_en_h).mean()) if y_en_h.size else nan - en_fe_h_corr, en_fe_h_err = _fusion_events(y_en_h, p_en_h, p_en_h_img, p_en_h_cd) - cl_auc_h = cl_acc_h = nan - cl_auc_h_img = cl_acc_h_img = cl_auc_h_md = cl_acc_h_md = nan - cl_fe_h_corr = cl_fe_h_err = 0 - else: - cl_auc_h = cl_acc_h = en_auc_h = en_acc_h = nan - cl_auc_h_img = cl_acc_h_img = cl_auc_h_md = cl_acc_h_md = nan - en_auc_h_img = en_acc_h_img = en_auc_h_md = en_acc_h_md = nan - cl_fe_h_corr = cl_fe_h_err = en_fe_h_corr = en_fe_h_err = 0 - if run_bilat: - y_bi_h, p_bi_h, _, _ = collect_probs_bilateral_components( - bilateral, holdout_loader, device - ) - _, bi_auc_h, _ = _score_arrays(y_bi_h, p_bi_h, num_classes) - bi_acc_h = float((p_bi_h.argmax(1) == y_bi_h).mean()) if y_bi_h.size else nan - else: - bi_auc_h = bi_acc_h = nan - else: - cl_auc_h = cl_acc_h = en_auc_h = en_acc_h = bi_auc_h = bi_acc_h = nan - cl_auc_h_img = cl_acc_h_img = cl_auc_h_md = cl_acc_h_md = nan - en_auc_h_img = en_acc_h_img = en_auc_h_md = en_acc_h_md = nan - cl_fe_h_corr = cl_fe_h_err = en_fe_h_corr = en_fe_h_err = 0 - - # --- fusion-event helpers for val sets ---------------------- - cl_fe_corr, cl_fe_err = _fusion_events(y_cl, p_cl, p_cl_img, p_cl_cd) if y_cl.size else (0, 0) - en_fe_corr, en_fe_err = _fusion_events(y_en, p_en, p_en_img, p_en_cd) if y_en.size else (0, 0) - bi_fe_corr, bi_fe_err = (0, 0) # bilateral components not separated the same way - - # --- train eval pass (eval mode, all 3 heads) ---------------- - tr_auc_f = tr_acc_f = tr_auc_i = tr_acc_i = tr_auc_m = tr_acc_m = nan - tr_fe_corr = tr_fe_err = tr_n = 0 - y_tr = np.array([], dtype=np.int64) - p_tr_f = p_tr_i = p_tr_m = np.zeros((0, num_classes), dtype=np.float32) - if run_single and train_eval_loader is not None and not _skip_val_eval: - y_tr, p_tr_f, p_tr_i, p_tr_m, tr_ids = collect_probs_eye_level( - single, train_eval_loader, device, return_ids=True - ) - if y_tr.size: - _, tr_auc_f, _ = _score_arrays(y_tr, p_tr_f, num_classes) - tr_acc_f = float((p_tr_f.argmax(1) == y_tr).mean()) - _, tr_auc_i, _ = _score_arrays(y_tr, p_tr_i, num_classes) - tr_acc_i = float((p_tr_i.argmax(1) == y_tr).mean()) - _, tr_auc_m, _ = _score_arrays(y_tr, p_tr_m, num_classes) - tr_acc_m = float((p_tr_m.argmax(1) == y_tr).mean()) - tr_fe_corr, tr_fe_err = _fusion_events(y_tr, p_tr_f, p_tr_i, p_tr_m) - tr_n = int(y_tr.size) - # accumulate for npy tensors - _epoch_train_pf.append(p_tr_f) - _epoch_train_pi.append(p_tr_i) - _epoch_train_pm.append(p_tr_m) - _epoch_train_ids.append(tr_ids) - _epoch_train_y.append(y_tr) - # record into PredictionStore - if pred_store is not None: - if tower_mode in ("single", "classic"): - pred_store.record(fold, epoch, tr_ids, "fused", p_tr_f) - pred_store.record(fold, epoch, tr_ids, "img", p_tr_i) - pred_store.record(fold, epoch, tr_ids, "md", p_tr_m) - else: # ensemble: separate OD and OS by eye suffix - od_mask = np.array([str(i).endswith("OD") for i in tr_ids]) - os_mask = ~od_mask - od_pids = [str(i)[:-2] for i in tr_ids[od_mask]] - os_pids = [str(i)[:-2] for i in tr_ids[os_mask]] - pred_store.record(fold, epoch, od_pids, "od_fused", p_tr_f[od_mask]) - pred_store.record(fold, epoch, od_pids, "od_img", p_tr_i[od_mask]) - pred_store.record(fold, epoch, od_pids, "od_md", p_tr_m[od_mask]) - pred_store.record(fold, epoch, os_pids, "os_fused", p_tr_f[os_mask]) - pred_store.record(fold, epoch, os_pids, "os_img", p_tr_i[os_mask]) - pred_store.record(fold, epoch, os_pids, "os_md", p_tr_m[os_mask]) - - # accumulate val for npy tensors - if run_single and tower_mode == "ensemble" and y_en.size: - _epoch_val_pf_od.append(_p_en_f_od) - _epoch_val_pi_od.append(_p_en_i_od) - _epoch_val_pm_od.append(_p_en_m_od) - _epoch_val_pf_os.append(_p_en_f_os) - _epoch_val_pi_os.append(_p_en_i_os) - _epoch_val_pm_os.append(_p_en_m_os) - _epoch_val_y.append(y_en) - _epoch_val_ids.append(_en_pat_ids) - elif run_single and tower_mode == "single" and y_cl.size: - # single mode: no per-eye split, reuse same array for both slots - _epoch_val_pf_od.append(p_cl) - _epoch_val_pi_od.append(p_cl_img) - _epoch_val_pm_od.append(p_cl_cd) - _epoch_val_pf_os.append(p_cl) - _epoch_val_pi_os.append(p_cl_img) - _epoch_val_pm_os.append(p_cl_cd) - _epoch_val_y.append(y_cl) - - # record val into PredictionStore - if pred_store is not None: - if run_single and tower_mode == "ensemble" and y_en.size: - pred_store.record(fold, epoch, _en_pat_ids, "od_fused", _p_en_f_od) - pred_store.record(fold, epoch, _en_pat_ids, "od_img", _p_en_i_od) - pred_store.record(fold, epoch, _en_pat_ids, "od_md", _p_en_m_od) - pred_store.record(fold, epoch, _en_pat_ids, "os_fused", _p_en_f_os) - pred_store.record(fold, epoch, _en_pat_ids, "os_img", _p_en_i_os) - pred_store.record(fold, epoch, _en_pat_ids, "os_md", _p_en_m_os) - elif run_single and tower_mode == "single" and y_cl.size: - # val in single mode: collect_probs_single_components(aggregate_patient=False) - # returns interleaved [all_OD, all_OS] per batch — IDs not tracked here yet - pass # single-mode val IDs not currently available; train IDs are sufficient - - # Best-epoch checks (restricted to main phase). - target_single_auc = cl_auc if tower_mode == "single" else en_auc - target_holdout_single_auc = cl_auc_h if tower_mode == "single" else en_auc_h - single_ckpt_eligible = run_single and (phase_single == "main") - is_best_single = ( - single_ckpt_eligible - and (not np.isnan(target_single_auc)) - and (target_single_auc > best_single_auc) - ) - if is_best_single: - best_single_auc = target_single_auc - best_epoch_single = epoch + 1 - best_single_state = copy.deepcopy(single.state_dict()) - if tower_mode == "single": - snap_cl, _, _, _ = _tune_and_snap(y_cl, p_cl, cl_acc, num_classes, args, args.ece_bins) - snap_classic = snap_cl - else: - snap_en, _, _, _ = _tune_and_snap(y_en, p_en, en_acc, num_classes, args, args.ece_bins) - snap_ensemble = snap_en - snap_holdout_single = { - "auc": float(target_holdout_single_auc), - "acc": float(cl_acc_h if tower_mode == "single" else en_acc_h), - } - - is_best_holdout_single = ( - holdout_loader is not None - and single_ckpt_eligible - and (not np.isnan(target_holdout_single_auc)) - and (target_holdout_single_auc > best_holdout_single_auc) - ) - if is_best_holdout_single: - best_holdout_single_auc = target_holdout_single_auc - best_epoch_holdout_single = epoch + 1 - best_holdout_single_state = copy.deepcopy(single.state_dict()) - - bilat_ckpt_eligible = run_bilat and (phase_bilat == "main") - is_best_bilat = ( - bilat_ckpt_eligible - and (not np.isnan(bi_auc)) - and (bi_auc > best_bilat_auc) - ) - if is_best_bilat: - best_bilat_auc = bi_auc - best_epoch_bilat = epoch + 1 - best_bilat_state = copy.deepcopy(bilateral.state_dict()) - snap_bi, _, _, _ = _tune_and_snap(y_bi, p_bi, bi_acc, num_classes, args, args.ece_bins) - snap_bilat = snap_bi - snap_holdout_bilat = {"auc": float(bi_auc_h), "acc": float(bi_acc_h)} - - is_best_holdout_bilat = ( - holdout_loader is not None - and bilat_ckpt_eligible - and (not np.isnan(bi_auc_h)) - and (bi_auc_h > best_holdout_bilat_auc) - ) - if is_best_holdout_bilat: - best_holdout_bilat_auc = bi_auc_h - best_epoch_holdout_bilat = epoch + 1 - best_holdout_bilat_state = copy.deepcopy(bilateral.state_dict()) - - # --- confusion matrix cells per split × head ------------------- - def _prefixed_cm(prefix: str, y: np.ndarray, pf: np.ndarray, - pi: np.ndarray, pm: np.ndarray) -> dict: - out: dict = {} - for head, p in (("fused", pf), ("img", pi), ("md", pm)): - for k, v in _cm_cells(y, p, num_classes).items(): - out[f"{prefix}_{head}_{k}"] = v - return out - - cm_row: dict = {} - cm_row.update(_prefixed_cm("classic_val", y_cl, p_cl, p_cl_img, p_cl_cd)) - cm_row.update(_prefixed_cm("ensemble_val", y_en, p_en, p_en_img, p_en_cd)) - cm_row.update(_prefixed_cm("classic_holdout", y_cl_h, p_cl_h, p_cl_h_img, p_cl_h_cd)) - cm_row.update(_prefixed_cm("ensemble_holdout", y_en_h, p_en_h, p_en_h_img, p_en_h_cd)) - cm_row.update(_prefixed_cm("train", y_tr, p_tr_f, p_tr_i, p_tr_m)) - - fold_logger.write_epoch_row({ - "fold": fold, "epoch": epoch + 1, - "phase_single": phase_single, "phase_bilat": phase_bilat, - "main_epoch_single": main_epoch_single, "main_epoch_bilat": main_epoch_bilat, - "single_active": int(single_active), "bilat_active": int(bilat_active), - "single_train_loss": _f(sl_loss), "single_train_acc": _f(sl_acc), - # val — fused - "classic_val_auc": _f(cl_auc), "classic_val_acc": _f(cl_acc), "classic_val_n": cl_n, - "ensemble_val_auc": _f(en_auc), "ensemble_val_acc": _f(en_acc), "ensemble_val_n": en_n, - "bilat_train_loss": _f(bl_loss), "bilat_train_acc": _f(bl_acc), - "bilat_val_auc": _f(bi_auc), "bilat_val_acc": _f(bi_acc), "bilat_val_n": bi_n, - # val — img/md + fusion events - "classic_val_auc_img": _f(cl_auc_img), "classic_val_acc_img": _f(cl_acc_img), - "classic_val_auc_cd": _f(cl_auc_cd), "classic_val_acc_cd": _f(cl_acc_cd), - "classic_val_fe_corr": cl_fe_corr, "classic_val_fe_err": cl_fe_err, - "ensemble_val_auc_img": _f(en_auc_img), "ensemble_val_acc_img": _f(en_acc_img), - "ensemble_val_auc_cd": _f(en_auc_cd), "ensemble_val_acc_cd": _f(en_acc_cd), - "ensemble_val_fe_corr": en_fe_corr, "ensemble_val_fe_err": en_fe_err, - "bilat_val_auc_img": _f(bi_auc_img), "bilat_val_acc_img": _f(bi_acc_img), - "bilat_val_auc_cd": _f(bi_auc_cd), "bilat_val_acc_cd": _f(bi_acc_cd), - "bilat_val_fe_corr": bi_fe_corr, "bilat_val_fe_err": bi_fe_err, - # holdout — fused - "classic_holdout_auc": _f(cl_auc_h), "classic_holdout_acc": _f(cl_acc_h), - "ensemble_holdout_auc": _f(en_auc_h), "ensemble_holdout_acc": _f(en_acc_h), - "bilat_holdout_auc": _f(bi_auc_h), "bilat_holdout_acc": _f(bi_acc_h), - # holdout — img/md + fusion events - "classic_holdout_auc_img": _f(cl_auc_h_img), "classic_holdout_acc_img": _f(cl_acc_h_img), - "classic_holdout_auc_cd": _f(cl_auc_h_md), "classic_holdout_acc_cd": _f(cl_acc_h_md), - "classic_holdout_fe_corr": cl_fe_h_corr, "classic_holdout_fe_err": cl_fe_h_err, - "ensemble_holdout_auc_img": _f(en_auc_h_img), "ensemble_holdout_acc_img": _f(en_acc_h_img), - "ensemble_holdout_auc_cd": _f(en_auc_h_md), "ensemble_holdout_acc_cd": _f(en_acc_h_md), - "ensemble_holdout_fe_corr": en_fe_h_corr, "ensemble_holdout_fe_err": en_fe_h_err, - # train eval pass - "train_auc_fused": _f(tr_auc_f), "train_acc_fused": _f(tr_acc_f), - "train_auc_img": _f(tr_auc_i), "train_acc_img": _f(tr_acc_i), - "train_auc_cd": _f(tr_auc_m), "train_acc_cd": _f(tr_acc_m), - "train_fe_corr": tr_fe_corr, "train_fe_err": tr_fe_err, - "train_n": tr_n, - "is_best_single": int(is_best_single), - "is_best_bilat": int(is_best_bilat), - "is_best_holdout_single": int(is_best_holdout_single), - "is_best_holdout_bilat": int(is_best_holdout_bilat), - **cm_row, - }, optional_cols=epoch_fields) - - # ---- md_warmup progress bar (replaces per-epoch print) -------- - if phase_single == "cd_warmup": - _bar_w = 30 - _filled = int(_bar_w * (epoch + 1) / single_warmup_cd) - _bar = "#" * _filled + "-" * (_bar_w - _filled) - _bar_msg = ( - f" [fold {fold+1}] md_warmup [{_bar}] " - f"{epoch + 1}/{single_warmup_cd} loss={sl_loss:.4f}" - ) - print(f"\r{_bar_msg}", end="", flush=True) - fold_logger.info(_bar_msg) - _prev_phase_single = phase_single - continue # skip normal log block entirely - - if _prev_phase_single == "cd_warmup": - print() # seal the progress bar line - - if args.log_every > 0 and (epoch + 1) % args.log_every == 0: - _epoch_secs = time.time() - _epoch_t0 - hld_auc = target_holdout_single_auc if run_single else bi_auc_h - hld_suffix = f" hld_auc={hld_auc:.4f}" if holdout_loader is not None else "" - - # Human-readable phase progress for console logs. - if phase_single == "tower_warmup": - single_phase_epoch = epoch - single_warmup_cd + 1 - single_phase_total = single_warmup_tower - elif phase_single == "fused_warmup": - single_phase_epoch = epoch - single_warmup_cd - single_warmup_tower + 1 - single_phase_total = single_warmup_fused - else: - single_phase_epoch = main_epoch_single - single_phase_total = main_epochs - - if phase_bilat == "tower_warmup": - bilat_phase_epoch = epoch + 1 - bilat_phase_total = bilat_warmup_tower - elif phase_bilat == "fused_warmup": - bilat_phase_epoch = epoch - bilat_warmup_tower + 1 - bilat_phase_total = bilat_warmup_fused - else: - bilat_phase_epoch = main_epoch_bilat - bilat_phase_total = main_epochs - - if run_single: - if tower_mode == "single": - msg = ( - f" ep {epoch+1:>3}/{total_epochs} ({_epoch_secs:.1f}s) " - f"[single:{phase_single} {single_phase_epoch}/{single_phase_total}] " - f"fused(acc={cl_acc:.4f},auc={cl_auc:.4f}) " - f"img(acc={cl_acc_img:.4f},auc={cl_auc_img:.4f}) " - f"md(acc={cl_acc_cd:.4f},auc={cl_auc_cd:.4f}) " - f"(best_fused={best_single_auc:.4f} @ep{best_epoch_single})" - f"{hld_suffix}" - ) - else: - msg = ( - f" ep {epoch+1:>3}/{total_epochs} ({_epoch_secs:.1f}s) " - f"[single:{phase_single} {single_phase_epoch}/{single_phase_total}] " - f"fused(acc={en_acc:.4f},auc={en_auc:.4f}) " - f"img(acc={en_acc_img:.4f},auc={en_auc_img:.4f}) " - f"md(acc={en_acc_cd:.4f},auc={en_auc_cd:.4f}) " - f"(best_fused={best_single_auc:.4f} @ep{best_epoch_single})" - f"{hld_suffix}" - ) - else: - msg = ( - f" ep {epoch+1:>3}/{total_epochs} ({_epoch_secs:.1f}s) " - f"[bilat:{phase_bilat} {bilat_phase_epoch}/{bilat_phase_total}] " - f"fused(acc={bi_acc:.4f},auc={bi_auc:.4f}) " - f"img(acc={bi_acc_img:.4f},auc={bi_auc_img:.4f}) " - f"md(acc={bi_acc_cd:.4f},auc={bi_auc_cd:.4f}) " - f"(best_bilat={best_bilat_auc:.4f} @ep{best_epoch_bilat})" - f"{hld_suffix}" - ) - print(msg, flush=True) - fold_logger.info(msg) - - _prev_phase_single = phase_single - - fold_logger.close() - - # Override: use final epoch state instead of best-AUC checkpoint - if getattr(args, "use_last_epoch", False): - best_single_state = copy.deepcopy(single.state_dict()) - if run_bilat: - best_bilat_state = copy.deepcopy(bilat.state_dict()) - - if args.save_checkpoints: - if best_single_state is not None: - torch.save(best_single_state, fold_dir / "best_single.pt") - if best_bilat_state is not None: - torch.save(best_bilat_state, fold_dir / "best_bilateral.pt") - if best_holdout_single_state is not None: - torch.save(best_holdout_single_state, fold_dir / "best_holdout_single.pt") - if best_holdout_bilat_state is not None: - torch.save(best_holdout_bilat_state, fold_dir / "best_holdout_bilateral.pt") - - # ---- Save per-epoch per-patient npy tensors ------------------------- - if _epoch_train_pf: - # Use the order from the first epoch (consistent since loader is non-shuffled) - ids_ref = _epoch_train_ids[0] - y_ref = _epoch_train_y[0] - np.save(fold_dir / "train_patient_ids.npy", ids_ref) - np.save(fold_dir / "train_y_true.npy", y_ref) - np.save(fold_dir / "train_probs_fused.npy", np.stack(_epoch_train_pf)) # (n_ep, n_pts, n_cls) - np.save(fold_dir / "train_probs_img.npy", np.stack(_epoch_train_pi)) - np.save(fold_dir / "train_probs_cd.npy", np.stack(_epoch_train_pm)) - if _epoch_val_pf_od: - np.save(fold_dir / "val_y_true_epochs.npy", np.stack(_epoch_val_y)) # (n_ep, N) - np.save(fold_dir / "val_probs_fused_od_epochs.npy", np.stack(_epoch_val_pf_od)) # (n_ep, N, C) - np.save(fold_dir / "val_probs_img_od_epochs.npy", np.stack(_epoch_val_pi_od)) - np.save(fold_dir / "val_probs_cd_od_epochs.npy", np.stack(_epoch_val_pm_od)) - np.save(fold_dir / "val_probs_fused_os_epochs.npy", np.stack(_epoch_val_pf_os)) - np.save(fold_dir / "val_probs_img_os_epochs.npy", np.stack(_epoch_val_pi_os)) - np.save(fold_dir / "val_probs_cd_os_epochs.npy", np.stack(_epoch_val_pm_os)) - if _epoch_val_ids: # only ensemble mode populates this - np.save(fold_dir / "val_patient_ids.npy", _epoch_val_ids[0]) - - # ---- Phase 2: fused head training (ensemble + --fused-head only) ---- - best_fused_auc = -1.0 - best_fused_state: Optional[dict] = None - best_holdout_fused_auc = -1.0 - - if run_fused and best_single_state is not None: - # Revert base to its best val checkpoint, then freeze it. - single.load_state_dict(best_single_state) - for p in single.parameters(): - p.requires_grad_(False) - - fused = FusedEnsembleHT(single, num_classes).to(device) - opt_fused = torch.optim.Adam(fused.eye_scorer.parameters(), lr=args.lr) - fusion_epochs = int(getattr(args, "fusion_epochs", 10)) - - print( - f" [fold {fold+1}] Phase 2: training fusion head " - f"bilat_train_n={len(bilat_train)} fusion_epochs={fusion_epochs}", - flush=True, - ) - - _val_pids_for_store = [str(s["id_1"]) for s in bilat_val] - for fep in range(fusion_epochs): - fu_loss, fu_acc = train_fusion_epoch(fused, train_bilat_loader, opt_fused, device) - y_fu, p_fu = collect_probs_fused(fused, val_loader, device) - fu_auc = _score_arrays(y_fu, p_fu, num_classes)[1] - if pred_store is not None and y_fu.size: - _store_ep = total_single_epochs + fep - pred_store.record(fold, _store_ep, _val_pids_for_store, "bilat_fused", p_fu) - - # Holdout eval (if available) - fu_hld_auc = nan - if holdout_loader is not None: - y_fu_h, p_fu_h = collect_probs_fused(fused, holdout_loader, device) - fu_hld_auc = _score_arrays(y_fu_h, p_fu_h, num_classes)[1] - - is_best_fused = not np.isnan(fu_auc) and fu_auc > best_fused_auc - if is_best_fused: - best_fused_auc = fu_auc - best_fused_state = copy.deepcopy(fused.state_dict()) - fu_acc_val = _score_arrays(y_fu, p_fu, num_classes)[0] - snap_fused, _, _, _ = _tune_and_snap(y_fu, p_fu, fu_acc_val, num_classes, args, args.ece_bins) - - is_best_hld_fused = not np.isnan(fu_hld_auc) and fu_hld_auc > best_holdout_fused_auc - if is_best_hld_fused: - best_holdout_fused_auc = fu_hld_auc - snap_holdout_fused = {"auc": fu_hld_auc, "acc": _score_arrays(y_fu_h, p_fu_h, num_classes)[0]} - - if (fep + 1) % max(1, getattr(args, "log_every", 1)) == 0: - print( - f" [fold {fold+1}] fusion ep{fep+1:>3} " - f"loss={fu_loss:.4f} train_acc={fu_acc:.4f} " - f"val_auc={fu_auc:.4f} hld_auc={fu_hld_auc:.4f}" - f"{' *' if is_best_fused else ''}", - flush=True, - ) - - if best_fused_state is not None: - if args.save_checkpoints: - torch.save(best_fused_state, fold_dir / "best_fused.pt") - print( - f" [fold {fold+1}] BEST " - f"fused_head(acc={snap_fused.get('acc', nan):.4f}," - f"auc={snap_fused.get('auc', nan):.4f}) " - f"kappa={snap_fused.get('kappa', nan):.4f} " - f"F1={snap_fused.get('macro_f1', nan):.4f} " - f"ECE={snap_fused.get('ece', nan):.4f} " - f"holdout_auc={snap_holdout_fused.get('auc', nan):.4f}", - flush=True, - ) - - if run_single: - if tower_mode == "single": - print( - f" [fold {fold+1}] BEST " - f"fused(acc={snap_classic.get('acc', nan):.4f},auc={snap_classic.get('auc', nan):.4f}) " - f"kappa={snap_classic.get('kappa', nan):.4f} " - f"F1={snap_classic.get('macro_f1', nan):.4f} " - f"ECE={snap_classic.get('ece', nan):.4f} @ep{best_epoch_single}", - flush=True, - ) - else: - print( - f" [fold {fold+1}] BEST " - f"ensemble(acc={snap_ensemble.get('acc', nan):.4f},auc={snap_ensemble.get('auc', nan):.4f}) " - f"kappa={snap_ensemble.get('kappa', nan):.4f} " - f"F1={snap_ensemble.get('macro_f1', nan):.4f} " - f"ECE={snap_ensemble.get('ece', nan):.4f} @ep{best_epoch_single}", - flush=True, - ) - else: - print( - f" [fold {fold+1}] BEST " - f"fused(acc={snap_bilat.get('acc', nan):.4f},auc={snap_bilat.get('auc', nan):.4f}) " - f"kappa={snap_bilat.get('kappa', nan):.4f} " - f"F1={snap_bilat.get('macro_f1', nan):.4f} " - f"ECE={snap_bilat.get('ece', nan):.4f} @ep{best_epoch_bilat}", - flush=True, - ) - - # Export best-epoch prediction artifacts. - if run_single and best_single_state is not None: - single.load_state_dict(best_single_state) - if run_bilat and best_bilat_state is not None: - bilateral.load_state_dict(best_bilat_state) - y_en_pe_best = p_en_pe_best = p_en_pe_best_img = p_en_pe_best_md = None - l_en_best = l_en_best_img = l_en_best_md = None - l_cl_best = l_cl_best_img = l_cl_best_md = None - l_en_pe_best = l_en_pe_best_img = l_en_pe_best_md = None - if run_single and tower_mode == "single": - y_cl_best, p_cl_best, p_cl_best_img, p_cl_best_md, \ - l_cl_best, l_cl_best_img, l_cl_best_md = collect_probs_single_components( - single, val_loader, device, aggregate_patient=False, return_logits=True - ) - y_en_best = p_en_best = p_en_best_img = p_en_best_md = None - elif run_single and tower_mode == "ensemble": - y_en_best, p_en_best, p_en_best_img, p_en_best_md, \ - l_en_best, l_en_best_img, l_en_best_md = collect_probs_single_components( - single, val_loader, device, aggregate_patient=True, return_logits=True - ) - y_en_pe_best, p_en_pe_best, p_en_pe_best_img, p_en_pe_best_md, \ - l_en_pe_best, l_en_pe_best_img, l_en_pe_best_md = collect_probs_single_components( - single, val_loader, device, aggregate_patient=False, return_logits=True - ) - y_cl_best = p_cl_best = p_cl_best_img = p_cl_best_md = None - else: - y_cl_best = y_en_best = None - p_cl_best = p_en_best = p_en_best_img = p_en_best_md = None - p_cl_best_img = p_cl_best_md = None - if run_bilat: - y_bi_best, p_bi_best = collect_probs_bilateral(bilateral, val_loader, device) - else: - y_bi_best = p_bi_best = None - - y_fu_best = p_fu_best = None - if run_fused and best_fused_state is not None: - fused.load_state_dict(best_fused_state) - y_fu_best, p_fu_best = collect_probs_fused(fused, val_loader, device) - - return FoldResult( - mode=mode, fold=fold, - best_epoch_single=best_epoch_single, best_epoch_bilat=best_epoch_bilat, - classic_val_auc=snap_classic.get("auc", nan), - classic_val_acc=snap_classic.get("acc", nan), - classic_val_kappa=snap_classic.get("kappa", nan), - classic_val_mcc=snap_classic.get("mcc", nan), - classic_val_f1=snap_classic.get("macro_f1", nan), - classic_val_recall=_sv(snap_classic.get("per_class_recall")), - classic_val_ece=snap_classic.get("ece", nan), - classic_val_threshold=snap_classic.get("threshold", nan), - classic_val_bias=_svf(snap_classic.get("bias")), - classic_val_n=snap_classic.get("n", 0), - ensemble_val_auc=snap_ensemble.get("auc", nan), - ensemble_val_acc=snap_ensemble.get("acc", nan), - ensemble_val_kappa=snap_ensemble.get("kappa", nan), - ensemble_val_mcc=snap_ensemble.get("mcc", nan), - ensemble_val_f1=snap_ensemble.get("macro_f1", nan), - ensemble_val_recall=_sv(snap_ensemble.get("per_class_recall")), - ensemble_val_ece=snap_ensemble.get("ece", nan), - ensemble_val_threshold=snap_ensemble.get("threshold", nan), - ensemble_val_bias=_svf(snap_ensemble.get("bias")), - ensemble_val_n=snap_ensemble.get("n", 0), - bilat_val_auc=snap_bilat.get("auc", nan), - bilat_val_acc=snap_bilat.get("acc", nan), - bilat_val_kappa=snap_bilat.get("kappa", nan), - bilat_val_mcc=snap_bilat.get("mcc", nan), - bilat_val_f1=snap_bilat.get("macro_f1", nan), - bilat_val_recall=_sv(snap_bilat.get("per_class_recall")), - bilat_val_ece=snap_bilat.get("ece", nan), - bilat_val_threshold=snap_bilat.get("threshold", nan), - bilat_val_bias=_svf(snap_bilat.get("bias")), - bilat_val_n=snap_bilat.get("n", 0), - classic_holdout_auc=snap_holdout_single.get("auc", nan) if tower_mode == "single" else nan, - classic_holdout_acc=snap_holdout_single.get("acc", nan) if tower_mode == "single" else nan, - ensemble_holdout_auc=snap_holdout_single.get("auc", nan) if tower_mode == "ensemble" else nan, - ensemble_holdout_acc=snap_holdout_single.get("acc", nan) if tower_mode == "ensemble" else nan, - bilat_holdout_auc=snap_holdout_bilat.get("auc", nan), - bilat_holdout_acc=snap_holdout_bilat.get("acc", nan), - holdout_n=len(holdout_bilat), - single_train_n=len(eye_train), - bilat_train_n=len(bilat_train), - fused_val_auc=snap_fused.get("auc", nan), - fused_val_acc=snap_fused.get("acc", nan), - fused_val_kappa=snap_fused.get("kappa", nan), - fused_val_mcc=snap_fused.get("mcc", nan), - fused_val_f1=snap_fused.get("macro_f1", nan), - fused_val_recall=_sv(snap_fused.get("per_class_recall")), - fused_val_ece=snap_fused.get("ece", nan), - fused_val_threshold=snap_fused.get("threshold", nan), - fused_val_bias=_svf(snap_fused.get("bias")), - fused_val_n=snap_fused.get("n", 0), - fused_holdout_auc=snap_holdout_fused.get("auc", nan), - fused_holdout_acc=snap_holdout_fused.get("acc", nan), - ), FoldArtifacts( - y_true_classic=y_cl_best, probs_classic=p_cl_best, - y_true_ensemble=y_en_best, probs_ensemble=p_en_best, - y_true_bilat=y_bi_best, probs_bilat=p_bi_best, - y_true_fused=y_fu_best, probs_fused=p_fu_best, - probs_ensemble_img=p_en_best_img, - probs_ensemble_md=p_en_best_md, - probs_classic_img=p_cl_best_img, - probs_classic_md=p_cl_best_md, - y_true_ensemble_pereye=y_en_pe_best, - probs_ensemble_pereye=p_en_pe_best, - probs_ensemble_img_pereye=p_en_pe_best_img, - probs_ensemble_md_pereye=p_en_pe_best_md, - logits_ensemble=l_en_best, - logits_ensemble_img=l_en_best_img, - logits_ensemble_md=l_en_best_md, - logits_classic=l_cl_best, - logits_classic_img=l_cl_best_img, - logits_classic_md=l_cl_best_md, - logits_ensemble_pereye=l_en_pe_best, - logits_ensemble_img_pereye=l_en_pe_best_img, - logits_ensemble_md_pereye=l_en_pe_best_md, - ) - - # ------------------------------------------------------------------ - # Summary helpers - # ------------------------------------------------------------------ - - @staticmethod - def _summary(results: list[FoldResult]) -> dict: - def _ms(vals): - v = np.array( - [x for x in vals if x is not None and not np.isnan(float(x))], dtype=float - ) - return (float(np.mean(v)) if v.size else None, float(np.std(v)) if v.size else None) - - out = {} - for label, prefix in [ - ("classic_best_val", "classic_val"), - ("ensemble_best_val", "ensemble_val"), - ("bilat_best_val", "bilat_val"), - ("fused_best_val", "fused_val"), - ]: - sub = {} - for m in ["auc", "acc", "kappa", "mcc", "f1", "ece", "threshold"]: - vals = [getattr(r, f"{prefix}_{m}") for r in results] - mean, std = _ms(vals) - sub[f"{m}_mean"] = mean - if m in ("auc", "f1", "kappa"): - sub[f"{m}_std"] = std - out[label] = sub - - for label, prefix in [ - ("classic_holdout", "classic_holdout"), - ("ensemble_holdout", "ensemble_holdout"), - ("bilat_holdout", "bilat_holdout"), - ("fused_holdout", "fused_holdout"), - ]: - sub = {} - for m in ["auc", "acc"]: - vals = [getattr(r, f"{prefix}_{m}") for r in results] - mean, std = _ms(vals) - sub[f"{m}_mean"] = mean - if m == "auc": - sub[f"{m}_std"] = std - out[label] = sub - - for delta_label, prefix_a, prefix_b in [ - ("delta_ensemble_vs_classic", "classic_val", "ensemble_val"), - ("delta_bilat_vs_ensemble", "ensemble_val", "bilat_val"), - ("delta_fused_vs_ensemble", "ensemble_val", "fused_val"), - ]: - delta = {} - for m in ["auc", "f1", "kappa"]: - pairs = [ - getattr(r, f"{prefix_b}_{m}") - getattr(r, f"{prefix_a}_{m}") - for r in results - if not np.isnan(float(getattr(r, f"{prefix_a}_{m}"))) - and not np.isnan(float(getattr(r, f"{prefix_b}_{m}"))) - ] - delta[f"{m}_mean"] = float(np.mean(pairs)) if pairs else None - delta[f"{m}_std"] = float(np.std(pairs)) if pairs else None - out[delta_label] = delta - - out["n_folds_completed"] = len(results) - out["single_train_mode"] = "eye-level (all OD+OS samples)" - out["bilat_train_mode"] = "patient-level (bilateral only)" - out["eval_note"] = ( - "classic=eye-level SingleEyeHT; " - "ensemble=patient-level SingleEyeHT (OD+OS averaged); " - "bilateral=patient-level BilateralHT; " - "fused=ensemble base + learned logit-level fusion head" - ) - return out - - @staticmethod - def _print_summary(mode: str, s: dict, tower_mode: str | None = None) -> None: - def f(v): - return "nan" if v is None else f"{v:.4f}" - - cv = s["classic_best_val"] - ev = s["ensemble_best_val"] - bv = s["bilat_best_val"] - fv = s["fused_best_val"] - d1 = s["delta_ensemble_vs_classic"] - d2 = s["delta_bilat_vs_ensemble"] - d3 = s["delta_fused_vs_ensemble"] - has_fused = fv["auc_mean"] is not None - - print(f"\n=== Summary [{mode}] — best-epoch val ===") - print(f" {'':26s} {'AUC':>8} {'ACC':>8} {'Kappa':>8} {'F1-mac':>8} {'ECE':>8}") - if tower_mode == "single": - rows = [("single (eye-lvl eval)", cv)] - elif tower_mode == "ensemble": - rows = [("ensemble (pat-lvl eval)", ev)] - if has_fused: - rows.append(("fused_head(pat-lvl eval)", fv)) - elif tower_mode == "bilateral": - rows = [("bilateral (bilat eval)", bv)] - else: - rows = [ - ("classic (eye-lvl eval)", cv), - ("ensemble (pat-lvl eval)", ev), - ("bilateral (bilat eval)", bv), - ] - if has_fused: - rows.append(("fused_head(pat-lvl eval)", fv)) - for label, d in rows: - print( - f" {label:26s} " - f"{f(d['auc_mean']):>8} {f(d['acc_mean']):>8} " - f"{f(d['kappa_mean']):>8} {f(d['f1_mean']):>8} {f(d['ece_mean']):>8}" - ) - if tower_mode is None: - print( - f" {'Δ ensemble−classic':26s} " - f"{f(d1['auc_mean']):>8} {'':>8} " - f"{f(d1['kappa_mean']):>8} {f(d1['f1_mean']):>8}" - ) - print( - f" {'Δ bilateral−ensemble':26s} " - f"{f(d2['auc_mean']):>8} {'':>8} " - f"{f(d2['kappa_mean']):>8} {f(d2['f1_mean']):>8}" - ) - if has_fused and tower_mode in ("ensemble", None): - print( - f" {'Δ fused−ensemble':26s} " - f"{f(d3['auc_mean']):>8} {'':>8} " - f"{f(d3['kappa_mean']):>8} {f(d3['f1_mean']):>8}" - ) - - -# --------------------------------------------------------------------------- -# V2ModeComparator — thin backward-compat shim -# --------------------------------------------------------------------------- - -class V2ModeComparator: - """Backward-compat shim used by run_multifold_v2_modes.py.""" - - @staticmethod - def build_parser() -> argparse.ArgumentParser: - return V2HyperTower.build_parser() - - @staticmethod - def run(cli_args=None) -> Path: - args = V2HyperTower.build_parser().parse_args(cli_args) - return V2HyperTower(args).run() - - -# --------------------------------------------------------------------------- -# Module-level aliases (kept for backward compat; use V2HyperTower directly) -# --------------------------------------------------------------------------- - -def build_parser() -> argparse.ArgumentParser: - return V2HyperTower.build_parser() - - -def run_mode(args) -> Path: - return V2HyperTower(args).run() diff --git a/v3/classes/v3_hypertower.py b/v3/classes/v3_hypertower.py index 0cb6785..cfbc837 100644 --- a/v3/classes/v3_hypertower.py +++ b/v3/classes/v3_hypertower.py @@ -29,7 +29,6 @@ from v3.classes.croppers import ( build_image_preprocessor_from_args, ) from v3.classes.image_loader import CachedImageLoader -from v3.classes.dataset import _ClinicalView # noqa: F401 from v3.classes.loader_factory import ( build_balanced_sampler, filter_bilateral_samples, @@ -37,7 +36,12 @@ from v3.classes.loader_factory import ( make_loader, ) from v3.classes.metrics import _score_arrays, _svf, _tune_and_snap -from v3.classes.models import ( +from v3.classes.bridges import Bridge +from v3.classes.towerbase import train_towers_epoch, collect_probs_towers +from v3.classes.image_towers import ImageTower +from v3.classes.clinical_towers import ClinicalDataTower +from v3.classes.geometry_towers import GeometryTower +from v3.classes.hypertower_models import ( BilateralHT, EmbeddingMLPEnsembleHT, FusedEnsembleHT, @@ -154,7 +158,7 @@ class V3HyperTower: ap.add_argument("--exclude-cols", nargs="*", default=[]) ap.add_argument("--eval-mode", choices=["binary", "multiclass"], default="binary") ap.add_argument( - "--tower-mode", choices=["single", "ensemble", "bilateral", "siamese", "classic"], + "--hypertower-mode", choices=["single", "ensemble", "bilateral", "siamese", "classic"], default="ensemble", ) ap.add_argument("--n-splits", type=int, default=5) @@ -252,6 +256,20 @@ class V3HyperTower: ap.add_argument("--geometry-source", default="gt", choices=["gt", "unet"], help="Source for geometry features: gt (GT contour annotations) or " "unet (U-Net segmentation). unet also requires --img-crop-weights.") + ap.add_argument("--geometry-tower", action="store_true", + help="Add a dedicated GeometryTower (disc/cup seg-map CNN) fused via the " + "bridge alongside ImageTower and ClinicalDataTower. Requires " + "--img-crop-manifest.") + ap.add_argument("--geometry-tower-backbone", default="resnet18", + choices=["resnet18", "resnet50", "efficientnet_b0"], + help="SegCNN backbone for GeometryTower (default: resnet18).") + ap.add_argument("--geometry-tower-in-channels", type=int, default=3, choices=[1, 3], + help="1 = single label map; 3 = one-hot disc/rim/cup (default: 3).") + ap.add_argument("--geometry-tower-frozen", action="store_true", + help="Freeze GeometryTower backbone throughout training.") + ap.add_argument("--geometry-tower-finetune-unet-epochs", type=int, default=0, + help="Epochs to fine-tune the U-Net per fold before seg-map extraction " + "(0 = disabled; only applies when --geometry-source unet).") return ap def __init__(self, args) -> None: @@ -333,7 +351,7 @@ class V3HyperTower: out_dir.mkdir(parents=True, exist_ok=True) mode = args.eval_mode - tower_mode = "single" if args.tower_mode == "classic" else args.tower_mode + tower_mode = "single" if args.hypertower_mode == "classic" else args.hypertower_mode df_mode = self.data.df.copy() if args.exclude_mixed_patients: @@ -482,6 +500,300 @@ class V3HyperTower: return out_dir + def _run_fold_towers( + self, + *, + fold: int, + split, + mode: str, + data, + num_classes: int, + profile_eye, + profile_patient, + fold_dir: Path, + pred_store, + image_cache, + ): + """Modular TowerBase training path (used when --geometry-tower is set). + + Builds [ImageTower, ClinicalDataTower, GeometryTower], runs the full + fold lifecycle (prepare_fold → augment_samples → loader build → epoch + loop → eval), and returns (FoldResult, FoldArtifacts) with metrics in + the ensemble_val_* slots. + """ + args = self.args + device = self.device + nan = float("nan") + + # ------------------------------------------------------------------ samples + eye_train = filter_eye_samples(profile_eye.build_samples(df=split.train, clinical=data)) + bilat_train = filter_bilateral_samples(profile_patient.build_samples(df=split.train, clinical=data)) + bilat_val = filter_bilateral_samples(profile_patient.build_samples(df=split.val, clinical=data)) + bilat_test = filter_bilateral_samples(profile_patient.build_samples( + df=split.test, clinical=data)) if split.test is not None else [] + + # Old --geometry-dim path still applies (injects geometry into clinical stream) + if self.geometry_provider is not None: + eye_train = self._augment_geometry(eye_train) + bilat_train = self._augment_geometry(bilat_train) + bilat_val = self._augment_geometry(bilat_val) + bilat_test = self._augment_geometry(bilat_test) + + if len(bilat_val) == 0: + empty = FoldResult( + mode=mode, fold=fold, + best_epoch_single=0, best_epoch_bilat=0, + classic_val_auc=nan, classic_val_acc=nan, classic_val_kappa=nan, + classic_val_mcc=nan, classic_val_f1=nan, classic_val_recall=None, + classic_val_ece=nan, classic_val_threshold=nan, classic_val_bias=None, + classic_val_n=0, + ensemble_val_auc=nan, ensemble_val_acc=nan, ensemble_val_kappa=nan, + ensemble_val_mcc=nan, ensemble_val_f1=nan, ensemble_val_recall=None, + ensemble_val_ece=nan, ensemble_val_threshold=nan, ensemble_val_bias=None, + ensemble_val_n=0, + bilat_val_auc=nan, bilat_val_acc=nan, bilat_val_kappa=nan, + bilat_val_mcc=nan, bilat_val_f1=nan, bilat_val_recall=None, + bilat_val_ece=nan, bilat_val_threshold=nan, bilat_val_bias=None, + bilat_val_n=0, + single_train_n=len(eye_train), bilat_train_n=len(bilat_train), + ) + return empty, FoldArtifacts( + y_true_classic=None, probs_classic=None, + y_true_ensemble=None, probs_ensemble=None, + y_true_bilat=None, probs_bilat=None, + ) + + # ------------------------------------------------------------------ towers + img_tower = ImageTower( + backbone=args.backbone, + freeze_ratio=args.freeze_ratio, + augment=args.augment, + use_se=getattr(args, "se_img_tower", False), + ) + cd_tower = ClinicalDataTower( + clinical_data=data, + cd_hidden_dim=args.cd_hidden_dim, + cd_dropout=getattr(args, "cd_dropout", 0.1), + use_se=getattr(args, "se_cd_tower", False), + ) + + geom_tower = GeometryTower( + backbone=getattr(args, "geometry_tower_backbone", "resnet18"), + in_channels=getattr(args, "geometry_tower_in_channels", 3), + pretrained=not getattr(args, "no_pretrained", False), + frozen=getattr(args, "geometry_tower_frozen", False), + geometry_source=getattr(args, "geometry_source", "gt"), + manifest_path=getattr(args, "img_crop_manifest", None), + weights_path=getattr(args, "img_crop_weights", None), + unet_normalize=getattr(args, "img_crop_normalize", "per_image"), + unet_threshold=getattr(args, "img_crop_threshold", 0.5), + finetune_unet_epochs=getattr(args, "geometry_tower_finetune_unet_epochs", 0), + ) + + # GeometryTower.prepare_fold must run before augment_samples (precomputes seg maps) + geom_tower.prepare_fold( + eye_train=eye_train, bilat_train=bilat_train, + bilat_val=bilat_val, bilat_test=bilat_test, + image_preprocessor=self.image_preprocessor, + image_cache=image_cache, device=device, args=args, + ) + # Inject seg_map_1/seg_map_2 into all sample lists before loaders are built + for sample_list in (eye_train, bilat_train, bilat_val, bilat_test): + geom_tower.augment_samples(sample_list) + + # Now ImageTower.prepare_fold sees augmented samples → loader includes seg maps + img_tower.prepare_fold( + eye_train=eye_train, bilat_train=bilat_train, + bilat_val=bilat_val, bilat_test=bilat_test, + image_preprocessor=self.image_preprocessor, + image_cache=image_cache, device=device, args=args, + ) + cd_tower.prepare_fold( + eye_train=eye_train, bilat_train=bilat_train, + bilat_val=bilat_val, bilat_test=bilat_test, + image_preprocessor=self.image_preprocessor, + image_cache=image_cache, device=device, args=args, + ) + + towers = [img_tower, cd_tower, geom_tower] + + # ------------------------------------------------------------------ bridge + tower_dims = [] + for t in towers: + tower_dims.extend(t.embed_dims) + bridge = Bridge( + tower_dims=tower_dims, + num_classes=num_classes, + fusion_dim=args.fusion_dim, + mode=getattr(args, "bridge_mode", "fused"), + dropout=getattr(args, "bridge_dropout", 0.5), + ) + + # Move all nn.Modules to device + for t in towers: + if isinstance(t, torch.nn.Module): + t.to(device) + bridge.to(device) + + # ------------------------------------------------------------------ loaders + slots_patient = profile_patient.slot_descriptors() + _persistent = args.num_workers > 0 + loader_kw = dict( + batch_size=args.batch_size, num_workers=args.num_workers, + image_cache=image_cache, persistent_workers=_persistent, + ) + eval_transform = build_eval_transform(args.backbone) + val_loader = make_loader( + bilat_val, slots_patient, + image_transform=eval_transform, + image_preprocessor=self.image_preprocessor, + shuffle=False, **loader_kw, + ) + test_loader = None + if bilat_test: + test_loader = make_loader( + bilat_test, slots_patient, + image_transform=eval_transform, + image_preprocessor=self.image_preprocessor, + shuffle=False, **loader_kw, + ) + + train_loader = img_tower.train_loader + val_loader.dataset.prebuild_image_cache() + if test_loader is not None: + test_loader.dataset.prebuild_image_cache() + + # ------------------------------------------------------------------ optimizer + all_params = list(bridge.parameters()) + for t in towers: + if isinstance(t, torch.nn.Module): + all_params.extend(t.parameters()) + optimizer = torch.optim.AdamW( + [p for p in all_params if p.requires_grad], + lr=args.lr, + weight_decay=getattr(args, "weight_decay", 1e-4), + ) + + # ------------------------------------------------------------------ epoch loop + global_warmup_tower = getattr(args, "warmup_tower_epochs", None) + global_warmup_fused = getattr(args, "warmup_fused_epochs", None) + warmup_cd = int(getattr(args, "warmup_cd_epochs", 0)) + warmup_tower = int(getattr(args, "single_warmup_tower_epochs", None) or global_warmup_tower or 2) + warmup_fused = int(getattr(args, "single_warmup_fused_epochs", None) or global_warmup_fused or 2) + main_epochs = int(args.epochs) + + schedule = [] + if warmup_cd > 0: schedule.append(("cd_warmup", warmup_cd)) + if warmup_tower > 0: schedule.append(("tower_warmup", warmup_tower)) + if warmup_fused > 0: schedule.append(("fused_warmup", warmup_fused)) + schedule.append(("main", main_epochs)) + + best_val_auc = float("-inf") + best_epoch = 0 + best_tower_states = None + best_bridge_state = None + epoch_idx = 0 + + for phase, n_epochs in schedule: + for _ in range(n_epochs): + for t in towers: + if isinstance(t, torch.nn.Module): + t.train() + train_towers_epoch( + towers, bridge, train_loader, optimizer, device, + phase=phase, + bcd_prob=getattr(args, "bcd_prob", 0.5), + tower_loss_mode=getattr(args, "tower_loss_mode", "bcd"), + ) + y_v, p_v = collect_probs_towers(towers, bridge, val_loader, device, + tower_mode="ensemble") + _, val_auc, _ = _score_arrays(y_v, p_v, num_classes) + if not np.isnan(val_auc) and val_auc > best_val_auc: + best_val_auc = val_auc + best_epoch = epoch_idx + best_tower_states = [ + t.state_dict() if isinstance(t, torch.nn.Module) else None + for t in towers + ] + best_bridge_state = bridge.state_dict() + epoch_idx += 1 + + # Restore best + if best_bridge_state is not None: + bridge.load_state_dict(best_bridge_state) + if best_tower_states is not None: + for t, st in zip(towers, best_tower_states): + if isinstance(t, torch.nn.Module) and st is not None: + t.load_state_dict(st) + + # ------------------------------------------------------------------ eval + y_val, p_val = collect_probs_towers(towers, bridge, val_loader, device, tower_mode="ensemble") + acc_val, auc_val, n_val = _score_arrays(y_val, p_val, num_classes) + snap_val, _, thr_val, bias_val = _tune_and_snap( + y_val, p_val, acc_val, num_classes, args, n_bins=10 + ) + + y_test = p_test = None + test_auc = test_acc = nan + test_n = 0 + if test_loader is not None: + y_test, p_test = collect_probs_towers(towers, bridge, test_loader, device, + tower_mode="ensemble") + test_acc, test_auc, test_n = _score_arrays(y_test, p_test, num_classes) + + result = FoldResult( + mode=mode, fold=fold, + best_epoch_single=best_epoch, best_epoch_bilat=0, + classic_val_auc=nan, classic_val_acc=nan, classic_val_kappa=nan, + classic_val_mcc=nan, classic_val_f1=nan, classic_val_recall=None, + classic_val_ece=nan, classic_val_threshold=nan, classic_val_bias=None, + classic_val_n=0, + ensemble_val_auc=snap_val["auc"], ensemble_val_acc=snap_val["acc"], + ensemble_val_kappa=snap_val["kappa"], ensemble_val_mcc=snap_val["mcc"], + ensemble_val_f1=snap_val["macro_f1"], + ensemble_val_recall=_sv(snap_val["per_class_recall"]), + ensemble_val_ece=snap_val["ece"], + ensemble_val_threshold=snap_val["threshold"], + ensemble_val_bias=_svf(bias_val), + ensemble_val_n=snap_val["n"], + bilat_val_auc=nan, bilat_val_acc=nan, bilat_val_kappa=nan, + bilat_val_mcc=nan, bilat_val_f1=nan, bilat_val_recall=None, + bilat_val_ece=nan, bilat_val_threshold=nan, bilat_val_bias=None, + bilat_val_n=0, + ensemble_test_auc=test_auc, ensemble_test_acc=test_acc, + test_n=test_n, + single_train_n=len(eye_train), bilat_train_n=len(bilat_train), + ) + artifacts = FoldArtifacts( + y_true_classic=None, probs_classic=None, + y_true_ensemble=y_val, probs_ensemble=p_val, + y_true_bilat=None, probs_bilat=None, + y_true_test=y_test, probs_test=p_test, + ) + return result, artifacts + + def _augment_geometry_slot(self, samples: list) -> list: + """Add geom_1/geom_2 keys to each sample dict (geometry tower mode). + + Unlike _augment_geometry, this does NOT touch matrix_1/matrix_2 — the + geometry vector lives in its own slot so ImageTower and ClinicalDataTower + each receive only their own modality. + """ + if self.geometry_provider is None: + return samples + geom_dim = int(getattr(self.args, "geometry_dim", 0)) or 5 + for s in samples: + for img_slot, geom_slot in (("image_1", "geom_1"), ("image_2", "geom_2")): + img_path = s.get(img_slot) + if img_path is None: + continue + vec = self.geometry_provider.geometry_for_image(img_path) + if vec is not None and len(vec) >= geom_dim: + s[geom_slot] = vec[:geom_dim].astype(np.float32) + else: + s[geom_slot] = np.zeros(geom_dim, dtype=np.float32) + return samples + def _augment_geometry(self, samples: list) -> list: """Append geometry features to matrix_1/matrix_2 in each sample dict.""" if self.geometry_provider is None: @@ -517,6 +829,16 @@ class V3HyperTower: image_cache, ): args = self.args + + # Modular tower path — bypasses the legacy single/bilat/siamese code entirely + if getattr(args, "geometry_tower", False): + return self._run_fold_towers( + fold=fold, split=split, mode=mode, data=data, + num_classes=num_classes, profile_eye=profile_eye, + profile_patient=profile_patient, fold_dir=fold_dir, + pred_store=pred_store, image_cache=image_cache, + ) + device = self.device image_preprocessor = self.image_preprocessor nan = float("nan") diff --git a/v3/distributed/server.py b/v3/distributed/server.py index 0e324fc..1292d16 100644 --- a/v3/distributed/server.py +++ b/v3/distributed/server.py @@ -10,6 +10,7 @@ Usage: Environment: HT_TOKEN — fallback if --token is not passed """ + from __future__ import annotations import argparse @@ -44,8 +45,8 @@ from .protocol import ( _TOKEN: str = "" _DB_PATH: Path = Path("v3/distributed/jobs.db") -_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 +_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 _clients: dict[str, ClientInfo] = {} _clients_lock = threading.Lock() @@ -60,11 +61,15 @@ def _reap_stale_clients(): # Step 1: evict timed-out clients from registry with _clients_lock: stale = [ - cid for cid, c in _clients.items() + cid + for cid, c in _clients.items() if datetime.fromisoformat(c.last_seen).timestamp() < cutoff ] for cid in stale: - print(f"[server] reaped stale client {cid} ({_clients[cid].hostname})", flush=True) + print( + f"[server] reaped stale client {cid} ({_clients[cid].hostname})", + flush=True, + ) del _clients[cid] known_ids = set(_clients.keys()) @@ -78,15 +83,20 @@ def _reap_stale_clients(): conn.execute( "UPDATE jobs SET state='pending', assigned_to=NULL, started_at=NULL " "WHERE job_id=?", - (row["job_id"],) + (row["job_id"],), ) - print(f"[server] re-queued job {row['job_id']} " - f"(client {row['assigned_to']} unknown)", flush=True) + print( + f"[server] re-queued job {row['job_id']} " + f"(client {row['assigned_to']} unknown)", + flush=True, + ) + # ────────────────────────────────────────────────────────────── # Database helpers # ────────────────────────────────────────────────────────────── + @contextmanager def _db(): conn = sqlite3.connect(str(_DB_PATH)) @@ -101,7 +111,8 @@ def _db(): def _init_db(): _DB_PATH.parent.mkdir(parents=True, exist_ok=True) with _db() as conn: - conn.execute(""" + conn.execute( + """ CREATE TABLE IF NOT EXISTS jobs ( job_id TEXT PRIMARY KEY, run_name TEXT NOT NULL, @@ -117,10 +128,16 @@ def _init_db(): error_msg TEXT, attempts INTEGER NOT NULL DEFAULT 0 ) - """) + """ + ) + conn.execute("CREATE INDEX IF NOT EXISTS idx_jobs_state ON jobs(state)") + conn.execute("CREATE INDEX IF NOT EXISTS idx_jobs_priority ON jobs(priority)") + conn.execute("CREATE INDEX IF NOT EXISTS idx_jobs_run_name ON jobs(run_name)") # Add attempts column to existing DBs that predate this field try: - conn.execute("ALTER TABLE jobs ADD COLUMN attempts INTEGER NOT NULL DEFAULT 0") + conn.execute( + "ALTER TABLE jobs ADD COLUMN attempts INTEGER NOT NULL DEFAULT 0" + ) except Exception: pass # column already exists # Note: running jobs are NOT reset on startup — active clients will re-register @@ -147,6 +164,7 @@ def _ensure_client(client_id: str, hostname: str = "", gpu_info: str = "") -> bo return True return False + # ────────────────────────────────────────────────────────────── # FastAPI app # ────────────────────────────────────────────────────────────── @@ -161,12 +179,20 @@ def _check_token(x_token: str = Header(...)): # ── Registration ────────────────────────────────────────────── -@app.post("/register", response_model=RegisterResponse, - dependencies=[Depends(_check_token)]) + +@app.post( + "/register", response_model=RegisterResponse, dependencies=[Depends(_check_token)] +) def register(req: RegisterRequest, reuse_id: Optional[str] = None): with _clients_lock: - client_id = reuse_id if (reuse_id and reuse_id in _clients) else str(uuid.uuid4())[:8] - existing_status = _clients[client_id].status if client_id in _clients else StatusPush(state="idle") + client_id = ( + reuse_id if (reuse_id and reuse_id in _clients) else str(uuid.uuid4())[:8] + ) + existing_status = ( + _clients[client_id].status + if client_id in _clients + else StatusPush(state="idle") + ) _clients[client_id] = ClientInfo( client_id=client_id, hostname=req.hostname, @@ -175,14 +201,16 @@ def register(req: RegisterRequest, reuse_id: Optional[str] = None): last_seen=_now(), ) action = "re-registered" if reuse_id else "registered" - print(f"[server] {action} {client_id} ({req.hostname} | {req.gpu_info})", flush=True) + print( + f"[server] {action} {client_id} ({req.hostname} | {req.gpu_info})", flush=True + ) return RegisterResponse(client_id=client_id) # ── Job polling ─────────────────────────────────────────────── -@app.post("/poll", response_model=PollResponse, - dependencies=[Depends(_check_token)]) + +@app.post("/poll", response_model=PollResponse, dependencies=[Depends(_check_token)]) def poll(client_id: str): with _clients_lock: needs_reregister = _ensure_client(client_id) @@ -206,7 +234,10 @@ def poll(client_id: str): (client_id, job_id), ) if cur.rowcount: - print(f"[server] reset {cur.rowcount} orphaned running job(s) for {client_id}", flush=True) + print( + f"[server] reset {cur.rowcount} orphaned running job(s) for {client_id}", + flush=True, + ) conn.execute( "UPDATE jobs SET state='running', assigned_to=?, started_at=? WHERE job_id=?", (client_id, _now(), job_id), @@ -231,6 +262,7 @@ def poll(client_id: str): # ── Status ──────────────────────────────────────────────────── + @app.post("/status/{client_id}", dependencies=[Depends(_check_token)]) def push_status(client_id: str, status: StatusPush): with _clients_lock: @@ -256,6 +288,7 @@ def get_client(client_id: str): # ── Job completion ──────────────────────────────────────────── + @app.post("/complete", dependencies=[Depends(_check_token)]) def complete(result: JobResult): with _db() as conn: @@ -274,14 +307,17 @@ def complete(result: JobResult): run_name = run_row["run_name"] remaining = conn.execute( "SELECT COUNT(*) FROM jobs WHERE run_name=? AND state != 'done'", - (run_name,) + (run_name,), ).fetchone()[0] if remaining == 0: total = conn.execute( "SELECT COUNT(*) FROM jobs WHERE run_name=?", (run_name,) ).fetchone()[0] conn.execute("DELETE FROM jobs WHERE run_name=?", (run_name,)) - print(f"[server] run '{run_name}' complete ({total} jobs) — cleared", flush=True) + print( + f"[server] run '{run_name}' complete ({total} jobs) — cleared", + flush=True, + ) else: row = conn.execute( "SELECT attempts FROM jobs WHERE job_id=?", (result.job_id,) @@ -293,21 +329,28 @@ def complete(result: JobResult): "attempts=?, error_msg=? WHERE job_id=?", (attempts, result.error_msg, result.job_id), ) - print(f"[server] job {result.job_id} failed (attempt {attempts}/{_MAX_ATTEMPTS}), " - f"re-queuing", flush=True) + print( + f"[server] job {result.job_id} failed (attempt {attempts}/{_MAX_ATTEMPTS}), " + f"re-queuing", + flush=True, + ) else: conn.execute( "UPDATE jobs SET state='failed', completed_at=?, attempts=?, error_msg=? " "WHERE job_id=?", (_now(), attempts, result.error_msg, result.job_id), ) - print(f"[server] job {result.job_id} failed permanently after " - f"{attempts} attempts", flush=True) + print( + f"[server] job {result.job_id} failed permanently after " + f"{attempts} attempts", + flush=True, + ) return {"ok": True} # ── Job queue management ────────────────────────────────────── + @app.post("/jobs", dependencies=[Depends(_check_token)]) def submit_job(job: JobSubmit): job_id = str(uuid.uuid4())[:12] @@ -316,8 +359,15 @@ def submit_job(job: JobSubmit): "INSERT INTO jobs " "(job_id, run_name, module, args, output_dir, priority, created_at) " "VALUES (?,?,?,?,?,?,?)", - (job_id, job.run_name, job.module, json.dumps(job.args), - job.output_dir, job.priority, _now()), + ( + job_id, + job.run_name, + job.module, + json.dumps(job.args), + job.output_dir, + job.priority, + _now(), + ), ) print(f"[server] queued {job_id} ({job.run_name})", flush=True) return {"job_id": job_id} @@ -347,7 +397,9 @@ def clear_jobs(body: dict): else: states = body.get("states", ["done", "failed", "cancelled"]) placeholders = ",".join("?" * len(states)) - cur = conn.execute(f"DELETE FROM jobs WHERE state IN ({placeholders})", states) + cur = conn.execute( + f"DELETE FROM jobs WHERE state IN ({placeholders})", states + ) print(f"[server] cleared {cur.rowcount} jobs", flush=True) return {"cleared": cur.rowcount} @@ -366,19 +418,33 @@ def cancel_job(job_id: str): # Entry point # ────────────────────────────────────────────────────────────── + def main(): - ap = argparse.ArgumentParser(description=__doc__, - formatter_class=argparse.RawDescriptionHelpFormatter) + ap = argparse.ArgumentParser( + description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter + ) ap.add_argument("--port", type=int, default=8765) ap.add_argument("--host", default="0.0.0.0") - ap.add_argument("--token", default=os.environ.get("HT_TOKEN", ""), - help="Shared secret (or set HT_TOKEN env var)") - ap.add_argument("--db", default="v3/distributed/jobs.db", - help="Path to SQLite job database") - ap.add_argument("--client-ttl", type=int, default=120, - help="Seconds of silence before a client is reaped (default: 120)") - ap.add_argument("--max-attempts", type=int, default=3, - help="Max times a failed job is retried before being left as failed (default: 3)") + ap.add_argument( + "--token", + default=os.environ.get("HT_TOKEN", ""), + help="Shared secret (or set HT_TOKEN env var)", + ) + ap.add_argument( + "--db", default="v3/distributed/jobs.db", help="Path to SQLite job database" + ) + ap.add_argument( + "--client-ttl", + type=int, + default=120, + help="Seconds of silence before a client is reaped (default: 120)", + ) + ap.add_argument( + "--max-attempts", + type=int, + default=3, + help="Max times a failed job is retried before being left as failed (default: 3)", + ) args = ap.parse_args() if not args.token: @@ -394,7 +460,10 @@ def main(): reaper = threading.Thread(target=_reap_stale_clients, daemon=True) reaper.start() - print(f"[server] listening on {args.host}:{args.port} client_ttl={_CLIENT_TTL}s", flush=True) + print( + f"[server] listening on {args.host}:{args.port} client_ttl={_CLIENT_TTL}s", + flush=True, + ) uvicorn.run(app, host=args.host, port=args.port, log_level="warning") diff --git a/v3/scripts/development/analyze_emergent_features.py b/v3/scripts/development/analyze_emergent_features.py new file mode 100644 index 0000000..9593f22 --- /dev/null +++ b/v3/scripts/development/analyze_emergent_features.py @@ -0,0 +1,256 @@ +import sys +from pathlib import Path + +# Make sure we can import from v3 classes when running directly +sys.path.insert(0, str(Path(__file__).resolve().parents[3])) + +import numpy as np +import torch +import torch.nn.functional as F +import matplotlib.pyplot as plt + +from v3.classes.hypertower_models import SingleEyeHT +from v3.classes.papila_builders import build_papila_data +from v3.classes.profiles import build_papila_profile +from v3.classes.loader_factory import make_loader, filter_eye_samples +from v3.classes.transforms import build_eval_transform +from v3.classes.utils import choose_device + + +def extract_emergent_features(model, loader, device, num_classes=2): + """ + Finds samples where the fused bridge is correct, but both individual + towers are wrong, and extracts the driving features from the fusion_dim. + + Assumes `model` is a SingleEyeHT. Can be adapted for NTowerHT. + """ + model.eval() + + emergent_samples = [] + + # Access the final linear layer weights in the HTClassifier + # HTClassifier head is: Sequential(ReLU(), Dropout(), Linear()) + # Index 2 is the Linear layer + final_linear = model.bridge.classifier_fused.head[2] + final_weights = ( + final_linear.weight.detach().cpu() + ) # Shape: [num_classes, fusion_dim] + final_bias = final_linear.bias.detach().cpu() + + with torch.no_grad(): + for batch in loader: + x = batch.get("image_1") + m = batch.get("matrix_1") + y = batch.get("label_1") + + if not (torch.is_tensor(x) and torch.is_tensor(m) and torch.is_tensor(y)): + continue + + x, m, y = x.to(device), m.to(device), y.to(device) + + # 1. Get Tower Embeddings + img_feats = model.img_tower(x) + md_feats = model.cd_tower(m) + + # 2. Get Independent Tower Predictions + logits_i = model.bridge.aux_heads[0](img_feats) + logits_m = model.bridge.aux_heads[1](md_feats) + pi = logits_i.argmax(dim=1) + pm = logits_m.argmax(dim=1) + + # 3. Get Fused Representation & Prediction + # _compute_fused applies the Hadamard product and optional SE gate + z_fused = model.bridge._compute_fused([img_feats, md_feats]) + logits_fused = model.bridge.classifier_fused(z_fused) + pf = logits_fused.argmax(dim=1) + + # 4. Find the "Aha!" Moments and "Corrections" + for i in range(len(y)): + yi = y[i].cpu().item() + is_correct = pf[i] == yi + img_wrong = pi[i] != yi + md_wrong = pm[i] != yi + + if not is_correct: + continue + + is_aha = img_wrong and md_wrong + is_correction = img_wrong or md_wrong + + if is_correction: + category = ( + "Aha!" + if is_aha + else ("Corrected Image" if img_wrong else "Corrected Clinical") + ) + # Apply the ReLU that happens inside HTClassifier before the Linear layer + z_act = F.relu(z_fused[i]).cpu() + + # Calculate how much each feature contributed to the correct class logit + feature_contributions = z_act * final_weights[yi] + + emergent_samples.append( + { + "patient_id": ( + batch.get("id_1", ["Unknown"])[i] + if "id_1" in batch + else "Unknown" + ), + "target_class": yi, + "category": category, + "z_activated": z_act.numpy(), + "contributions": feature_contributions.numpy(), + "total_logit": logits_fused[i, yi].cpu().item(), + } + ) + + return emergent_samples + + +def plot_top_emergent_features(emergent_samples, top_k=10): + """ + Plots the top K contributing dimensions across all emergent success samples. + """ + if not emergent_samples: + print("No emergent success or correction samples found in this pass.") + return + + ahas = [s for s in emergent_samples if s["category"] == "Aha!"] + corrected_img = [s for s in emergent_samples if s["category"] == "Corrected Image"] + corrected_clin = [ + s for s in emergent_samples if s["category"] == "Corrected Clinical" + ] + + print(f"\nFound {len(emergent_samples)} total events of interest:") + print(f" - 'Aha!' Moments (Both wrong, Fused right): {len(ahas)}") + print(f" - Corrected Image (Image wrong, Fused right): {len(corrected_img)}") + print( + f" - Corrected Clinical (Clinical wrong, Fused right): {len(corrected_clin)}" + ) + + # Prioritize true Aha moments if they exist, otherwise use corrections + plot_samples = ahas if ahas else emergent_samples + plot_title = ( + "Aha! Moments (Both Wrong, Fused Right)" + if ahas + else "Fusion Corrections (At least one wrong)" + ) + + # Average the feature contributions across samples + all_contribs = np.stack([s["contributions"] for s in plot_samples]) + mean_contribs = all_contribs.mean(axis=0) + + # Get indices of the top K features with the highest absolute contribution + top_indices = np.argsort(np.abs(mean_contribs))[-top_k:][::-1] + + top_values = mean_contribs[top_indices] + labels = [f"Dim {idx}" for idx in top_indices] + + plt.figure(figsize=(10, 6)) + colors = ["green" if v > 0 else "red" for v in top_values] + + plt.barh(np.arange(top_k), top_values[::-1], color=colors[::-1]) + plt.yticks(np.arange(top_k), labels[::-1]) + plt.xlabel("Mean Contribution to Correct Logit") + plt.title(f"Top {top_k} Features Driving {plot_title}") + plt.tight_layout() + plt.show() + + print(f"\nAnalyzed {len(plot_samples)} samples for this plot.") + print("\nTop Feature Breakdown:") + for idx, val in zip(top_indices, top_values): + print(f"Dimension {idx:3d}: {val:+.4f} average logit push") + + +def main(): + device = choose_device("auto") + print(f"Using device: {device}") + + repo_root = Path(__file__).resolve().parents[3] + image_dir = repo_root / "Papila" / "FundusImages" + clinical_dir = repo_root / "Papila" / "ClinicalData" + + print("Loading PAPILA data with Phase 5 settings...") + data = build_papila_data( + image_dir=str(image_dir), + clinical_dir=str(clinical_dir), + label_col="Diagnosis", + cat_cols=["Gender", "Phakic/Pseudophakic"], + iop_corr_method="ratio", + iop_drop_raw=True, + exclude_cols=["Axial_Length"], + ) + # Filter to binary + data.df = data.df[data.df["Diagnosis"].isin([0, 1])].reset_index(drop=True) + + print("Building dataloader (All samples)...") + profile_eye = build_papila_profile( + patient_col="Patient ID", label_col="Diagnosis", sample_mode="eye" + ) + eye_samples = filter_eye_samples( + profile_eye.build_samples(df=data.df, clinical=data) + ) + + loader = make_loader( + eye_samples, + profile_eye.slot_descriptors(), + image_transform=build_eval_transform("refugelike"), + batch_size=16, + shuffle=False, + num_workers=4, + ) + + print("Building SingleEyeHT model...") + model = SingleEyeHT( + backbone="refugelike", + freeze_ratio=0.0, + augment=False, + clinical_data=data, + num_classes=2, + cd_hidden_dim=128, + fusion_dim=256, + bridge_mode="fused", + ).to(device) + + base_ckpt_dir = repo_root / "v3" / "results" / "phase5" / "logit_mlp_head_ckpt" + checkpoints = sorted(list(base_ckpt_dir.rglob("best_single.pt"))) + + if not checkpoints: + print(f"\n[!] No checkpoints found in {base_ckpt_dir}") + print("Please ensure you ran the jobs with the --save-checkpoints flag.") + return + + all_emergent_data = [] + + for checkpoint_path in checkpoints: + print(f"\nProcessing {checkpoint_path.relative_to(repo_root)}...") + state_dict = torch.load(checkpoint_path, map_location=device) + + # Backward compatibility for checkpoints saved before the N-tower bridge refactor + new_state_dict = {} + for k, v in state_dict.items(): + k = k.replace("bridge.W_img.", "bridge.W.0.") + k = k.replace("bridge.W_md.", "bridge.W.1.") + k = k.replace("bridge.ln_img.", "bridge.ln.0.") + k = k.replace("bridge.ln_md.", "bridge.ln.1.") + k = k.replace("bridge.classifier_img.", "bridge.aux_heads.0.") + k = k.replace("bridge.classifier_cd.", "bridge.aux_heads.1.") + k = k.replace( + "bridge.classifier_fused.2.", "bridge.classifier_fused.head.2." + ) + new_state_dict[k] = v + + model.load_state_dict(new_state_dict) + + emergent_data = extract_emergent_features(model, loader, device) + all_emergent_data.extend(emergent_data) + print(f" -> Found {len(emergent_data)} events of interest.") + + print( + f"\nTotal aggregated events across {len(checkpoints)} checkpoints: {len(all_emergent_data)}" + ) + plot_top_emergent_features(all_emergent_data, top_k=15) + + +if __name__ == "__main__": + main() diff --git a/v3/scripts/development/stage_seg_cnn/run_seg_cnn.py b/v3/scripts/development/stage_seg_cnn/run_seg_cnn.py index 1d8da23..2a49ba2 100644 --- a/v3/scripts/development/stage_seg_cnn/run_seg_cnn.py +++ b/v3/scripts/development/stage_seg_cnn/run_seg_cnn.py @@ -61,7 +61,7 @@ from tqdm import tqdm sys.path.insert(0, str(Path(__file__).resolve().parents[4])) -from v3.classes.seg_cnn import ( +from v3.classes.geometry_towers import ( SegCNN, SegMapDataset, SegMapRecord, UNetFineTuneDataset, precompute_unet_seg_maps, ) diff --git a/v3/scripts/main/phase3/epoch_grid.json b/v3/scripts/main/phase3/epoch_grid.json index b2d8d5f..3a6cf17 100644 --- a/v3/scripts/main/phase3/epoch_grid.json +++ b/v3/scripts/main/phase3/epoch_grid.json @@ -7,7 +7,7 @@ "common_args": [ "--eval-mode", "binary", - "--tower-mode", "single", + "--hypertower-mode", "single", "--in-memory-cache", "--augment", "--tune-binary-threshold", diff --git a/v3/scripts/main/phase3/experiment_grid.json b/v3/scripts/main/phase3/experiment_grid.json index f7a87a1..f557e85 100644 --- a/v3/scripts/main/phase3/experiment_grid.json +++ b/v3/scripts/main/phase3/experiment_grid.json @@ -12,7 +12,7 @@ "common_args": [ "--eval-mode", "binary", - "--tower-mode", + "--hypertower-mode", "single", "--epochs", "30", diff --git a/v3/scripts/main/phase3/phase35_grid.json b/v3/scripts/main/phase3/phase35_grid.json index 4c0d594..40a2b18 100644 --- a/v3/scripts/main/phase3/phase35_grid.json +++ b/v3/scripts/main/phase3/phase35_grid.json @@ -8,7 +8,7 @@ "common_args": [ "--eval-mode", "binary", - "--tower-mode", "single", + "--hypertower-mode", "single", "--epochs", "30", "--in-memory-cache", "--augment", diff --git a/v3/scripts/main/phase4/experiment_grid.json b/v3/scripts/main/phase4/experiment_grid.json index 02fd347..4042f48 100644 --- a/v3/scripts/main/phase4/experiment_grid.json +++ b/v3/scripts/main/phase4/experiment_grid.json @@ -31,7 +31,7 @@ "baseline": { "run_name": "phase4/single", "description": "Single-eye baseline with best phase 3 image settings. Direct comparison point for bilateral modes.", - "extra_args": ["--tower-mode", "single"] + "extra_args": ["--hypertower-mode", "single"] }, "groups": [ @@ -42,17 +42,17 @@ { "run_name": "phase4/ensemble", "description": "Ensemble: two independent single-eye forward passes, patient-level average of OD+OS scores.", - "extra_args": ["--tower-mode", "ensemble"] + "extra_args": ["--hypertower-mode", "ensemble"] }, { "run_name": "phase4/bilateral", "description": "BilateralHT: shared towers, concat OD+OS → learned joint projection MLP → classifier.", - "extra_args": ["--tower-mode", "bilateral"] + "extra_args": ["--hypertower-mode", "bilateral"] }, { "run_name": "phase4/siamese", "description": "SiameseHT: shared backbone, mean+delta (asymmetry) representation → classifier.", - "extra_args": ["--tower-mode", "siamese"] + "extra_args": ["--hypertower-mode", "siamese"] } ] }, @@ -64,12 +64,12 @@ { "run_name": "phase4/bilateral_loss_all", "description": "BilateralHT with all-losses mode — joint training dynamics may differ from single-eye.", - "extra_args": ["--tower-mode", "bilateral", "--tower-loss-mode", "all"] + "extra_args": ["--hypertower-mode", "bilateral", "--tower-loss-mode", "all"] }, { "run_name": "phase4/siamese_loss_all", "description": "SiameseHT with all-losses mode.", - "extra_args": ["--tower-mode", "siamese", "--tower-loss-mode", "all"] + "extra_args": ["--hypertower-mode", "siamese", "--tower-loss-mode", "all"] } ] } diff --git a/v3/scripts/main/phase5/dispatch_logit_mlp_ckpt.py b/v3/scripts/main/phase5/dispatch_logit_mlp_ckpt.py index 27609b9..410c6b3 100644 --- a/v3/scripts/main/phase5/dispatch_logit_mlp_ckpt.py +++ b/v3/scripts/main/phase5/dispatch_logit_mlp_ckpt.py @@ -38,7 +38,7 @@ N_REPS = 10 RUN_ARGS = [ "--eval-mode", "binary", "--bridge-mode", "fused", - "--tower-mode", "ensemble", + "--hypertower-mode", "ensemble", "--fused-head", "--head-type", "logit_mlp", "--epochs", "30", diff --git a/v3/scripts/main/phase5/experiment_grid.json b/v3/scripts/main/phase5/experiment_grid.json index 8338234..b7d4e1b 100644 --- a/v3/scripts/main/phase5/experiment_grid.json +++ b/v3/scripts/main/phase5/experiment_grid.json @@ -36,7 +36,7 @@ "baseline": { "run_name": "phase5/single_fused", "description": "Single-eye + clinical data — phase 3 best config, re-run as direct comparison baseline for phase 5.", - "extra_args": ["--tower-mode", "single"] + "extra_args": ["--hypertower-mode", "single"] }, "groups": [ @@ -47,17 +47,17 @@ { "run_name": "phase5/ensemble_fused", "description": "Ensemble (independent OD+OS) + clinical data via fused bridge.", - "extra_args": ["--tower-mode", "ensemble"] + "extra_args": ["--hypertower-mode", "ensemble"] }, { "run_name": "phase5/bilateral_fused", "description": "BilateralHT + clinical data — full canonical HyperTower.", - "extra_args": ["--tower-mode", "bilateral"] + "extra_args": ["--hypertower-mode", "bilateral"] }, { "run_name": "phase5/siamese_fused", "description": "SiameseHT + clinical data — siamese mean+delta with fused clinical bridge.", - "extra_args": ["--tower-mode", "siamese"] + "extra_args": ["--hypertower-mode", "siamese"] } ] }, @@ -69,17 +69,17 @@ { "run_name": "phase5/ensemble_fused_head", "description": "Ensemble + clinical data + attention scorer head (Linear(C→1) per eye, softmax-weighted average).", - "extra_args": ["--tower-mode", "ensemble", "--fused-head"] + "extra_args": ["--hypertower-mode", "ensemble", "--fused-head"] }, { "run_name": "phase5/logit_mlp_head", "description": "Ensemble + clinical data + logit-level MLP head (cat([logit_od, logit_os]) → FC(64) → FC(C)).", - "extra_args": ["--tower-mode", "ensemble", "--fused-head", "--head-type", "logit_mlp"] + "extra_args": ["--hypertower-mode", "ensemble", "--fused-head", "--head-type", "logit_mlp"] }, { "run_name": "phase5/embedding_mlp_head", "description": "Ensemble + clinical data + embedding-level MLP head (cat([z_od, z_os]) → FC(256) → FC(C)).", - "extra_args": ["--tower-mode", "ensemble", "--fused-head", "--head-type", "embedding_mlp"] + "extra_args": ["--hypertower-mode", "ensemble", "--fused-head", "--head-type", "embedding_mlp"] } ] } diff --git a/v3/scripts/main/phase6/experiment_grid.json b/v3/scripts/main/phase6/experiment_grid.json index fb0dab1..434ac7b 100644 --- a/v3/scripts/main/phase6/experiment_grid.json +++ b/v3/scripts/main/phase6/experiment_grid.json @@ -47,7 +47,7 @@ "baseline": { "run_name": "phase6/single_no_geom", "description": "Single-eye + clinical data, no geometry — needed as Phase 6 baseline since no prior phase ran single with all tuned hyperparameters (iop_ratio_drop_raw, bcd_p05, refugelike).", - "extra_args": ["--tower-mode", "single"] + "extra_args": ["--hypertower-mode", "single"] }, "groups": [ @@ -58,17 +58,17 @@ { "run_name": "phase6/vec_gt_single", "description": "Single-eye + clinical + GT geometry vector appended to clinical stream.", - "extra_args": ["--tower-mode", "single", "--geometry-dim", "5", "--geometry-source", "gt"] + "extra_args": ["--hypertower-mode", "single", "--geometry-dim", "5", "--geometry-source", "gt"] }, { "run_name": "phase6/vec_gt_ensemble", "description": "Ensemble + clinical + GT geometry vector (per-eye geometry, independent OD+OS mean).", - "extra_args": ["--tower-mode", "ensemble", "--geometry-dim", "5", "--geometry-source", "gt"] + "extra_args": ["--hypertower-mode", "ensemble", "--geometry-dim", "5", "--geometry-source", "gt"] }, { "run_name": "phase6/vec_gt_fused_head", "description": "Ensemble + fused head + clinical + GT geometry vector.", - "extra_args": ["--tower-mode", "ensemble", "--fused-head", "--geometry-dim", "5", "--geometry-source", "gt"] + "extra_args": ["--hypertower-mode", "ensemble", "--fused-head", "--geometry-dim", "5", "--geometry-source", "gt"] } ] }, @@ -80,17 +80,17 @@ { "run_name": "phase6/vec_unet_single", "description": "Single-eye + clinical + U-Net geometry vector.", - "extra_args": ["--tower-mode", "single", "--geometry-dim", "5", "--geometry-source", "unet"] + "extra_args": ["--hypertower-mode", "single", "--geometry-dim", "5", "--geometry-source", "unet"] }, { "run_name": "phase6/vec_unet_ensemble", "description": "Ensemble + clinical + U-Net geometry vector.", - "extra_args": ["--tower-mode", "ensemble", "--geometry-dim", "5", "--geometry-source", "unet"] + "extra_args": ["--hypertower-mode", "ensemble", "--geometry-dim", "5", "--geometry-source", "unet"] }, { "run_name": "phase6/vec_unet_fused_head", "description": "Ensemble + fused head + clinical + U-Net geometry vector.", - "extra_args": ["--tower-mode", "ensemble", "--fused-head", "--geometry-dim", "5", "--geometry-source", "unet"] + "extra_args": ["--hypertower-mode", "ensemble", "--fused-head", "--geometry-dim", "5", "--geometry-source", "unet"] } ] }, @@ -103,17 +103,17 @@ { "run_name": "phase6/tower_gt_single", "description": "Single-eye + clinical tower + dedicated GT geometry tower.", - "extra_args": ["--tower-mode", "single", "--geometry-tower", "--geometry-source", "gt"] + "extra_args": ["--hypertower-mode", "single", "--geometry-tower", "--geometry-source", "gt"] }, { "run_name": "phase6/tower_gt_ensemble", "description": "Ensemble + clinical tower + dedicated GT geometry tower.", - "extra_args": ["--tower-mode", "ensemble", "--geometry-tower", "--geometry-source", "gt"] + "extra_args": ["--hypertower-mode", "ensemble", "--geometry-tower", "--geometry-source", "gt"] }, { "run_name": "phase6/tower_gt_fused_head", "description": "Ensemble + fused head + clinical tower + dedicated GT geometry tower.", - "extra_args": ["--tower-mode", "ensemble", "--fused-head", "--geometry-tower", "--geometry-source", "gt"] + "extra_args": ["--hypertower-mode", "ensemble", "--fused-head", "--geometry-tower", "--geometry-source", "gt"] } ] }, @@ -126,17 +126,17 @@ { "run_name": "phase6/tower_unet_single", "description": "Single-eye + clinical tower + dedicated U-Net geometry tower.", - "extra_args": ["--tower-mode", "single", "--geometry-tower", "--geometry-source", "unet"] + "extra_args": ["--hypertower-mode", "single", "--geometry-tower", "--geometry-source", "unet"] }, { "run_name": "phase6/tower_unet_ensemble", "description": "Ensemble + clinical tower + dedicated U-Net geometry tower.", - "extra_args": ["--tower-mode", "ensemble", "--geometry-tower", "--geometry-source", "unet"] + "extra_args": ["--hypertower-mode", "ensemble", "--geometry-tower", "--geometry-source", "unet"] }, { "run_name": "phase6/tower_unet_fused_head", "description": "Ensemble + fused head + clinical tower + dedicated U-Net geometry tower.", - "extra_args": ["--tower-mode", "ensemble", "--fused-head", "--geometry-tower", "--geometry-source", "unet"] + "extra_args": ["--hypertower-mode", "ensemble", "--fused-head", "--geometry-tower", "--geometry-source", "unet"] } ] } diff --git a/v3/scripts/main/run_cv.py b/v3/scripts/main/run_cv.py index f9811ec..7ebe127 100644 --- a/v3/scripts/main/run_cv.py +++ b/v3/scripts/main/run_cv.py @@ -9,7 +9,7 @@ Usage (single fold-seed, 5-fold, binary, ensemble): python -m v3.scripts.main.run_cv \ --run-name my_run \ --eval-mode binary \ - --tower-mode ensemble \ + --hypertower-mode ensemble \ --epochs 40 \ --augment \ --tune-binary-threshold \ @@ -22,7 +22,7 @@ Usage (10x5 rep-CV, seeds 100..1000): --rep-seed-start 100 \ --rep-seed-step 100 \ --eval-mode binary \ - --tower-mode ensemble \ + --hypertower-mode ensemble \ --epochs 40 \ --augment \ --tune-binary-threshold \ diff --git a/v3/scripts/main/run_ntower_cv.py b/v3/scripts/main/run_ntower_cv.py new file mode 100644 index 0000000..8dff92f --- /dev/null +++ b/v3/scripts/main/run_ntower_cv.py @@ -0,0 +1,533 @@ +#!/usr/bin/env python +""" +NTowerHT + HyperBridge ensemble cross-validation runner. + +Reproduces phase5/embedding_mlp_head using the new module architecture: + + Stage 1 — per-eye NTowerHT (eye-level samples, BCD training) + img_tower + cd_tower → Bridge → z_fused [B, fusion_dim] + + Stage 2 — HyperBridge(embedding_mlp) (patient-level bilateral samples) + cat([z_od, z_os]) → Linear(2*fusion_dim → hidden_dim) → ReLU → Dropout → Linear → logits + +Training mirrors v3_hypertower ensemble+fused_head: + - Warmup phases for NTowerHT (tower_warmup → fused_warmup → main) + - NTowerHT frozen; HyperBridge trained on bilateral samples + +Usage (phase5/embedding_mlp_head equivalent): + python -m v3.scripts.main.run_ntower_cv \\ + --run-name ntower/ensemble_fused \\ + --eval-mode binary \\ + --epochs 30 \\ + --fusion-epochs 10 \\ + --in-memory-cache \\ + --augment \\ + --tune-binary-threshold \\ + --backbone refugelike \\ + --iop-corr-method ratio \\ + --iop-drop-raw \\ + --exclude-cols Axial_Length +""" +from __future__ import annotations + +import argparse +import json +import sys +import time +from pathlib import Path + +import numpy as np +import torch +import torch.nn.functional as F + +sys.path.insert(0, str(Path(__file__).resolve().parents[3])) + +from v3.classes.hypertower_models import NTowerHT, train_ntower_epoch, collect_probs_ntower +from v3.classes.towerbase import _to_label_tensor +from v3.classes.bridges import HyperBridge +from v3.classes.image_towers import ImageEncoder +from v3.classes.clinical_towers import ClinicalEncoder +from v3.classes.papila_builders import build_papila_data +from v3.classes.profiles import build_papila_profile +from v3.classes.split_manager import PatientFirstSplitManager +from types import SimpleNamespace +from v3.classes.loader_factory import ( + filter_eye_samples, + filter_bilateral_samples, + make_loader, + build_balanced_sampler, +) +from v3.classes.metrics import _score_arrays, compute_extended_metrics, tune_binary_threshold +from v3.classes.transforms import build_eval_transform +from v3.classes.utils import seed_everything, choose_device +from v3.classes.croppers import build_image_preprocessor_from_args +from v3.classes.image_loader import CachedImageLoader + +REPO_ROOT = Path(__file__).resolve().parents[3] +IMAGE_DIR = REPO_ROOT / "Papila" / "FundusImages" +CLINICAL_DIR = REPO_ROOT / "Papila" / "ClinicalData" + +# Batch key mapping for per-eye NTowerHT training +EYE_KEY_MAP = {"img": "image_1", "cd": "matrix_1"} + + +# --------------------------------------------------------------------------- +# CLI +# --------------------------------------------------------------------------- + +def build_parser() -> argparse.ArgumentParser: + ap = argparse.ArgumentParser( + description=__doc__, + formatter_class=argparse.RawDescriptionHelpFormatter, + ) + ap.add_argument("--run-name", default="ntower/ensemble_fused") + ap.add_argument("--output-root", default="v3/results") + ap.add_argument("--eval-mode", default="binary", choices=["binary", "multiclass"]) + ap.add_argument("--epochs", type=int, default=30, + help="NTowerHT main-phase epochs") + ap.add_argument("--fusion-epochs", type=int, default=10, + help="HyperBridge training epochs (after NTowerHT is frozen)") + ap.add_argument("--warmup-cd-epochs", type=int, default=40, + help="Pre-train cd tower + aux head only (no image tower)") + ap.add_argument("--folds", type=int, default=5) + ap.add_argument("--fold-seed", type=int, default=100, + help="Seed for patient splits (rep00=100, rep01=200, ...)") + ap.add_argument("--seed", type=int, default=1234, + help="Seed for model init / per-fold RNG (matches V3HyperTower default)") + ap.add_argument("--batch-size", type=int, default=16) + ap.add_argument("--lr", type=float, default=1e-4) + ap.add_argument("--backbone", default="refugelike") + ap.add_argument("--freeze-ratio", type=float, default=0.0) + ap.add_argument("--augment", action="store_true") + ap.add_argument("--fusion-dim", type=int, default=256) + ap.add_argument("--hyper-hidden-dim", type=int, default=256, + help="HyperBridge hidden dim (default matches EmbeddingMLPEnsembleHT)") + ap.add_argument("--cd-hidden-dim", type=int, default=128) + ap.add_argument("--bcd-prob", type=float, default=0.5) + ap.add_argument("--warmup-tower-epochs", type=int, default=3) + ap.add_argument("--warmup-fused-epochs", type=int, default=3) + ap.add_argument("--label-col", default="Diagnosis") + ap.add_argument("--iop-corr-method", default="ratio") + ap.add_argument("--iop-drop-raw", action="store_true") + ap.add_argument("--exclude-cols", nargs="*", default=[]) + ap.add_argument("--num-workers", type=int, default=0) + ap.add_argument("--in-memory-cache", action="store_true") + ap.add_argument("--tune-binary-threshold", action="store_true") + ap.add_argument("--device", default=None) + return ap + + +# --------------------------------------------------------------------------- +# Data +# --------------------------------------------------------------------------- + +def load_data(args): + return build_papila_data( + image_dir=str(IMAGE_DIR), + clinical_dir=str(CLINICAL_DIR), + label_col=args.label_col, + cat_cols=["Gender", "Phakic/Pseudophakic"], + iop_corr_method=args.iop_corr_method, + iop_drop_raw=args.iop_drop_raw, + exclude_cols=args.exclude_cols or [], + ) + + +# --------------------------------------------------------------------------- +# Model factories +# --------------------------------------------------------------------------- + +def build_nt(data, num_classes: int, args) -> NTowerHT: + """Per-eye NTowerHT: image + clinical → bridge.""" + img_enc = ImageEncoder(backbone=args.backbone, freeze_ratio=args.freeze_ratio, augment=args.augment) + cd_enc = ClinicalEncoder(clinical_data=data, hidden_dim=args.cd_hidden_dim) + return NTowerHT( + towers={"img": img_enc, "cd": cd_enc}, + num_classes=num_classes, + fusion_dim=args.fusion_dim, + ) + + +def build_hb(num_classes: int, args) -> HyperBridge: + """HyperBridge(embedding_mlp): cat([z_od, z_os]) → MLP → logits.""" + return HyperBridge( + input_dims={"od": args.fusion_dim, "os": args.fusion_dim}, + num_classes=num_classes, + hidden_dim=args.hyper_hidden_dim, + mode="embedding_mlp", + ) + + + + +# --------------------------------------------------------------------------- +# Per-tower pre-warmup + HyperBridge training/inference +# --------------------------------------------------------------------------- + +def train_tower_pre_warmup( + nt: NTowerHT, + tower_idx: int, + loader, + batch_key: str, + opt, + device: torch.device, +) -> tuple[float, float]: + """Pre-train a single tower (by index) + its bridge aux head only. + + Everything else is frozen. Caller is responsible for passing a loader + that omits unnecessary slots (e.g. cd_only_loader strips image_1). + """ + tower_name = list(nt.towers.keys())[tower_idx] + + for p in nt.parameters(): + p.requires_grad_(False) + for p in nt.towers[tower_name].parameters(): + p.requires_grad_(True) + for p in nt.bridge.aux_heads[tower_idx].parameters(): + p.requires_grad_(True) + + nt.train() + total_loss = total_correct = total_n = 0 + for batch in loader: + x = batch.get(batch_key) + y = batch.get("label_1") + if not torch.is_tensor(x): + continue + y_t = _to_label_tensor(y, device) + z = nt.towers[tower_name](x.to(device)) + logits = nt.bridge.aux_heads[tower_idx](z) + loss = F.cross_entropy(logits, y_t) + opt.zero_grad(); loss.backward(); opt.step() + total_loss += loss.item() * len(y_t) + total_correct += int((logits.argmax(1) == y_t).sum()) + total_n += len(y_t) + + for p in nt.parameters(): + p.requires_grad_(True) + + return ( + total_loss / total_n if total_n else float("nan"), + total_correct / total_n if total_n else float("nan"), + ) + + +def _encode_eye(nt: NTowerHT, batch: dict, batch_key_map: dict[str, str], device) -> torch.Tensor: + """Run all towers from a single-eye batch dict and return z_fused.""" + embeddings = {name: nt.towers[name](batch[key].to(device)) + for name, key in batch_key_map.items()} + return nt.encode(embeddings) + + +def train_hb_epoch( + nt: NTowerHT, + hb: HyperBridge, + loader, + opt, + device: torch.device, +) -> tuple[float, float]: + """Train HyperBridge with NTowerHT frozen. + + Each bilateral batch provides both eyes; we encode each through NTowerHT + to get z_od, z_os, then train HyperBridge to fuse them. + """ + nt.eval() + hb.train() + total_loss = total_correct = total_n = 0 + for batch in loader: + x1 = batch.get("image_1"); m1 = batch.get("matrix_1") + x2 = batch.get("image_2"); m2 = batch.get("matrix_2") + y = batch.get("label_1") + if not all(torch.is_tensor(t) for t in (x1, m1, x2, m2)): + continue + y_t = _to_label_tensor(y, device) + # Build per-eye batch dicts (keyed by batch key, as _encode_eye expects) + od_batch = {"image_1": x1, "matrix_1": m1} + os_batch = {"image_1": x2, "matrix_1": m2} + with torch.no_grad(): + z_od = _encode_eye(nt, od_batch, EYE_KEY_MAP, device) + z_os = _encode_eye(nt, os_batch, EYE_KEY_MAP, device) + logits, _ = hb({"od": z_od, "os": z_os}) + loss = F.cross_entropy(logits, y_t) + opt.zero_grad(); loss.backward(); opt.step() + total_loss += loss.item() * len(y_t) + total_correct += int((logits.argmax(1) == y_t).sum()) + total_n += len(y_t) + return ( + total_loss / total_n if total_n else float("nan"), + total_correct / total_n if total_n else float("nan"), + ) + + +def collect_probs_hb( + nt: NTowerHT, + hb: HyperBridge, + loader, + device: torch.device, +) -> tuple[np.ndarray, np.ndarray]: + """Collect HyperBridge predictions (patient-level).""" + nt.eval(); hb.eval() + y_all, p_all = [], [] + with torch.no_grad(): + for batch in loader: + x1 = batch.get("image_1"); m1 = batch.get("matrix_1") + x2 = batch.get("image_2"); m2 = batch.get("matrix_2") + y = batch.get("label_1") + if not all(torch.is_tensor(t) for t in (x1, m1, x2, m2)): + continue + y_t = _to_label_tensor(y, device) + z_od = _encode_eye(nt, {"image_1": x1, "matrix_1": m1}, EYE_KEY_MAP, device) + z_os = _encode_eye(nt, {"image_1": x2, "matrix_1": m2}, EYE_KEY_MAP, device) + logits, _ = hb({"od": z_od, "os": z_os}) + y_all.append(y_t.cpu().numpy()) + p_all.append(F.softmax(logits, dim=1).cpu().numpy()) + if not y_all: + return np.zeros(0, dtype=np.int64), np.zeros((0, 0), dtype=np.float32) + return np.concatenate(y_all), np.concatenate(p_all, axis=0) + + +# --------------------------------------------------------------------------- +# Fold runner +# --------------------------------------------------------------------------- + +def phase_for_epoch(epoch: int, warmup_tower: int, warmup_fused: int) -> str: + if epoch < warmup_tower: + return "tower_warmup" + if epoch < warmup_tower + warmup_fused: + return "fused_warmup" + return "main" + + +def run_fold(fold: int, plans, data, num_classes: int, device, args, + profile_eye, profile_patient, image_preprocessor, image_cache) -> dict: + nan = float("nan") + seed_everything(args.seed + fold * 100) + + split = plans[fold] # PatientSplit with .train/.val/.test DataFrames + + # Eye-level splits (for NTowerHT training) + eye_train = filter_eye_samples(profile_eye.build_samples(df=split.train, clinical=data)) + eye_val = filter_eye_samples(profile_eye.build_samples(df=split.val, clinical=data)) + + # Patient-level splits (for HyperBridge training/eval) + bilat_train = filter_bilateral_samples(profile_patient.build_samples(df=split.train, clinical=data)) + bilat_val = filter_bilateral_samples(profile_patient.build_samples(df=split.val, clinical=data)) + bilat_test = filter_bilateral_samples(profile_patient.build_samples( + df=split.test, clinical=data)) if split.test is not None else [] + + if not bilat_val: + print(f" fold{fold+1}: no bilateral val samples, skipping.", flush=True) + return {"fold": fold, "val_auc": nan, "val_acc": nan, "val_n": 0, + "val_kappa": nan, "val_mcc": nan, "val_f1": nan, + "val_threshold": 0.5, "test_auc": nan, "test_acc": nan, "test_n": nan} + + # Build Stage 1 model only — HyperBridge is built after Stage 1 completes, + # matching phase5 where EmbeddingMLPEnsembleHT is constructed at Phase 2 start. + # Building hb here would consume random state and shift all subsequent dropout ops. + nt = build_nt(data, num_classes, args).to(device) + opt_nt = torch.optim.Adam(nt.parameters(), lr=args.lr) + + slots_eye = profile_eye.slot_descriptors() + slots_patient = profile_patient.slot_descriptors() + + loader_kw = dict(batch_size=args.batch_size, num_workers=args.num_workers, + image_cache=image_cache, persistent_workers=args.num_workers > 0) + eval_transform = build_eval_transform(args.backbone) + + # Unified eye-level loader (all slots, shuffle=True — matches phase5 exactly) + train_eye_loader = make_loader( + eye_train, slots_eye, image_transform=nt.transform, + image_preprocessor=image_preprocessor, shuffle=True, **loader_kw, + ) + # cd-only loader for warmup: strips image_1 so image decoding is skipped entirely + slots_cd_only = {k: v for k, v in slots_eye.items() if k != "image_1"} + cd_warmup_loader = make_loader( + eye_train, slots_cd_only, image_transform=None, + image_preprocessor=None, shuffle=True, + sampler=build_balanced_sampler(eye_train), **loader_kw, + ) + + val_eye_loader = make_loader( + eye_val, slots_eye, image_transform=eval_transform, + image_preprocessor=image_preprocessor, shuffle=False, **loader_kw, + ) + train_bilat_loader = make_loader( + bilat_train, slots_patient, image_transform=nt.transform, + image_preprocessor=image_preprocessor, shuffle=True, **loader_kw, + ) + val_bilat_loader = make_loader( + bilat_val, slots_patient, image_transform=eval_transform, + image_preprocessor=image_preprocessor, shuffle=False, **loader_kw, + ) + test_bilat_loader = make_loader( + bilat_test, slots_patient, image_transform=eval_transform, + image_preprocessor=image_preprocessor, shuffle=False, **loader_kw, + ) if bilat_test else None + + # ── Stage 1: train NTowerHT ─────────────────────────────────────────── + tower_names = list(nt.towers.keys()) + tower_keys = list(EYE_KEY_MAP.values()) + + # Per-tower pre-warmup using the cd-only loader (no image loading overhead) + pre_warmup_epochs = [args.warmup_cd_epochs if name == "cd" else 0 + for name in tower_names] + warmup_loaders = {"cd": cd_warmup_loader} + for idx, (name, n_epochs) in enumerate(zip(tower_names, pre_warmup_epochs)): + if n_epochs == 0: + continue + for epoch in range(n_epochs): + tr_loss, tr_acc = train_tower_pre_warmup( + nt, idx, warmup_loaders[name], tower_keys[idx], opt_nt, device, + ) + print( + f" fold{fold+1} [NT] ep{epoch+1:03d}/{n_epochs} [{name}_warmup ]" + f" loss={tr_loss:.4f} tr_acc={tr_acc:.3f}", + flush=True, + ) + + total_nt_epochs = args.warmup_tower_epochs + args.warmup_fused_epochs + args.epochs + + for epoch in range(total_nt_epochs): + phase = phase_for_epoch(epoch, args.warmup_tower_epochs, args.warmup_fused_epochs) + tr_loss, tr_acc = train_ntower_epoch( + nt, train_eye_loader, opt_nt, device, + batch_key_map=EYE_KEY_MAP, phase=phase, bcd_prob=args.bcd_prob, + ) + y_v, p_v = collect_probs_ntower(nt, val_eye_loader, device, batch_key_map=EYE_KEY_MAP) + _, val_auc, _ = _score_arrays(y_v, p_v, num_classes) + print( + f" fold{fold+1} [NT] ep{epoch+1:03d}/{total_nt_epochs} [{phase:14s}]" + f" loss={tr_loss:.4f} tr_acc={tr_acc:.3f} val_auc={val_auc:.4f}", + flush=True, + ) + + # Freeze NTowerHT (final-epoch weights, matching phase5 — no best-checkpoint restore) + for p in nt.parameters(): + p.requires_grad_(False) + + # ── Stage 2: train HyperBridge ──────────────────────────────────────── + # Build here (not at fold start) to match phase5 random-state sequence + hb = build_hb(num_classes, args).to(device) + opt_hb = torch.optim.Adam(hb.parameters(), lr=args.lr) + + for epoch in range(args.fusion_epochs): + tr_loss, tr_acc = train_hb_epoch(nt, hb, train_bilat_loader, opt_hb, device) + y_v, p_v = collect_probs_hb(nt, hb, val_bilat_loader, device) + _, val_auc, _ = _score_arrays(y_v, p_v, num_classes) + print( + f" fold{fold+1} [HB] ep{epoch+1:02d}/{args.fusion_epochs} [fusion ]" + f" loss={tr_loss:.4f} tr_acc={tr_acc:.3f} val_auc={val_auc:.4f}", + flush=True, + ) + # Final-epoch weights used (no best-checkpoint restore, matching phase5) + + # ── Final eval ──────────────────────────────────────────────────────── + y_val, p_val = collect_probs_hb(nt, hb, val_bilat_loader, device) + val_acc, val_auc, val_n = _score_arrays(y_val, p_val, num_classes) + ext = compute_extended_metrics(y_val, p_val, num_classes) if y_val.size else {} + + val_threshold = 0.5 + if args.tune_binary_threshold and num_classes == 2 and y_val.size >= 2: + val_threshold = tune_binary_threshold(y_val, p_val[:, 1]) + + test_auc = test_acc = test_n = nan + if test_bilat_loader is not None: + y_te, p_te = collect_probs_hb(nt, hb, test_bilat_loader, device) + test_acc, test_auc, test_n = _score_arrays(y_te, p_te, num_classes) + + return { + "fold": fold, + "val_auc": val_auc, + "val_acc": val_acc, + "val_n": val_n, + "val_kappa": ext.get("kappa", nan), + "val_mcc": ext.get("mcc", nan), + "val_f1": ext.get("macro_f1", nan), + "val_threshold": val_threshold, + "test_auc": test_auc, + "test_acc": test_acc, + "test_n": test_n, + } + + +# --------------------------------------------------------------------------- +# Main +# --------------------------------------------------------------------------- + +def main(): + ap = build_parser() + args = ap.parse_args() + + num_classes = 2 if args.eval_mode == "binary" else 3 + device = choose_device(args.device) + print(f"Device: {device}", flush=True) + + print("Loading data ...", flush=True) + data = load_data(args) + print(f" feature_dim={data.feature_dim}", flush=True) + + image_cache = CachedImageLoader() if args.in_memory_cache else None + image_preprocessor = build_image_preprocessor_from_args(args) + + profile_eye = build_papila_profile(patient_col="Patient ID", label_col=args.label_col, sample_mode="eye") + profile_patient = build_papila_profile(patient_col="Patient ID", label_col=args.label_col, sample_mode="patient") + + # Filter to binary labels before splitting (mirrors V3HyperTower) + df_mode = data.df.copy() + if args.eval_mode == "binary": + df_mode = df_mode[df_mode[args.label_col].isin([0, 1])].reset_index(drop=True) + + split_mgr = PatientFirstSplitManager(patient_col="Patient ID", label_col=args.label_col) + split_args = SimpleNamespace(eval_mode=args.eval_mode, n_splits=args.folds, fold_seed=args.fold_seed) + clinical_ns = SimpleNamespace(df=df_mode, label_col=args.label_col) + plans = split_mgr.build_plans(clinical=clinical_ns, args=split_args, profile=None) + + out_dir = REPO_ROOT / args.output_root / args.run_name / "binary" / "ntower" + out_dir.mkdir(parents=True, exist_ok=True) + + fold_results = [] + t0 = time.time() + + for fold in range(args.folds): + split = plans[fold] + bilat_val = filter_bilateral_samples(profile_patient.build_samples(df=split.val, clinical=data)) + print( + f"\n── fold {fold+1}/{args.folds}" + f" train_patients={split.train['Patient ID'].nunique()}" + f" val={len(bilat_val)} ──", + flush=True, + ) + result = run_fold(fold, plans, data, num_classes, device, args, + profile_eye, profile_patient, image_preprocessor, image_cache) + fold_results.append(result) + print( + f" fold{fold+1} DONE val_auc={result['val_auc']:.4f}" + f" test_auc={result['test_auc']:.4f}", + flush=True, + ) + + if fold_results: + val_aucs = [r["val_auc"] for r in fold_results if not np.isnan(r["val_auc"])] + test_aucs = [r["test_auc"] for r in fold_results if not np.isnan(r["test_auc"])] + summary = { + "run_name": args.run_name, + "backbone": args.backbone, + "nt_epochs": args.epochs, + "fusion_epochs": args.fusion_epochs, + "folds": args.folds, + "mean_val_auc": float(np.mean(val_aucs)) if val_aucs else float("nan"), + "std_val_auc": float(np.std(val_aucs)) if val_aucs else float("nan"), + "mean_test_auc": float(np.mean(test_aucs)) if test_aucs else float("nan"), + "std_test_auc": float(np.std(test_aucs)) if test_aucs else float("nan"), + "elapsed_s": round(time.time() - t0, 1), + "fold_results": fold_results, + } + summary_path = out_dir / "summary.json" + summary_path.write_text(json.dumps(summary, indent=2)) + print(f"\n{'='*60}", flush=True) + print(f"Val AUC: {summary['mean_val_auc']:.4f} ± {summary['std_val_auc']:.4f}", flush=True) + print(f"Test AUC: {summary['mean_test_auc']:.4f} ± {summary['std_test_auc']:.4f}", flush=True) + print(f"Saved: {summary_path}", flush=True) + + +if __name__ == "__main__": + main() diff --git a/v3/scripts/output_analysis/explainability/gradcam_phase5.py b/v3/scripts/output_analysis/explainability/gradcam_phase5.py index 156fb56..1b5879e 100644 --- a/v3/scripts/output_analysis/explainability/gradcam_phase5.py +++ b/v3/scripts/output_analysis/explainability/gradcam_phase5.py @@ -239,7 +239,7 @@ def make_disc_attention_detail( def build_model(ckpt_path: Path, device: torch.device): """Reconstruct SingleEyeHT from checkpoint and load weights.""" from types import SimpleNamespace - from v3.classes.models import SingleEyeHT + from v3.classes.hypertower_models import SingleEyeHT sd = torch.load(ckpt_path, map_location="cpu") # ClinicalTower only reads clinical_data.feature_dim at init time cd_in = sd["cd_tower.block0.0.weight"].shape[1] diff --git a/v3/scripts/output_analysis/explainability/permutation_importance_phase5.py b/v3/scripts/output_analysis/explainability/permutation_importance_phase5.py index a4c4601..4828501 100644 --- a/v3/scripts/output_analysis/explainability/permutation_importance_phase5.py +++ b/v3/scripts/output_analysis/explainability/permutation_importance_phase5.py @@ -44,7 +44,7 @@ SEED = 0 # ── Model ───────────────────────────────────────────────────────────────────── def build_model(ckpt_path: Path, device: torch.device): - from v3.classes.models import SingleEyeHT + from v3.classes.hypertower_models import SingleEyeHT sd = torch.load(ckpt_path, map_location="cpu") cd_in = sd["cd_tower.block0.0.weight"].shape[1] model = SingleEyeHT( @@ -169,7 +169,7 @@ def run_fold(rep_idx: int, fold_idx: int, model, data, device: torch.device, # Baseline AUC with torch.no_grad(): md_feats = model.cd_tower(meta_all.to(device)) - out_f, _, _ = model.bridge(img_feats, md_feats) + out_f, _ = model.bridge.fuse([img_feats, md_feats]) probs_base = F.softmax(out_f, dim=1)[:, 1].cpu().numpy() baseline_auc = roc_auc_score(y_true, probs_base) print(f" fold{fold_idx}: baseline AUC={baseline_auc:.4f} N={len(y_true)}") @@ -186,7 +186,7 @@ def run_fold(rep_idx: int, fold_idx: int, model, data, device: torch.device, meta_perm[:, dims] = meta_perm[perm_idx][:, dims] with torch.no_grad(): md_p = model.cd_tower(meta_perm.to(device)) - out_p, _, _ = model.bridge(img_feats, md_p) + out_p, _ = model.bridge.fuse([img_feats, md_p]) probs_p = F.softmax(out_p, dim=1)[:, 1].cpu().numpy() try: drops.append(baseline_auc - roc_auc_score(y_true, probs_p)) diff --git a/v4/__init__.py b/v4/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/v4/classes/__init__.py b/v4/classes/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/v4/classes/accessory/__init__.py b/v4/classes/accessory/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/v4/classes/accessory/backbones.py b/v4/classes/accessory/backbones.py new file mode 100644 index 0000000..a90bcc7 --- /dev/null +++ b/v4/classes/accessory/backbones.py @@ -0,0 +1,217 @@ +"""backbones — backbone registry and builder for v4 image towers.""" +from __future__ import annotations + +import math +from dataclasses import dataclass +from pathlib import Path +from typing import Callable, Dict, List + +import torch +from torch import nn +from torchvision import models, transforms + +_REPO_ROOT = Path(__file__).resolve().parents[3] + +REFUGELIKE_BACKBONE_PATH = _REPO_ROOT / "models/v2/refuge/refugelike_backbone.pt" +REFUGE_DENSENET_PATH = _REPO_ROOT / "models/refuge/classifier/refuge_densenet_backbone.pt" +REFUGE_EFFICIENT_B0_PATH = _REPO_ROOT / "models/refuge/classifier/refuge_efficient_b0_backbone.pt" +REFUGE_EFFICIENT_B7_PATH = _REPO_ROOT / "models/refuge/classifier/refuge_efficient_b7_backbone.pt" + + +@dataclass(frozen=True) +class BackboneSpec: + ctor: Callable + weights_default: object + strip: Callable[[nn.Module], tuple] + blocks: Callable[[nn.Module], List[nn.Module]] + + +# --------------------------------------------------------------------------- +# Strip helpers +# --------------------------------------------------------------------------- + +def _strip_efficientnet(m: models.EfficientNet): + out_dim = m.classifier[1].in_features + m.classifier = nn.Identity() + return out_dim, m + +def _strip_resnet(m: models.ResNet): + out_dim = m.fc.in_features + m.fc = nn.Identity() + return out_dim, m + +def _strip_densenet(m: models.DenseNet): + out_dim = m.classifier.in_features + m.classifier = nn.Identity() + return out_dim, m + +def _strip_vgg(m: models.VGG): + out_dim = m.classifier[0].in_features + m.classifier = nn.Identity() + return out_dim, m + +def _strip_mobilenet_v2(m: models.MobileNetV2): + out_dim = m.classifier[1].in_features + m.classifier = nn.Identity() + return out_dim, m + +def _strip_inception_v3(m: models.Inception3): + out_dim = m.fc.in_features + m.fc = nn.Identity() + m.aux_logits = False + m.AuxLogits = None + return out_dim, m + + +# --------------------------------------------------------------------------- +# Block splitters for ratio-based freezing +# --------------------------------------------------------------------------- + +def _blocks_efficientnet(m: models.EfficientNet): + return list(m.features) + +def _blocks_resnet(m: models.ResNet): + stem = nn.Sequential(m.conv1, m.bn1, m.relu, m.maxpool) + return [stem, m.layer1, m.layer2, m.layer3, m.layer4] + +def _blocks_densenet(m: models.DenseNet): + f = m.features + stem = nn.Sequential(f.conv0, f.norm0, f.relu0, f.pool0) + return [stem, f.denseblock1, f.transition1, f.denseblock2, f.transition2, + f.denseblock3, f.transition3, f.denseblock4, f.norm5] + +def _blocks_vgg(m: models.VGG): + stages, cur = [], [] + for mod in m.features: + cur.append(mod) + if isinstance(mod, nn.MaxPool2d): + stages.append(nn.Sequential(*cur)); cur = [] + if cur: + stages.append(nn.Sequential(*cur)) + return stages + +def _blocks_mobilenet_v2(m: models.MobileNetV2): + return list(m.features) + +def _blocks_inception_v3(m: models.Inception3): + return [child for name, child in m.named_children() + if name not in ("fc", "AuxLogits")] + + +# --------------------------------------------------------------------------- +# Registry +# --------------------------------------------------------------------------- + +BACKBONES: Dict[str, BackboneSpec] = { + "efficientnet_b0": BackboneSpec( + ctor=models.efficientnet_b0, + weights_default=models.EfficientNet_B0_Weights.DEFAULT, + strip=_strip_efficientnet, + blocks=_blocks_efficientnet, + ), + "resnet50": BackboneSpec( + ctor=models.resnet50, + weights_default=models.ResNet50_Weights.DEFAULT, + strip=_strip_resnet, + blocks=_blocks_resnet, + ), + "densenet121": BackboneSpec( + ctor=models.densenet121, + weights_default=models.DenseNet121_Weights.DEFAULT, + strip=_strip_densenet, + blocks=_blocks_densenet, + ), + "vgg16": BackboneSpec( + ctor=models.vgg16, + weights_default=models.VGG16_Weights.DEFAULT, + strip=_strip_vgg, + blocks=_blocks_vgg, + ), + "mobilenet_v2": BackboneSpec( + ctor=models.mobilenet_v2, + weights_default=models.MobileNet_V2_Weights.DEFAULT, + strip=_strip_mobilenet_v2, + blocks=_blocks_mobilenet_v2, + ), + "inception_v3": BackboneSpec( + ctor=models.inception_v3, + weights_default=models.Inception_V3_Weights.DEFAULT, + strip=_strip_inception_v3, + blocks=_blocks_inception_v3, + ), + "refugelike": BackboneSpec( + ctor=models.resnet50, + weights_default=None, + strip=_strip_resnet, + blocks=_blocks_resnet, + ), + "refuge_densenet": BackboneSpec( + ctor=models.densenet121, + weights_default=None, + strip=_strip_densenet, + blocks=_blocks_densenet, + ), + "refuge_efficient_b0": BackboneSpec( + ctor=models.efficientnet_b0, + weights_default=None, + strip=_strip_efficientnet, + blocks=_blocks_efficientnet, + ), + "refuge_efficient_b7": BackboneSpec( + ctor=models.efficientnet_b7, + weights_default=None, + strip=_strip_efficientnet, + blocks=_blocks_efficientnet, + ), +} + + +def list_names() -> List[str]: + return list(BACKBONES.keys()) + + +def load_backbone_weights(key: str, model: nn.Module) -> None: + paths = { + "refugelike": REFUGELIKE_BACKBONE_PATH, + "refuge_densenet": REFUGE_DENSENET_PATH, + "refuge_efficient_b0": REFUGE_EFFICIENT_B0_PATH, + "refuge_efficient_b7": REFUGE_EFFICIENT_B7_PATH, + } + path = paths.get(key) + if path is None: + return + if not path.exists(): + raise FileNotFoundError( + f"Custom backbone weights not found at {path}. " + "Export them via refuge_build.py --export-backbone first." + ) + state = torch.load(path, map_location="cpu") + model.load_state_dict(state, strict=False) + + +def build_backbone(name: str, freeze_ratio: float = 0.0) -> tuple[nn.Module, int, list]: + """Instantiate a backbone, strip its classifier head, apply freeze ratio. + + Returns (model, out_dim, blocks) where blocks is the ordered list of + freezable units — callers use it to dynamically adjust freeze_ratio later. + """ + key = (name or "").lower() + if key not in BACKBONES: + raise ValueError(f"Unknown backbone '{name}'. Available: {list_names()}") + + spec = BACKBONES[key] + if spec.weights_default is not None: + m = spec.ctor(weights=spec.weights_default) + else: + m = spec.ctor(weights=None) + out_dim, m = spec.strip(m) + load_backbone_weights(key, m) + + blocks = spec.blocks(m) + fr = max(0.0, min(1.0, float(freeze_ratio))) + n_freeze = int(math.floor(len(blocks) * fr)) + for b in blocks[:n_freeze]: + for p in b.parameters(): + p.requires_grad = False + + return m, out_dim, blocks diff --git a/v4/classes/accessory/se_block.py b/v4/classes/accessory/se_block.py new file mode 100644 index 0000000..9cb5c57 --- /dev/null +++ b/v4/classes/accessory/se_block.py @@ -0,0 +1,126 @@ +"""se_block — Squeeze-and-Excitation channel gating.""" +from __future__ import annotations + +import torch +import torch.nn as nn + + +class SEBlock(nn.Module): + """SE-style channel gating for vectors, feature maps, and sequences. + + Input shapes: + [N, C] (vector) — squeeze is identity + [N, C, H, W] (image map) — squeeze over H, W + [N, T, C] (sequence) — squeeze over T + + Gate modes: + residual (default): gate = 1 + tanh(MLP(s)) in (0, 2), identity at init + plain: gate = sigmoid(MLP(s)) in (0, 1) + """ + + def __init__( + self, + dim: int, + reduction: int = 16, + residual: bool = True, + identity_init: bool = True, + ): + super().__init__() + hid = max(1, dim // max(1, reduction)) + self.fc1 = nn.Linear(dim, hid, bias=True) + self.act = nn.ReLU(inplace=True) + self.fc2 = nn.Linear(hid, dim, bias=True) + self.residual = residual + + if residual and identity_init: + nn.init.zeros_(self.fc2.weight) + nn.init.zeros_(self.fc2.bias) + + def _squeeze(self, x: torch.Tensor) -> torch.Tensor: + if x.dim() == 2: + return x + if x.dim() == 4: + return x.mean(dim=(2, 3)) + if x.dim() == 3: + return x.mean(dim=1) + return x.view(x.size(0), -1) + + def _broadcast(self, gate: torch.Tensor, like: torch.Tensor) -> torch.Tensor: + if like.dim() == 2: + return gate + if like.dim() == 3: + return gate.unsqueeze(1) + if like.dim() == 4: + return gate.unsqueeze(-1).unsqueeze(-1) + return gate.view_as(like) + + def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + s = self._squeeze(x) + u = self.fc2(self.act(self.fc1(s))) + gate = (1.0 + torch.tanh(u)) if self.residual else torch.sigmoid(u) + return x * self._broadcast(gate, x), gate + + +class SEGateLogger: + """Running stats over SE gate activations across batches.""" + + def __init__( + self, + enabled: bool = True, + track_channels: bool = False, + dim: int | None = None, + ): + self.enabled = enabled + self.track_channels = track_channels + self.dim = dim + self.reset() + + def reset(self) -> None: + self._n = 0 + self._sum = 0.0 + self._sum2 = 0.0 + self._lt02 = 0 + self._gt08 = 0 + self._ch_sum = None + self._ch_count = 0 + if self.track_channels and self.dim is not None: + self._ch_sum = torch.zeros(self.dim, dtype=torch.float32) + + @torch.no_grad() + def accumulate(self, gates: torch.Tensor) -> None: + if not self.enabled: + return + if gates.dim() == 4: + g = gates.mean(dim=(2, 3)) + elif gates.dim() == 3: + g = gates.mean(dim=1) + elif gates.dim() == 2: + g = gates + else: + g = gates.view(gates.size(0), -1) + g = g.detach() + self._n += g.numel() + self._sum += g.sum().item() + self._sum2 += (g * g).sum().item() + self._lt02 += (g < 0.2).sum().item() + self._gt08 += (g > 0.8).sum().item() + if self._ch_sum is not None: + self._ch_sum += g.sum(dim=0).cpu() + self._ch_count += g.size(0) + + def get(self, reset: bool = True) -> dict | None: + if self._n == 0: + return None + mean = self._sum / self._n + var = max(0.0, self._sum2 / self._n - mean * mean) + out = { + "mean": mean, + "std": var ** 0.5, + "pct_lt_0.2": self._lt02 / self._n, + "pct_gt_0.8": self._gt08 / self._n, + } + if self._ch_sum is not None and self._ch_count > 0: + out["channel_mean"] = (self._ch_sum / float(self._ch_count)).tolist() + if reset: + self.reset() + return out diff --git a/v4/classes/accessory/transforms.py b/v4/classes/accessory/transforms.py new file mode 100644 index 0000000..d24b651 --- /dev/null +++ b/v4/classes/accessory/transforms.py @@ -0,0 +1,66 @@ +"""transforms — image transform utilities for v4 towers.""" +from __future__ import annotations + +from dataclasses import dataclass +from typing import Tuple + +from torchvision import transforms + +from v4.classes.accessory.backbones import BACKBONES + +IMAGENET_MEAN: Tuple[float, float, float] = (0.485, 0.456, 0.406) +IMAGENET_STD: Tuple[float, float, float] = (0.229, 0.224, 0.225) + + +@dataclass +class ImageTransformConfig: + crop_size: int = 224 + resize_size: int = 256 + mean: Tuple[float, float, float] = IMAGENET_MEAN + std: Tuple[float, float, float] = IMAGENET_STD + augment: bool = True + rotation_deg: int = 15 + color_jitter: Tuple[float, float, float, float] = (0.1, 0.1, 0.1, 0.05) + hflip: bool = True + vflip: bool = True + + def build(self) -> transforms.Compose: + ops = [ + transforms.Resize(self.resize_size), + transforms.CenterCrop(self.crop_size), + ] + if self.augment: + if self.hflip: + ops.append(transforms.RandomHorizontalFlip()) + if self.vflip: + ops.append(transforms.RandomVerticalFlip()) + if self.rotation_deg: + ops.append(transforms.RandomRotation(self.rotation_deg)) + if self.color_jitter: + ops.append(transforms.ColorJitter(*self.color_jitter)) + ops += [ + transforms.ToTensor(), + transforms.Normalize(mean=self.mean, std=self.std), + ] + return transforms.Compose(ops) + + +def backbone_transform_config(backbone_name: str, augment: bool = True) -> ImageTransformConfig: + """Build an ImageTransformConfig using the backbone's default normalisation stats.""" + key = (backbone_name or "").lower() + if key not in BACKBONES: + raise ValueError(f"Unknown backbone '{backbone_name}'.") + spec = BACKBONES[key] + mean = getattr(spec.weights_default, "meta", {}).get("mean", IMAGENET_MEAN) + std = getattr(spec.weights_default, "meta", {}).get("std", IMAGENET_STD) + crop = 299 if key == "inception_v3" else 224 + return ImageTransformConfig(crop_size=crop, mean=mean, std=std, augment=augment) + + +def build_backbone_transform(backbone_name: str, augment: bool = True) -> transforms.Compose: + return backbone_transform_config(backbone_name, augment=augment).build() + + +def build_eval_transform(backbone_name: str) -> transforms.Compose: + """Deterministic eval transform — no augmentation, backbone-matched normalisation.""" + return build_backbone_transform(backbone_name, augment=False) diff --git a/v4/classes/bridges/__init__.py b/v4/classes/bridges/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/v4/classes/bridges/fusion_bridge.py b/v4/classes/bridges/fusion_bridge.py new file mode 100644 index 0000000..c23a907 --- /dev/null +++ b/v4/classes/bridges/fusion_bridge.py @@ -0,0 +1,57 @@ +"""fusion_bridge — FusionBridge: N-input Hadamard-product fusion, embedding output only.""" +from __future__ import annotations + +import torch +import torch.nn as nn + +from v4.classes.accessory.se_block import SEBlock + + +class FusionBridge(nn.Module): + """Project N input embeddings to a shared dim, fuse via element-wise product. + + Pure embedding producer — no classification head. Attach a head stage in + the pipeline config to produce logits. + + Parameters + ---------- + input_dims : ordered list of input embedding dims + fusion_dim : projection / output dimension + use_se : SE gate on the fused vector + se_reduction : SE reduction factor + se_pre_norm : LayerNorm before each projection; else Identity + """ + + def __init__( + self, + input_dims: list[int], + fusion_dim: int = 256, + use_se: bool = True, + se_reduction: int = 16, + se_pre_norm: bool = True, + ): + super().__init__() + self.out_dim = fusion_dim + self.W = nn.ModuleList([nn.Linear(d, fusion_dim) for d in input_dims]) + self.ln = nn.ModuleList( + [nn.LayerNorm(fusion_dim) if se_pre_norm else nn.Identity() + for _ in input_dims] + ) + self.se = SEBlock(fusion_dim, reduction=se_reduction, residual=True) if use_se else None + + def forward(self, embeddings: list[torch.Tensor]) -> torch.Tensor: + assert len(embeddings) == len(self.W), ( + f"FusionBridge expects {len(self.W)} inputs, got {len(embeddings)}" + ) + h = self.ln[0](self.W[0](embeddings[0])) + for i in range(1, len(embeddings)): + h = h * self.ln[i](self.W[i](embeddings[i])) + if self.se is not None: + h, _ = self.se(h) + return h + + def set_phase(self, phase: str) -> None: + """Freeze bridge during tower_warmup; trainable otherwise.""" + enabled = phase not in ("tower_warmup", "cd_warmup") + for p in self.parameters(): + p.requires_grad_(enabled) diff --git a/v4/classes/bridges/hyperbridge.py b/v4/classes/bridges/hyperbridge.py new file mode 100644 index 0000000..83eaee6 --- /dev/null +++ b/v4/classes/bridges/hyperbridge.py @@ -0,0 +1,56 @@ +"""hyperbridge — HyperBridge: bilateral fusion over paired embeddings, embedding output only.""" +from __future__ import annotations + +import torch +import torch.nn as nn + + +class HyperBridge(nn.Module): + """Fuse side embeddings (e.g. two z_fused vectors) into a single embedding. + + Pure embedding producer — no classification head. Attach a head stage in + the pipeline config to produce logits. + + Modes + ----- + embedding_mlp (default) + Linear projection of concatenated inputs → hidden_dim embedding. + + classic_bridge + Per-side projection → Hadamard product → hidden_dim embedding. + + Parameters + ---------- + input_dims : {side_key: dim} — e.g. {"a": 256, "b": 256} + hidden_dim : output embedding dimension + mode : "embedding_mlp" | "classic_bridge" + """ + + def __init__( + self, + input_dims: dict[str, int], + hidden_dim: int = 256, + mode: str = "embedding_mlp", + ): + super().__init__() + self.input_names = list(input_dims.keys()) + self.mode = mode + self.out_dim = hidden_dim + dims = list(input_dims.values()) + + if mode == "embedding_mlp": + self.proj = nn.Linear(sum(dims), hidden_dim) + elif mode == "classic_bridge": + self.W = nn.ModuleList([nn.Linear(d, hidden_dim) for d in dims]) + self.ln = nn.ModuleList([nn.LayerNorm(hidden_dim) for _ in dims]) + else: + raise ValueError(f"Unknown HyperBridge mode: {mode!r}") + + def forward(self, inputs: dict[str, torch.Tensor]) -> torch.Tensor: + ordered = [inputs[name] for name in self.input_names] + if self.mode == "embedding_mlp": + return self.proj(torch.cat(ordered, dim=1)) + h = self.ln[0](self.W[0](ordered[0])) + for i in range(1, len(ordered)): + h = h * self.ln[i](self.W[i](ordered[i])) + return h diff --git a/v4/classes/dataset.py b/v4/classes/dataset.py new file mode 100644 index 0000000..1c5e82d --- /dev/null +++ b/v4/classes/dataset.py @@ -0,0 +1,248 @@ +"""dataset — data packaging for v4: shells, DataBundle, HTDataset. + +ShellEntry / LoaderShell + Minimal, data-free structures representing *who* to sample and in what + order. entity_id is opaque to the orchestrator; towers interpret it. + +DataBundle + Accumulates per-eye DataFrames and derives scalar/categorical stats used + by ClinicalDataView. No kfold, no vectorization — those live in the + profile and orchestrator respectively. + +HTDataset / ht_collate + PyTorch Dataset that delegates sample retrieval to towers via get_sample. +""" +from __future__ import annotations + +from dataclasses import dataclass, field +from pathlib import Path +from typing import Any, Dict, Iterable, List, Optional + +import numpy as np +import pandas as pd +import torch +from torch.utils.data import Dataset + + +# --------------------------------------------------------------------------- +# Shell +# --------------------------------------------------------------------------- + +@dataclass +class ShellEntry: + """One sample slot in a LoaderShell. + + entity_id : opaque — defined by the profile, interpreted by towers. + label : integer class label. + meta : per-entry context a profile wants to pass through. + """ + entity_id: Any + label: int + meta: dict = field(default_factory=dict) + + +@dataclass +class LoaderShell: + """An ordered sequence of ShellEntry objects for one split/fold.""" + entries: list[ShellEntry] + + def __len__(self) -> int: + return len(self.entries) + + def __iter__(self): + return iter(self.entries) + + +# --------------------------------------------------------------------------- +# DataBundle +# --------------------------------------------------------------------------- + +class DataBundle: + """Accumulates per-eye DataFrames and derives feature metadata. + + Responsibilities: column type inference, scalar stats (min/max/median), + categorical index maps, and feature dim. Everything else (splits, + vectorization, image paths) lives in the profile that uses this bundle. + """ + + def __init__( + self, + *, + image_dir: str, + clinical_dir: Optional[str] = None, + label_col: str, + patient_col: str = "Patient ID", + cat_cols: Optional[Iterable[str]] = None, + max_unique_for_cat: int = 4, + n_splits: int = 5, + random_seed: int = 42, + filename_template: str = "RET{pid:03d}{eye}.jpg", + ) -> None: + self.image_dir = Path(image_dir) + self.label_col = label_col + self.patient_col = patient_col + self.max_unique_for_cat = max_unique_for_cat + self.filename_template = filename_template + self.clinical_dir = Path(clinical_dir) if clinical_dir else None + + self.frames: List[pd.DataFrame] = [] + self.df: pd.DataFrame = pd.DataFrame() + self.scalar_cols: List[str] = [] + self.cat_cols: List[str] = list(cat_cols) if cat_cols else [] + self.scalar_stats: Dict[str, Dict[str, float]] = {} + self.cat_maps: Dict[str, Dict[object, int]] = {} + self.feature_dim: int = 0 + + def add_df( + self, + df: pd.DataFrame, + *, + id_column: Optional[str] = None, + exclude_cols: Optional[Iterable[str]] = None, + ) -> None: + df = df.copy() + self._ensure_patient_id(df, id_column) + if self.label_col not in df.columns: + raise ValueError(f"label_col '{self.label_col}' not found in added dataframe") + self.frames.append(df) + self._refresh_master_df(exclude_cols=exclude_cols) + self._infer_or_validate_feature_types(exclude_cols=exclude_cols) + self._compute_numeric_stats() + self._build_cat_maps() + self._compute_feature_dim() + + # ── Internal ───────────────────────────────────────────────────────────── + + def _ensure_patient_id(self, df: pd.DataFrame, id_column: Optional[str]) -> None: + if self.patient_col in df.columns: + return + if id_column and id_column in df.columns: + df.rename(columns={id_column: self.patient_col}, inplace=True) + return + candidates = [ + c for c in df.columns + if c.lower().replace(" ", "") in {"patientid", "patient", "pid"} + ] + if len(candidates) == 1: + df.rename(columns={candidates[0]: self.patient_col}, inplace=True) + return + raise ValueError( + f"A '{self.patient_col}' column is required; " + f"provide id_column=... if it has a different name." + ) + + def _refresh_master_df(self, exclude_cols: Optional[Iterable[str]] = None) -> None: + self.df = pd.concat(self.frames, axis=0, ignore_index=True) + if exclude_cols: + self.df = self.df.drop(columns=[c for c in exclude_cols if c in self.df.columns]) + + def _infer_or_validate_feature_types(self, exclude_cols: Optional[Iterable[str]] = None) -> None: + excluded = set(exclude_cols or []) | {self.label_col, self.patient_col} + candidates = [c for c in self.df.columns if c not in excluded] + cats = set(self.cat_cols) + scalars: set[str] = set() + for c in candidates: + if c in cats: + continue + s = self.df[c] + as_num = pd.to_numeric(s, errors="coerce") + n_uniq = s.dropna().nunique() + if as_num.notna().any() and as_num.isna().mean() < 1.0 and n_uniq > self.max_unique_for_cat: + scalars.add(c) + else: + cats.add(c) + self.cat_cols = sorted(cats) + self.scalar_cols = sorted(scalars) + + def _compute_numeric_stats(self) -> None: + self.scalar_stats.clear() + for col in self.scalar_cols: + vals = pd.to_numeric(self.df[col], errors="coerce").dropna().astype(float).values + if vals.size == 0: + lo, hi, med = 0.0, 1.0, 0.0 + else: + lo, hi = float(np.min(vals)), float(np.max(vals)) + med = float(np.median(vals)) + if hi <= lo: + hi = lo + 1.0 + self.scalar_stats[col] = {"min": lo, "max": hi, "median": med} + + def _build_cat_maps(self) -> None: + self.cat_maps.clear() + for col in self.cat_cols: + cats = [v for v in self.df[col].dropna().unique().tolist()] + try: + cats = sorted(cats) + except Exception: + pass + mapping: Dict[object, int] = {"": 0} + for i, v in enumerate(cats, start=1): + mapping[v] = i + self.cat_maps[col] = mapping + + def _compute_feature_dim(self) -> None: + self.feature_dim = ( + len(self.scalar_cols) + + sum(len(m) for m in self.cat_maps.values()) + + len(self.scalar_cols) + ) + + +# --------------------------------------------------------------------------- +# HTDataset / ht_collate +# --------------------------------------------------------------------------- + +class HTDataset(Dataset): + """PyTorch Dataset backed by a LoaderShell. + + Delegates sample retrieval to each tower's get_sample(entry). + Batch keys are tower names plus "label". + """ + + def __init__(self, shell: LoaderShell, towers: dict) -> None: + self.entries = shell.entries + self.towers = towers + + def __len__(self) -> int: + return len(self.entries) + + def __getitem__(self, idx: int) -> dict[str, Any]: + entry = self.entries[idx] + sample = { + "label": torch.tensor(entry.label, dtype=torch.long), + "entity_id": entry.entity_id, + } + for name, tower in self.towers.items(): + sample[name] = tower.get_sample(entry) + return sample + + +def to_label_tensor(labels, device: torch.device) -> torch.Tensor: + """Normalise a batch of labels (tensor or list) to a long tensor on device.""" + if torch.is_tensor(labels): + return labels.to(device=device, dtype=torch.long) + return torch.as_tensor(labels, dtype=torch.long, device=device) + + +def ht_collate(batch: list[dict[str, Any]]) -> dict[str, Any]: + """Collate HTDataset samples. + + Tensor values are stacked; dict values (side dicts from patient-level + shells) are stacked per inner key; everything else becomes a list. + """ + if not batch: + return {} + result: dict[str, Any] = {} + for key in batch[0]: + vals = [s[key] for s in batch] + first = vals[0] + if isinstance(first, torch.Tensor): + result[key] = torch.stack(vals, dim=0) + elif isinstance(first, dict): + result[key] = { + side: torch.stack([v[side] for v in vals], dim=0) + for side in first + } + else: + result[key] = vals + return result diff --git a/v4/classes/heads/__init__.py b/v4/classes/heads/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/v4/classes/heads/classifier.py b/v4/classes/heads/classifier.py new file mode 100644 index 0000000..5060ae1 --- /dev/null +++ b/v4/classes/heads/classifier.py @@ -0,0 +1,24 @@ +"""classifier — ClassificationHead output head.""" +from __future__ import annotations + +import torch +import torch.nn as nn + + +class ClassificationHead(nn.Module): + """Minimal classification head: ReLU → Dropout → Linear(in_dim → num_classes). + + Used as the output stage of bridges and any module that needs a reusable, + swappable task head producing class logits. + """ + + def __init__(self, in_dim: int, num_classes: int, dropout: float = 0.5): + super().__init__() + self.head = nn.Sequential( + nn.ReLU(), + nn.Dropout(dropout), + nn.Linear(in_dim, num_classes), + ) + + def forward(self, z: torch.Tensor) -> torch.Tensor: + return self.head(z) diff --git a/v4/classes/loaders/__init__.py b/v4/classes/loaders/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/v4/classes/loaders/image_loader.py b/v4/classes/loaders/image_loader.py new file mode 100644 index 0000000..6a5ca84 --- /dev/null +++ b/v4/classes/loaders/image_loader.py @@ -0,0 +1,132 @@ +"""image_loader — CachedImageLoader: disk I/O with optional in-memory cache. + +A single instance is shared across all towers/datasets for a run so that +images are decoded from disk at most once. The optional preprocessor runs +at cache-fill time (resize, deterministic crop, etc.) so that per-batch +transforms in the image tower only need to apply stochastic augmentations. + +Usage +----- + loader = CachedImageLoader(enabled=True, workers=4) + loader.warm(paths, preprocessor=resize_fn) # optional parallel pre-fill + img = loader.load(path, preprocessor=resize_fn) # returns PIL Image +""" +from __future__ import annotations + +from concurrent.futures import ThreadPoolExecutor, as_completed +from pathlib import Path +from typing import Callable, Iterable, Optional + +import numpy as np +from PIL import Image + + +class CachedImageLoader: + """Loads PIL Images from disk with an optional shared in-memory cache. + + The cache stores decoded, pre-preprocessed images as uint8 numpy arrays + (RGB, HWC). Storing after the preprocessor runs means the deterministic + resize/crop step executes only once per image across all folds and epochs. + + Parameters + ---------- + enabled : bool + When False the cache is bypassed and every call hits disk. + workers : int + Thread count for ``warm()``. 0 or 1 → single-threaded. + """ + + def __init__(self, *, enabled: bool = True, workers: int = 4) -> None: + self._cache: dict[str, np.ndarray] | None = {} if enabled else None + self._workers = workers + + @property + def enabled(self) -> bool: + return self._cache is not None + + def __len__(self) -> int: + return len(self._cache) if self._cache is not None else 0 + + def load( + self, + path: str | Path, + preprocessor: Optional[Callable[..., Image.Image]] = None, + ) -> Image.Image: + """Return a PIL Image for *path*, using the cache when enabled.""" + key = str(path) + if self._cache is not None: + cached = self._cache.get(key) + if cached is not None: + return Image.fromarray(cached, mode="RGB") + + img = Image.open(path).convert("RGB") + if preprocessor is not None: + img = call_preprocessor(preprocessor, img, path) + + if self._cache is not None: + self._cache[key] = np.asarray(img, dtype=np.uint8) + + return img + + def warm( + self, + paths: Iterable[str | Path], + preprocessor: Optional[Callable[..., Image.Image]] = None, + ) -> None: + """Pre-populate the cache for all *paths* (no-op when disabled). + + Already-cached paths are skipped, so calling warm() multiple times + (e.g. once per fold) is safe. + """ + if self._cache is None: + return + paths = list(paths) + to_warm = [str(p) for p in paths if str(p) not in self._cache] + if not to_warm: + return + already = len(paths) - len(to_warm) + print( + f"[image_cache] warming {len(to_warm)} images" + + (f" ({already} already cached)" if already else ""), + flush=True, + ) + + def _warm_one(path_str: str) -> None: + if path_str in self._cache: + return + img = Image.open(path_str).convert("RGB") + if preprocessor is not None: + img = call_preprocessor(preprocessor, img, Path(path_str)) + self._cache[path_str] = np.asarray(img, dtype=np.uint8) + + try: + from tqdm import tqdm + except ImportError: + tqdm = None + + if self._workers <= 1: + it = tqdm(to_warm, desc="Warm image cache", unit="img") if tqdm else to_warm + for p in it: + _warm_one(p) + else: + with ThreadPoolExecutor(max_workers=self._workers) as ex: + futures = {ex.submit(_warm_one, p): p for p in to_warm} + it = ( + tqdm(as_completed(futures), total=len(futures), + desc="Warm image cache", unit="img") + if tqdm else as_completed(futures) + ) + for fut in it: + fut.result() + + +def call_preprocessor( + fn: Callable[..., Image.Image], + img: Image.Image, + path: Path | str, +) -> Image.Image: + """Call preprocessor with (img, path) or just (img) depending on arity.""" + try: + return fn(img, path) + except TypeError: + return fn(img) diff --git a/v4/classes/logging/__init__.py b/v4/classes/logging/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/v4/classes/logging/prediction_store.py b/v4/classes/logging/prediction_store.py new file mode 100644 index 0000000..f98ccf3 --- /dev/null +++ b/v4/classes/logging/prediction_store.py @@ -0,0 +1,394 @@ +"""prediction_store — per-epoch logit and embedding recording across folds and phases. + +PredictionStore — records logits for every head, fold, phase, and epoch. +FeatureStore — records embeddings (opt-in); same structure but per-head + tensors since embedding dims vary across heads. + +HDF5 layout — PredictionStore +------------------------------ +/{phase}/logits float32 (n_folds, n_epochs, n_samples, n_heads, n_classes) +/{phase}/head_names str (n_heads,) +/{phase}/y_true int64 (n_samples,) +/{phase}/entity_id_{k} int64|str (n_samples,) — one dataset per id component +/{phase}/split str (n_folds, n_samples) +/{phase}/loss float32 (n_folds, n_epochs) + +HDF5 layout — FeatureStore +--------------------------- +/{phase}/{head_name} float32 (n_folds, n_epochs, n_samples, embedding_dim) +/{phase}/y_true int64 (n_samples,) +/{phase}/entity_id_{k} int64|str (n_samples,) +/{phase}/split str (n_folds, n_samples) +""" +from __future__ import annotations + +from pathlib import Path +from typing import Sequence + +import numpy as np + +try: + import h5py +except ImportError as e: + raise ImportError("PredictionStore requires h5py: pip install h5py") from e + +_STR_DT = h5py.string_dtype() + + +# --------------------------------------------------------------------------- +# Internal phase buffer +# --------------------------------------------------------------------------- + +class _PhaseBuffer: + def __init__( + self, + entity_ids: list[tuple], + y_true: np.ndarray, + head_names: list[str], + n_epochs: int, + n_folds: int, + n_classes: int, + ): + n_s = len(entity_ids) + n_h = len(head_names) + self.entity_ids = list(entity_ids) + self.y_true = np.asarray(y_true, dtype=np.int64) + self.head_names = list(head_names) + self.n_epochs = n_epochs + self.logits = np.full((n_folds, n_epochs, n_s, n_h, n_classes), np.nan, dtype=np.float32) + self.split = np.full((n_folds, n_s), "", dtype=object) + self.loss = np.full((n_folds, n_epochs), np.nan, dtype=np.float32) + self._sid = {str(eid): i for i, eid in enumerate(entity_ids)} + self._hid = {h: i for i, h in enumerate(head_names)} + + +# --------------------------------------------------------------------------- +# Internal feature buffer (per-head, variable embedding_dim) +# --------------------------------------------------------------------------- + +class _FeaturePhaseBuffer: + def __init__( + self, + entity_ids: list[tuple], + y_true: np.ndarray, + n_folds: int, + ): + self.entity_ids = list(entity_ids) + self.y_true = np.asarray(y_true, dtype=np.int64) + self.split = np.full((n_folds, len(entity_ids)), "", dtype=object) + self._sid = {str(eid): i for i, eid in enumerate(entity_ids)} + # head_name → (buffer array, n_epochs) + self._heads: dict[str, tuple[np.ndarray, int]] = {} + + def register_head(self, head: str, n_epochs: int, embedding_dim: int, n_folds: int) -> None: + n_s = len(self.entity_ids) + self._heads[head] = ( + np.full((n_folds, n_epochs, n_s, embedding_dim), np.nan, dtype=np.float32), + n_epochs, + ) + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + +def _write_entity_ids(grp: h5py.Group, entity_ids: list[tuple]) -> None: + if not entity_ids: + return + n_components = max(len(eid) for eid in entity_ids) + for k in range(n_components): + vals = [eid[k] if k < len(eid) else "" for eid in entity_ids] + if all(isinstance(v, (int, np.integer)) for v in vals): + grp.create_dataset(f"entity_id_{k}", data=np.array(vals, dtype=np.int64)) + else: + grp.create_dataset(f"entity_id_{k}", data=np.array(vals, dtype=object), dtype=_STR_DT) + + +def _read_entity_ids(grp: h5py.Group, n_samples: int) -> list[tuple]: + k, components = 0, [] + while f"entity_id_{k}" in grp: + arr = grp[f"entity_id_{k}"][:] + if arr.dtype.kind in ("S", "O", "U"): + arr = np.array([v.decode() if isinstance(v, bytes) else str(v) for v in arr]) + components.append(arr) + k += 1 + if not components: + return [() for _ in range(n_samples)] + return [tuple(c[i] for c in components) for i in range(n_samples)] + + +def _decode_str_array(arr: np.ndarray) -> list[str]: + return [v.decode() if isinstance(v, bytes) else str(v) for v in arr.flat] + + +# --------------------------------------------------------------------------- +# PredictionStore +# --------------------------------------------------------------------------- + +class PredictionStore: + """Records per-epoch logits across all folds and phases, saves to HDF5. + + The store is generic — it knows nothing about what heads or phases exist. + The orchestrator registers phases and records whatever heads it builds. + """ + + def __init__(self, n_folds: int, n_classes: int) -> None: + self.n_folds = n_folds + self.n_classes = n_classes + self._phases: dict[str, _PhaseBuffer] = {} + + def register_phase( + self, + phase: str, + entity_ids: list[tuple], + y_true: Sequence[int], + head_names: list[str], + n_epochs: int, + ) -> None: + """Register a training phase before recording begins.""" + self._phases[phase] = _PhaseBuffer( + entity_ids=list(entity_ids), + y_true=np.asarray(y_true, dtype=np.int64), + head_names=list(head_names), + n_epochs=n_epochs, + n_folds=self.n_folds, + n_classes=self.n_classes, + ) + + def record( + self, + phase: str, + fold: int, + epoch: int, + entity_ids: Sequence[tuple], + head: str, + logits: np.ndarray, + ) -> None: + """Record a batch of logits for one head at one epoch.""" + buf = self._phases[phase] + hidx = buf._hid.get(head) + if hidx is None: + return + for i, eid in enumerate(entity_ids): + sidx = buf._sid.get(str(eid)) + if sidx is not None: + buf.logits[fold, epoch, sidx, hidx, :] = logits[i] + + def record_loss(self, phase: str, fold: int, epoch: int, loss: float) -> None: + self._phases[phase].loss[fold, epoch] = float(loss) + + def set_split( + self, + phase: str, + fold: int, + entity_ids: Sequence[tuple], + label: str, + ) -> None: + """Mark samples as 'train', 'val', or 'test' for a fold.""" + buf = self._phases[phase] + for eid in entity_ids: + sidx = buf._sid.get(str(eid)) + if sidx is not None: + buf.split[fold, sidx] = label + + # ------------------------------------------------------------------ + # Persistence + # ------------------------------------------------------------------ + + def save(self, path: str | Path) -> None: + path = Path(path) + path.parent.mkdir(parents=True, exist_ok=True) + with h5py.File(path, "w") as f: + for phase, buf in self._phases.items(): + grp = f.create_group(phase) + grp.create_dataset("logits", data=buf.logits, compression="gzip", compression_opts=4) + 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) + _write_entity_ids(grp, buf.entity_ids) + + @classmethod + def load(cls, path: str | Path) -> "PredictionStore": + """Load all phases into memory.""" + with h5py.File(path, "r") as f: + first = next(iter(f.values())) + n_folds, _, _, _, n_classes = first["logits"].shape + store = cls(n_folds=n_folds, n_classes=n_classes) + for phase in f: + grp = f[phase] + logits = grp["logits"][:] + n_folds_, n_epochs, n_samples, n_heads, _ = logits.shape + head_names = _decode_str_array(grp["head_names"][:]) + entity_ids = _read_entity_ids(grp, n_samples) + buf = _PhaseBuffer( + entity_ids=entity_ids, + y_true=grp["y_true"][:], + head_names=head_names, + n_epochs=n_epochs, + n_folds=n_folds_, + n_classes=n_classes, + ) + buf.logits = logits + buf.loss = grp["loss"][:] + split_raw = grp["split"][:] + buf.split = np.array( + [[v.decode() if isinstance(v, bytes) else str(v) for v in row] + for row in split_raw], + dtype=object, + ) + store._phases[phase] = buf + return store + + # ------------------------------------------------------------------ + # Query helpers + # ------------------------------------------------------------------ + + @property + def phases(self) -> list[str]: + return list(self._phases.keys()) + + def head_names(self, phase: str) -> list[str]: + return self._phases[phase].head_names + + def entity_ids(self, phase: str) -> list[tuple]: + return self._phases[phase].entity_ids + + def get_logits( + self, + phase: str, + head: str, + fold: int | None = None, + epoch: int | None = None, + ) -> np.ndarray: + """Slice logits for one head. Unspecified dims return the full axis. + + Returns shape (folds, epochs, samples, classes) by default, + with leading dims dropped for each specified index. + """ + buf = self._phases[phase] + hidx = buf._hid[head] + data = buf.logits[:, :, :, hidx, :] # (folds, epochs, samples, classes) + if fold is not None: data = data[fold] # (epochs, samples, classes) + if epoch is not None: data = data[..., epoch, :, :] if fold is None else data[epoch] + return data + + def get_split(self, phase: str, fold: int) -> dict[str, list[tuple]]: + """Return {'train': [...], 'val': [...], 'test': [...]} entity_id lists.""" + buf = self._phases[phase] + labels = buf.split[fold] + out: dict[str, list[tuple]] = {} + for eid, lbl in zip(buf.entity_ids, labels): + out.setdefault(lbl, []).append(eid) + return out + + +# --------------------------------------------------------------------------- +# FeatureStore +# --------------------------------------------------------------------------- + +class FeatureStore: + """Records per-epoch embeddings (variable dim per head), saves to HDF5. + + Opt-in companion to PredictionStore. Typically written only on checkpoint + runs where you want to do dimensionality reduction or cluster analysis. + """ + + def __init__(self, n_folds: int) -> None: + self.n_folds = n_folds + self._phases: dict[str, _FeaturePhaseBuffer] = {} + + def register_phase( + self, + phase: str, + entity_ids: list[tuple], + y_true: Sequence[int], + ) -> None: + self._phases[phase] = _FeaturePhaseBuffer( + entity_ids=list(entity_ids), + y_true=np.asarray(y_true, dtype=np.int64), + n_folds=self.n_folds, + ) + + def register_head( + self, + phase: str, + head: str, + n_epochs: int, + embedding_dim: int, + ) -> None: + self._phases[phase].register_head(head, n_epochs, embedding_dim, self.n_folds) + + def record( + self, + phase: str, + fold: int, + epoch: int, + entity_ids: Sequence[tuple], + head: str, + embeddings: np.ndarray, + ) -> None: + buf = self._phases[phase] + arr, _= buf._heads[head] + for i, eid in enumerate(entity_ids): + sidx = buf._sid.get(str(eid)) + if sidx is not None: + arr[fold, epoch, sidx, :] = embeddings[i] + + def set_split( + self, + phase: str, + fold: int, + entity_ids: Sequence[tuple], + label: str, + ) -> None: + buf = self._phases[phase] + for eid in entity_ids: + sidx = buf._sid.get(str(eid)) + if sidx is not None: + buf.split[fold, sidx] = label + + def save(self, path: str | Path) -> None: + path = Path(path) + path.parent.mkdir(parents=True, exist_ok=True) + with h5py.File(path, "w") as f: + 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) + _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) + + @classmethod + def load(cls, path: str | Path) -> "FeatureStore": + with h5py.File(path, "r") as f: + n_folds = next( + arr.shape[0] + for grp in f.values() + for k, arr in grp.items() + if k not in ("y_true", "split") and not k.startswith("entity_id_") + ) + store = cls(n_folds=n_folds) + _meta = {"y_true", "split"} + for phase in f: + grp = f[phase] + n_samples = grp["y_true"].shape[0] + entity_ids = _read_entity_ids(grp, n_samples) + buf = _FeaturePhaseBuffer( + entity_ids=entity_ids, + y_true=grp["y_true"][:], + n_folds=n_folds, + ) + buf.split = np.array( + [[v.decode() if isinstance(v, bytes) else str(v) for v in row] + for row in grp["split"][:]], + dtype=object, + ) + for key in grp: + if key in _meta or key.startswith("entity_id_"): + continue + arr = grp[key][:] + buf._heads[key] = (arr, arr.shape[1]) + store._phases[phase] = buf + return store diff --git a/v4/classes/metrics.py b/v4/classes/metrics.py new file mode 100644 index 0000000..efabd28 --- /dev/null +++ b/v4/classes/metrics.py @@ -0,0 +1,175 @@ +"""metrics — loss, scoring, calibration, and threshold/bias tuning.""" +from __future__ import annotations + +from typing import Optional + +import numpy as np +import torch +import torch.nn.functional as F +from sklearn.metrics import ( + cohen_kappa_score, + f1_score, + matthews_corrcoef, + recall_score, + roc_auc_score, + roc_curve, +) + + +# --------------------------------------------------------------------------- +# Loss +# --------------------------------------------------------------------------- + +def focal_loss( + logits: torch.Tensor, + targets: torch.Tensor, + gamma: float = 0.0, + weight: Optional[torch.Tensor] = None, + reduction: str = "mean", +) -> torch.Tensor: + """Focal loss; reduces to cross-entropy when gamma=0.""" + if gamma <= 0: + return F.cross_entropy(logits, targets, weight=weight, reduction=reduction) + log_probs = F.log_softmax(logits, dim=1) + probs = log_probs.exp() + targets = targets.long().view(-1, 1) + logpt = log_probs.gather(1, targets) + pt = probs.gather(1, targets) + loss = -(((1.0 - pt).clamp_min(0.0) ** gamma) * logpt) + if weight is not None: + loss = loss * weight.gather(0, targets.view(-1)).view(-1, 1) + loss = loss.view(-1) + if reduction == "sum": return loss.sum() + if reduction == "mean": return loss.mean() + return loss + + +# --------------------------------------------------------------------------- +# Basic array scoring +# --------------------------------------------------------------------------- + +def score_arrays(y_true: np.ndarray, probs: np.ndarray, num_classes: int): + """Return (acc, auc, n).""" + if y_true.size == 0: + return float("nan"), float("nan"), 0 + acc = float((probs.argmax(1) == y_true).mean()) + try: + auc = ( + float(roc_auc_score(y_true, probs[:, 1])) + if num_classes == 2 + else float(roc_auc_score(y_true, probs, multi_class="ovr", average="macro")) + ) + except Exception: + auc = float("nan") + return acc, auc, int(len(y_true)) + + +# --------------------------------------------------------------------------- +# Calibration +# --------------------------------------------------------------------------- + +def compute_ece(y_true: np.ndarray, probs: np.ndarray, n_bins: int = 10) -> float: + """Expected Calibration Error: weighted mean |confidence − accuracy| per bin.""" + if y_true.size == 0: + return float("nan") + confidences = probs.max(axis=1) + predictions = probs.argmax(axis=1) + bin_edges = np.linspace(0.0, 1.0, n_bins + 1) + ece = 0.0 + n = len(y_true) + for i, (lo, hi) in enumerate(zip(bin_edges[:-1], bin_edges[1:])): + mask = (confidences >= lo) & ( + confidences <= hi if i == n_bins - 1 else confidences < hi + ) + if not mask.any(): + continue + ece += float(mask.sum()) / n * abs( + float(confidences[mask].mean()) - float((predictions[mask] == y_true[mask]).mean()) + ) + return float(ece) + + +def compute_extended_metrics( + y_true: np.ndarray, + probs: np.ndarray, + num_classes: int, + n_bins: int = 10, + preds_override: Optional[np.ndarray] = None, +) -> dict: + nan = float("nan") + if y_true.size == 0: + return dict( + kappa=nan, mcc=nan, macro_f1=nan, + per_class_recall=np.full(num_classes, nan), ece=nan, + ) + preds = preds_override if preds_override is not None else probs.argmax(axis=1) + try: kappa = float(cohen_kappa_score(y_true, preds)) + except: kappa = nan + try: mcc = float(matthews_corrcoef(y_true, preds)) + except: mcc = nan + try: macro_f1 = float(f1_score(y_true, preds, average="macro", zero_division=0)) + except: macro_f1 = nan + try: + pcr = recall_score( + y_true, preds, average=None, + labels=list(range(num_classes)), zero_division=0, + ).astype(float) + except: + pcr = np.full(num_classes, nan) + return dict( + kappa=kappa, mcc=mcc, macro_f1=macro_f1, + per_class_recall=pcr, ece=compute_ece(y_true, probs, n_bins=n_bins), + ) + + +# --------------------------------------------------------------------------- +# Threshold / bias tuning +# --------------------------------------------------------------------------- + +def tune_binary_threshold(y_true: np.ndarray, p1: np.ndarray) -> float: + """Pick threshold via Youden's J (sensitivity + specificity − 1). + + Class-distribution independent; falls back to 0.5 if fewer than two + classes are present in y_true. + """ + if y_true.size == 0 or len(np.unique(y_true)) < 2: + return 0.5 + fpr, tpr, thresholds = roc_curve(y_true, p1) + return float(thresholds[np.argmax(tpr + (1.0 - fpr) - 1.0)]) + + +def multiclass_acc_with_bias( + y_true: np.ndarray, probs: np.ndarray, bias: np.ndarray +) -> float: + """Balanced accuracy (mean per-class recall) after applying log-space bias.""" + if y_true.size == 0: + return float("nan") + logits = np.log(np.clip(probs, 1e-8, 1.0)) + bias.reshape(1, -1) + preds = np.argmax(logits, axis=1) + classes = np.unique(y_true) + return float(np.mean([(preds[y_true == c] == c).mean() for c in classes])) + + +def tune_multiclass_bias( + y_true: np.ndarray, probs: np.ndarray, *, iters: int = 2 +) -> np.ndarray: + """Grid-search per-class log-space bias to maximise balanced accuracy.""" + if y_true.size == 0 or probs.size == 0: + return np.zeros((0,), dtype=float) + c = probs.shape[1] + bias = np.zeros((c,), dtype=float) + grid = np.linspace(-1.0, 1.0, 41) + for _ in range(iters): + for k in range(c): + best_v = bias[k] + best_acc = multiclass_acc_with_bias(y_true, probs, bias) + old = bias[k] + for v in grid: + bias[k] = float(v) + acc = multiclass_acc_with_bias(y_true, probs, bias) + if acc > best_acc or (acc == best_acc and abs(v) < abs(best_v)): + best_acc, best_v = acc, float(v) + bias[k] = best_v + if np.isnan(best_acc): + bias[k] = old + return bias diff --git a/v4/classes/profiles/__init__.py b/v4/classes/profiles/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/v4/classes/profiles/v4papila.py b/v4/classes/profiles/v4papila.py new file mode 100644 index 0000000..7936815 --- /dev/null +++ b/v4/classes/profiles/v4papila.py @@ -0,0 +1,663 @@ +"""v4papila — self-contained PAPILA data module. + +Public contract (v4 orchestrator interface) +------------------------------------------- + bundle = build_data(args: dict) -> PapilaBundle + +PapilaBundle exposes: + .df full preprocessed DataFrame (for split building) + .label_col, .patient_col, .feature_dim + .id_names tuple of semantic names for each entity_id slot + e.g. ("patient_id", "eye") — used by orchestrator for logging + .matrix ClinicalDataView (all eyes) + .matrix.od / .matrix.os scoped views (OD or OS only) + .image ImageDataView (all eyes) + .image.od / .image.os scoped views + .build_shells(df, *, level) -> LoaderShell + +DataView interface (consumed by towers' get_sample) +---------------------------------------------------- +Both views accept positional id slots (*ids) matching entity_id tuple positions. +Semantic names for each position are in view.id_names. + +ClinicalDataView: + .feature_dim + .id_names e.g. ("patient_id", "eye") + .vectorize_entity(*ids) -> np.ndarray + .side_map -> dict mapping generic keys {"a", "b"} to id_1 values + .od, .os -> scoped ClinicalDataView + +ImageDataView: + .id_names e.g. ("patient_id", "eye") + .get_image_path(*ids) -> Path + .load_image(*ids) -> PIL.Image + .side_map -> dict mapping generic keys {"a", "b"} to id_1 values + .od, .os -> scoped ImageDataView +""" +from __future__ import annotations + +import sys +from functools import cached_property +from pathlib import Path +from typing import Callable, Dict, List, Optional + +import numpy as np +import pandas as pd +from PIL import Image + +_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.loaders.image_loader import CachedImageLoader, call_preprocessor +from v4.classes.dataset import DataBundle, LoaderShell, ShellEntry + + +# --------------------------------------------------------------------------- +# Pachymetry → IOP correction (PAPILA Table 3) +# --------------------------------------------------------------------------- + +_PACHY_TABLE: Dict[int, int] = { + 475: +5, 485: +4, 495: +4, 505: +3, 515: +2, + 525: +1, 535: +1, 545: 0, 555: -1, 565: -1, + 575: -2, 585: -3, 595: -4, 605: -4, 615: -5, +} +_PACHY_KEYS = np.array(sorted(_PACHY_TABLE.keys())) + + +def _nearest_pachy_key(x: float) -> int: + return int(_PACHY_KEYS[int(np.argmin(np.abs(_PACHY_KEYS - float(x))))]) + + +def _fit_perkins_converter( + frames: List[pd.DataFrame], method: str +) -> Callable[[float, Optional[float]], float]: + combined = pd.concat(frames, ignore_index=True) + paired = combined.dropna(subset=["Pneumatic", "Perkins"]) + if len(paired) == 0: + raise ValueError("No paired Pneumatic+Perkins rows; cannot fit converter.") + pneumatic = paired["Pneumatic"].values.astype(float) + perkins = paired["Perkins"].values.astype(float) + + if method == "ratio": + ratio = float((pneumatic / perkins).mean()) + def _conv(p: float, pachy: Optional[float] = None) -> float: + return p * ratio + return _conv + + elif method == "ols": + from scipy import stats as _stats + slope, intercept, *_ = _stats.linregress(perkins, pneumatic) + slope, intercept = float(slope), float(intercept) + def _conv(p: float, pachy: Optional[float] = None) -> float: + return p * slope + intercept + return _conv + + elif method == "lad": + from scipy import stats as _stats + from scipy.optimize import minimize as _minimize + slope0, intercept0, *_ = _stats.linregress(perkins, pneumatic) + def _lad_loss(params): + a, b = params + return np.abs(pneumatic - (a * perkins + b)).mean() + res = _minimize(_lad_loss, x0=[slope0, intercept0], method="Nelder-Mead") + slope, intercept = float(res.x[0]), float(res.x[1]) + def _conv(p: float, pachy: Optional[float] = None) -> float: + return p * slope + intercept + return _conv + + elif method == "multi": + from numpy.linalg import lstsq as _lstsq + pm = combined.dropna(subset=["Pneumatic", "Perkins", "Pachymetry"]) + if len(pm) == 0: + raise ValueError("No Pneumatic+Perkins+Pachymetry rows; cannot fit multi.") + pneu = pm["Pneumatic"].values.astype(float) + perk = pm["Perkins"].values.astype(float) + pv = pm["Pachymetry"].values.astype(float) + X = np.column_stack([perk, pv, np.ones(len(perk))]) + coeffs, *_ = _lstsq(X, pneu, rcond=None) + slope, pachy_coef, intercept = float(coeffs[0]), float(coeffs[1]), float(coeffs[2]) + fallback = float(pv.mean()) + def _conv(p: float, pachy: Optional[float] = None) -> float: + pval = pachy if (pachy is not None and not np.isnan(pachy)) else fallback + return p * slope + pachy_coef * pval + intercept + return _conv + + else: + raise ValueError(f"Unknown iop_corr_method: {method!r}. Choose ratio/ols/lad/multi.") + + +def _pick_iop(row: pd.Series, converter: Callable) -> float: + pneumatic = row.get("Pneumatic", np.nan) + if not pd.isna(pneumatic): + return float(pneumatic) + perkins = row.get("Perkins", np.nan) + if pd.isna(perkins): + return np.nan + pachy = row.get("Pachymetry", np.nan) + return converter(float(perkins), None if pd.isna(pachy) else float(pachy)) + + +def _correct_iop(raw_iop: float, pachy: float) -> float: + if pd.isna(raw_iop): + return np.nan + if pd.isna(pachy): + return float(raw_iop) + key = _nearest_pachy_key(float(pachy)) + return float(raw_iop) + float(_PACHY_TABLE[key]) + + +def _apply_iop_and_drop_md( + df: pd.DataFrame, converter: Callable, drop_raw: bool = False +) -> pd.DataFrame: + df["IOP_raw"] = df.apply(lambda row: _pick_iop(row, converter), axis=1) + pachy = df.get("Pachymetry", pd.Series(np.nan, index=df.index)) + df["IOP_corr"] = [ + _correct_iop(r, p) for r, p in zip(df["IOP_raw"].values, pachy.values) + ] + drop = [c for c in ("Pneumatic", "Perkins", "VF_MD") if c in df.columns] + if drop_raw: + drop.append("IOP_raw") + if drop: + df.drop(columns=drop, inplace=True) + return df + + +def _canonicalize_eye_column(df: pd.DataFrame) -> None: + if "eyeID" in df.columns: + src = "eyeID" + else: + src = next((c for c in df.columns if "eye" in c.lower()), None) + if src is None: + df["eyeID"] = "OS" + return + + def norm(v): + if pd.isna(v): + return None + x = str(v).strip().upper() + if x in {"OS", "L", "LEFT", "0"}: return "OS" + if x in {"OD", "R", "RIGHT", "1"}: return "OD" + try: + num = int(float(x)) + return "OD" if num % 2 == 1 else "OS" + except Exception: + return None + + mapped = df[src].map(norm) + uniq = {u for u in mapped.dropna().unique().tolist()} + if not uniq.issubset({"OS", "OD"}): + raise ValueError(f"eyeID must be binary; found {sorted(uniq)}") + df["eyeID"] = mapped.fillna("OS") + + +# --------------------------------------------------------------------------- +# DataView classes +# --------------------------------------------------------------------------- + +class ClinicalDataView: + """Tabular feature view over a (possibly eye-filtered) clinical DataFrame. + + Exposes feature_dim and vectorize_entity so towers can retrieve + feature vectors by entity identity without knowing about the DataFrame. + + Scoped views (OD or OS only) are accessed via .od and .os properties. + + id_names gives semantic labels for each positional slot in an entity_id tuple, + e.g. ("patient_id", "eye"). The orchestrator uses this for logging without + needing to know PAPILA-specific field names itself. + """ + + # PAPILA canonical side keys used in ShellEntry entity_ids + SIDE_A = "OD" + SIDE_B = "OS" + + # Semantic name for each entity_id position (id_0, id_1, ...) + id_names: tuple[str, ...] = ("patient_id", "eye") + + def __init__( + self, + df: pd.DataFrame, + patient_col: str, + scalar_cols: list[str], + cat_cols: list[str], + scalar_stats: dict, + cat_maps: dict, + *, + eye_filter: str | None = None, # "OD", "OS", or None (all eyes) + ): + self._df = df + self.patient_col = patient_col + self.scalar_cols = scalar_cols + self.cat_cols = cat_cols + self.scalar_stats = scalar_stats + self.cat_maps = cat_maps + self._eye_filter = eye_filter + + # Build a (patient_id, eyeID) → row index for fast lookup + if "eyeID" in df.columns: + self._idx = df.set_index([patient_col, "eyeID"]) + else: + self._idx = df.set_index(patient_col) + + # ── Public API ────────────────────────────────────────────────────────── + + @property + def feature_dim(self) -> int: + n_scalar = len(self.scalar_cols) + n_cat = sum(len(m) for m in self.cat_maps.values()) + return n_scalar + n_cat + n_scalar # scalars + one-hots + missing flags + + def vectorize_entity(self, *ids) -> np.ndarray: + """Return the feature vector for an entity identified by positional ids. + + Positional slots match entity_id tuple positions (see id_names). + For PAPILA: ids = (id_0, id_1) = (patient_id, eye). + """ + try: + row = self._idx.loc[ids if len(ids) > 1 else ids[0]].copy() + except KeyError: + names = self.id_names[:len(ids)] + raise KeyError( + f"ClinicalDataView: no row found for {dict(zip(names, ids))}" + ) + # set_index removes index-level columns from the row; restore any that + # _vectorize_row needs (e.g. eyeID is a cat feature AND an index level) + idx_names = (self._idx.index.names + if hasattr(self._idx.index, 'names') + else [self._idx.index.name]) + for name, val in zip(idx_names, ids if len(ids) > 1 else [ids[0]]): + if name not in row.index: + row[name] = val + return self._vectorize_row(row) + + # ── Scoped views ───────────────────────────────────────────────────────── + + @property + def side_map(self) -> dict[str, str]: + """Generic side-key → dataset side string. Towers use this for patient-level shells.""" + return {"a": self.SIDE_A, "b": self.SIDE_B} + + @cached_property + def od(self) -> "ClinicalDataView": + return self._scoped(self.SIDE_A) + + @cached_property + def os(self) -> "ClinicalDataView": + return self._scoped(self.SIDE_B) + + def _scoped(self, eye: str) -> "ClinicalDataView": + sub = self._df[self._df["eyeID"] == eye].reset_index(drop=True) + return ClinicalDataView( + df=sub, + patient_col=self.patient_col, + scalar_cols=self.scalar_cols, + cat_cols=self.cat_cols, + scalar_stats=self.scalar_stats, + cat_maps=self.cat_maps, + eye_filter=eye, + ) + + # ── Internal ───────────────────────────────────────────────────────────── + + def _vectorize_row(self, row: pd.Series) -> np.ndarray: + feats: list[float] = [] + miss: list[float] = [] + for col in self.scalar_cols: + v = pd.to_numeric(row.get(col), errors="coerce") + if pd.isna(v): + miss.append(1.0) + v = self.scalar_stats[col]["median"] + else: + miss.append(0.0) + lo = self.scalar_stats[col]["min"] + hi = self.scalar_stats[col]["max"] + feats.append((float(v) - lo) / (hi - lo) if hi > lo else 0.0) + for col in self.cat_cols: + mapping = self.cat_maps[col] + one = [0.0] * len(mapping) + key = row.get(col) + one[mapping.get(key, 0)] = 1.0 + feats.extend(one) + feats.extend(miss) + return np.asarray(feats, dtype=np.float32) + + +class ImageDataView: + """Image path and loading view over a (possibly eye-filtered) DataFrame. + + Provides get_image_path and load_image keyed by positional id slots. + Scoped views (.od, .os) are available for single-side towers. + The optional image_cache is a shared CachedImageLoader for the run. + + id_names gives semantic labels for each positional slot in an entity_id tuple, + e.g. ("patient_id", "eye"). + """ + + SIDE_A = "OD" + SIDE_B = "OS" + + id_names: tuple[str, ...] = ("patient_id", "eye") + + def __init__( + self, + df: pd.DataFrame, + patient_col: str, + image_dir: str, + filename_template: str, + preprocessor: Callable | None = None, + image_cache: CachedImageLoader | None = None, + *, + eye_filter: str | None = None, + ): + self._df = df + self.patient_col = patient_col + self.image_dir = Path(image_dir) + self.filename_template = filename_template + self.preprocessor = preprocessor + self.image_cache = image_cache + self._eye_filter = eye_filter + + if "eyeID" in df.columns: + self._idx = df.set_index([patient_col, "eyeID"]) + else: + self._idx = df.set_index(patient_col) + + # ── Public API ────────────────────────────────────────────────────────── + + def get_image_path(self, *ids) -> Path: + """Return the image path for an entity identified by positional ids. + + For PAPILA: ids = (id_0, id_1) = (patient_id, eye). + """ + id_0, id_1 = ids # PAPILA always uses two slots + try: + row = self._idx.loc[(id_0, id_1)] + except KeyError: + names = self.id_names[:len(ids)] + raise KeyError( + f"ImageDataView: no row for {dict(zip(names, ids))}" + ) + pid = int(row[self.patient_col]) if self.patient_col in row.index else int(id_0) + return self.image_dir / self.filename_template.format(pid=pid, eye=id_1) + + def load_image(self, *ids) -> Image.Image: + """Load image for an entity identified by positional ids.""" + path = self.get_image_path(*ids) + if self.image_cache is not None: + return self.image_cache.load(path, preprocessor=self.preprocessor) + img = Image.open(path).convert("RGB") + if self.preprocessor is not None: + img = call_preprocessor(self.preprocessor, img, path) + return img + + # ── Scoped views ───────────────────────────────────────────────────────── + + @cached_property + def od(self) -> "ImageDataView": + return self._scoped(self.SIDE_A) + + @cached_property + def os(self) -> "ImageDataView": + return self._scoped(self.SIDE_B) + + @property + def side_map(self) -> dict[str, str]: + """Generic side-key → dataset side string. Towers use this for patient-level shells.""" + return {"a": self.SIDE_A, "b": self.SIDE_B} + + def _scoped(self, eye: str) -> "ImageDataView": + sub = self._df[self._df["eyeID"] == eye].reset_index(drop=True) + return ImageDataView( + df=sub, + patient_col=self.patient_col, + image_dir=str(self.image_dir), + filename_template=self.filename_template, + preprocessor=self.preprocessor, + image_cache=self.image_cache, + eye_filter=eye, + ) + + +# --------------------------------------------------------------------------- +# PapilaBundle — the v4 DataBundle returned by build_data +# --------------------------------------------------------------------------- + +class PapilaBundle: + """V4 DataBundle for PAPILA. + + Wraps the v3 DataBundle for backward compatibility (df, feature_dim, + vectorize_row, get_image_path, patient_col, label_col) while adding + the v4 DataView interface and build_shells(). + """ + + def __init__( + self, + bundle: DataBundle, + image_dir: str, + preprocessor: Callable | None = None, + image_cache: CachedImageLoader | None = None, + ): + self._bundle = bundle + self._image_dir = image_dir + + # ── ClinicalDataView (all eyes) ────────────────────────────────────── + self.matrix = ClinicalDataView( + df=bundle.df, + patient_col=bundle.patient_col, + scalar_cols=bundle.scalar_cols, + cat_cols=bundle.cat_cols, + scalar_stats=bundle.scalar_stats, + cat_maps=bundle.cat_maps, + ) + + # ── ImageDataView (all eyes) ───────────────────────────────────────── + self.image = ImageDataView( + df=bundle.df, + patient_col=bundle.patient_col, + image_dir=image_dir, + filename_template=bundle.filename_template, + preprocessor=preprocessor, + image_cache=image_cache, + ) + + # ── Entity-id metadata (for orchestrator logging) ──────────────────────── + + @property + def id_names(self) -> tuple[str, ...]: + """Semantic names for each entity_id position, e.g. ('patient_id', 'eye'). + + Orchestrators use this to decode entity_ids for logging without + hardcoding dataset-specific field names. + """ + return self.matrix.id_names # both views share the same structure + + # ── Identity column registry (for orchestrator split_identity_level) ──────── + + @property + def identity_cols(self) -> list[str]: + """Ordered list of grouping columns, one per identity level. + + identity_level=1 → identity_cols[0] → patient column (group by patient) + identity_level=2 → identity_cols[1] → eye column (group by patient+eye) + """ + return [self._bundle.patient_col, "eyeID"] + + # ── Backward-compat delegates ──────────────────────────────────────────── + + @property + def df(self) -> pd.DataFrame: + return self._bundle.df + + @property + def label_col(self) -> str: + return self._bundle.label_col + + @property + def patient_col(self) -> str: + return self._bundle.patient_col + + @property + def feature_dim(self) -> int: + return self._bundle.feature_dim + + def vectorize_row(self, row: pd.Series) -> np.ndarray: + return self._bundle.vectorize_row(row) + + def get_image_path(self, row: pd.Series): + return self._bundle.get_image_path(row) + + # ── Shell building ─────────────────────────────────────────────────────── + + def build_shells( + self, + df: pd.DataFrame, + *, + level: str = "eye", + label_filter: list[int] | None = None, + ) -> LoaderShell: + """Build a LoaderShell from a split DataFrame. + + level="eye" — one ShellEntry per eye row. + entity_id = (patient_id, side_key) + e.g. (42, "OD") or (42, "OS") + + level="patient" — one ShellEntry per patient. + entity_id = (patient_id,) — 1-tuple + Towers that need both sides use data_view.side_map + to assemble them in get_sample. + Patients missing either eye are excluded. + """ + pc = self.patient_col + lc = self.label_col + + if label_filter is not None: + df = df[df[lc].isin(label_filter)] + + entries: list[ShellEntry] = [] + + if level == "eye": + for _, row in df.iterrows(): + pid = int(row[pc]) + label = int(row[lc]) + side = str(row.get("eyeID", "OD")) + entries.append(ShellEntry(entity_id=(pid, side), label=label)) + + elif level == "patient": + for pid, grp in df.groupby(pc): + if "eyeID" in grp.columns: + eyes = set(grp["eyeID"].unique()) + if "OD" not in eyes or "OS" not in eyes: + continue + label_mode = grp[lc].mode() + label = int(label_mode.iloc[0]) if not label_mode.empty else int(grp[lc].iloc[0]) + entries.append(ShellEntry(entity_id=(int(pid),), label=label)) + + else: + raise ValueError(f"Unknown shell level: {level!r}. Choose 'eye' or 'patient'.") + + return LoaderShell(entries=entries) + + +# --------------------------------------------------------------------------- +# Resolve helper — used by the orchestrator to inject DataView into towers +# --------------------------------------------------------------------------- + +def resolve_data_source(bundle: PapilaBundle, path: str): + """Resolve a dot-path data source string against a PapilaBundle. + + Examples + -------- + "matrix" → bundle.matrix + "matrix.od" → bundle.matrix.od + "image" → bundle.image + "image.os" → bundle.image.os + """ + obj = bundle + for part in path.split("."): + obj = getattr(obj, part) + return obj + + +# --------------------------------------------------------------------------- +# Public contract: build_data(args: dict) -> PapilaBundle +# --------------------------------------------------------------------------- + +_DEFAULT_CAT_COLS = ["Gender", "Phakic/Pseudophakic"] + + +def build_data(args: dict) -> PapilaBundle: + """Build and return a PapilaBundle. + + args keys + --------- + image_dir (required) + clinical_dir (required) + label_col (default: "Diagnosis") + iop_corr_method (default: "ratio") + iop_drop_raw (default: False) + exclude_cols (default: []) + cat_cols (default: ["Gender", "Phakic/Pseudophakic"]) + n_splits (default: 5) + random_seed (default: 42) + in_memory_cache (default: False) — enable shared image cache for the run + """ + image_dir = args["image_dir"] + clinical_dir = args["clinical_dir"] + label_col = args.get("label_col", "Diagnosis") + iop_method = args.get("iop_corr_method", "ratio") + iop_drop_raw = bool(args.get("iop_drop_raw", False)) + exclude_cols = list(args.get("exclude_cols", [])) + cat_cols = list(args.get("cat_cols", _DEFAULT_CAT_COLS)) + n_splits = int(args.get("n_splits", 5)) + random_seed = int(args.get("random_seed", 42)) + use_cache = bool(args.get("in_memory_cache", False)) + + effective_cat = [c for c in cat_cols if c not in exclude_cols] + + bundle = DataBundle( + image_dir=image_dir, + clinical_dir=clinical_dir, + label_col=label_col, + patient_col="Patient ID", + cat_cols=effective_cat, + n_splits=n_splits, + random_seed=random_seed, + filename_template="RET{pid:03d}{eye}.jpg", + ) + + od = pd.read_excel(f"{clinical_dir}/patient_data_od.xlsx", header=1) + od["eyeID"] = "OD" + os_ = pd.read_excel(f"{clinical_dir}/patient_data_os.xlsx", header=1) + os_["eyeID"] = "OS" + + for frame in (od, os_): + if "Patient ID" not in frame.columns and "ID" in frame.columns: + frame.rename(columns={"ID": "Patient ID"}, inplace=True) + frame["Patient ID"] = ( + frame["Patient ID"].astype(str).str.extract(r"(\d+)")[0].astype(int) + ) + _canonicalize_eye_column(frame) + + bundle.add_df(od, id_column="ID", exclude_cols=exclude_cols or None) + bundle.add_df(os_, id_column="ID", exclude_cols=exclude_cols or None) + + converter = _fit_perkins_converter(bundle.frames, method=iop_method) + for i in range(len(bundle.frames)): + bundle.frames[i] = _apply_iop_and_drop_md( + bundle.frames[i], converter=converter, drop_raw=iop_drop_raw + ) + + bundle._refresh_master_df(exclude_cols=exclude_cols or None) + bundle._infer_or_validate_feature_types(exclude_cols=exclude_cols or None) + bundle._compute_numeric_stats() + bundle._build_cat_maps() + bundle._compute_feature_dim() + + image_cache = CachedImageLoader() if use_cache else None + + return PapilaBundle( + bundle=bundle, + image_dir=image_dir, + image_cache=image_cache, + ) diff --git a/v4/classes/split_manager.py b/v4/classes/split_manager.py new file mode 100644 index 0000000..1ab221f --- /dev/null +++ b/v4/classes/split_manager.py @@ -0,0 +1,177 @@ +"""split_manager — generic stratified k-fold splitter. + +SplitManager splits a DataFrame into train/val/test folds using an +outer/inner k-fold scheme. The grouping identity is controlled by +``group_col``: + + group_col=None — row-level splits (each row is its own identity) + group_col="Patient ID" — group-level splits (all rows sharing a group + key land in the same fold) + +The translation from a conceptual "identity level" to a concrete column +name belongs in the caller (typically the orchestrator), which has access +to the data bundle's column schema. +""" +from __future__ import annotations + +from dataclasses import dataclass +from typing import Any, Iterable, Optional + +import numpy as np +import pandas as pd +from sklearn.model_selection import KFold, StratifiedKFold + + +# --------------------------------------------------------------------------- +# Data structures +# --------------------------------------------------------------------------- + +@dataclass(frozen=True) +class SplitPlan: + train_ids: set[Any] + val_ids: set[Any] + test_ids: set[Any] + + +@dataclass +class Split: + """One fold's train/val/test DataFrames.""" + train: pd.DataFrame + val: pd.DataFrame + test: Optional[pd.DataFrame] = None + + +# --------------------------------------------------------------------------- +# Core splitter +# --------------------------------------------------------------------------- + +def _can_stratify(labels: np.ndarray, n_splits: int) -> bool: + if labels.size == 0: + return False + unique, counts = np.unique(labels, return_counts=True) + return len(unique) >= 2 and bool(np.all(counts >= n_splits)) + + +def _build_split_plans( + ids: np.ndarray, + labels: np.ndarray, + n_splits: int, + seed: int, +) -> list[SplitPlan]: + """Outer/inner k-fold: test=fold k, val=fold (k+1)%n, train=remaining.""" + if ids.size == 0: + raise ValueError("No samples available for splitting") + if len(set(ids.tolist())) != ids.size: + raise ValueError("ids must be unique") + if n_splits < 3: + raise ValueError("n_splits must be >= 3 for outer/inner k-fold") + + if _can_stratify(labels, n_splits): + splitter = StratifiedKFold(n_splits=n_splits, shuffle=True, random_state=seed) + folds = list(splitter.split(ids, labels)) + else: + splitter = KFold(n_splits=n_splits, shuffle=True, random_state=seed) + folds = list(splitter.split(ids)) + + fold_sets = [set(ids[test_idx].tolist()) for _, test_idx in folds] + + plans = [] + for k in range(n_splits): + test_ids = fold_sets[k] + val_ids = fold_sets[(k + 1) % n_splits] + train_ids = set().union(*(fold_sets[j] for j in range(n_splits) + if j != k and j != (k + 1) % n_splits)) + plans.append(SplitPlan(train_ids=train_ids, val_ids=val_ids, test_ids=test_ids)) + return plans + + +# --------------------------------------------------------------------------- +# SplitManager +# --------------------------------------------------------------------------- + +class SplitManager: + """Generic stratified k-fold split manager. + + Parameters + ---------- + group_col : str | None + Column whose values define the grouping identity for fold assignment. + ``None`` splits on individual rows (no grouping). + label_col : str | None + Column used for stratification. Resolved from the DataFrame at + ``build_plans`` time if not provided here. + """ + + def __init__( + self, + group_col: Optional[str] = None, + label_col: Optional[str] = None, + ) -> None: + self.group_col = group_col + self.label_col = label_col + + def build_plans( + self, + df: pd.DataFrame, + *, + n_splits: int = 5, + seed: int = 42, + label_col: Optional[str] = None, + ) -> list[Split]: + """Build n_splits fold plans from df. + + Parameters + ---------- + df : full DataFrame (pre-filtered to the desired eval mode) + n_splits : number of folds (must be >= 3) + seed : random seed for reproducibility + label_col : override for stratification column (falls back to + ``self.label_col``, then raises) + """ + lc = label_col or self.label_col + if lc is None: + raise ValueError("label_col must be provided to build_plans or SplitManager") + if lc not in df.columns: + raise ValueError(f"label_col {lc!r} not found in DataFrame") + + df = df.copy().reset_index(drop=True) + + if self.group_col is None: + # Row-level: each row is its own identity + ids = df.index.to_numpy() + labels = df[lc].to_numpy() + plans = _build_split_plans(ids, labels, n_splits, seed) + return [ + Split( + train=df[df.index.isin(p.train_ids)].reset_index(drop=True), + val =df[df.index.isin(p.val_ids) ].reset_index(drop=True), + test =df[df.index.isin(p.test_ids) ].reset_index(drop=True), + ) + for p in plans + ] + else: + gc = self.group_col + if gc not in df.columns: + raise ValueError(f"group_col {gc!r} not found in DataFrame") + # Group-level: collapse to one row per group, then split + group_table = ( + df.groupby(gc, as_index=False)[lc] + .agg(lambda s: s.mode().iloc[0] if not s.mode().empty else s.iloc[0]) + .rename(columns={lc: "_label"}) + .sort_values(gc) + .reset_index(drop=True) + ) + plans = _build_split_plans( + group_table[gc].to_numpy(), + group_table["_label"].to_numpy(), + n_splits, + seed, + ) + return [ + Split( + train=df[df[gc].isin(p.train_ids)].reset_index(drop=True), + val =df[df[gc].isin(p.val_ids) ].reset_index(drop=True), + test =df[df[gc].isin(p.test_ids) ].reset_index(drop=True), + ) + for p in plans + ] diff --git a/v4/classes/stages/__init__.py b/v4/classes/stages/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/v4/classes/stages/fusion.py b/v4/classes/stages/fusion.py new file mode 100644 index 0000000..b8dbb47 --- /dev/null +++ b/v4/classes/stages/fusion.py @@ -0,0 +1,273 @@ +"""stages/fusion — fusion stage runner: trains a bridge + associated head stages.""" +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 ( + encode_embedding, get_out_dim, resolve_input_dims, phase_for_epoch, +) + + +def collect_probs( + bridge, + primary_head, + stage_cfg: dict, + towers: dict, + stage_models: dict, + cfg_stages: list[dict], + loader, + device, + num_classes: int, +) -> tuple[np.ndarray, np.ndarray]: + """Eval pass for one fusion stage; returns (y_true, softmax_probs).""" + 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 = [], [] + + with torch.no_grad(): + for batch in loader: + y = batch.get("label") + if not torch.is_tensor(y): + continue + + if is_bilateral: + side_embs = { + side: encode_embedding(src, batch, side, towers, stage_models, cfg_stages, device) + for side, src in inputs.items() + } + z = bridge(side_embs) + else: + embs = [ + encode_embedding(n, batch, None, towers, stage_models, cfg_stages, device) + for n in inputs + ] + z = bridge(embs) + + logits = primary_head(z) + y_all.append(to_label_tensor(y, device).cpu().numpy()) + p_all.append(F.softmax(logits, dim=1).cpu().numpy()) + + 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) + + +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, +) -> tuple[dict, dict]: + """Train one fusion stage + its associated head stages. + + Returns (updated_stage_models, metrics_dict). + """ + nan = float("nan") + name = stage_cfg["name"] + level = stage_cfg["level"] + epochs = stage_cfg["epochs"] + 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) + 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) + return stage_models, {} + + bs = cfg["training"]["batch_size"] + 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) + + # ── Bridge ─────────────────────────────────────────────────────────────── + input_dims = resolve_input_dims(inputs, towers, stage_models) + bmod = importlib.import_module(stage_cfg["module"]) + bridge = getattr(bmod, stage_cfg["class"])(input_dims, **stage_cfg.get("args", {})).to(device) + + # ── Head stages ────────────────────────────────────────────────────────── + 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")) + head_models[hs["name"]] = h_cls(h_dim, num_classes).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)] + + if primary_hs_cfg is None: + print(f" WARNING: no primary head for stage {name!r}; skipping.", flush=True) + return stage_models, {} + + primary_head = head_models[primary_hs_cfg["name"]] + + # ── Freeze prior stages ─────────────────────────────────────────────────── + for m in stage_models.values(): + for p in m.parameters(): + p.requires_grad_(False) + m.eval() + + # ── Optimizer ──────────────────────────────────────────────────────────── + train_towers = stage_cfg.get("train_towers", False) + opt_params = ( + ([p for t in towers.values() for p in t.parameters()] if train_towers else []) + + 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 = stage_cfg.get("warmup", {}) + wt = 0 if is_bilateral else warmup_cfg.get("tower_epochs", 0) + wf = 0 if is_bilateral else warmup_cfg.get("fused_epochs", 0) + bcd_prob = cfg["training"].get("bcd_prob", 0.5) + + # ── Epoch loop ──────────────────────────────────────────────────────────── + for epoch in range(epochs): + bridge.train() + for h in head_models.values(): + h.train() + for t in towers.values(): + if train_towers: + t.train() + else: + t.eval() + + if not is_bilateral: + phase = phase_for_epoch(epoch, wt, wf) + if hasattr(bridge, "set_phase"): + bridge.set_phase(phase) + for t in towers.values(): + if hasattr(t, "set_phase"): + t.set_phase(phase) + else: + phase = "fusion" + + total_loss = total_correct = total_n = 0 + + for batch in 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 + + if is_bilateral: + side_embs = { + side: encode_embedding(src, batch, side, towers, stage_models, cfg_stages, device) + for side, src in inputs.items() + } + local_embs = {name: bridge(side_embs)} + else: + 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"]: head_models[hs["name"]](local_embs[hs["input"]]) + for hs in head_stage_cfgs + if hs["input"] in local_embs + } + + 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) + for hs in bcd_head_cfgs if hs["name"] in head_logits] + if not losses: + continue + loss = sum(losses) / len(losses) + opt.zero_grad(); loss.backward(); opt.step() + total_loss += loss.item() * len(y_t) + total_n += len(y_t) + continue + else: + if bcd_head_cfgs and _random() < bcd_prob: + logits = head_logits.get(choice(bcd_head_cfgs)["name"]) + else: + logits = head_logits.get(primary_hs_cfg["name"]) + + if logits is None: + continue + loss = F.cross_entropy(logits, y_t) + opt.zero_grad(); loss.backward(); 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, 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}]" + f" loss={tr_loss:.4f} acc={tr_acc:.3f} val_auc={val_auc:.4f}", + flush=True, + ) + + # ── Final eval ──────────────────────────────────────────────────────────── + y_val, p_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 {} + + 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]) + + 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) + test_acc, test_auc, test_n = (score_arrays(y_te, p_te, num_classes) + if y_te.size else (nan, nan, nan)) + + updated = dict(stage_models) + updated[name] = bridge + updated.update(head_models) + + metrics = { + 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, + } + return updated, metrics diff --git a/v4/classes/stages/helpers.py b/v4/classes/stages/helpers.py new file mode 100644 index 0000000..b93122e --- /dev/null +++ b/v4/classes/stages/helpers.py @@ -0,0 +1,66 @@ +"""stages/helpers — shared utilities for stage runners.""" +from __future__ import annotations + +import torch + + +def get_out_dim(name: str, towers: dict, stage_models: dict) -> int: + if name in towers: + return towers[name].out_dim + if name in stage_models: + return stage_models[name].out_dim + raise KeyError(f"No out_dim for input {name!r}") + + +def resolve_input_dims(inputs, towers: dict, stage_models: dict): + """Return list[int] for list inputs, dict[str, int] for dict inputs.""" + if isinstance(inputs, list): + return [get_out_dim(n, towers, stage_models) for n in inputs] + return {k: get_out_dim(v, towers, stage_models) for k, v in inputs.items()} + + +def encode_embedding( + name: str, + batch: dict, + side: str | None, + towers: dict, + stage_models: dict, + cfg_stages: list[dict], + device, +) -> torch.Tensor: + """Return the embedding for a named tower or prior frozen fusion stage. + + For bilateral batches, *side* selects which side dict entry to use. + Fusion stages are re-encoded recursively from their own inputs. + Runs under torch.no_grad() — only used for frozen passes. + """ + if name in towers: + t = batch[name] + if side is not None and isinstance(t, dict): + t = t[side] + with torch.no_grad(): + return towers[name](t.to(device)) + + s_cfg = next(s for s in cfg_stages if s["name"] == name) + model = stage_models[name] + inputs = s_cfg["inputs"] + + if isinstance(inputs, list): + embeddings = [ + encode_embedding(n, batch, side, towers, stage_models, cfg_stages, device) + for n in inputs + ] + with torch.no_grad(): + return model(embeddings) + + raise NotImplementedError( + f"Stage {name!r} uses dict inputs and cannot be used as a fusion input." + ) + + +def phase_for_epoch(epoch: int, warmup_tower: int, warmup_fused: int) -> str: + if epoch < warmup_tower: + return "tower_warmup" + if epoch < warmup_tower + warmup_fused: + return "fused_warmup" + return "main" diff --git a/v4/classes/stages/warm.py b/v4/classes/stages/warm.py new file mode 100644 index 0000000..2d2003a --- /dev/null +++ b/v4/classes/stages/warm.py @@ -0,0 +1,73 @@ +"""stages/warm — warm stage runner: pre-trains a single tower with a temporary probe.""" +from __future__ import annotations + +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 + + +def run( + stage_cfg: dict, + towers: dict, + data, + split, + label_filter, + cfg: dict, + num_classes: int, + device, + 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"] + + if n_epochs == 0: + return + + s_train = data.build_shells(split.train, level=level, label_filter=label_filter) + bs = cfg["training"]["batch_size"] + loader = _make_loader( + s_train, {tower_name: towers[tower_name]}, + batch_size=bs, shuffle=False, + sampler=_balanced_sampler(s_train), + ) + + for n, t in towers.items(): + for p in t.parameters(): + p.requires_grad_(n == tower_name) + + 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"], + ) + + towers[tower_name].train() + for epoch in range(n_epochs): + total_loss = total_correct = total_n = 0 + for batch in 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 = probe(towers[tower_name](x.to(device))) + loss = F.cross_entropy(logits, y_t) + opt.zero_grad(); loss.backward(); opt.step() + total_loss += loss.item() * len(y_t) + total_correct += int((logits.argmax(1) == y_t).sum()) + total_n += len(y_t) + print( + f" fold{fold+1} [warm/{tower_name}] ep{epoch+1:03d}/{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) diff --git a/v4/classes/towerbase.py b/v4/classes/towerbase.py new file mode 100644 index 0000000..7d1585a --- /dev/null +++ b/v4/classes/towerbase.py @@ -0,0 +1,57 @@ +"""towerbase — v4 TowerBase ABC. + +All v4 towers inherit from TowerBase. The only required interface is: + + out_dim : int property — embedding dimensionality + _side_map : dict property — generic side keys {"a","b"} → dataset-specific ids + _get(*ids) : retrieve and transform one sample by entity_id slots + +get_sample handles the eye-level / patient-level dispatch automatically: + len(entity_id) > 1 — full key; calls _get(*entity_id) directly + len(entity_id) == 1 — patient key; builds {"a": _get(...), "b": _get(...)} + using _side_map to expand the missing slot + +Towers may optionally implement early_pass(context) for cross-tower +communication before loaders are built. +""" +from __future__ import annotations + +from abc import ABC, abstractmethod + +import torch +from torch import nn + +from v4.classes.dataset import ShellEntry + + +class TowerBase(nn.Module, ABC): + + @property + @abstractmethod + def out_dim(self) -> int: ... + + @property + @abstractmethod + def _side_map(self) -> dict[str, str]: ... + + @abstractmethod + def _get(self, *ids) -> torch.Tensor: ... + + def get_sample(self, entry: ShellEntry) -> torch.Tensor | dict[str, torch.Tensor]: + eid = entry.entity_id + if len(eid) > 1: + return self._get(*eid) + return {key: self._get(eid[0], id_1) for key, id_1 in self._side_map.items()} + + def set_phase(self, phase: str) -> None: + """Called by the orchestrator at the start of each training epoch. + + Default: freeze all parameters during fused_warmup, train otherwise. + Override to implement tower-specific phase behaviour. + """ + trainable = phase != "fused_warmup" + for p in self.parameters(): + p.requires_grad_(trainable) + + def early_pass(self, context) -> None: + pass diff --git a/v4/classes/towers/__init__.py b/v4/classes/towers/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/v4/classes/towers/clinical_tower.py b/v4/classes/towers/clinical_tower.py new file mode 100644 index 0000000..eb9782d --- /dev/null +++ b/v4/classes/towers/clinical_tower.py @@ -0,0 +1,94 @@ +"""clinical_tower — ClinicalEncoder for v4. + +Self-contained: no v3 dependencies. +Inherits get_sample dispatch from TowerBase. +""" +from __future__ import annotations + +import numpy as np +import torch +from torch import nn + +from v4.classes.towerbase import TowerBase +from v4.classes.accessory.se_block import SEBlock + + +class ClinicalEncoder(TowerBase): + """MLP over tabular clinical features. + + clinical_data : ClinicalDataView — provides feature_dim, vectorize_entity, side_map + hidden_dim : output embedding dimensionality + dropout : applied after the first linear block + use_se : wrap output with SEBlock channel gating + se_reduction : SEBlock bottleneck factor + se_pre_norm : apply LayerNorm before SEBlock + """ + + 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, + ): + super().__init__() + self.clinical_data = clinical_data + self._out_dim = hidden_dim + feature_dim = clinical_data.feature_dim + + self.block0 = nn.Sequential( + nn.Linear(feature_dim, hidden_dim), + nn.LayerNorm(hidden_dim), + nn.ReLU(inplace=True), + nn.Dropout(dropout), + ) + self.block1 = nn.Sequential( + nn.Linear(hidden_dim, hidden_dim), + nn.ReLU(inplace=True), + ) + self.net = nn.Sequential(self.block0, self.block1) + + self.tower_ln = nn.LayerNorm(hidden_dim) if se_pre_norm else nn.Identity() + self.tower_se = SEBlock(hidden_dim, reduction=se_reduction, residual=True) if use_se else None + + # ── TowerBase interface ────────────────────────────────────────────────── + + @property + def out_dim(self) -> int: + return self._out_dim + + @property + def _side_map(self) -> dict[str, str]: + return self.clinical_data.side_map + + def _get(self, *ids) -> torch.Tensor: + arr = self.clinical_data.vectorize_entity(*ids) + return torch.from_numpy(arr.astype(np.float32, copy=False)) + + # ── nn.Module forward ──────────────────────────────────────────────────── + + def forward(self, x) -> torch.Tensor: + if not isinstance(x, torch.Tensor): + x = torch.as_tensor(x, dtype=torch.float32) + h = self.net(x) + if self.tower_se is not None: + h, _ = self.tower_se(self.tower_ln(h)) + return h + + # ── Utilities ──────────────────────────────────────────────────────────── + + def set_freeze_ratio(self, ratio: float) -> None: + """Freeze the earliest MLP block proportionally.""" + r = max(0.0, min(1.0, float(ratio))) + for p in self.block0.parameters(): + p.requires_grad = True + for p in self.block1.parameters(): + p.requires_grad = True + if r >= 0.5: + for p in self.block0.parameters(): + p.requires_grad = False + if r >= 1.0: + for p in self.block1.parameters(): + p.requires_grad = False diff --git a/v4/classes/towers/image_tower.py b/v4/classes/towers/image_tower.py new file mode 100644 index 0000000..595a165 --- /dev/null +++ b/v4/classes/towers/image_tower.py @@ -0,0 +1,83 @@ +"""image_tower — ImageEncoder for v4. + +Self-contained: no v3 dependencies. +Inherits get_sample dispatch from TowerBase. +""" +from __future__ import annotations + +import math + +import torch +from torch import nn + +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 + + +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 + """ + + 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, + ): + 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.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 + + # ── 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: + img = self.image_data.load_image(*ids) + t = self.transform if self.training else self.eval_transform + return t(img) + + # ── nn.Module forward ──────────────────────────────────────────────────── + + def forward(self, x: torch.Tensor) -> torch.Tensor: + y = self.backbone(x) + if self.tower_se is not None: + y, _ = self.tower_se(self.tower_ln(y)) + return y + + # ── Utilities ──────────────────────────────────────────────────────────── + + def set_freeze_ratio(self, ratio: float) -> None: + """Dynamically freeze the earliest floor(N * ratio) backbone blocks.""" + 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 diff --git a/v4/classes/utils.py b/v4/classes/utils.py new file mode 100644 index 0000000..2b5d5b5 --- /dev/null +++ b/v4/classes/utils.py @@ -0,0 +1,59 @@ +"""utils — general-purpose pipeline utilities.""" +from __future__ import annotations + +import random as pyrandom + +import numpy as np +import pandas as pd +import torch + + +def seed_everything(seed: int) -> None: + pyrandom.seed(seed) + np.random.seed(seed) + torch.manual_seed(seed) + if torch.cuda.is_available(): + torch.cuda.manual_seed_all(seed) + + +def choose_device(device_arg: str | None) -> torch.device: + if device_arg and device_arg != "auto": + return torch.device(device_arg) + return torch.device("cuda" if torch.cuda.is_available() else "cpu") + + +def drop_mixed_label_patients( + df: pd.DataFrame, *, patient_col: str, label_col: str +): + """Remove patients whose rows carry conflicting labels. + + Returns (clean_df, mixed_pids). + """ + per_patient = ( + df.groupby(patient_col)[label_col] + .agg(lambda s: set(pd.to_numeric(s, errors="coerce").dropna().astype(int).tolist())) + ) + mixed = [pid for pid, labels in per_patient.items() if len(labels) > 1] + if not mixed: + return df, [] + return df[~df[patient_col].isin(mixed)].reset_index(drop=True), mixed + + +def relabel_mixed_patients_to_max( + df: pd.DataFrame, *, patient_col: str, label_col: str +): + """Set all rows for each patient to that patient's max observed label. + + Returns (df, changed_rows, still_mixed_pids). + """ + out = df.copy() + labels = pd.to_numeric(out[label_col], errors="coerce") + patient_max = labels.groupby(out[patient_col]).transform("max") + changed_rows = int((labels != patient_max).fillna(False).sum()) + out[label_col] = patient_max.astype(int) + still_mixed = ( + out.groupby(patient_col)[label_col] + .nunique(dropna=True) + .pipe(lambda s: s[s > 1].index.tolist()) + ) + return out.reset_index(drop=True), changed_rows, still_mixed diff --git a/v4/classes/v4_hypertower.py b/v4/classes/v4_hypertower.py new file mode 100644 index 0000000..e36857c --- /dev/null +++ b/v4/classes/v4_hypertower.py @@ -0,0 +1,285 @@ +#!/usr/bin/env python +""" +V4 HyperTower orchestrator — config-driven stage pipeline. + +Stage logic lives in v4/classes/stages/: + warm.py — pre-trains a single tower with a temporary linear probe + fusion.py — trains a bridge + associated head stages, freezes for downstream use + helpers.py — encode_embedding, resolve_input_dims, phase_for_epoch, etc. + +Usage: + python -m v4.classes.v4_hypertower --config v4/configs/ensemble_fused.json +""" +from __future__ import annotations + +import argparse +import importlib +import json +import sys +import time +from pathlib import Path +from typing import Any + +import numpy as np +from torch.utils.data import DataLoader + +REPO_ROOT = Path(__file__).resolve().parents[2] +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 + + +# --------------------------------------------------------------------------- +# Early-pass protocol +# --------------------------------------------------------------------------- + +class EarlyPassContext: + def __init__(self) -> None: + self._store: dict[str, Any] = {} + + def put(self, key: str, value: Any) -> None: + self._store[key] = value + + def get(self, key: str, default: Any = None) -> Any: + return self._store.get(key, default) + + def require(self, key: str) -> Any: + if key not in self._store: + raise KeyError( + f"EarlyPassContext: required key '{key}' not present. " + f"Available: {sorted(self._store.keys())}" + ) + return self._store[key] + + def keys(self) -> set[str]: + return set(self._store.keys()) + + +def _validate_epc_requests(towers_cfg: list[dict], provided_keys: set[str]) -> None: + available = set(provided_keys) + for t in towers_cfg: + for req in t.get("epc_requests", []): + if req not in available: + raise ValueError( + f"Tower '{t['name']}' requests EPC key '{req}' " + f"but no supplier provides it. Available: {sorted(available)}" + ) + available.update(t.get("epc_supplies", [])) + + +# --------------------------------------------------------------------------- +# Data + tower helpers +# --------------------------------------------------------------------------- + +cfg_ref: dict = {} + + +def load_data(cfg: dict): + data_cfg = cfg["data"] + args = dict(data_cfg.get("args", {})) + for key in ("image_dir", "clinical_dir"): + if key in args: + p = Path(args[key]) + if not p.is_absolute(): + args[key] = str(REPO_ROOT / p) + mod = importlib.import_module(data_cfg["module"]) + return mod.build_data(args) + + +def build_towers(towers_cfg: list[dict], data) -> dict: + def _resolve(path: str): + mod_name = cfg_ref.get("data", {}).get("module", "") + if mod_name: + try: + m = importlib.import_module(mod_name) + if hasattr(m, "resolve_data_source"): + return m.resolve_data_source(data, path) + except Exception: + pass + obj = data + for part in path.split("."): + obj = getattr(obj, part) + return obj + + towers = {} + for t in towers_cfg: + mod = importlib.import_module(t["module"]) + cls = getattr(mod, t["class"]) + kwargs = dict(t.get("args", {})) + if "data_source" in t: + towers[t["name"]] = cls(_resolve(t["data_source"]), **kwargs) + elif "data_arg" in t: + kwargs[t["data_arg"]] = data + towers[t["name"]] = cls(**kwargs) + else: + towers[t["name"]] = cls(**kwargs) + return towers + + +def _make_loader(shell: LoaderShell, towers: dict, *, batch_size: int, + shuffle: bool, sampler=None) -> DataLoader: + dataset = HTDataset(shell, towers) + return DataLoader( + dataset, + batch_size=batch_size, + shuffle=(shuffle and sampler is None), + sampler=sampler, + collate_fn=ht_collate, + num_workers=0, + persistent_workers=False, + ) + + +def _balanced_sampler(shell: LoaderShell): + from torch.utils.data import WeightedRandomSampler + labels = [e.label for e in shell.entries] + counts = {} + for l in labels: + counts[l] = counts.get(l, 0) + 1 + weights = [1.0 / counts[l] for l in labels] + return WeightedRandomSampler(weights, num_samples=len(weights), replacement=True) + + +# --------------------------------------------------------------------------- +# Fold runner +# --------------------------------------------------------------------------- + +def run_fold(fold: int, splits, cfg: dict, data, num_classes: int, device) -> dict: + seed_everything(cfg["seed"] + fold * 100) + split = splits[fold] + label_filter = cfg.get("label_filter", None) + cfg_stages = cfg["stages"] + + towers = build_towers(cfg["towers"], data) + for t in towers.values(): + t.to(device) + + context = EarlyPassContext() + context.put("device", device) + context.put("data", data) + context.put("split", split) + context.put("label_filter", label_filter) + _validate_epc_requests(cfg["towers"], context.keys()) + for tower in towers.values(): + if hasattr(tower, "early_pass"): + tower.early_pass(context) + + stage_models: dict = {} + fold_result = {"fold": fold} + + 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) + + elif stype == "fusion": + stage_models, metrics = fusion.run( + stage_cfg, cfg, towers, stage_models, data, split, + label_filter, num_classes, device, fold, cfg_stages, + _make_loader, + ) + fold_result.update(metrics) + + # head stages are handled inside fusion.run + + return fold_result + + +# --------------------------------------------------------------------------- +# Main +# --------------------------------------------------------------------------- + +def main(): + global cfg_ref + + ap = argparse.ArgumentParser(description=__doc__, + formatter_class=argparse.RawDescriptionHelpFormatter) + ap.add_argument("--config", required=True) + ap.add_argument("--device", default=None) + args = ap.parse_args() + + with open(args.config) as f: + cfg = json.load(f) + cfg_ref = cfg + + device = choose_device(args.device or cfg.get("device")) + num_classes = cfg.get("num_classes", 2) + label_filter = cfg.get("label_filter", None) + print(f"Device: {device}", flush=True) + + print("Loading data ...", flush=True) + data = load_data(cfg) + print(f" feature_dim={data.feature_dim}", flush=True) + + label_col = cfg["data"]["args"].get("label_col", "Diagnosis") + df_mode = data.df.copy() + if label_filter is not None: + df_mode = df_mode[df_mode[label_col].isin(label_filter)].reset_index(drop=True) + + identity_level = cfg.get("split_identity_level", 1) + identity_cols = getattr(data, "identity_cols", []) + group_col = identity_cols[identity_level - 1] if identity_level and identity_cols else None + + splits = SplitManager(group_col=group_col).build_plans( + df_mode, + label_col=label_col, + n_splits=cfg.get("folds", 5), + seed=cfg.get("fold_seed", 100), + ) + + out_dir_tags = cfg.get("out_dir_tags", []) + out_dir = REPO_ROOT / cfg.get("output_root", "v4/results") / cfg["run_name"] + for tag in out_dir_tags: + out_dir = out_dir / tag + out_dir.mkdir(parents=True, exist_ok=True) + + eval_stage = cfg.get("eval_stage", "hb") + fold_results = [] + t0 = time.time() + + for fold in range(cfg.get("folds", 5)): + split = splits[fold] + 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) + fold_results.append(result) + print( + f" fold{fold+1} DONE" + f" val_auc={result.get(f'{eval_stage}_val_auc', float('nan')):.4f}" + f" test_auc={result.get(f'{eval_stage}_test_auc', float('nan')):.4f}", + flush=True, + ) + + if fold_results: + val_aucs = [r.get(f"{eval_stage}_val_auc", float("nan")) for r in fold_results] + test_aucs = [r.get(f"{eval_stage}_test_auc", float("nan")) for r in fold_results] + val_aucs = [v for v in val_aucs if not np.isnan(v)] + test_aucs = [v for v in test_aucs if not np.isnan(v)] + summary = { + "run_name": cfg["run_name"], + "eval_stage": eval_stage, + "config": cfg, + "mean_val_auc": float(np.mean(val_aucs)) if val_aucs else float("nan"), + "std_val_auc": float(np.std(val_aucs)) if val_aucs else float("nan"), + "mean_test_auc": float(np.mean(test_aucs)) if test_aucs else float("nan"), + "std_test_auc": float(np.std(test_aucs)) if test_aucs else float("nan"), + "elapsed_s": round(time.time() - t0, 1), + "fold_results": fold_results, + } + summary_path = out_dir / "summary.json" + summary_path.write_text(json.dumps(summary, indent=2)) + print(f"\n{'='*60}", flush=True) + print(f"Val AUC: {summary['mean_val_auc']:.4f} ± {summary['std_val_auc']:.4f}", flush=True) + print(f"Test AUC: {summary['mean_test_auc']:.4f} ± {summary['std_test_auc']:.4f}", flush=True) + print(f"Saved: {summary_path}", flush=True) + + +if __name__ == "__main__": + main() diff --git a/v4/classes/vehicles-deprecated/__init__.py b/v4/classes/vehicles-deprecated/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/v4/classes/vehicles-deprecated/htbase.py b/v4/classes/vehicles-deprecated/htbase.py new file mode 100644 index 0000000..6a92a8b --- /dev/null +++ b/v4/classes/vehicles-deprecated/htbase.py @@ -0,0 +1,40 @@ +"""htbase — HTBase: abstract base for all v4 vehicle classes.""" +from __future__ import annotations + +from abc import ABC, abstractmethod + +import torch +import torch.nn as nn + + +class HTBase(nn.Module, ABC): + """Shared interface for all HyperTower vehicles. + + Subclasses must implement ``encode`` and ``forward``. + + ``transform`` walks ``self.towers`` (if present) and returns the transform + from the first tower that exposes one — used by data loaders. + + ``forward`` contract: returns ``(logits, aux_dict)`` where + ``aux_dict`` maps a name or index to per-component logits. + HTMono returns an empty dict to keep the signature uniform. + """ + + @abstractmethod + def encode(self, inputs) -> torch.Tensor: + """Return the pre-classifier embedding.""" + + @abstractmethod + def forward(self, inputs) -> tuple[torch.Tensor, dict]: + """Return (logits, aux_dict).""" + + @property + def transform(self): + towers = getattr(self, "towers", None) or {} + for t in (towers.values() if hasattr(towers, "values") else []): + if hasattr(t, "transform"): + return t.transform + encoder = getattr(self, "encoder", None) + if encoder is not None: + return getattr(encoder, "transform", None) + return None diff --git a/v4/classes/vehicles-deprecated/htfusion.py b/v4/classes/vehicles-deprecated/htfusion.py new file mode 100644 index 0000000..69f01a6 --- /dev/null +++ b/v4/classes/vehicles-deprecated/htfusion.py @@ -0,0 +1,61 @@ +"""htfusion — HTFusion: N named towers fused through a FusionBridge.""" +from __future__ import annotations + +import torch +import torch.nn as nn + +from v4.classes.bridges.fusion_bridge import FusionBridge +from v4.classes.vehicles.htbase import HTBase + + +class HTFusion(HTBase): + """General N-tower fusion vehicle. + + Each named encoder is registered as a submodule; the FusionBridge + projects and Hadamard-fuses their embeddings. + + Parameters + ---------- + towers : ordered dict ``{name: encoder}``. Each encoder must + expose ``.out_dim``. + num_classes : output classes + fusion_dim : bridge projection dimensionality + dropout : bridge dropout + use_se : SE gate on the fused vector + + Forward contract + ---------------- + ``forward(embeddings)`` takes a ``dict[str, Tensor]`` of pre-computed + per-tower embeddings and returns ``(logits_fused, aux_dict)`` where + ``aux_dict`` maps each tower name to its auxiliary head logits. + """ + + def __init__( + self, + towers: dict[str, nn.Module], + num_classes: int, + fusion_dim: int = 256, + dropout: float = 0.5, + use_se: bool = False, + ): + super().__init__() + self.towers = nn.ModuleDict(towers) + self.bridge = FusionBridge( + tower_dims=[t.out_dim for t in self.towers.values()], + num_classes=num_classes, + fusion_dim=fusion_dim, + dropout=dropout, + use_se=use_se, + ) + + def encode(self, embeddings: dict[str, torch.Tensor]) -> torch.Tensor: + """Return z_fused (pre-classifier) from a dict of per-tower embeddings.""" + return self.bridge.encode([embeddings[name] for name in self.towers]) + + def forward( + self, + embeddings: dict[str, torch.Tensor], + ) -> tuple[torch.Tensor, dict[str, torch.Tensor]]: + ordered = [embeddings[name] for name in self.towers] + logits, aux = self.bridge.fuse(ordered) + return logits, {name: aux[i] for i, name in enumerate(self.towers)} diff --git a/v4/classes/vehicles-deprecated/htlateral.py b/v4/classes/vehicles-deprecated/htlateral.py new file mode 100644 index 0000000..2f7f0ae --- /dev/null +++ b/v4/classes/vehicles-deprecated/htlateral.py @@ -0,0 +1,67 @@ +"""htlateral — HTLateral: shared encoder over N same-type inputs.""" +from __future__ import annotations + +import torch +import torch.nn as nn + +from v4.classes.vehicles.htbase import HTBase + + +class HTLateral(HTBase): + """N same-type inputs through a shared encoder, jointly compressed, then classified. + + All inputs share the same encoder weights (one forward pass per input). + The joint MLP compresses the concatenated embeddings before classification. + + Aux heads provide per-input logits before the joint MLP — useful for + BCD-style training. + + Parameters + ---------- + encoder : shared encoder module with ``.out_dim`` + input_names : ordered slot names (e.g. ``["od", "os"]``) + num_classes : output classes + fusion_dim : joint MLP hidden dim + dropout : dropout in MLP and classifier + """ + + def __init__( + self, + encoder: nn.Module, + input_names: list[str], + num_classes: int, + fusion_dim: int = 256, + dropout: float = 0.5, + ): + super().__init__() + self.encoder = encoder + self.input_names = list(input_names) + n = len(input_names) + in_dim: int = encoder.out_dim # type: ignore[assignment] + + self.joint = nn.Sequential( + nn.Linear(n * in_dim, fusion_dim), nn.LayerNorm(fusion_dim), + nn.ReLU(), nn.Dropout(dropout), nn.Linear(fusion_dim, in_dim), + ) + self.aux_heads = nn.ModuleList([ + nn.Linear(in_dim, num_classes) for _ in range(n) + ]) + self.head = nn.Sequential( + nn.ReLU(), nn.Dropout(dropout), nn.Linear(in_dim, num_classes), + ) + + def encode(self, inputs: dict[str, torch.Tensor]) -> torch.Tensor: + """Return joint embedding (post-MLP, pre-classifier).""" + zs = [self.encoder(inputs[name]) for name in self.input_names] + return self.joint(torch.cat(zs, dim=1)) + + def forward( + self, + inputs: dict[str, torch.Tensor], + ) -> tuple[torch.Tensor, dict[str, torch.Tensor]]: + zs = [self.encoder(inputs[name]) for name in self.input_names] + z_joint = self.joint(torch.cat(zs, dim=1)) + logits = self.head(z_joint) + aux = {name: head(z) + for name, head, z in zip(self.input_names, self.aux_heads, zs)} + return logits, aux diff --git a/v4/classes/vehicles-deprecated/htmono.py b/v4/classes/vehicles-deprecated/htmono.py new file mode 100644 index 0000000..7161c75 --- /dev/null +++ b/v4/classes/vehicles-deprecated/htmono.py @@ -0,0 +1,44 @@ +"""htmono — HTMono: single tower + ClassificationHead, no bridge.""" +from __future__ import annotations + +import torch +import torch.nn as nn + +from v4.classes.heads.classifier import ClassificationHead +from v4.classes.vehicles.htbase import HTBase + + +class HTMono(HTBase): + """Single-tower vehicle: tower embedding fed directly into a ClassificationHead. + + No bridge or projection — the tower's output goes straight to + ReLU → Dropout → Linear. Returns ``(logits, {})`` from ``forward`` + to match the HTFusion / HTLateral interface. + + Parameters + ---------- + tower : encoder module with ``.out_dim`` + num_classes : output classes + dropout : dropout before the output linear layer + """ + + def __init__( + self, + tower: nn.Module, + num_classes: int, + dropout: float = 0.5, + ): + super().__init__() + self.tower = tower + self.head = ClassificationHead(tower.out_dim, num_classes, dropout) # type: ignore[arg-type] + + def encode(self, inputs) -> torch.Tensor: + """Return tower embedding (pre-classifier).""" + if isinstance(inputs, dict): + # single-entry dict from HTDataset eye-level pass + (z,) = inputs.values() + return self.tower(z) if torch.is_tensor(z) else self.tower(*z.values()) + return self.tower(inputs) + + def forward(self, inputs) -> tuple[torch.Tensor, dict]: + return self.head(self.encode(inputs)), {} diff --git a/v4/configs/ensemble_fused.json b/v4/configs/ensemble_fused.json new file mode 100644 index 0000000..ccbf5a0 --- /dev/null +++ b/v4/configs/ensemble_fused.json @@ -0,0 +1,121 @@ +{ + "_notes": [ + "V4 stage-pipeline: warm → fusion stages with parallel head stages.", + "Bridges are pure embedding producers; heads are separate swappable stages.", + "BCD-eligible heads are sampled during tower_warmup and main phases.", + "eval_stage names the fusion stage whose primary head is used for final metrics." + ], + + "run_name": "v4/ensemble_fused", + "num_classes": 2, + "label_filter": [0, 1], + "split_identity_level": 1, + "eval_stage": "hb", + "save_predictions": false, + "seed": 1234, + "folds": 5, + "fold_seed": 100, + "output_root": "v4/results", + "out_dir_tags": ["binary", "ntower"], + + "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", + "args": { + "backbone": "refugelike", + "freeze_ratio": 0.0, + "augment": true + } + }, + { + "name": "cd", + "module": "v4.classes.towers.clinical_tower", + "class": "ClinicalEncoder", + "data_source": "matrix", + "args": { + "hidden_dim": 128 + } + } + ], + + "stages": [ + { + "name": "cd_warm", + "type": "warm", + "tower": "cd", + "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" + } + ], + + "training": { + "lr": 1e-4, + "batch_size": 16, + "bcd_prob": 0.5, + "tune_binary_threshold": true + } +} diff --git a/v4/results/v4/ensemble_fused/binary/ntower/gold_standard.json b/v4/results/v4/ensemble_fused/binary/ntower/gold_standard.json new file mode 100644 index 0000000..cf58369 --- /dev/null +++ b/v4/results/v4/ensemble_fused/binary/ntower/gold_standard.json @@ -0,0 +1,146 @@ +{ + "run_name": "v4/ensemble_fused", + "config": { + "_notes": [ + "V4 ensemble_fused: img + cd towers, HTFusion Stage 1, HyperBridge Stage 2.", + "Matches phase5/embedding_mlp_head setup for direct comparison.", + "data_source resolves against the PapilaBundle returned by build_data.", + "image_dir / clinical_dir are project-relative; orchestrator resolves against REPO_ROOT." + ], + "run_name": "v4/ensemble_fused", + "eval_mode": "binary", + "split_identity_level": 1, + "epochs": 30, + "fusion_epochs": 10, + "folds": 5, + "fold_seed": 100, + "seed": 1234, + "output_root": "v4/results", + "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", + "args": { + "backbone": "refugelike", + "freeze_ratio": 0.0, + "augment": true + }, + "warmup_epochs": 0 + }, + { + "name": "cd", + "module": "v4.classes.towers.clinical_tower", + "class": "ClinicalEncoder", + "data_source": "matrix", + "args": { + "hidden_dim": 128 + }, + "warmup_epochs": 40, + "warmup_exclude_towers": [ + "img" + ] + } + ], + "bridge": { + "mode": "embedding_mlp", + "fusion_dim": 256, + "hidden_dim": 256 + }, + "training": { + "lr": 0.0001, + "batch_size": 16, + "bcd_prob": 0.5, + "warmup_tower_epochs": 3, + "warmup_fused_epochs": 3, + "tune_binary_threshold": true + } + }, + "mean_val_auc": 0.9036764705882353, + "std_val_auc": 0.02318218054475656, + "mean_test_auc": 0.9007352941176471, + "std_test_auc": 0.04178903599524356, + "elapsed_s": 1563.1, + "fold_results": [ + { + "fold": 0, + "val_auc": 0.9044117647058824, + "val_acc": 0.8571428571428571, + "val_n": 42, + "val_kappa": 0.42727272727272725, + "val_mcc": 0.4622975667767223, + "val_f1": 0.7083333333333333, + "val_threshold": 0.1659889668226242, + "test_auc": 0.8566176470588236, + "test_acc": 0.9285714285714286, + "test_n": 42 + }, + { + "fold": 1, + "val_auc": 0.8970588235294118, + "val_acc": 0.8809523809523809, + "val_n": 42, + "val_kappa": 0.5945945945945946, + "val_mcc": 0.5965587590013045, + "val_f1": 0.7971014492753623, + "val_threshold": 0.017083797603845596, + "test_auc": 0.9044117647058824, + "test_acc": 0.8095238095238095, + "test_n": 42 + }, + { + "fold": 2, + "val_auc": 0.863970588235294, + "val_acc": 0.8333333333333334, + "val_n": 42, + "val_kappa": 0.43243243243243246, + "val_mcc": 0.4338609156373123, + "val_f1": 0.7159420289855072, + "val_threshold": 0.09843172132968903, + "test_auc": 0.9044117647058824, + "test_acc": 0.8571428571428571, + "test_n": 42 + }, + { + "fold": 3, + "val_auc": 0.9227941176470588, + "val_acc": 0.8809523809523809, + "val_n": 42, + "val_kappa": 0.631578947368421, + "val_mcc": 0.6333004963811236, + "val_f1": 0.8156277436347674, + "val_threshold": 0.01531070377677679, + "test_auc": 0.8639705882352942, + "test_acc": 0.8571428571428571, + "test_n": 42 + }, + { + "fold": 4, + "val_auc": 0.9301470588235294, + "val_acc": 0.9047619047619048, + "val_n": 42, + "val_kappa": 0.6181818181818182, + "val_mcc": 0.6688560540599386, + "val_f1": 0.8055555555555556, + "val_threshold": 0.19587548077106476, + "test_auc": 0.9742647058823529, + "test_acc": 0.9285714285714286, + "test_n": 42 + } + ] +} \ No newline at end of file diff --git a/v4/results/v4/ensemble_fused/binary/ntower/summary.json b/v4/results/v4/ensemble_fused/binary/ntower/summary.json new file mode 100644 index 0000000..7533a70 --- /dev/null +++ b/v4/results/v4/ensemble_fused/binary/ntower/summary.json @@ -0,0 +1,264 @@ +{ + "run_name": "v4/ensemble_fused", + "eval_stage": "hb", + "config": { + "_notes": [ + "V4 stage-pipeline: warm \u2192 fusion stages with parallel head stages.", + "Bridges are pure embedding producers; heads are separate swappable stages.", + "BCD-eligible heads are sampled during tower_warmup and main phases.", + "eval_stage names the fusion stage whose primary head is used for final metrics." + ], + "run_name": "v4/ensemble_fused", + "num_classes": 2, + "label_filter": [ + 0, + 1 + ], + "split_identity_level": 1, + "eval_stage": "hb", + "save_predictions": false, + "seed": 1234, + "folds": 5, + "fold_seed": 100, + "output_root": "v4/results", + "out_dir_tags": [ + "binary", + "ntower" + ], + "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", + "args": { + "backbone": "refugelike", + "freeze_ratio": 0.0, + "augment": true + } + }, + { + "name": "cd", + "module": "v4.classes.towers.clinical_tower", + "class": "ClinicalEncoder", + "data_source": "matrix", + "args": { + "hidden_dim": 128 + } + } + ], + "stages": [ + { + "name": "cd_warm", + "type": "warm", + "tower": "cd", + "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" + } + ], + "training": { + "lr": 0.0001, + "batch_size": 16, + "bcd_prob": 0.5, + "tune_binary_threshold": true + } + }, + "mean_val_auc": 0.8926470588235293, + "std_val_auc": 0.06105155516072669, + "mean_test_auc": 0.9022058823529413, + "std_test_auc": 0.035001853633564395, + "elapsed_s": 1525.9, + "fold_results": [ + { + "fold": 0, + "nt_val_auc": 0.858647936786655, + "nt_val_acc": 0.8928571428571429, + "nt_val_n": 84, + "nt_val_kappa": 0.6272189349112426, + "nt_val_mcc": 0.6411186083279721, + "nt_val_f1": 0.8124534854874721, + "nt_val_threshold": 0.012987074442207813, + "nt_test_auc": 0.9157155399473222, + "nt_test_acc": 0.9047619047619048, + "nt_test_n": 84, + "hb_val_auc": 0.9080882352941178, + "hb_val_acc": 0.8333333333333334, + "hb_val_n": 42, + "hb_val_kappa": 0.2898550724637682, + "hb_val_mcc": 0.33633639699815626, + "hb_val_f1": 0.6338729763387297, + "hb_val_threshold": 0.011891158297657967, + "hb_test_auc": 0.9522058823529411, + "hb_test_acc": 0.9523809523809523, + "hb_test_n": 42 + }, + { + "fold": 1, + "nt_val_auc": 0.8282828282828283, + "nt_val_acc": 0.8928571428571429, + "nt_val_n": 84, + "nt_val_kappa": 0.6111111111111112, + "nt_val_mcc": 0.6633249580710799, + "nt_val_f1": 0.801418439716312, + "nt_val_threshold": 0.008082838729023933, + "nt_test_auc": 0.8726953467954346, + "nt_test_acc": 0.8690476190476191, + "nt_test_n": 84, + "hb_val_auc": 0.7904411764705882, + "hb_val_acc": 0.9047619047619048, + "hb_val_n": 42, + "hb_val_kappa": 0.6181818181818182, + "hb_val_mcc": 0.6688560540599386, + "hb_val_f1": 0.8055555555555556, + "hb_val_threshold": 0.013159404508769512, + "hb_test_auc": 0.8970588235294117, + "hb_test_acc": 0.8571428571428571, + "hb_test_n": 42 + }, + { + "fold": 2, + "nt_val_auc": 0.8906882591093117, + "nt_val_acc": 0.8452380952380952, + "nt_val_n": 84, + "nt_val_kappa": 0.5125, + "nt_val_mcc": 0.5217535056401378, + "nt_val_f1": 0.7548821548821549, + "nt_val_threshold": 0.052708517760038376, + "nt_test_auc": 0.856060606060606, + "nt_test_acc": 0.8333333333333334, + "nt_test_n": 84, + "hb_val_auc": 0.9117647058823529, + "hb_val_acc": 0.8333333333333334, + "hb_val_n": 42, + "hb_val_kappa": 0.36909871244635195, + "hb_val_mcc": 0.3833788364965519, + "hb_val_f1": 0.6814734561213435, + "hb_val_threshold": 0.0401872955262661, + "hb_test_auc": 0.8566176470588236, + "hb_test_acc": 0.8571428571428571, + "hb_test_n": 42 + }, + { + "fold": 3, + "nt_val_auc": 0.921875, + "nt_val_acc": 0.8928571428571429, + "nt_val_n": 84, + "nt_val_kappa": 0.6758147512864494, + "nt_val_mcc": 0.6797955088067001, + "nt_val_f1": 0.837593984962406, + "nt_val_threshold": 0.708580732345581, + "nt_test_auc": 0.8631578947368421, + "nt_test_acc": 0.8452380952380952, + "nt_test_n": 84, + "hb_val_auc": 0.9779411764705882, + "hb_val_acc": 0.9285714285714286, + "hb_val_n": 42, + "hb_val_kappa": 0.7567567567567568, + "hb_val_mcc": 0.7592566023652966, + "hb_val_f1": 0.8782608695652174, + "hb_val_threshold": 0.04653067886829376, + "hb_test_auc": 0.875, + "hb_test_acc": 0.8809523809523809, + "hb_test_n": 42 + }, + { + "fold": 4, + "nt_val_auc": 0.8279192273924496, + "nt_val_acc": 0.8571428571428571, + "nt_val_n": 84, + "nt_val_kappa": 0.46325878594249204, + "nt_val_mcc": 0.49610717544581684, + "nt_val_f1": 0.7269772481040087, + "nt_val_threshold": 0.06419441103935242, + "nt_test_auc": 0.8961397058823529, + "nt_test_acc": 0.8928571428571429, + "nt_test_n": 84, + "hb_val_auc": 0.875, + "hb_val_acc": 0.9047619047619048, + "hb_val_n": 42, + "hb_val_kappa": 0.6181818181818182, + "hb_val_mcc": 0.6688560540599386, + "hb_val_f1": 0.8055555555555556, + "hb_val_threshold": 0.23509082198143005, + "hb_test_auc": 0.9301470588235294, + "hb_test_acc": 0.8809523809523809, + "hb_test_n": 42 + } + ] +} \ No newline at end of file diff --git a/v4/scripts/debug_splits.py b/v4/scripts/debug_splits.py new file mode 100644 index 0000000..0f7e4bb --- /dev/null +++ b/v4/scripts/debug_splits.py @@ -0,0 +1,65 @@ +"""Compare fold assignments between v3 PatientFirstSplitManager and v4 SplitManager.""" +import sys +from pathlib import Path +from types import SimpleNamespace + +sys.path.insert(0, str(Path(__file__).resolve().parents[2])) + +from v3.classes.split_manager import PatientFirstSplitManager +from v4.classes.split_manager import SplitManager +from v4.classes.profiles.v4papila import build_data + +args = { + "image_dir": "Papila/FundusImages", + "clinical_dir": "Papila/ClinicalData", + "label_col": "Diagnosis", + "iop_corr_method": "ratio", + "iop_drop_raw": True, + "exclude_cols": ["Axial_Length"], +} +# Resolve relative paths +root = Path(__file__).resolve().parents[2] +args["image_dir"] = str(root / args["image_dir"]) +args["clinical_dir"] = str(root / args["clinical_dir"]) + +data = build_data(args) + +label_col = data.label_col +patient_col = data.patient_col +df_mode = data.df[data.df[label_col].isin([0, 1])].reset_index(drop=True) + +# ── v3 splits ──────────────────────────────────────────────────────────────── +split_mgr_v3 = PatientFirstSplitManager(patient_col=patient_col, label_col=label_col) +split_args_v3 = SimpleNamespace(eval_mode="binary", n_splits=5, fold_seed=100) +clinical_ns = SimpleNamespace(df=df_mode, label_col=label_col) +splits_v3 = split_mgr_v3.build_plans(clinical=clinical_ns, args=split_args_v3, profile=None) + +# ── v4 splits ──────────────────────────────────────────────────────────────── +splits_v4 = SplitManager(group_col=patient_col).build_plans( + df_mode, label_col=label_col, n_splits=5, seed=100, +) + +# ── Compare ────────────────────────────────────────────────────────────────── +print(f"{'Fold':<6} {'Set':<6} {'v3 patients':<8} {'v4 patients':<8} {'Match'}") +print("-" * 50) + +all_match = True +for fold in range(5): + s3, s4 = splits_v3[fold], splits_v4[fold] + for label, df3, df4 in [ + ("train", s3.train, s4.train), + ("val", s3.val, s4.val), + ("test", s3.test, s4.test), + ]: + ids3 = set(df3[patient_col].unique()) if df3 is not None else set() + ids4 = set(df4[patient_col].unique()) if df4 is not None else set() + match = ids3 == ids4 + if not match: + all_match = False + print(f"{fold+1:<6} {label:<6} {len(ids3):<8} {len(ids4):<8} {'✓' if match else '✗ DIFF'}") + if not match: + print(f" only in v3: {sorted(ids3 - ids4)[:10]}") + print(f" only in v4: {sorted(ids4 - ids3)[:10]}") + +print() +print("All folds match!" if all_match else "SPLITS DIFFER — fold assignments changed.")