ryan/hypertower
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

v4 update

Browse Source
This commit is contained in:
rpotter6298
2026-04-20 18:01:31 +02:00
parent 13290575d5
commit 4dea45df78
71 changed files with 8316 additions and 4112 deletions
Download Patch File Download Diff File Expand all files Collapse all files
TODO.md
+23
View File
@@ -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.
v3/PLAN_v4_hypertower.md
+179
View File
@@ -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
v3/classes/__init__.py
+61 -31
View File
@@ -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",
]
v3/classes/bridges.py
+254 -44
View File
@@ -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)
v3/classes/clinical_towers.py
+162
View File
@@ -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
v3/classes/config_builder.py
-276
View File
@@ -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
v3/classes/dataset.py
-115
View File
@@ -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])
v3/classes/filters.py
-119
View File
@@ -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
v3/classes/seg_cnn.py → v3/classes/geometry_towers.py
+349 -92
View File
@@ -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
v3/classes/hypertower_models.py
+1187
View File
File diff suppressed because it is too large Load Diff
v3/classes/image_towers.py
+252
View File
@@ -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,
)
v3/classes/loader_factory.py
+2 -2
View File
@@ -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(
v3/classes/model_builder.py
-148
View File
@@ -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())
v3/classes/models.py
-1074
View File
File diff suppressed because it is too large Load Diff
v3/classes/towerbase.py
+455
View File
@@ -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)
v3/classes/towers.py
-279
View File
@@ -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
v3/classes/v2_hypertower.py
-1853
View File
File diff suppressed because it is too large Load Diff
v3/classes/v3_hypertower.py
+326 -4
View File
@@ -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")
v3/distributed/server.py
+107 -38
View File
@@ -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")
v3/scripts/development/analyze_emergent_features.py
+256
View File
@@ -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()
v3/scripts/development/stage_seg_cnn/run_seg_cnn.py
+1 -1
View File
@@ -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,
)
v3/scripts/main/phase3/epoch_grid.json
+1 -1
View File
@@ -7,7 +7,7 @@
"common_args": [
"--eval-mode", "binary",
"--tower-mode", "single",
"--hypertower-mode", "single",
"--in-memory-cache",
"--augment",
"--tune-binary-threshold",
v3/scripts/main/phase3/experiment_grid.json
+1 -1
View File
@@ -12,7 +12,7 @@
"common_args": [
"--eval-mode",
"binary",
"--tower-mode",
"--hypertower-mode",
"single",
"--epochs",
"30",
v3/scripts/main/phase3/phase35_grid.json
+1 -1
View File
@@ -8,7 +8,7 @@
"common_args": [
"--eval-mode", "binary",
"--tower-mode", "single",
"--hypertower-mode", "single",
"--epochs", "30",
"--in-memory-cache",
"--augment",
v3/scripts/main/phase4/experiment_grid.json
+6 -6
View File
@@ -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"]
}
]
}
v3/scripts/main/phase5/dispatch_logit_mlp_ckpt.py
+1 -1
View File
@@ -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",
v3/scripts/main/phase5/experiment_grid.json
+7 -7
View File
@@ -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"]
}
]
}
v3/scripts/main/phase6/experiment_grid.json
+13 -13
View File
@@ -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"]
}
]
}
v3/scripts/main/run_cv.py
+2 -2
View File
@@ -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 \
v3/scripts/main/run_ntower_cv.py
+533
View File
@@ -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()
v3/scripts/output_analysis/explainability/gradcam_phase5.py
+1 -1
View File
@@ -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]
v3/scripts/output_analysis/explainability/permutation_importance_phase5.py
+3 -3
View File
@@ -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))
v4/__init__.py
View File
v4/classes/__init__.py
View File
v4/classes/accessory/__init__.py
View File
v4/classes/accessory/backbones.py
+217
View File
@@ -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
v4/classes/accessory/se_block.py
+126
View File
@@ -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
v4/classes/accessory/transforms.py
+66
View File
@@ -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)
v4/classes/bridges/__init__.py
View File
v4/classes/bridges/fusion_bridge.py
+57
View File
@@ -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)
v4/classes/bridges/hyperbridge.py
+56
View File
@@ -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
v4/classes/dataset.py
+248
View File
@@ -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] = {"<UNK>": 0}
for i, v in enumerate(cats, start=1):
mapping[v] = i
self.cat_maps[col] = mapping
def _compute_feature_dim(self) -> None:
self.feature_dim = (
len(self.scalar_cols)
+ sum(len(m) for m in self.cat_maps.values())
+ len(self.scalar_cols)
)
# ---------------------------------------------------------------------------
# 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
v4/classes/heads/__init__.py
View File
v4/classes/heads/classifier.py
+24
View File
@@ -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)
v4/classes/loaders/__init__.py
View File
v4/classes/loaders/image_loader.py
+132
View File
@@ -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)
v4/classes/logging/__init__.py
View File
v4/classes/logging/prediction_store.py
+394
View File
@@ -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
v4/classes/metrics.py
+175
View File
@@ -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
v4/classes/profiles/__init__.py
View File
v4/classes/profiles/v4papila.py
+663
View File
@@ -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,
)
v4/classes/split_manager.py
+177
View File
@@ -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
]
v4/classes/stages/__init__.py
View File
v4/classes/stages/fusion.py
+273
View File
@@ -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
v4/classes/stages/helpers.py
+66
View File
@@ -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"
v4/classes/stages/warm.py
+73
View File
@@ -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)
v4/classes/towerbase.py
+57
View File
@@ -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
v4/classes/towers/__init__.py
View File
v4/classes/towers/clinical_tower.py
+94
View File
@@ -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
v4/classes/towers/image_tower.py
+83
View File
@@ -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
v4/classes/utils.py
+59
View File
@@ -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
v4/classes/v4_hypertower.py
+285
View File
@@ -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()
v4/classes/vehicles-deprecated/__init__.py
View File
v4/classes/vehicles-deprecated/htbase.py
+40
View File
@@ -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
v4/classes/vehicles-deprecated/htfusion.py
+61
View File
@@ -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)}
v4/classes/vehicles-deprecated/htlateral.py
+67
View File
@@ -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
v4/classes/vehicles-deprecated/htmono.py
+44
View File
@@ -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)), {}
v4/configs/ensemble_fused.json
+121
View File
@@ -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
}
}
v4/results/v4/ensemble_fused/binary/ntower/gold_standard.json
+146
View File
@@ -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
}
]
}
v4/results/v4/ensemble_fused/binary/ntower/summary.json
+264
View File
@@ -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
}
]
}
v4/scripts/debug_splits.py
+65
View File
@@ -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.")