pre-refactor 041426
This commit is contained in:
rpotter6298
+19
-16
@@ -14,6 +14,7 @@ class Bridge(nn.Module):
|
||||
num_classes,
|
||||
fusion_dim=256,
|
||||
mode="fused",
|
||||
dropout: float = 0.5,
|
||||
use_se: bool = True,
|
||||
se_reduction: int = 16,
|
||||
se_pre_norm: bool = True,
|
||||
@@ -37,7 +38,7 @@ class Bridge(nn.Module):
|
||||
# heads
|
||||
self.classifier_fused = nn.Sequential(
|
||||
nn.ReLU(),
|
||||
nn.Dropout(0.5),
|
||||
nn.Dropout(dropout),
|
||||
nn.Linear(fusion_dim, num_classes),
|
||||
)
|
||||
self.classifier_img = nn.Linear(img_dim, num_classes)
|
||||
@@ -54,26 +55,28 @@ class Bridge(nn.Module):
|
||||
return self.se_log.get(reset=reset)
|
||||
return None
|
||||
|
||||
def _compute_fused(self, img_feats, md_feats):
|
||||
"""Return z_fused embedding (before classifier_fused). Used by encode() and forward()."""
|
||||
hi = self.ln_img(self.W_img(img_feats))
|
||||
hm = self.ln_md(self.W_md(md_feats))
|
||||
fused = hi * hm
|
||||
if self.se is not None:
|
||||
fused, gates = self.se(fused)
|
||||
if self.se_log.enabled:
|
||||
self.se_log.accumulate(gates)
|
||||
return fused
|
||||
|
||||
def encode(self, img_feats, md_feats) -> torch.Tensor:
|
||||
"""Return z_fused embedding without applying the classifier head."""
|
||||
assert self.mode == "fused", "encode() only valid in fused mode"
|
||||
return self._compute_fused(img_feats, md_feats)
|
||||
|
||||
def forward(self, img_feats, md_feats):
|
||||
out_img = None if self.mode == "clinical_only" else self.classifier_img(img_feats)
|
||||
out_md = None if self.mode == "image_only" else self.classifier_cd(md_feats)
|
||||
|
||||
if self.mode == "fused":
|
||||
hi = self.ln_img(self.W_img(img_feats)) # image features
|
||||
hm = self.ln_md(self.W_md(md_feats)) # clinical data features
|
||||
fused = hi * hm # elementwise product
|
||||
# apply SE gates
|
||||
if self.se is not None:
|
||||
fused, gates = self.se(fused)
|
||||
if self.se_log.enabled:
|
||||
self.se_log.accumulate(gates)
|
||||
|
||||
if self.se is not None and self.training and self.se_log.enabled:
|
||||
if not hasattr(self, "_dbg_seen"):
|
||||
self._dbg_seen = 0
|
||||
if self._dbg_seen < 3: # print only a few times
|
||||
print("[SE] gate mean this batch:", gates.mean().item())
|
||||
self._dbg_seen += 1
|
||||
fused = self._compute_fused(img_feats, md_feats)
|
||||
out_f = self.classifier_fused(fused)
|
||||
return out_f, out_img, out_md
|
||||
# if ablation modes:
|
||||
|
||||
@@ -189,6 +189,31 @@ class UNetImageCropper:
|
||||
return None
|
||||
return np.asarray(features, dtype=np.float32)
|
||||
|
||||
def precompute_geometry(self, image_paths) -> None:
|
||||
"""Pre-compute geometry features for all image_paths into an in-memory cache.
|
||||
Safe to call in the main process; geometry_for_image() can then be called
|
||||
from DataLoader workers without touching CUDA.
|
||||
"""
|
||||
self._geometry_cache: Dict[str, Optional[np.ndarray]] = {}
|
||||
paths = list(image_paths)
|
||||
print(f"[UNetImageCropper] pre-computing geometry for {len(paths)} images...", flush=True)
|
||||
for img_path in paths:
|
||||
key = str(Path(img_path).resolve())
|
||||
try:
|
||||
img = Image.open(img_path).convert("RGB")
|
||||
self._geometry_cache[key] = self.geometry_features(img, img_path)
|
||||
except Exception:
|
||||
self._geometry_cache[key] = None
|
||||
n_ok = sum(1 for v in self._geometry_cache.values() if v is not None)
|
||||
print(f"[UNetImageCropper] {n_ok}/{len(paths)} geometry vectors computed", flush=True)
|
||||
|
||||
def geometry_for_image(self, image_path) -> Optional[np.ndarray]:
|
||||
"""Return pre-computed geometry vector for image_path (call precompute_geometry first)."""
|
||||
cache = getattr(self, "_geometry_cache", None)
|
||||
if cache is None:
|
||||
raise RuntimeError("Call precompute_geometry() before geometry_for_image()")
|
||||
return cache.get(str(Path(image_path).resolve()))
|
||||
|
||||
|
||||
class ManifestImageCropper:
|
||||
def __init__(
|
||||
@@ -371,6 +396,35 @@ class ManifestImageCropper:
|
||||
return None
|
||||
return np.asarray(features, dtype=np.float32)
|
||||
|
||||
def precompute_geometry(self, image_paths) -> None:
|
||||
"""Pre-compute geometry features for all image_paths into an in-memory cache.
|
||||
Safe to call in the main process; geometry_for_image() can then be called
|
||||
without re-opening images or re-loading annotations.
|
||||
"""
|
||||
self._geometry_cache: Dict[str, Optional[np.ndarray]] = {}
|
||||
paths = list(image_paths)
|
||||
print(f"[ManifestImageCropper] pre-computing geometry for {len(paths)} images...", flush=True)
|
||||
for img_path in paths:
|
||||
key = str(Path(img_path).resolve())
|
||||
entry = self.entries.get(key)
|
||||
if entry is None:
|
||||
self._geometry_cache[key] = None
|
||||
continue
|
||||
try:
|
||||
img = Image.open(img_path).convert("RGB")
|
||||
self._geometry_cache[key] = self.geometry_features(img, img_path)
|
||||
except Exception:
|
||||
self._geometry_cache[key] = None
|
||||
n_ok = sum(1 for v in self._geometry_cache.values() if v is not None)
|
||||
print(f"[ManifestImageCropper] {n_ok}/{len(paths)} geometry vectors computed", flush=True)
|
||||
|
||||
def geometry_for_image(self, image_path) -> Optional[np.ndarray]:
|
||||
"""Return pre-computed geometry vector for image_path (call precompute_geometry first)."""
|
||||
cache = getattr(self, "_geometry_cache", None)
|
||||
if cache is None:
|
||||
raise RuntimeError("Call precompute_geometry() before geometry_for_image()")
|
||||
return cache.get(str(Path(image_path).resolve()))
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Factory
|
||||
|
||||
@@ -211,6 +211,7 @@ def make_loader(
|
||||
shuffle: bool,
|
||||
num_workers: int,
|
||||
sampler: Optional[WeightedRandomSampler] = None,
|
||||
persistent_workers: bool = False,
|
||||
) -> DataLoader:
|
||||
ds = SlotDataset(
|
||||
samples,
|
||||
@@ -226,6 +227,7 @@ def make_loader(
|
||||
sampler=sampler,
|
||||
num_workers=num_workers,
|
||||
collate_fn=slot_collate,
|
||||
persistent_workers=(persistent_workers and num_workers > 0),
|
||||
)
|
||||
|
||||
|
||||
|
||||
+222
-6
@@ -11,7 +11,7 @@ from torch import nn
|
||||
from torch.utils.data import DataLoader
|
||||
|
||||
from v3.classes.bridges import Bridge
|
||||
from v3.classes.towers import ImageTower, ClinicalTower
|
||||
from v3.classes.towers import ImageTower, ClinicalTower, SiameseImageTower
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
@@ -35,18 +35,24 @@ class SingleEyeHT(nn.Module):
|
||||
cd_hidden_dim: int = 128,
|
||||
fusion_dim: int = 256,
|
||||
bridge_mode: str = "fused",
|
||||
bridge_dropout: float = 0.5,
|
||||
cd_dropout: float = 0.1,
|
||||
se_img_tower: bool = False,
|
||||
se_cd_tower: bool = False,
|
||||
se_bridge: bool = False,
|
||||
):
|
||||
super().__init__()
|
||||
self.img_tower = ImageTower(
|
||||
backbone=backbone,
|
||||
freeze_ratio=freeze_ratio,
|
||||
augment=augment,
|
||||
use_se=False,
|
||||
use_se=se_img_tower,
|
||||
)
|
||||
self.cd_tower = ClinicalTower(
|
||||
clinical_data=clinical_data,
|
||||
hidden_dim=cd_hidden_dim,
|
||||
use_se=False,
|
||||
dropout=cd_dropout,
|
||||
use_se=se_cd_tower,
|
||||
)
|
||||
self.bridge = Bridge(
|
||||
img_dim=self.img_tower.out_dim,
|
||||
@@ -54,13 +60,20 @@ class SingleEyeHT(nn.Module):
|
||||
num_classes=num_classes,
|
||||
fusion_dim=fusion_dim,
|
||||
mode=bridge_mode,
|
||||
use_se=False,
|
||||
dropout=bridge_dropout,
|
||||
use_se=se_bridge,
|
||||
)
|
||||
|
||||
@property
|
||||
def transform(self):
|
||||
return self.img_tower.transform
|
||||
|
||||
def encode(self, x: torch.Tensor, meta: torch.Tensor) -> torch.Tensor:
|
||||
"""Return z_fused embedding (fusion_dim) without applying the classifier head."""
|
||||
img_feats = self.img_tower(x)
|
||||
md_feats = self.cd_tower(meta)
|
||||
return self.bridge.encode(img_feats, md_feats)
|
||||
|
||||
def forward(self, x: torch.Tensor, meta: torch.Tensor) -> torch.Tensor:
|
||||
img_feats = None if self.bridge.mode == "clinical_only" else self.img_tower(x)
|
||||
md_feats = None if self.bridge.mode == "image_only" else self.cd_tower(meta)
|
||||
@@ -158,6 +171,59 @@ class BilateralHT(nn.Module):
|
||||
return out_f
|
||||
|
||||
|
||||
class SiameseHT(nn.Module):
|
||||
"""
|
||||
Bilateral model using a shared-weight SiameseImageTower (mean+delta).
|
||||
|
||||
Both eyes pass through the same backbone; features are combined as
|
||||
cat([mean(f_od, f_os), f_od - f_os]) giving the model both a shared
|
||||
bilateral representation and an asymmetry signal.
|
||||
|
||||
Uses the same bilateral loader and training loop as BilateralHT.
|
||||
The bridge operates in image_only mode (no clinical data in phase 4).
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
backbone: str,
|
||||
freeze_ratio: float,
|
||||
augment: bool,
|
||||
num_classes: int,
|
||||
fusion_dim: int = 256,
|
||||
bridge_mode: str = "image_only",
|
||||
):
|
||||
super().__init__()
|
||||
self.img_tower = SiameseImageTower(
|
||||
backbone=backbone,
|
||||
freeze_ratio=freeze_ratio,
|
||||
augment=augment,
|
||||
use_se=False,
|
||||
)
|
||||
img_dim = self.img_tower.out_dim # 2 * backbone_out_dim
|
||||
self.bridge = Bridge(
|
||||
img_dim=img_dim,
|
||||
meta_dim=1, # dummy — not used in image_only mode
|
||||
num_classes=num_classes,
|
||||
fusion_dim=fusion_dim,
|
||||
mode="image_only",
|
||||
use_se=False,
|
||||
)
|
||||
self.aux_img = nn.Linear(img_dim, num_classes)
|
||||
|
||||
@property
|
||||
def transform(self):
|
||||
return self.img_tower.transform
|
||||
|
||||
def encode(self, x_od: torch.Tensor, x_os: torch.Tensor) -> torch.Tensor:
|
||||
return self.img_tower(x_od, x_os)
|
||||
|
||||
def forward(self, x_od: torch.Tensor, x_os: torch.Tensor) -> torch.Tensor:
|
||||
feats = self.encode(x_od, x_os)
|
||||
out_f, _, _ = self.bridge(feats, None)
|
||||
return out_f
|
||||
|
||||
|
||||
class FusedEnsembleHT(nn.Module):
|
||||
"""
|
||||
SingleEyeHT base with a per-eye attention scorer for bilateral fusion.
|
||||
@@ -190,6 +256,10 @@ class FusedEnsembleHT(nn.Module):
|
||||
# Learns the GC-direction in logit space from bilateral labels.
|
||||
self.eye_scorer = nn.Linear(num_classes, 1, bias=True)
|
||||
|
||||
@property
|
||||
def head(self) -> nn.Module:
|
||||
return self.eye_scorer
|
||||
|
||||
def forward(
|
||||
self,
|
||||
x_od: torch.Tensor,
|
||||
@@ -205,6 +275,78 @@ class FusedEnsembleHT(nn.Module):
|
||||
return alpha[:, 0:1] * logit_od + alpha[:, 1:2] * logit_os # [B, C]
|
||||
|
||||
|
||||
class LogitMLPEnsembleHT(nn.Module):
|
||||
"""
|
||||
MLP head trained on concatenated per-eye logits.
|
||||
|
||||
Both eyes pass through the frozen base independently, producing per-eye
|
||||
logit vectors. These are concatenated and fed through a small MLP:
|
||||
|
||||
cat([logit_od, logit_os]) [B, 2C]
|
||||
→ Linear(2C, hidden) → ReLU → Dropout → Linear(hidden, C)
|
||||
|
||||
Permutation-variant by design: the model can learn left/right asymmetries
|
||||
directly from the concatenated pair, at the cost of needing a consistent
|
||||
OD-first input ordering.
|
||||
"""
|
||||
|
||||
def __init__(self, base: SingleEyeHT, num_classes: int, hidden: int = 64):
|
||||
super().__init__()
|
||||
self.base = base
|
||||
self.head = nn.Sequential(
|
||||
nn.Linear(2 * num_classes, hidden),
|
||||
nn.ReLU(),
|
||||
nn.Dropout(0.3),
|
||||
nn.Linear(hidden, num_classes),
|
||||
)
|
||||
|
||||
def forward(
|
||||
self,
|
||||
x_od: torch.Tensor, meta_od: torch.Tensor,
|
||||
x_os: torch.Tensor, meta_os: torch.Tensor,
|
||||
) -> torch.Tensor:
|
||||
logit_od = self.base(x_od, meta_od)
|
||||
logit_os = self.base(x_os, meta_os)
|
||||
return self.head(torch.cat([logit_od, logit_os], dim=1))
|
||||
|
||||
|
||||
class EmbeddingMLPEnsembleHT(nn.Module):
|
||||
"""
|
||||
MLP head trained on concatenated per-eye z_fused embeddings.
|
||||
|
||||
Both eyes pass through the frozen base independently, and their bridge
|
||||
embeddings (pre-classifier, shape [B, fusion_dim]) are concatenated and
|
||||
fed through an MLP:
|
||||
|
||||
cat([z_od, z_os]) [B, 2 * fusion_dim]
|
||||
→ Linear(2*fusion_dim, hidden) → ReLU → Dropout → Linear(hidden, C)
|
||||
|
||||
Richer than logit-level: the head sees pre-softmax feature vectors rather
|
||||
than the compressed C-dimensional output, giving it more signal to work
|
||||
with when fusion_dim >> C.
|
||||
"""
|
||||
|
||||
def __init__(self, base: SingleEyeHT, num_classes: int, hidden: int = 256):
|
||||
super().__init__()
|
||||
self.base = base
|
||||
fusion_dim = base.bridge.W_img.out_features
|
||||
self.head = nn.Sequential(
|
||||
nn.Linear(2 * fusion_dim, hidden),
|
||||
nn.ReLU(),
|
||||
nn.Dropout(0.3),
|
||||
nn.Linear(hidden, num_classes),
|
||||
)
|
||||
|
||||
def forward(
|
||||
self,
|
||||
x_od: torch.Tensor, meta_od: torch.Tensor,
|
||||
x_os: torch.Tensor, meta_os: torch.Tensor,
|
||||
) -> torch.Tensor:
|
||||
z_od = self.base.encode(x_od, meta_od)
|
||||
z_os = self.base.encode(x_os, meta_os)
|
||||
return self.head(torch.cat([z_od, z_os], dim=1))
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Phase control
|
||||
# ---------------------------------------------------------------------------
|
||||
@@ -429,15 +571,89 @@ def train_bilateral_epoch(
|
||||
)
|
||||
|
||||
|
||||
def train_siamese_epoch(
|
||||
model: SiameseHT,
|
||||
loader: DataLoader,
|
||||
opt,
|
||||
device: torch.device,
|
||||
*,
|
||||
bcd_prob: float = 0.5,
|
||||
tower_loss_mode: str = "bcd",
|
||||
) -> tuple[float, float]:
|
||||
"""Train one epoch of SiameseHT on bilateral (patient-level) samples."""
|
||||
model.train()
|
||||
total_loss = total_correct = total_n = 0
|
||||
for batch in loader:
|
||||
x1 = batch.get("image_1")
|
||||
x2 = batch.get("image_2")
|
||||
y = batch.get("label_1")
|
||||
if not (torch.is_tensor(x1) and torch.is_tensor(x2)):
|
||||
continue
|
||||
x1 = x1.to(device); x2 = x2.to(device)
|
||||
y = _to_label_tensor(y, device)
|
||||
feats = model.encode(x1, x2)
|
||||
|
||||
if tower_loss_mode == "all":
|
||||
logits_i = model.aux_img(feats)
|
||||
out_f, _, _ = model.bridge(feats, None)
|
||||
loss = F.cross_entropy(out_f, y) + F.cross_entropy(logits_i, y)
|
||||
logits = out_f
|
||||
elif random() < bcd_prob:
|
||||
logits = model.aux_img(feats)
|
||||
loss = F.cross_entropy(logits, y)
|
||||
else:
|
||||
logits, _, _ = model.bridge(feats, None)
|
||||
loss = F.cross_entropy(logits, y)
|
||||
|
||||
opt.zero_grad()
|
||||
loss.backward()
|
||||
opt.step()
|
||||
bs = y.shape[0]
|
||||
total_loss += float(loss.item()) * bs
|
||||
total_correct += int((logits.argmax(1) == y).sum())
|
||||
total_n += bs
|
||||
return (
|
||||
total_loss / total_n if total_n else float("nan"),
|
||||
total_correct / total_n if total_n else float("nan"),
|
||||
)
|
||||
|
||||
|
||||
def collect_probs_siamese(
|
||||
model: SiameseHT,
|
||||
loader: DataLoader,
|
||||
device: torch.device,
|
||||
) -> tuple[np.ndarray, np.ndarray]:
|
||||
"""Patient-level probs from a bilateral loader using SiameseHT."""
|
||||
model.eval()
|
||||
y_c, p_c = [], []
|
||||
with torch.no_grad():
|
||||
for batch in loader:
|
||||
x1 = batch.get("image_1")
|
||||
x2 = batch.get("image_2")
|
||||
y = batch.get("label_1")
|
||||
if not (torch.is_tensor(x1) and torch.is_tensor(x2)):
|
||||
continue
|
||||
x1 = x1.to(device); x2 = x2.to(device)
|
||||
feats = model.encode(x1, x2)
|
||||
logits, _, _ = model.bridge(feats, None)
|
||||
probs = torch.softmax(logits, dim=1)
|
||||
y_c.append(np.array(y) if not torch.is_tensor(y) else y.cpu().numpy())
|
||||
p_c.append(probs.cpu().numpy())
|
||||
return (
|
||||
np.concatenate(y_c, axis=0),
|
||||
np.concatenate(p_c, axis=0),
|
||||
)
|
||||
|
||||
|
||||
def train_fusion_epoch(
|
||||
model: FusedEnsembleHT,
|
||||
model, # FusedEnsembleHT | LogitMLPEnsembleHT | EmbeddingMLPEnsembleHT
|
||||
loader: DataLoader,
|
||||
opt,
|
||||
device: torch.device,
|
||||
) -> tuple[float, float]:
|
||||
"""Train only the fusion head; the base SingleEyeHT is frozen in eval mode."""
|
||||
model.base.eval()
|
||||
model.eye_scorer.train()
|
||||
model.head.train()
|
||||
total_loss = total_correct = total_n = 0
|
||||
for batch in loader:
|
||||
x1 = batch.get("image_1"); m1 = batch.get("matrix_1")
|
||||
|
||||
@@ -31,6 +31,8 @@ def head_names_for_mode(tower_mode: str, *, fused_head: bool = False) -> list[st
|
||||
return names + ["bilat_fused"] if fused_head else names
|
||||
if tower_mode == "bilateral":
|
||||
return ["fused", "img_joint", "md_joint"]
|
||||
if tower_mode == "siamese":
|
||||
return ["fused"]
|
||||
raise ValueError(f"Unknown tower_mode: {tower_mode!r}")
|
||||
|
||||
|
||||
|
||||
@@ -0,0 +1,521 @@
|
||||
"""Segmentation-map CNN for glaucoma grading.
|
||||
|
||||
Trains a CNN on combined disc/cup segmentation maps — pixel values
|
||||
0 = background, 1 = disc (rim only), 2 = cup
|
||||
— instead of raw RGB fundus images, forcing the model to learn purely
|
||||
from optic nerve head geometry (CDR, rim width, cup location, etc.).
|
||||
|
||||
Two segmentation sources are supported:
|
||||
gt – rasterise expert contour/mask annotations directly (pure NumPy/PIL,
|
||||
no CUDA — safe in DataLoader worker processes)
|
||||
unet – run a trained UNetSegmenter on the raw fundus image
|
||||
|
||||
Usage (import from training script):
|
||||
from v3.classes.seg_cnn import SegMapRecord, SegMapDataset, SegCNN, seg_map_to_tensor
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
from typing import List, Optional, Tuple
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
from PIL import Image, ImageDraw
|
||||
from PIL.Image import Resampling
|
||||
from torch.utils.data import Dataset
|
||||
from torchvision import models, transforms
|
||||
from tqdm import tqdm
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Data record
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
@dataclass
|
||||
class SegMapRecord:
|
||||
"""One labelled eye sample for the seg-map CNN."""
|
||||
sample_id: str
|
||||
image_path: Path # original fundus image (used by unet mode)
|
||||
annotation_disc: Path # contour (.txt) or mask (.bmp/.png)
|
||||
annotation_cup: Path
|
||||
annotation_type_disc: str # "contour" or "mask"
|
||||
annotation_type_cup: str
|
||||
patient_id: int # for group-CV: keep both eyes of a patient together
|
||||
eye: str # "OD" or "OS"
|
||||
label: int # 0 = Normal, 1 = Glaucoma
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Seg-map utilities
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def _combine_masks(disc_mask: np.ndarray, cup_mask: np.ndarray) -> np.ndarray:
|
||||
"""
|
||||
Combine binary disc and cup masks into a 3-class label map.
|
||||
|
||||
Returns a uint8 array with values:
|
||||
0 — background
|
||||
1 — optic disc rim (disc but not cup)
|
||||
2 — optic cup
|
||||
"""
|
||||
disc = (disc_mask > 0).astype(np.uint8)
|
||||
cup = (cup_mask > 0).astype(np.uint8)
|
||||
cup = (cup & disc) # structural prior: cup must be inside disc
|
||||
seg = disc + cup # 0, 1 (rim), or 2 (cup)
|
||||
return seg.astype(np.uint8)
|
||||
|
||||
|
||||
def crop_to_disc(seg_map: np.ndarray) -> np.ndarray:
|
||||
"""
|
||||
Crop a seg map tightly to the disc bounding box.
|
||||
|
||||
The disc is anywhere seg_map > 0 (i.e. rim or cup).
|
||||
Returns the original array unchanged if no disc is found.
|
||||
"""
|
||||
rows = np.any(seg_map > 0, axis=1)
|
||||
cols = np.any(seg_map > 0, axis=0)
|
||||
if not rows.any():
|
||||
return seg_map
|
||||
r0, r1 = int(np.argmax(rows)), int(len(rows) - 1 - np.argmax(rows[::-1]))
|
||||
c0, c1 = int(np.argmax(cols)), int(len(cols) - 1 - np.argmax(cols[::-1]))
|
||||
return seg_map[r0:r1 + 1, c0:c1 + 1]
|
||||
|
||||
|
||||
def seg_map_to_tensor(
|
||||
seg_map: np.ndarray,
|
||||
channels: int,
|
||||
target_size: int,
|
||||
) -> torch.Tensor:
|
||||
"""
|
||||
Convert an (H, W) seg map with values {0, 1, 2} to a float tensor.
|
||||
|
||||
channels=1 → (1, H, W) float in [0, 1] (values 0/0.5/1.0)
|
||||
channels=3 → (3, H, W) one-hot binary channels [bg, disc_rim, cup]
|
||||
"""
|
||||
pil = Image.fromarray(seg_map.astype(np.uint8), mode="L")
|
||||
pil = pil.resize((target_size, target_size), Image.NEAREST)
|
||||
seg = np.array(pil, dtype=np.uint8)
|
||||
|
||||
if channels == 1:
|
||||
arr = seg.astype(np.float32) / 2.0 # {0, 0.5, 1.0}
|
||||
return torch.from_numpy(arr).unsqueeze(0)
|
||||
|
||||
if channels == 3:
|
||||
bg = (seg == 0).astype(np.float32)
|
||||
disc_rim = (seg == 1).astype(np.float32)
|
||||
cup = (seg == 2).astype(np.float32)
|
||||
return torch.from_numpy(np.stack([bg, disc_rim, cup], axis=0))
|
||||
|
||||
raise ValueError(f"channels must be 1 or 3, got {channels}")
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# GT mask loading (pure NumPy / PIL — no CUDA, safe in DataLoader workers)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def _load_contour(path: Path) -> np.ndarray:
|
||||
"""Load x,y contour pairs from a whitespace- or comma-delimited text file."""
|
||||
for delimiter in (",", None):
|
||||
try:
|
||||
arr = np.loadtxt(str(path), delimiter=delimiter, comments="#", dtype=np.float32)
|
||||
if arr.size > 0:
|
||||
break
|
||||
except Exception:
|
||||
arr = np.zeros((0, 2), dtype=np.float32)
|
||||
if arr.size == 0 or arr.ndim == 1:
|
||||
return np.zeros((0, 2), dtype=np.float32)
|
||||
if arr.shape[1] < 2:
|
||||
return np.zeros((0, 2), dtype=np.float32)
|
||||
return arr[:, :2]
|
||||
|
||||
|
||||
def _contour_to_mask(
|
||||
coords: np.ndarray, image_size: Tuple[int, int], target_size: int
|
||||
) -> np.ndarray:
|
||||
"""
|
||||
Rasterise a polygon defined by (x, y) coords into a binary mask.
|
||||
|
||||
image_size is the (width, height) of the original fundus image — the
|
||||
coordinate space the contour was annotated in. The mask is drawn at
|
||||
that resolution then resized to target_size, matching UNetSegmenter's
|
||||
behaviour and avoiding off-canvas clipping.
|
||||
"""
|
||||
if coords is None or len(coords) < 3:
|
||||
return np.zeros((target_size, target_size), dtype=np.uint8)
|
||||
points = [tuple(map(float, pt)) for pt in coords]
|
||||
img = Image.new("L", image_size, 0)
|
||||
ImageDraw.Draw(img).polygon(points, outline=1, fill=1)
|
||||
img = img.resize((target_size, target_size), Resampling.NEAREST)
|
||||
return (np.array(img, dtype=np.uint8) > 0).astype(np.uint8)
|
||||
|
||||
|
||||
def _extract_masks_from_image(
|
||||
mask_path: Path, target_size: int
|
||||
) -> Tuple[np.ndarray, np.ndarray]:
|
||||
"""
|
||||
Extract disc and cup binary masks from a segmentation image file.
|
||||
|
||||
Handles both grayscale label images (e.g. REFUGE .bmp) and
|
||||
RGB colour-coded masks. Returns (disc_mask, cup_mask) both at
|
||||
target_size × target_size.
|
||||
"""
|
||||
raw = Image.open(mask_path)
|
||||
arr = np.array(raw)
|
||||
|
||||
if arr.ndim == 2:
|
||||
# Grayscale: identify background from edge statistics
|
||||
edges = np.concatenate([arr[0], arr[-1], arr[:, 0], arr[:, -1]])
|
||||
bg_val = int(np.argmax(np.bincount(edges.astype(np.int64).clip(0, 255), minlength=256)))
|
||||
disc_arr = (arr != bg_val).astype(np.uint8)
|
||||
vals = np.unique(arr)
|
||||
non_bg = vals[vals != bg_val]
|
||||
cup_arr: np.ndarray
|
||||
if non_bg.size > 1:
|
||||
cup_val = int(non_bg.min())
|
||||
cup_arr = (arr == cup_val).astype(np.uint8)
|
||||
else:
|
||||
cup_arr = np.zeros_like(disc_arr, dtype=np.uint8)
|
||||
else:
|
||||
img_rgb = raw.convert("RGB")
|
||||
arr = np.array(img_rgb)
|
||||
h, w, c = arr.shape
|
||||
edges_rgb = np.concatenate(
|
||||
[arr[0], arr[-1], arr[:, 0], arr[:, -1]], axis=0
|
||||
)
|
||||
edge_colors, edge_counts = np.unique(edges_rgb.reshape(-1, c), axis=0, return_counts=True)
|
||||
bg_color = edge_colors[int(np.argmax(edge_counts))]
|
||||
colors, counts = np.unique(arr.reshape(-1, c), axis=0, return_counts=True)
|
||||
not_bg = np.any(colors != bg_color.reshape(1, -1), axis=1)
|
||||
colors, counts = colors[not_bg], counts[not_bg]
|
||||
disc_arr = np.zeros((h, w), dtype=np.uint8)
|
||||
cup_arr = np.zeros((h, w), dtype=np.uint8)
|
||||
if colors.shape[0] >= 1:
|
||||
order = np.argsort(-counts)
|
||||
disc_color = colors[order[0]]
|
||||
disc_arr[np.all(arr == disc_color, axis=-1)] = 1
|
||||
if colors.shape[0] >= 2:
|
||||
cup_color = colors[order[1]]
|
||||
cup_arr[np.all(arr == cup_color, axis=-1)] = 1
|
||||
|
||||
# Resize to target_size with nearest-neighbour to preserve binary values
|
||||
def _resize(m: np.ndarray) -> np.ndarray:
|
||||
pil = Image.fromarray((m > 0).astype(np.uint8) * 255)
|
||||
pil = pil.resize((target_size, target_size), Resampling.NEAREST)
|
||||
return (np.array(pil) > 0).astype(np.uint8)
|
||||
|
||||
return _resize(disc_arr), _resize(cup_arr)
|
||||
|
||||
|
||||
def load_gt_masks(rec: "SegMapRecord", target_size: int) -> Tuple[np.ndarray, np.ndarray]:
|
||||
"""
|
||||
Load GT disc + cup masks for one record.
|
||||
|
||||
Handles annotation_type "contour" (x,y text file) and "mask" (image file).
|
||||
Returns (disc_mask, cup_mask) as uint8 arrays of shape (target_size, target_size).
|
||||
"""
|
||||
disc_mask: Optional[np.ndarray] = None
|
||||
cup_mask: Optional[np.ndarray] = None
|
||||
|
||||
# Get original image size so contour coordinates are drawn in the right space
|
||||
with Image.open(rec.image_path) as _img:
|
||||
image_size = _img.size # (width, height)
|
||||
|
||||
# ---- Disc ----
|
||||
if rec.annotation_type_disc == "mask":
|
||||
disc_mask, cup_from_disc = _extract_masks_from_image(rec.annotation_disc, target_size)
|
||||
if cup_from_disc.any():
|
||||
cup_mask = cup_from_disc
|
||||
else: # contour
|
||||
coords = _load_contour(rec.annotation_disc)
|
||||
disc_mask = _contour_to_mask(coords, image_size, target_size)
|
||||
|
||||
# ---- Cup ----
|
||||
if cup_mask is None:
|
||||
if rec.annotation_type_cup == "mask":
|
||||
_, cup_from_cup = _extract_masks_from_image(rec.annotation_cup, target_size)
|
||||
cup_mask = cup_from_cup
|
||||
else: # contour
|
||||
coords = _load_contour(rec.annotation_cup)
|
||||
cup_mask = _contour_to_mask(coords, image_size, target_size)
|
||||
|
||||
if disc_mask is None:
|
||||
disc_mask = np.zeros((target_size, target_size), dtype=np.uint8)
|
||||
if cup_mask is None:
|
||||
cup_mask = np.zeros((target_size, target_size), dtype=np.uint8)
|
||||
|
||||
# Structural prior: cup must lie within disc
|
||||
cup_mask = (cup_mask > 0) & (disc_mask > 0)
|
||||
return disc_mask.astype(np.uint8), cup_mask.astype(np.uint8)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# U-Net fine-tuning dataset
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
class UNetFineTuneDataset(Dataset):
|
||||
"""
|
||||
Loads (image_tensor, mask_tensor) pairs for fine-tuning the U-Net on
|
||||
PAPILA GT annotations. Uses the same preprocessing as UNetSegmenter
|
||||
so the fine-tuned weights are compatible with inference.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
records: List[SegMapRecord],
|
||||
target_size: int = 512,
|
||||
normalize: str = "per_image",
|
||||
) -> None:
|
||||
self.records = records
|
||||
self.target_size = target_size
|
||||
self.normalize = normalize
|
||||
self.to_tensor = transforms.ToTensor()
|
||||
|
||||
def __len__(self) -> int:
|
||||
return len(self.records)
|
||||
|
||||
def _normalize(self, tensor: torch.Tensor) -> torch.Tensor:
|
||||
if self.normalize == "per_image":
|
||||
mean = tensor.mean(dim=(1, 2), keepdim=True)
|
||||
std = tensor.std(dim=(1, 2), keepdim=True).clamp(min=1e-6)
|
||||
return (tensor - mean) / std
|
||||
if self.normalize == "imagenet":
|
||||
mean = torch.tensor([0.485, 0.456, 0.406]).view(-1, 1, 1)
|
||||
std = torch.tensor([0.229, 0.224, 0.225]).view(-1, 1, 1)
|
||||
return (tensor - mean) / std
|
||||
return tensor
|
||||
|
||||
def __getitem__(self, idx: int):
|
||||
rec = self.records[idx]
|
||||
image = Image.open(rec.image_path).convert("RGB")
|
||||
image = image.resize((self.target_size, self.target_size), Resampling.BILINEAR)
|
||||
img_tensor = self._normalize(self.to_tensor(image))
|
||||
|
||||
disc_mask, cup_mask = load_gt_masks(rec, self.target_size)
|
||||
mask_tensor = torch.from_numpy(
|
||||
np.stack([disc_mask, cup_mask], axis=0).astype(np.float32)
|
||||
)
|
||||
return img_tensor, mask_tensor
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# U-Net precomputation (run once per full record list, not per fold)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def precompute_unet_seg_maps(
|
||||
records: List["SegMapRecord"],
|
||||
segmenter,
|
||||
threshold: float = 0.5,
|
||||
) -> List[np.ndarray]:
|
||||
"""
|
||||
Run the U-Net on every record and return a list of combined seg maps.
|
||||
|
||||
Call this once before the CV loop and pass the results to each fold's
|
||||
SegMapDataset via precomputed_seg_maps, so the U-Net isn't re-run per fold.
|
||||
"""
|
||||
to_tensor = transforms.ToTensor()
|
||||
seg_maps = []
|
||||
for rec in tqdm(records, desc="U-Net inference", unit="img", leave=False):
|
||||
image = Image.open(rec.image_path).convert("RGB")
|
||||
resized = segmenter.preprocess_image(image)
|
||||
tensor = segmenter._normalize_tensor(
|
||||
to_tensor(resized).to(segmenter.device)
|
||||
).unsqueeze(0)
|
||||
with torch.no_grad():
|
||||
logits = segmenter.model(tensor)
|
||||
probs = torch.sigmoid(logits)[0].cpu().numpy()
|
||||
disc = (probs[0] > threshold).astype(np.uint8)
|
||||
cup = (probs[1] > threshold).astype(np.uint8)
|
||||
cup = (cup & disc)
|
||||
seg_maps.append(_combine_masks(disc, cup.astype(np.uint8)))
|
||||
return seg_maps
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Dataset
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
class SegMapDataset(Dataset):
|
||||
"""
|
||||
PyTorch Dataset that yields (seg_tensor, label) pairs.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
records : list of SegMapRecord
|
||||
target_size : CNN input spatial size (images are resized to this)
|
||||
channels : 1 = single-channel label map; 3 = one-hot three channels
|
||||
augment : apply random flips + rotation (for training set)
|
||||
unet_segmenter : if provided, use U-Net predictions instead of GT masks;
|
||||
must be a loaded UNetSegmenter with model weights set
|
||||
unet_threshold : threshold for U-Net logit → binary mask
|
||||
seg_target_size: resolution at which GT masks are rasterised (or U-Net
|
||||
output size). Default 512 matches UNetSegmenter default.
|
||||
crop_to_disc : crop the seg map tightly to the disc bounding box before
|
||||
resizing to target_size (default True — eliminates the
|
||||
background zeros that make up most of the full image)
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
records: List[SegMapRecord],
|
||||
target_size: int = 224,
|
||||
channels: int = 3,
|
||||
augment: bool = False,
|
||||
unet_segmenter=None,
|
||||
unet_threshold: float = 0.5,
|
||||
seg_target_size: int = 512,
|
||||
crop_to_disc: bool = True,
|
||||
precomputed_seg_maps: Optional[List[np.ndarray]] = None,
|
||||
) -> None:
|
||||
self.records = records
|
||||
self.target_size = target_size
|
||||
self.channels = channels
|
||||
self.augment = augment
|
||||
self.seg_target_size = seg_target_size
|
||||
self.crop_to_disc = crop_to_disc
|
||||
|
||||
if precomputed_seg_maps is not None:
|
||||
self._seg_maps = precomputed_seg_maps
|
||||
elif unet_segmenter is not None:
|
||||
self._seg_maps = precompute_unet_seg_maps(
|
||||
records, unet_segmenter, unet_threshold
|
||||
)
|
||||
else:
|
||||
self._seg_maps = None
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
def __len__(self) -> int:
|
||||
return len(self.records)
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
def _augment(self, seg_map: np.ndarray) -> np.ndarray:
|
||||
"""Random flips + 90° rotations (label-safe since NEAREST resize)."""
|
||||
if np.random.rand() < 0.5:
|
||||
seg_map = np.fliplr(seg_map)
|
||||
if np.random.rand() < 0.5:
|
||||
seg_map = np.flipud(seg_map)
|
||||
k = np.random.randint(0, 4)
|
||||
if k:
|
||||
seg_map = np.rot90(seg_map, k=k)
|
||||
return np.ascontiguousarray(seg_map)
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
def __getitem__(self, idx: int):
|
||||
rec = self.records[idx]
|
||||
|
||||
if self._seg_maps is not None:
|
||||
seg_map = self._seg_maps[idx]
|
||||
else:
|
||||
disc_mask, cup_mask = load_gt_masks(rec, self.seg_target_size)
|
||||
seg_map = _combine_masks(disc_mask, cup_mask)
|
||||
|
||||
if self.crop_to_disc:
|
||||
seg_map = crop_to_disc(seg_map)
|
||||
|
||||
if self.augment:
|
||||
seg_map = self._augment(seg_map)
|
||||
|
||||
tensor = seg_map_to_tensor(seg_map, self.channels, self.target_size)
|
||||
return tensor, rec.label
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Model
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
class SegCNN(nn.Module):
|
||||
"""
|
||||
Pretrained CNN backbone adapted for segmentation-map input.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
num_classes : output classes (2 for binary glaucoma grading)
|
||||
backbone : "resnet18" | "resnet50" | "efficientnet_b0"
|
||||
pretrained : initialise with ImageNet weights (recommended even for
|
||||
non-RGB input — transfer generalises across domains)
|
||||
in_channels : 1 (single label map) or 3 (one-hot channels)
|
||||
dropout : dropout rate before the final classifier head
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
num_classes: int = 2,
|
||||
backbone: str = "resnet18",
|
||||
pretrained: bool = True,
|
||||
in_channels: int = 3,
|
||||
dropout: float = 0.3,
|
||||
) -> None:
|
||||
super().__init__()
|
||||
|
||||
weights_arg = "DEFAULT" if pretrained else None
|
||||
|
||||
if backbone == "resnet18":
|
||||
base = models.resnet18(weights=weights_arg)
|
||||
feat_dim = base.fc.in_features
|
||||
base.fc = nn.Identity()
|
||||
elif backbone == "resnet50":
|
||||
base = models.resnet50(weights=weights_arg)
|
||||
feat_dim = base.fc.in_features
|
||||
base.fc = nn.Identity()
|
||||
elif backbone == "efficientnet_b0":
|
||||
base = models.efficientnet_b0(weights=weights_arg)
|
||||
feat_dim = base.classifier[1].in_features
|
||||
base.classifier = nn.Identity()
|
||||
else:
|
||||
raise ValueError(f"Unknown backbone: {backbone!r}")
|
||||
|
||||
# Adapt first conv layer if in_channels ≠ 3
|
||||
if in_channels != 3:
|
||||
first_conv = self._find_first_conv(base)
|
||||
new_conv = nn.Conv2d(
|
||||
in_channels,
|
||||
first_conv.out_channels,
|
||||
kernel_size=first_conv.kernel_size,
|
||||
stride=first_conv.stride,
|
||||
padding=first_conv.padding,
|
||||
bias=first_conv.bias is not None,
|
||||
)
|
||||
if pretrained:
|
||||
# Average pretrained RGB weights across channel dim
|
||||
with torch.no_grad():
|
||||
new_conv.weight.copy_(
|
||||
first_conv.weight.mean(dim=1, keepdim=True).expand_as(new_conv.weight)
|
||||
)
|
||||
self._replace_first_conv(base, new_conv)
|
||||
|
||||
self.backbone = base
|
||||
self.head = nn.Sequential(
|
||||
nn.Dropout(p=dropout),
|
||||
nn.Linear(feat_dim, num_classes),
|
||||
)
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
@staticmethod
|
||||
def _find_first_conv(module: nn.Module) -> nn.Conv2d:
|
||||
for m in module.modules():
|
||||
if isinstance(m, nn.Conv2d):
|
||||
return m
|
||||
raise RuntimeError("No Conv2d found in backbone")
|
||||
|
||||
@staticmethod
|
||||
def _replace_first_conv(module: nn.Module, new_conv: nn.Conv2d) -> None:
|
||||
"""Replace the first Conv2d in-place (handles resnet and efficientnet)."""
|
||||
for name, child in module.named_children():
|
||||
if isinstance(child, nn.Conv2d):
|
||||
setattr(module, name, new_conv)
|
||||
return
|
||||
try:
|
||||
SegCNN._replace_first_conv(child, new_conv)
|
||||
return
|
||||
except RuntimeError:
|
||||
pass
|
||||
raise RuntimeError("Could not replace first Conv2d")
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||||
feats = self.backbone(x)
|
||||
if feats.dim() > 2:
|
||||
feats = feats.flatten(1)
|
||||
return self.head(feats)
|
||||
@@ -83,6 +83,65 @@ def build_patient_split_plans(
|
||||
return plans
|
||||
|
||||
|
||||
class EyeLevelSplitManager:
|
||||
"""
|
||||
Eye-level (leaky) splitter — splits on individual eye rows, ignoring
|
||||
patient grouping. Same patient's eyes can appear in different folds.
|
||||
Used to demonstrate the effect of data leakage.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
patient_col: str = "Patient ID",
|
||||
label_col: Optional[str] = None,
|
||||
) -> None:
|
||||
self.patient_col = patient_col
|
||||
self.label_col = label_col
|
||||
|
||||
def build_plans(
|
||||
self,
|
||||
*,
|
||||
clinical: Any,
|
||||
args: Any,
|
||||
profile: Optional[Any] = None,
|
||||
) -> list[PatientSplit]:
|
||||
profile_label_col = getattr(profile, "label_col", None) if profile is not None else None
|
||||
label_col = self.label_col or profile_label_col or getattr(clinical, "label_col", None)
|
||||
if label_col is None:
|
||||
raise ValueError("Could not resolve label column")
|
||||
|
||||
if not hasattr(clinical, "df"):
|
||||
raise ValueError("Clinical object must expose a dataframe at .df")
|
||||
df_full = clinical.df.copy().reset_index(drop=True)
|
||||
|
||||
eval_mode = str(getattr(args, "eval_mode", "multiclass")).lower()
|
||||
if eval_mode == "binary":
|
||||
df_full = df_full[df_full[label_col].isin([0, 1])].reset_index(drop=True)
|
||||
|
||||
n_splits = int(getattr(args, "n_splits", 5))
|
||||
fold_seed = int(getattr(args, "fold_seed", 42))
|
||||
|
||||
labels = df_full[label_col].to_numpy()
|
||||
eye_ids = df_full.index.to_numpy()
|
||||
|
||||
# Reuse build_patient_split_plans with eye-row IDs as the "patients"
|
||||
plans = build_patient_split_plans(
|
||||
patient_ids=eye_ids,
|
||||
patient_labels=labels,
|
||||
n_splits=n_splits,
|
||||
seed=fold_seed,
|
||||
)
|
||||
|
||||
out: list[PatientSplit] = []
|
||||
for plan in plans:
|
||||
train_df = df_full[df_full.index.isin(plan.train_patient_ids)].reset_index(drop=True)
|
||||
val_df = df_full[df_full.index.isin(plan.val_patient_ids)].reset_index(drop=True)
|
||||
test_df = df_full[df_full.index.isin(plan.test_patient_ids)].reset_index(drop=True)
|
||||
out.append(PatientSplit(train=train_df, val=val_df, test=test_df))
|
||||
return out
|
||||
|
||||
|
||||
class PatientFirstSplitManager:
|
||||
"""Patient-level splitter for V3. Outer/inner k-fold, no holdout."""
|
||||
|
||||
|
||||
+223
-51
@@ -23,7 +23,11 @@ from typing import Optional
|
||||
import numpy as np
|
||||
import torch
|
||||
|
||||
from v3.classes.croppers import build_image_preprocessor_from_args
|
||||
from v3.classes.croppers import (
|
||||
ManifestImageCropper,
|
||||
UNetImageCropper,
|
||||
build_image_preprocessor_from_args,
|
||||
)
|
||||
from v3.classes.image_loader import CachedImageLoader
|
||||
from v3.classes.dataset import _ClinicalView # noqa: F401
|
||||
from v3.classes.loader_factory import (
|
||||
@@ -35,10 +39,14 @@ from v3.classes.loader_factory import (
|
||||
from v3.classes.metrics import _score_arrays, _svf, _tune_and_snap
|
||||
from v3.classes.models import (
|
||||
BilateralHT,
|
||||
EmbeddingMLPEnsembleHT,
|
||||
FusedEnsembleHT,
|
||||
LogitMLPEnsembleHT,
|
||||
SiameseHT,
|
||||
SingleEyeHT,
|
||||
V2ModeComparisonOps,
|
||||
collect_probs_bilateral,
|
||||
collect_probs_siamese,
|
||||
collect_probs_bilateral_components,
|
||||
collect_probs_classic,
|
||||
collect_probs_ensemble,
|
||||
@@ -47,6 +55,7 @@ from v3.classes.models import (
|
||||
collect_probs_fused,
|
||||
collect_probs_single_components,
|
||||
train_bilateral_epoch,
|
||||
train_siamese_epoch,
|
||||
train_fusion_epoch,
|
||||
train_single_epoch,
|
||||
)
|
||||
@@ -54,7 +63,7 @@ from v3.classes.papila_builders import build_papila_data
|
||||
from v3.classes.predictions import PredictionStore, head_names_for_mode
|
||||
from v3.classes.profiles import build_papila_profile
|
||||
from v3.classes.results import FoldArtifacts, FoldResult, _f, _nan, _sv
|
||||
from v3.classes.split_manager import PatientFirstSplitManager
|
||||
from v3.classes.split_manager import EyeLevelSplitManager, PatientFirstSplitManager
|
||||
from v3.classes.transforms import build_eval_transform
|
||||
from v3.classes.utils import (
|
||||
_drop_mixed_label_patients,
|
||||
@@ -145,11 +154,14 @@ class V3HyperTower:
|
||||
ap.add_argument("--exclude-cols", nargs="*", default=[])
|
||||
ap.add_argument("--eval-mode", choices=["binary", "multiclass"], default="binary")
|
||||
ap.add_argument(
|
||||
"--tower-mode", choices=["single", "ensemble", "bilateral", "classic"],
|
||||
"--tower-mode", choices=["single", "ensemble", "bilateral", "siamese", "classic"],
|
||||
default="ensemble",
|
||||
)
|
||||
ap.add_argument("--n-splits", type=int, default=5)
|
||||
ap.add_argument("--fold-seed", type=int, default=42)
|
||||
ap.add_argument("--leaky-cv", action="store_true",
|
||||
help="Split at eye level (leaky: same patient can span folds). "
|
||||
"Used to demonstrate data-leakage effect.")
|
||||
ap.add_argument(
|
||||
"--folds", type=int, default=None,
|
||||
help="Optional cap on number of folds to run.",
|
||||
@@ -157,9 +169,9 @@ class V3HyperTower:
|
||||
ap.add_argument("--epochs", type=int, default=40)
|
||||
ap.add_argument("--warmup-tower-epochs", type=int, default=None)
|
||||
ap.add_argument("--warmup-fused-epochs", type=int, default=None)
|
||||
ap.add_argument("--single-warmup-tower-epochs", type=int, default=None)
|
||||
ap.add_argument("--single-warmup-fused-epochs", type=int, default=None)
|
||||
ap.add_argument("--warmup-cd-epochs", type=int, default=0)
|
||||
ap.add_argument("--single-warmup-tower-epochs", type=int, default=3)
|
||||
ap.add_argument("--single-warmup-fused-epochs", type=int, default=3)
|
||||
ap.add_argument("--warmup-cd-epochs", type=int, default=40)
|
||||
ap.add_argument("--bilat-warmup-tower-epochs", type=int, default=None)
|
||||
ap.add_argument("--bilat-warmup-fused-epochs", type=int, default=None)
|
||||
ap.add_argument("--batch-size", type=int, default=8)
|
||||
@@ -170,7 +182,7 @@ class V3HyperTower:
|
||||
ap.add_argument("--freeze-ratio", type=float, default=0.0)
|
||||
ap.add_argument("--augment", action="store_true")
|
||||
ap.add_argument("--balanced-sampling", action="store_true")
|
||||
ap.add_argument("--num-workers", type=int, default=4)
|
||||
ap.add_argument("--num-workers", type=int, default=8)
|
||||
ap.add_argument("--in-memory-cache", action="store_true", default=True)
|
||||
ap.add_argument("--no-in-memory-cache", action="store_false", dest="in_memory_cache")
|
||||
ap.add_argument("--cache-workers", type=int, default=4)
|
||||
@@ -191,10 +203,20 @@ class V3HyperTower:
|
||||
ap.add_argument("--img-crop-cache", type=str, default="cache_data/hypertower_crops")
|
||||
ap.add_argument("--persist-img-crop-cache", action="store_true")
|
||||
# Architecture
|
||||
ap.add_argument("--cd-hidden-dim", type=int, default=128)
|
||||
ap.add_argument("--fusion-dim", type=int, default=256)
|
||||
ap.add_argument("--bridge-mode", default="fused",
|
||||
ap.add_argument("--cd-hidden-dim", type=int, default=128)
|
||||
ap.add_argument("--fusion-dim", type=int, default=256)
|
||||
ap.add_argument("--bridge-mode", default="fused",
|
||||
choices=["fused", "image_only", "clinical_only"])
|
||||
ap.add_argument("--bridge-dropout", type=float, default=0.5,
|
||||
help="Dropout in bridge classifier_fused (default: 0.5)")
|
||||
ap.add_argument("--cd-dropout", type=float, default=0.1,
|
||||
help="Dropout in clinical tower MLP (default: 0.1)")
|
||||
ap.add_argument("--se-img-tower", action="store_true",
|
||||
help="Enable SE gate on image tower output features")
|
||||
ap.add_argument("--se-cd-tower", action="store_true",
|
||||
help="Enable SE gate on clinical tower output features")
|
||||
ap.add_argument("--se-bridge", action="store_true",
|
||||
help="Enable SE gate on fused vector inside the bridge")
|
||||
# Mixed patients
|
||||
ap.add_argument("--exclude-mixed-patients", dest="exclude_mixed_patients",
|
||||
action="store_true")
|
||||
@@ -217,6 +239,19 @@ class V3HyperTower:
|
||||
# Fused head
|
||||
ap.add_argument("--fused-head", action="store_true")
|
||||
ap.add_argument("--fusion-epochs", type=int, default=10)
|
||||
ap.add_argument("--head-type",
|
||||
choices=["attention", "logit_mlp", "embedding_mlp"],
|
||||
default="attention",
|
||||
help="Which bilateral head to train on top of frozen ensemble base")
|
||||
ap.add_argument("--save-checkpoints", action="store_true",
|
||||
help="Save best_single.pt per fold for explainability / GradCAM")
|
||||
# Geometry features
|
||||
ap.add_argument("--geometry-dim", type=int, default=0,
|
||||
help="Append N geometry features to clinical metadata (0=disabled, 5=all). "
|
||||
"Requires --img-crop-manifest.")
|
||||
ap.add_argument("--geometry-source", default="gt", choices=["gt", "unet"],
|
||||
help="Source for geometry features: gt (GT contour annotations) or "
|
||||
"unet (U-Net segmentation). unet also requires --img-crop-weights.")
|
||||
return ap
|
||||
|
||||
def __init__(self, args) -> None:
|
||||
@@ -246,6 +281,49 @@ class V3HyperTower:
|
||||
patient_col="Patient ID", label_col=args.label_col, sample_mode="patient"
|
||||
)
|
||||
|
||||
# Build geometry provider if requested, extend feature_dim to include geometry.
|
||||
# Both ManifestImageCropper and UNetImageCropper already have geometry_features()
|
||||
# and precompute_geometry() — we just pick the right one and pre-compute upfront.
|
||||
self.geometry_provider = None
|
||||
geom_dim = int(getattr(args, "geometry_dim", 0))
|
||||
if geom_dim > 0:
|
||||
source = getattr(args, "geometry_source", "gt")
|
||||
manifest = getattr(args, "img_crop_manifest", None)
|
||||
if not manifest:
|
||||
raise ValueError("--geometry-dim requires --img-crop-manifest")
|
||||
all_paths = [
|
||||
self.data.get_image_path(row)
|
||||
for _, row in self.data.df.iterrows()
|
||||
]
|
||||
if source == "gt":
|
||||
# Reuse image_preprocessor if it's already a ManifestImageCropper,
|
||||
# otherwise build a lightweight one just for geometry (no crop cache).
|
||||
if isinstance(self.image_preprocessor, ManifestImageCropper):
|
||||
provider = self.image_preprocessor
|
||||
else:
|
||||
provider = ManifestImageCropper(manifest_path=Path(manifest))
|
||||
print(f"[geometry] GT source — pre-computing geometry from {manifest}", flush=True)
|
||||
elif source == "unet":
|
||||
weights = getattr(args, "img_crop_weights", None)
|
||||
if not weights:
|
||||
raise ValueError("--geometry-source unet requires --img-crop-weights")
|
||||
if isinstance(self.image_preprocessor, UNetImageCropper):
|
||||
provider = self.image_preprocessor
|
||||
else:
|
||||
provider = UNetImageCropper(
|
||||
manifest_path=Path(manifest),
|
||||
weights_path=Path(weights),
|
||||
normalize=getattr(args, "img_crop_normalize", "per_image"),
|
||||
threshold=getattr(args, "img_crop_threshold", 0.5),
|
||||
)
|
||||
print(f"[geometry] UNet source — pre-computing geometry from {weights}", flush=True)
|
||||
else:
|
||||
raise ValueError(f"Unknown --geometry-source: {source!r}")
|
||||
provider.precompute_geometry(all_paths)
|
||||
self.geometry_provider = provider
|
||||
self.data.feature_dim += geom_dim
|
||||
print(f"[geometry] feature_dim extended to {self.data.feature_dim} (+{geom_dim} geometry)", flush=True)
|
||||
|
||||
def run(self) -> Path:
|
||||
"""Execute the full fold loop."""
|
||||
args = self.args
|
||||
@@ -276,7 +354,11 @@ class V3HyperTower:
|
||||
num_classes = 2 if mode == "binary" else int(df_mode[args.label_col].nunique())
|
||||
print(f"\n[{mode}] num_classes={num_classes} rows={len(df_mode)} patients={df_mode['Patient ID'].nunique()}", flush=True)
|
||||
|
||||
split_manager = PatientFirstSplitManager(patient_col="Patient ID", label_col=args.label_col)
|
||||
if getattr(args, "leaky_cv", False):
|
||||
split_manager = EyeLevelSplitManager(patient_col="Patient ID", label_col=args.label_col)
|
||||
print("[CV] WARNING: leaky-cv mode — eye-level splits, same patient can span folds.", flush=True)
|
||||
else:
|
||||
split_manager = PatientFirstSplitManager(patient_col="Patient ID", label_col=args.label_col)
|
||||
split_args = SimpleNamespace(
|
||||
eval_mode=mode,
|
||||
n_splits=args.n_splits,
|
||||
@@ -390,7 +472,7 @@ class V3HyperTower:
|
||||
_test_key = "classic_test"
|
||||
elif tower_mode == "ensemble":
|
||||
_test_key = "ensemble_test"
|
||||
elif tower_mode == "bilateral":
|
||||
elif tower_mode in ("bilateral", "siamese"):
|
||||
_test_key = "bilat_test"
|
||||
else:
|
||||
_test_key = "classic_test"
|
||||
@@ -400,6 +482,25 @@ class V3HyperTower:
|
||||
|
||||
return out_dir
|
||||
|
||||
def _augment_geometry(self, samples: list) -> list:
|
||||
"""Append geometry features to matrix_1/matrix_2 in each sample dict."""
|
||||
if self.geometry_provider is None:
|
||||
return samples
|
||||
geom_dim = int(getattr(self.args, "geometry_dim", 0))
|
||||
for s in samples:
|
||||
for img_slot, mat_slot in (("image_1", "matrix_1"), ("image_2", "matrix_2")):
|
||||
img_path = s.get(img_slot)
|
||||
mat = s.get(mat_slot)
|
||||
if img_path is None or mat is None:
|
||||
continue
|
||||
vec = self.geometry_provider.geometry_for_image(img_path)
|
||||
if vec is not None and len(vec) >= geom_dim:
|
||||
geom = vec[:geom_dim].astype(np.float32)
|
||||
else:
|
||||
geom = np.zeros(geom_dim, dtype=np.float32)
|
||||
s[mat_slot] = np.concatenate([np.asarray(mat, dtype=np.float32), geom])
|
||||
return samples
|
||||
|
||||
def _run_fold(
|
||||
self,
|
||||
*,
|
||||
@@ -420,9 +521,10 @@ class V3HyperTower:
|
||||
image_preprocessor = self.image_preprocessor
|
||||
nan = float("nan")
|
||||
|
||||
run_single = tower_mode in ("single", "ensemble")
|
||||
run_bilat = tower_mode == "bilateral"
|
||||
run_fused = tower_mode == "ensemble" and getattr(args, "fused_head", False)
|
||||
run_single = tower_mode in ("single", "ensemble")
|
||||
run_bilat = tower_mode == "bilateral"
|
||||
run_siamese = tower_mode == "siamese"
|
||||
run_fused = tower_mode == "ensemble" and getattr(args, "fused_head", False)
|
||||
|
||||
# ---- warmup schedule -------------------------------------------
|
||||
global_warmup_tower = getattr(args, "warmup_tower_epochs", None)
|
||||
@@ -442,7 +544,7 @@ class V3HyperTower:
|
||||
single_warmup_cd = int(getattr(args, "warmup_cd_epochs", 0)) if run_single else 0
|
||||
if not run_single:
|
||||
single_warmup_tower = single_warmup_fused = 0
|
||||
if not run_bilat:
|
||||
if not run_bilat and not run_siamese:
|
||||
bilat_warmup_tower = bilat_warmup_fused = 0
|
||||
# Warmup is meaningless in single-pathway modes — skip it entirely
|
||||
_bridge_mode = getattr(args, "bridge_mode", "fused")
|
||||
@@ -451,7 +553,7 @@ class V3HyperTower:
|
||||
bilat_warmup_tower = bilat_warmup_fused = 0
|
||||
main_epochs = int(args.epochs)
|
||||
total_single_epochs = (single_warmup_cd + single_warmup_tower + single_warmup_fused + main_epochs) if run_single else 0
|
||||
total_bilat_epochs = (bilat_warmup_tower + bilat_warmup_fused + main_epochs) if run_bilat else 0
|
||||
total_bilat_epochs = (bilat_warmup_tower + bilat_warmup_fused + main_epochs) if (run_bilat or run_siamese) else 0
|
||||
total_epochs = max(total_single_epochs, total_bilat_epochs)
|
||||
|
||||
# ---- samples ---------------------------------------------------
|
||||
@@ -460,6 +562,12 @@ class V3HyperTower:
|
||||
bilat_val = filter_bilateral_samples(profile_patient.build_samples(df=split.val, clinical=data))
|
||||
bilat_test = filter_bilateral_samples(profile_patient.build_samples(df=split.test, clinical=data)) if split.test is not None else []
|
||||
|
||||
if self.geometry_provider is not None:
|
||||
eye_train = self._augment_geometry(eye_train)
|
||||
bilat_train = self._augment_geometry(bilat_train)
|
||||
bilat_val = self._augment_geometry(bilat_val)
|
||||
bilat_test = self._augment_geometry(bilat_test)
|
||||
|
||||
if pred_store is not None:
|
||||
if tower_mode in ("single", "classic"):
|
||||
train_sids = [f"{s['id_1']}{s.get('eye_id_1','')}" for s in eye_train]
|
||||
@@ -502,27 +610,41 @@ class V3HyperTower:
|
||||
y_true_bilat=None, probs_bilat=None,
|
||||
)
|
||||
|
||||
# ---- models ----------------------------------------------------
|
||||
# ---- models (CPU for now — moved to device after workers spawn) ---
|
||||
single = None
|
||||
bilateral = None
|
||||
siamese = None
|
||||
if run_single:
|
||||
single = SingleEyeHT(
|
||||
backbone=args.backbone, freeze_ratio=args.freeze_ratio,
|
||||
augment=args.augment, clinical_data=data, num_classes=num_classes,
|
||||
cd_hidden_dim=args.cd_hidden_dim, fusion_dim=args.fusion_dim,
|
||||
bridge_mode=getattr(args, "bridge_mode", "fused"),
|
||||
).to(device)
|
||||
bridge_dropout=getattr(args, "bridge_dropout", 0.5),
|
||||
cd_dropout=getattr(args, "cd_dropout", 0.1),
|
||||
se_img_tower=getattr(args, "se_img_tower", False),
|
||||
se_cd_tower=getattr(args, "se_cd_tower", False),
|
||||
se_bridge=getattr(args, "se_bridge", False),
|
||||
)
|
||||
if run_bilat:
|
||||
bilateral = BilateralHT(
|
||||
backbone=args.backbone, freeze_ratio=args.freeze_ratio,
|
||||
augment=args.augment, clinical_data=data, num_classes=num_classes,
|
||||
cd_hidden_dim=args.cd_hidden_dim, fusion_dim=args.fusion_dim,
|
||||
).to(device)
|
||||
)
|
||||
if run_siamese:
|
||||
siamese = SiameseHT(
|
||||
backbone=args.backbone, freeze_ratio=args.freeze_ratio,
|
||||
augment=args.augment, num_classes=num_classes,
|
||||
fusion_dim=args.fusion_dim,
|
||||
)
|
||||
|
||||
slots_eye = profile_eye.slot_descriptors()
|
||||
slots_patient = profile_patient.slot_descriptors()
|
||||
_persistent_workers = args.num_workers > 0
|
||||
loader_kw = dict(batch_size=args.batch_size, num_workers=args.num_workers,
|
||||
image_cache=image_cache)
|
||||
image_cache=image_cache,
|
||||
persistent_workers=_persistent_workers)
|
||||
|
||||
# ---- loaders ---------------------------------------------------
|
||||
use_balanced = bool(getattr(args, "balanced_sampling", False))
|
||||
@@ -553,6 +675,13 @@ class V3HyperTower:
|
||||
image_preprocessor=image_preprocessor, shuffle=True,
|
||||
sampler=bilat_sampler, **loader_kw,
|
||||
)
|
||||
elif run_siamese:
|
||||
siamese_sampler = build_balanced_sampler(bilat_train) if use_balanced else None
|
||||
train_bilat_loader = make_loader(
|
||||
bilat_train, slots_patient, image_transform=siamese.transform,
|
||||
image_preprocessor=image_preprocessor, shuffle=True,
|
||||
sampler=siamese_sampler, **loader_kw,
|
||||
)
|
||||
elif run_fused:
|
||||
fused_sampler = build_balanced_sampler(bilat_train) if use_balanced else None
|
||||
train_bilat_loader = make_loader(
|
||||
@@ -581,8 +710,26 @@ class V3HyperTower:
|
||||
if _ldr is not None:
|
||||
_ldr.dataset.prebuild_image_cache()
|
||||
|
||||
opt_single = torch.optim.Adam(single.parameters(), lr=args.lr) if run_single else None
|
||||
opt_bilateral = torch.optim.Adam(bilateral.parameters(), lr=args.lr) if run_bilat else None
|
||||
# ---- spawn DataLoader workers BEFORE CUDA init -----------------
|
||||
# Workers fork here (clean process state, no CUDA context yet).
|
||||
# persistent_workers=True keeps them alive so the training loop
|
||||
# reuses them rather than re-forking after .to(device).
|
||||
if _persistent_workers:
|
||||
for _ldr in [train_single_loader, train_bilat_loader, val_loader, test_loader]:
|
||||
if _ldr is not None:
|
||||
_ = iter(_ldr) # triggers fork now, before CUDA
|
||||
|
||||
# ---- move models to device (CUDA init happens here) ------------
|
||||
if single is not None:
|
||||
single = single.to(device)
|
||||
if bilateral is not None:
|
||||
bilateral = bilateral.to(device)
|
||||
if siamese is not None:
|
||||
siamese = siamese.to(device)
|
||||
|
||||
opt_single = torch.optim.Adam(single.parameters(), lr=args.lr) if run_single else None
|
||||
opt_bilateral = torch.optim.Adam(bilateral.parameters(), lr=args.lr) if run_bilat else None
|
||||
opt_siamese = torch.optim.Adam(siamese.parameters(), lr=args.lr) if run_siamese else None
|
||||
|
||||
# ---- epoch log -------------------------------------------------
|
||||
epoch_fields = [
|
||||
@@ -679,7 +826,7 @@ class V3HyperTower:
|
||||
else:
|
||||
phase_single, main_epoch_single, single_active = "done", main_epochs, False
|
||||
|
||||
if not run_bilat:
|
||||
if not run_bilat and not run_siamese:
|
||||
phase_bilat, main_epoch_bilat, bilat_active = "inactive", 0, False
|
||||
elif epoch < bilat_warmup_tower:
|
||||
phase_bilat, main_epoch_bilat, bilat_active = "tower_warmup", 0, True
|
||||
@@ -709,6 +856,12 @@ class V3HyperTower:
|
||||
phase=phase_bilat, bcd_prob=float(args.bcd_prob),
|
||||
tower_loss_mode=args.tower_loss_mode,
|
||||
)
|
||||
elif run_siamese and bilat_active:
|
||||
bl_loss, bl_acc = train_siamese_epoch(
|
||||
siamese, train_bilat_loader, opt_siamese, device,
|
||||
bcd_prob=float(args.bcd_prob),
|
||||
tower_loss_mode=args.tower_loss_mode,
|
||||
)
|
||||
else:
|
||||
bl_loss, bl_acc = nan, nan
|
||||
|
||||
@@ -763,6 +916,10 @@ class V3HyperTower:
|
||||
bi_acc_cd = float((p_bi_cd.argmax(1) ==y_bi).mean()) if y_bi.size else nan
|
||||
_, bi_auc_img, _ = _score_arrays(y_bi, p_bi_img, num_classes)
|
||||
_, bi_auc_cd, _ = _score_arrays(y_bi, p_bi_cd, num_classes)
|
||||
elif run_siamese and not _skip_val_eval:
|
||||
y_bi, p_bi = collect_probs_siamese(siamese, val_loader, device)
|
||||
bi_acc, bi_auc, bi_n = _score_arrays(y_bi, p_bi, num_classes)
|
||||
bi_acc_img = bi_acc_cd = bi_auc_img = bi_auc_cd = nan
|
||||
else:
|
||||
y_bi = np.array([], dtype=np.int64)
|
||||
p_bi = np.zeros((0, 0), dtype=np.float32)
|
||||
@@ -905,7 +1062,7 @@ class V3HyperTower:
|
||||
_bar_w = 30
|
||||
_filled = int(_bar_w * (epoch + 1) / single_warmup_cd)
|
||||
_bar = "#" * _filled + "-" * (_bar_w - _filled)
|
||||
msg = f" [fold {fold+1}] md_warmup [{_bar}] {epoch+1}/{single_warmup_cd} loss={sl_loss:.4f}"
|
||||
msg = f" [fold {fold+1}] md_warmup [{_bar}] {epoch+1}/{single_warmup_cd} loss={sl_loss:.2f}"
|
||||
print(f"\r{msg}", end="", flush=True)
|
||||
fold_logger.info(msg)
|
||||
_prev_phase_single = phase_single
|
||||
@@ -917,13 +1074,13 @@ class V3HyperTower:
|
||||
if args.log_every > 0 and (epoch + 1) % args.log_every == 0:
|
||||
_epoch_secs = time.time() - _epoch_t0
|
||||
if tower_mode == "ensemble":
|
||||
msg = f" ep {epoch+1:>3}/{total_epochs} ({_epoch_secs:.1f}s) auc={en_auc:.4f} acc={en_acc:.4f}"
|
||||
elif tower_mode == "bilateral":
|
||||
msg = f" ep {epoch+1:>3}/{total_epochs} ({_epoch_secs:.1f}s) auc={bi_auc:.4f} acc={bi_acc:.4f}"
|
||||
_auc_v, _acc_v = en_auc, en_acc
|
||||
elif tower_mode in ("bilateral", "siamese"):
|
||||
_auc_v, _acc_v = bi_auc, bi_acc
|
||||
else:
|
||||
msg = f" ep {epoch+1:>3}/{total_epochs} ({_epoch_secs:.1f}s) auc={cl_auc:.4f} acc={cl_acc:.4f}"
|
||||
print(msg, flush=True)
|
||||
fold_logger.info(msg)
|
||||
_auc_v, _acc_v = cl_auc, cl_acc
|
||||
print(f" ep {epoch+1:>3}/{total_epochs} ({_epoch_secs:.1f}s) auc={_auc_v:.2f} acc={_acc_v:.2f}", flush=True)
|
||||
fold_logger.info(f" ep {epoch+1:>3}/{total_epochs} ({_epoch_secs:.1f}s) auc={_auc_v:.4f} acc={_acc_v:.4f}")
|
||||
|
||||
_prev_phase_single = phase_single
|
||||
|
||||
@@ -958,8 +1115,14 @@ class V3HyperTower:
|
||||
if run_fused and single is not None:
|
||||
for p in single.parameters():
|
||||
p.requires_grad_(False)
|
||||
fused = FusedEnsembleHT(single, num_classes).to(device)
|
||||
opt_fused = torch.optim.Adam(fused.eye_scorer.parameters(), lr=args.lr)
|
||||
head_type = getattr(args, "head_type", "attention")
|
||||
if head_type == "logit_mlp":
|
||||
fused = LogitMLPEnsembleHT(single, num_classes).to(device)
|
||||
elif head_type == "embedding_mlp":
|
||||
fused = EmbeddingMLPEnsembleHT(single, num_classes).to(device)
|
||||
else:
|
||||
fused = FusedEnsembleHT(single, num_classes).to(device)
|
||||
opt_fused = torch.optim.Adam(fused.head.parameters(), lr=args.lr)
|
||||
fusion_epochs = int(getattr(args, "fusion_epochs", 10))
|
||||
print(
|
||||
f" [fold {fold+1}] Phase 2: fusion head bilat_train_n={len(bilat_train)} epochs={fusion_epochs}",
|
||||
@@ -977,8 +1140,8 @@ class V3HyperTower:
|
||||
snap_fused, _, _, _ = _tune_and_snap(y_fu, p_fu, fu_acc_val, num_classes, args, args.ece_bins)
|
||||
if (fep + 1) % max(1, args.log_every) == 0:
|
||||
print(
|
||||
f" [fold {fold+1}] fusion ep{fep+1:>3} loss={fu_loss:.4f} "
|
||||
f"val_auc={fu_auc:.4f}",
|
||||
f" [fold {fold+1}] fusion ep{fep+1:>3} loss={fu_loss:.2f} "
|
||||
f"val_auc={fu_auc:.2f}",
|
||||
flush=True,
|
||||
)
|
||||
# No checkpoint saving for fused head either.
|
||||
@@ -1015,9 +1178,16 @@ class V3HyperTower:
|
||||
p_cl_best_img = p_cl_best_md = None
|
||||
if run_bilat:
|
||||
y_bi_best, p_bi_best = collect_probs_bilateral(bilateral, val_loader, device)
|
||||
elif run_siamese:
|
||||
y_bi_best, p_bi_best = collect_probs_siamese(siamese, val_loader, device)
|
||||
else:
|
||||
y_bi_best = p_bi_best = None
|
||||
|
||||
# Optional checkpoint saving (final-epoch weights for explainability)
|
||||
if getattr(args, "save_checkpoints", False) and run_single and single is not None:
|
||||
import torch as _torch
|
||||
_torch.save(single.state_dict(), fold_dir / "best_single.pt")
|
||||
|
||||
# Compute val snaps from final-epoch model state
|
||||
if run_single and tower_mode == "single" and y_cl_best is not None:
|
||||
snap_cl, _, _, _ = _tune_and_snap(y_cl_best, p_cl_best, float((p_cl_best.argmax(1) == y_cl_best).mean()), num_classes, args, args.ece_bins)
|
||||
@@ -1025,7 +1195,7 @@ class V3HyperTower:
|
||||
elif run_single and tower_mode == "ensemble" and y_en_best is not None:
|
||||
snap_en, _, _, _ = _tune_and_snap(y_en_best, p_en_best, float((p_en_best.argmax(1) == y_en_best).mean()), num_classes, args, args.ece_bins)
|
||||
snap_ensemble = snap_en
|
||||
if run_bilat and y_bi_best is not None:
|
||||
if (run_bilat or run_siamese) and y_bi_best is not None:
|
||||
snap_bi, _, _, _ = _tune_and_snap(y_bi_best, p_bi_best, float((p_bi_best.argmax(1) == y_bi_best).mean()), num_classes, args, args.ece_bins)
|
||||
snap_bilat = snap_bi
|
||||
|
||||
@@ -1050,6 +1220,8 @@ class V3HyperTower:
|
||||
y_test_out, p_test_out, _, _ = collect_probs_bilateral_components(
|
||||
bilateral, test_loader, device
|
||||
)
|
||||
elif run_siamese:
|
||||
y_test_out, p_test_out = collect_probs_siamese(siamese, test_loader, device)
|
||||
if y_test_out is not None and y_test_out.size:
|
||||
test_acc_raw = float((p_test_out.argmax(1) == y_test_out).mean())
|
||||
snap_test, _, _, _ = _tune_and_snap(
|
||||
@@ -1057,11 +1229,11 @@ class V3HyperTower:
|
||||
)
|
||||
print(
|
||||
f" [fold {fold+1}] TEST "
|
||||
f"auc={snap_test.get('auc', nan):.4f} "
|
||||
f"acc={snap_test.get('acc', nan):.4f} "
|
||||
f"kappa={snap_test.get('kappa', nan):.4f} "
|
||||
f"f1={snap_test.get('macro_f1', nan):.4f} "
|
||||
f"ece={snap_test.get('ece', nan):.4f} "
|
||||
f"auc={snap_test.get('auc', nan):.2f} "
|
||||
f"acc={snap_test.get('acc', nan):.2f} "
|
||||
f"kappa={snap_test.get('kappa', nan):.2f} "
|
||||
f"f1={snap_test.get('macro_f1', nan):.2f} "
|
||||
f"ece={snap_test.get('ece', nan):.2f} "
|
||||
f"n={snap_test.get('n', 0)}",
|
||||
flush=True,
|
||||
)
|
||||
@@ -1118,11 +1290,11 @@ class V3HyperTower:
|
||||
classic_test_kappa=snap_test.get("kappa", nan) if tower_mode == "single" else nan,
|
||||
classic_test_f1=snap_test.get("macro_f1", nan) if tower_mode == "single" else nan,
|
||||
classic_test_ece=snap_test.get("ece", nan) if tower_mode == "single" else nan,
|
||||
bilat_test_auc=snap_test.get("auc", nan) if tower_mode == "bilateral" else nan,
|
||||
bilat_test_acc=snap_test.get("acc", nan) if tower_mode == "bilateral" else nan,
|
||||
bilat_test_kappa=snap_test.get("kappa", nan) if tower_mode == "bilateral" else nan,
|
||||
bilat_test_f1=snap_test.get("macro_f1", nan) if tower_mode == "bilateral" else nan,
|
||||
bilat_test_ece=snap_test.get("ece", nan) if tower_mode == "bilateral" else nan,
|
||||
bilat_test_auc=snap_test.get("auc", nan) if tower_mode in ("bilateral", "siamese") else nan,
|
||||
bilat_test_acc=snap_test.get("acc", nan) if tower_mode in ("bilateral", "siamese") else nan,
|
||||
bilat_test_kappa=snap_test.get("kappa", nan) if tower_mode in ("bilateral", "siamese") else nan,
|
||||
bilat_test_f1=snap_test.get("macro_f1", nan) if tower_mode in ("bilateral", "siamese") else nan,
|
||||
bilat_test_ece=snap_test.get("ece", nan) if tower_mode in ("bilateral", "siamese") else nan,
|
||||
test_n=test_n,
|
||||
single_train_n=len(eye_train),
|
||||
bilat_train_n=len(bilat_train),
|
||||
@@ -1222,18 +1394,18 @@ class V3HyperTower:
|
||||
@staticmethod
|
||||
def _print_summary(mode: str, s: dict, tower_mode: str | None = None) -> None:
|
||||
def f(v):
|
||||
return " nan " if v is None else f"{v:.4f}"
|
||||
return " nan " if v is None else f"{v:.2f}"
|
||||
def fsd(mean, std):
|
||||
if mean is None: return " nan "
|
||||
if std is None: return f"{mean:.4f} "
|
||||
return f"{mean:.4f}±{std:.4f}"
|
||||
if mean is None: return " nan "
|
||||
if std is None: return f"{mean:.2f} "
|
||||
return f"{mean:.2f}±{std:.2f}"
|
||||
|
||||
# Resolve test key
|
||||
if tower_mode in ("single", "classic"):
|
||||
test_key = "classic_test"
|
||||
elif tower_mode == "ensemble":
|
||||
test_key = "ensemble_test"
|
||||
elif tower_mode == "bilateral":
|
||||
elif tower_mode in ("bilateral", "siamese"):
|
||||
test_key = "bilat_test"
|
||||
else:
|
||||
test_key = "classic_test"
|
||||
|
||||
Reference in New Issue
Block a user