v4 update
This commit is contained in:
rpotter6298
@@ -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
|
||||
@@ -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
|
||||
@@ -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)
|
||||
@@ -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)
|
||||
@@ -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
|
||||
@@ -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
|
||||
@@ -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)
|
||||
@@ -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)
|
||||
@@ -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
|
||||
@@ -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
|
||||
@@ -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,
|
||||
)
|
||||
@@ -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
|
||||
]
|
||||
@@ -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
|
||||
@@ -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"
|
||||
@@ -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)
|
||||
@@ -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
|
||||
@@ -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
|
||||
@@ -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
|
||||
@@ -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
|
||||
@@ -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()
|
||||
@@ -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
|
||||
@@ -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)}
|
||||
@@ -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
|
||||
@@ -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)), {}
|
||||
@@ -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
|
||||
}
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
]
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
]
|
||||
}
|
||||
@@ -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.")
|
||||
Reference in New Issue
Block a user