from __future__ import annotations import os import math import argparse import json import time import copy import csv from pathlib import Path import shutil import random as pyrandom from dataclasses import dataclass import pandas as pd from PIL import Image, ImageDraw import torch from torch import nn from torch.utils.data import DataLoader from torch.utils.data.sampler import WeightedRandomSampler from classes.v2.data_bundle import DataBundle from classes.v2.dataset import ClinicalDataset from classes.early_stop import EarlyStopper from classes.unet_segmenter import UNetSegmenter from classes.geometry_features import ( FEATURE_DIM, compute_geometry_features, disc_cup_from_mask_image, ) from classes.refuge_classification import _geometry_from_mask from classes.v2.bridges import Bridge, VoteBridge from classes.v2.hypertower_logger import HypertowerLogger from classes.v2.towers import ImageTower, MDTower from torchvision import transforms # from clinical_data import ClinicalData # from dataset import ClinicalDataset # from image_tower import ImageTower # from md_tower import MDTower # from bridge import Bridge, VoteBridge from random import random from sklearn.metrics import roc_auc_score, cohen_kappa_score, f1_score, matthews_corrcoef, recall_score import torch.nn.functional as F import numpy as np from sklearn.metrics import roc_curve, auc from sklearn.preprocessing import label_binarize from typing import Optional, Tuple from types import SimpleNamespace from classes.backbones import BACKBONES from classes.v2.papila_builders import build_papila_data from classes.v2.split_manager import PatientFirstSplitManager LOG_FIELDS = [ "epoch", "eval_loss", "acc_fused", "acc_img", "acc_md", "auc_fused", "auc_img", "auc_md", "top2_fused", "top2_img", "top2_md", "margin_fused", "margin_img", "margin_md", "agree_fused_img", "agree_fused_md", "pct_fused", "pct_img", "pct_md", "phase", "holdout_loss", "holdout_acc_fused", "holdout_acc_img", "holdout_acc_md", "holdout_auc_fused", "holdout_auc_img", "holdout_auc_md", ] def focal_loss( logits: torch.Tensor, targets: torch.Tensor, gamma: float = 0.0, weight: Optional[torch.Tensor] = None, reduction: str = "mean", ) -> torch.Tensor: """ Standard focal loss wrapper. When gamma=0 it reduces to cross entropy. weight should be per-class weights (same semantics as CrossEntropyLoss). """ 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) focal_factor = (1.0 - pt).clamp_min(0.0) ** gamma loss = -focal_factor * logpt if weight is not None: class_weight = weight.gather(0, targets.view(-1)) loss = loss * class_weight.view(-1, 1) loss = loss.view(-1) if reduction == "sum": return loss.sum() if reduction == "mean": return loss.mean() return loss class UNetImageCropper: def __init__( self, manifest_path: Path, weights_path: Path, normalize: str = "per_image", threshold: float = 0.5, tta: bool = False, scale: float = 2.5, target_size: int = 224, cache_dir: Optional[Path] = None, ) -> None: self.segmenter = UNetSegmenter( manifest_path=manifest_path, normalize=normalize, ) state = torch.load(weights_path, map_location=self.segmenter.device) state_dict = state.get("model", state) self.segmenter.model.load_state_dict(state_dict) self.segmenter.model.to(self.segmenter.device) self.segmenter.model.eval() self.threshold = threshold self.tta = tta self.scale = scale self.target_size = target_size self.cache_dir = Path(cache_dir) if cache_dir is not None else None if self.cache_dir is not None: self.cache_dir.mkdir(parents=True, exist_ok=True) self.to_tensor = transforms.ToTensor() def _cache_path(self, image_path: Path) -> Optional[Path]: if self.cache_dir is None: return None stem = image_path.stem return self.cache_dir / f"{stem}_s{int(self.scale * 100)}.npz" def _infer_masks(self, image: Image.Image) -> Optional[Tuple[np.ndarray, np.ndarray]]: resized = self.segmenter.preprocess_image(image) tensor = self.to_tensor(resized).unsqueeze(0).to(self.segmenter.device) with torch.no_grad(): logits = self.segmenter.model(tensor) if self.tta: t_h = torch.flip(tensor, dims=[3]) log_h = self.segmenter.model(t_h) log_h = torch.flip(log_h, dims=[3]) t_v = torch.flip(tensor, dims=[2]) log_v = self.segmenter.model(t_v) log_v = torch.flip(log_v, dims=[2]) logits = (logits + log_h + log_v) / 3.0 probs = torch.sigmoid(logits)[0].cpu().numpy() disc_pred = (probs[0] > self.threshold).astype(np.uint8) * 255 cup_pred = (probs[1] > self.threshold).astype(np.uint8) * 255 disc_img = Image.fromarray(disc_pred, mode="L").resize(image.size, Image.NEAREST) disc_mask = np.array(disc_img, dtype=np.uint8) cup_img = Image.fromarray(cup_pred, mode="L").resize(image.size, Image.NEAREST) cup_mask = (np.array(cup_img, dtype=np.uint8) > 0).astype(np.uint8) cup_mask = (cup_mask > 0) & (disc_mask > 0) cup_mask = cup_mask.astype(np.uint8) disc_mask = (disc_mask > 0).astype(np.uint8) return disc_mask, cup_mask def _compute_crop_info(self, image: Image.Image, image_path: Path) -> Optional[dict]: image_path = Path(image_path).resolve() cache_path = self._cache_path(image_path) cached_bounds = None if cache_path is not None and cache_path.exists(): data = np.load(cache_path, allow_pickle=False) try: cached_bounds = { "left": float(data["left"]), "upper": float(data["upper"]), "right": float(data["right"]), "lower": float(data["lower"]), } if "features" in data.files: cached_bounds["features"] = data["features"].astype(np.float32) return cached_bounds except KeyError: cached_bounds = None masks = self._infer_masks(image) if masks is None: return cached_bounds disc_mask, cup_mask = masks try: geom = _geometry_from_mask(disc_mask, self.scale) except Exception: return cached_bounds cx = geom["centre_x"] cy = geom["centre_y"] r = geom["crop_radius"] left = max(0.0, cx - r) upper = max(0.0, cy - r) right = min(float(image.width), cx + r) lower = min(float(image.height), cy + r) features = compute_geometry_features(disc_mask, cup_mask) info = { "left": left, "upper": upper, "right": right, "lower": lower, "features": features, } if cache_path is not None: np.savez( cache_path, left=left, upper=upper, right=right, lower=lower, width=float(image.width), height=float(image.height), scale=self.scale, target_size=self.target_size, features=features, ) return info def __call__(self, image: Image.Image, image_path: Path) -> Image.Image: info = self._compute_crop_info(image, image_path) if info is None: return image left = info["left"] upper = info["upper"] right = info["right"] lower = info["lower"] if right <= left or lower <= upper: return image crop = image.crop((left, upper, right, lower)) return crop.resize((self.target_size, self.target_size), Image.BILINEAR) def geometry_features(self, image: Image.Image, image_path: Path) -> Optional[np.ndarray]: info = self._compute_crop_info(image, image_path) if info is None: return None features = info.get("features") if features is None: return None return np.asarray(features, dtype=np.float32) class ManifestImageCropper: def __init__( self, manifest_path: Path, scale: float = 2.5, target_size: int = 224, cache_dir: Optional[Path] = None, ) -> None: self.scale = scale self.target_size = target_size self.cache_dir = Path(cache_dir) if cache_dir is not None else None if self.cache_dir is not None: self.cache_dir.mkdir(parents=True, exist_ok=True) df = pd.read_csv(manifest_path) self.entries: Dict[str, dict] = {} for _, row in df.iterrows(): img_path = Path(row["image_path"]).resolve() self.entries[str(img_path)] = { "annotation_disc": row.get("annotation_disc"), "annotation_cup": row.get("annotation_cup"), "annotation_type_disc": row.get("annotation_type_disc"), "annotation_type_cup": row.get("annotation_type_cup"), } def _cache_path(self, image_path: Path) -> Optional[Path]: if self.cache_dir is None: return None return self.cache_dir / f"{image_path.stem}_s{int(self.scale * 100)}.npz" @staticmethod def _load_contour(path: Path) -> np.ndarray: coords = np.loadtxt(path) if coords.ndim == 1: coords = coords.reshape(-1, 2) return coords @staticmethod def _contour_to_mask(coords: np.ndarray, size: tuple[int, int]) -> np.ndarray: if coords is None or coords.size == 0: return np.zeros((size[1], size[0]), dtype=np.uint8) img = Image.new("L", size, 0) draw = ImageDraw.Draw(img) points = [tuple(map(float, pt)) for pt in coords] draw.polygon(points, outline=1, fill=1) return np.array(img, dtype=np.uint8) def _load_masks(self, entry: dict, image: Image.Image) -> Optional[Tuple[np.ndarray, np.ndarray]]: disc_path = entry.get("annotation_disc") cup_path = entry.get("annotation_cup") disc_type = (entry.get("annotation_type_disc") or "").lower() cup_type = (entry.get("annotation_type_cup") or "").lower() disc_mask: Optional[np.ndarray] = None cup_mask: Optional[np.ndarray] = None if disc_path and not pd.isna(disc_path): disc_path = Path(disc_path) try: if disc_type == "mask": mask_img = Image.open(disc_path) mask_img = mask_img.resize(image.size, Image.NEAREST) disc_mask, cup_from_mask = disc_cup_from_mask_image(mask_img) if cup_from_mask.sum() > 0: cup_mask = cup_from_mask elif disc_type == "contour": coords = self._load_contour(disc_path) disc_mask = self._contour_to_mask(coords, image.size) except Exception: disc_mask = None if cup_mask is None and cup_path and not pd.isna(cup_path): cup_path = Path(cup_path) try: if cup_type == "mask": mask_img = Image.open(cup_path) mask_img = mask_img.resize(image.size, Image.NEAREST) _, cup_mask = disc_cup_from_mask_image(mask_img) elif cup_type == "contour": coords = self._load_contour(cup_path) cup_mask = self._contour_to_mask(coords, image.size) except Exception: cup_mask = None if disc_mask is None: return None disc_mask = (disc_mask > 0).astype(np.uint8) if cup_mask is None: cup_mask = np.zeros_like(disc_mask, dtype=np.uint8) cup_mask = ((cup_mask > 0) & (disc_mask > 0)).astype(np.uint8) return disc_mask, cup_mask def _compute_crop_info(self, image: Image.Image, image_path: Path) -> Optional[dict]: image_path = Path(image_path).resolve() entry = self.entries.get(str(image_path)) if entry is None: return None cache_path = self._cache_path(image_path) cached_bounds = None if cache_path is not None and cache_path.exists(): data = np.load(cache_path, allow_pickle=False) try: cached_bounds = { "left": float(data["left"]), "upper": float(data["upper"]), "right": float(data["right"]), "lower": float(data["lower"]), } if "features" in data.files: cached_bounds["features"] = data["features"].astype(np.float32) return cached_bounds except KeyError: cached_bounds = None masks = self._load_masks(entry, image) if masks is None: return cached_bounds disc_mask, cup_mask = masks try: geom = _geometry_from_mask(disc_mask, self.scale) except Exception: return cached_bounds cx = geom["centre_x"] cy = geom["centre_y"] r = geom["crop_radius"] left = max(0.0, cx - r) upper = max(0.0, cy - r) right = min(float(image.width), cx + r) lower = min(float(image.height), cy + r) features = compute_geometry_features(disc_mask, cup_mask) info = { "left": left, "upper": upper, "right": right, "lower": lower, "features": features, } if cache_path is not None: np.savez( cache_path, left=left, upper=upper, right=right, lower=lower, width=float(image.width), height=float(image.height), scale=self.scale, target_size=self.target_size, features=features, ) return info def __call__(self, image: Image.Image, image_path: Path) -> Image.Image: info = self._compute_crop_info(image, image_path) if info is None: return image left = info["left"] upper = info["upper"] right = info["right"] lower = info["lower"] if right <= left or lower <= upper: return image crop = image.crop((left, upper, right, lower)) return crop.resize((self.target_size, self.target_size), Image.BILINEAR) def geometry_features(self, image: Image.Image, image_path: Path) -> Optional[np.ndarray]: info = self._compute_crop_info(image, image_path) if info is None: return None features = info.get("features") if features is None: return None return np.asarray(features, dtype=np.float32) def _num(x): """float(x) or None if NaN/None/invalid.""" try: v = float(x) except Exception: return None return None if math.isnan(v) else v class _ClinicalView: """Minimal shim so ClinicalDataset can iterate an epoch-specific DataFrame while still delegating encoding/paths/labels to the DataBundle object.""" def __init__(self, base: DataBundle, df): self.base = base self.df = df @property def image_dir(self): return self.base.image_dir @property def clinical_dir(self): return self.base.clinical_dir @property def id_cols(self): return ("Patient ID", "eyeID") @property def label_col(self): return self.base.label_col @property def filename_template(self): # prefer whatever the dataset defined; otherwise fall back to RET{pid}{eye}.jpg style return getattr(self.base, "filename_template", "RET{pid:03d}{eye}.jpg") @property def dim(self): return self.base.feature_dim def encode_metadata(self, row): # DataBundle returns numpy; convert to torch here to keep towers torch-only import torch as _torch vec = self.base.vectorize_row(row) return _torch.as_tensor(vec, dtype=_torch.float32) def get_image_path(self, row): return self.base.get_image_path(row) def get_label(self, row): return int(row[self.base.label_col]) class HyperTower: def __init__(self, clinical: DataBundle, args): # Expects a fully built DataBundle (add_df handled upstream) self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu") print(f"Using device: {self.device}") self.clinical = clinical sample_mode = str(getattr(args, "sample_mode", "eye") or "eye").lower() tower_warmup = getattr(args, "warmup_tower_epochs", None) fused_warmup = getattr(args, "warmup_fused_epochs", None) if tower_warmup is None: tower_warmup = 8 if sample_mode == "patient" else 2 if fused_warmup is None: fused_warmup = 5 if sample_mode == "patient" else 3 self.warmup_tower_epochs = int(tower_warmup) self.warmup_fused_epochs = int(fused_warmup) self.args = args # cache for checkpointing self._early = None self.aux_img = getattr(args, "aux_img", 0.05) self.aux_md = getattr(args, "aux_md", 0.05) self.aux_detach = getattr(args, "aux_detach", True) # SE placement configuration se_where = getattr(args, "se_where", "bridge") # 'bridge'|'tower'|'both'|'none' use_se_tower = se_where in ("tower", "both") # Optional disc-centric cropping self.image_preprocessor = None crop_manifest = getattr(args, "img_crop_manifest", None) crop_weights = getattr(args, "img_crop_weights", None) use_gt = getattr(args, "img_crop_gt", False) if crop_manifest: crop_cache = getattr(args, "img_crop_cache", Path("analysis_data/hypertower_crops")) crop_cache = Path(crop_cache) if use_gt: self.image_preprocessor = ManifestImageCropper( manifest_path=Path(crop_manifest), scale=getattr(args, "img_crop_scale", 2.5), target_size=getattr(args, "img_crop_size", 224), cache_dir=crop_cache, ) print(f"[HyperTower] GT disc cropper enabled → cache at {crop_cache}") elif crop_weights: self.image_preprocessor = UNetImageCropper( manifest_path=Path(crop_manifest), weights_path=Path(crop_weights), normalize=getattr(args, "img_crop_normalize", "per_image"), threshold=getattr(args, "img_crop_threshold", 0.5), tta=getattr(args, "img_crop_tta", False), scale=getattr(args, "img_crop_scale", 2.5), target_size=getattr(args, "img_crop_size", 224), cache_dir=crop_cache, ) print(f"[HyperTower] UNet disc cropper enabled → cache at {crop_cache}") else: print("[HyperTower] img_crop_manifest provided but no weights/gt flag; skipping cropping") self.use_geometry_features = bool(getattr(args, "img_geometry_features", False)) if self.use_geometry_features and self.image_preprocessor is None: raise ValueError("img_geometry_features requires --img-crop-manifest with either --img-crop-weights or --img-crop-gt.") self.geometry_dim = FEATURE_DIM if self.use_geometry_features else 0 # Towers derive their dimensions from DataBundle self.img_tower = ImageTower( backbone=getattr(args, "backbone", "efficientnet_b0"), freeze_ratio=getattr(args, "freeze_ratio", 0.0), use_se=use_se_tower, se_reduction=getattr(args, "se_reduction_tower", getattr(args, "se_reduction", 16)), se_pre_norm=getattr(args, "se_pre_norm_tower", getattr(args, "se_pre_norm", True)), augment=getattr(args, "img_augment", True), geometry_dim=self.geometry_dim, ).to(self.device) self.md_tower = MDTower( self.clinical, use_se=use_se_tower, se_reduction=getattr(args, "se_reduction_tower", getattr(args, "se_reduction", 16)), se_pre_norm=getattr(args, "se_pre_norm_tower", getattr(args, "se_pre_norm", True)), ).to(self.device) self.mode = args.fusion_mode # cache # Training hyperparams / objects self.batch_size = args.batch_size # Main epochs exclude warmup; total training epochs adds warmup phases. self.epochs = int(args.epochs) self.total_epochs = int(self.warmup_tower_epochs + self.warmup_fused_epochs + self.epochs) print( f"[HyperTower] Warmup schedule: tower={self.warmup_tower_epochs}, " f"fused={self.warmup_fused_epochs}, main={self.epochs}, total={self.total_epochs}" ) self.lr = args.lr self.fold = args.fold self.criterion = nn.CrossEntropyLoss() self._ce_weight = self.criterion.weight.detach().clone() if self.criterion.weight is not None else None if self._ce_weight is not None: self._ce_weight = self._ce_weight.to(self.device) self._ce_reduction = self.criterion.reduction self.focal_gamma = float(getattr(args, "focal_gamma", 0.0) or 0.0) # Run and model directories come from the script self.run_dir = Path(getattr(args, "run_dir", ".")).resolve() self.run_dir.mkdir(parents=True, exist_ok=True) # Per-run log locations to avoid collisions across concurrent workers self.epoch_log_path = self.run_dir / "epoch_log.csv" self.train_log_path = self.run_dir / "train.log" self.models_dir = Path(getattr(args, "models_dir", "models")).resolve() self.models_dir.mkdir(parents=True, exist_ok=True) self.holdout_df = getattr(args, "holdout_df", None) self.holdout_loader = None if isinstance(self.holdout_df, pd.DataFrame) and not self.holdout_df.empty: if getattr(args, "eval_mode", "multiclass") == "binary": label_col = getattr(self.clinical, "label_col", None) if label_col and label_col in self.holdout_df.columns: self.holdout_df = self.holdout_df[self.holdout_df[label_col].isin([0, 1])].reset_index(drop=True) self.holdout_loader = self._make_loader_for_df(self.holdout_df, is_train=False) print(f"[HyperTower] Holdout loader prepared with {len(self.holdout_df)} samples") if self.mode == "vote": # Per-tower classification heads (logits) + vote combiner self.head_img = nn.Linear(self.img_tower.out_dim, args.num_classes).to(self.device) self.head_md = nn.Linear(self.md_tower.out_dim, args.num_classes).to(self.device) self.vote = VoteBridge(num_classes=args.num_classes).to(self.device) # Optimizer: towers + heads + vote self.optimizer = torch.optim.Adam( list(self.img_tower.parameters()) + list(self.md_tower.parameters()) + list(self.head_img.parameters()) + list(self.head_md.parameters()) + list(self.vote.parameters()), lr=self.lr, ) else: # Existing feature-fusion bridge self.bridge = Bridge( img_dim=self.img_tower.out_dim, meta_dim=self.md_tower.out_dim, num_classes=args.num_classes, fusion_dim=256, mode=args.fusion_mode, use_se=(se_where in ("bridge","both")) and getattr(args, "use_se", True), se_reduction=(getattr(args, "se_reduction", 16) if getattr(args, "use_se", True) else 0), se_pre_norm=getattr(args, "se_pre_norm", True) ).to(self.device) self.optimizer = torch.optim.Adam( list(self.img_tower.parameters()) + list(self.md_tower.parameters()) + list(self.bridge.parameters()), lr=self.lr, ) # EMA-forgiveness + BCD knobs; logging stays simple for now self.ema_alpha = args.ema_alpha self.bcd_prob = args.bcd_prob self.ema_fused_loss = None # Shared V2 logger for train.log + epoch_log.csv self.ht_logger = HypertowerLogger( run_dir=self.run_dir, train_log_path=self.train_log_path, epoch_log_path=self.epoch_log_path, ) self.logger = self.ht_logger.logger # Dynamic BCD configuration (epoch baselines + batch nudges) self.bcd_cfg = { "metric": getattr(args, "bcd_metric", "auc"), "p0": getattr(args, "bcd_p0", 0.20), "k": getattr(args, "bcd_k", 0.4), "pmin": getattr(args, "bcd_min", 0.05), "pmax": getattr(args, "bcd_max", 0.30), "alpha_batch": getattr(args, "bcd_alpha_batch", 0.2), "alpha_tower": getattr(args, "bcd_alpha_tower", 0.3), "explore_floor": getattr(args, "bcd_explore_floor", 0.15), "entropy_ema": getattr(args, "entropy_ema", 0.7), } # Track last epoch's eval metrics for baseline deficits; initialize safely self.last_eval = { "acc_fused": 0.0, "acc_img": 0.0, "acc_md": 0.0, "auc_fused": 0.0, "auc_img": 0.0, "auc_md": 0.0, } # Running EMA of per-head uncertainty (entropy in [0,1]) self.entropy_ema = {"fused": 0.5, "img": 0.5, "md": 0.5} # DataLoaders are (re)built each epoch from DataBundle's splits self.train_loader = None self.test_loader = None self._last_step_mix = None if not getattr(self.clinical, "folds", None): raise RuntimeError("DataBundle has no built folds. Did you call add_df(...) upstream?") @staticmethod def _confusion(preds, labels): # Quick TP/TN/FP/FN for binary debug (kept as-is) tp = ((preds == 1) & (labels == 1)).sum().item() tn = ((preds == 0) & (labels == 0)).sum().item() fp = ((preds == 1) & (labels == 0)).sum().item() fn = ((preds == 0) & (labels == 1)).sum().item() return {"tp": tp, "tn": tn, "fp": fp, "fn": fn} def _unpack_batch(self, batch): if len(batch) == 4: imgs, metas, geometry, labels = batch else: imgs, metas, labels = batch if self.geometry_dim > 0: geometry = torch.zeros(imgs.size(0), self.geometry_dim, dtype=torch.float32) else: geometry = None return imgs, metas, geometry, labels def _classification_loss(self, logits: torch.Tensor, labels: torch.Tensor) -> torch.Tensor: """ Shared classification loss (CrossEntropy or focal depending on gamma). """ weight = self._ce_weight if weight is not None and weight.device != logits.device: weight = weight.to(logits.device) return focal_loss( logits, labels, gamma=self.focal_gamma, weight=weight, reduction=self._ce_reduction, ) def _step_image_only(self, imgs, metas, geometry, labels): if self.mode == "vote": img_feats = self.img_tower(imgs, geometry) out_img = self.head_img(img_feats) loss_i = self._classification_loss(out_img, labels) preds_i = out_img.argmax(dim=1) acc_i = (preds_i == labels).float().mean().item() cm = self._confusion(preds_i, labels) self.optimizer.zero_grad() loss_i.backward() self.optimizer.step() return None, loss_i.item(), None, None, acc_i, None, {"img": cm} else:# Optimize the image head alone (used when BCD chooses image-only) img_feats = self.img_tower(imgs, geometry) out_img = self.bridge.classifier_img(img_feats) loss_i = self._classification_loss(out_img, labels) preds_i = out_img.argmax(dim=1) acc_i = (preds_i == labels).float().mean().item() cm = self._confusion(preds_i, labels) self.optimizer.zero_grad() loss_i.backward() self.optimizer.step() return None, loss_i.item(), None, None, acc_i, None, {"img": cm} def _step_meta_only(self, imgs, metas, geometry, labels): # Optimize the metadata head alone (used when BCD chooses meta-only) if self.mode == "vote": meta_feats = self.md_tower(metas) out_md = self.head_md(meta_feats) loss_m = self._classification_loss(out_md, labels) preds_m = out_md.argmax(dim=1) acc_m = (preds_m == labels).float().mean().item() cm = self._confusion(preds_m, labels) self.optimizer.zero_grad() loss_m.backward() self.optimizer.step() return None, None, loss_m.item(), None, None, acc_m, {"meta": cm} else: meta_feats = self.md_tower(metas) out_md = self.bridge.classifier_md(meta_feats) loss_m = self._classification_loss(out_md, labels) preds_m = out_md.argmax(dim=1) acc_m = (preds_m == labels).float().mean().item() cm = self._confusion(preds_m, labels) self.optimizer.zero_grad() loss_m.backward() self.optimizer.step() return None, None, loss_m.item(), None, None, acc_m, {"meta": cm} def _step_fused(self, imgs, metas, geometry, labels): if self.mode == "vote": img_feats = self.img_tower(imgs, geometry) md_feats = self.md_tower(metas) out_img = self.head_img(img_feats) out_md = self.head_md(md_feats) out_fused = self.vote(out_img, out_md) preds_f = out_fused.argmax(dim=1) acc_f = (preds_f == labels).float().mean().item() loss_f = self._classification_loss(out_fused, labels) loss_f_val = float(loss_f.detach().item()) # Optional: keep your EMA + small aux tower losses just like before if self.ema_fused_loss is None: self.ema_fused_loss = loss_f_val L_total = (1.0 - self.ema_alpha) * loss_f + self.ema_alpha * loss_f.new_tensor(self.ema_fused_loss) aux_terms = 0.0 if self.aux_img > 0: aux_terms = aux_terms + self.aux_img * self._classification_loss( out_img.detach() if self.aux_detach else out_img, labels, ) if self.aux_md > 0: aux_terms = aux_terms + self.aux_md * self._classification_loss( out_md.detach() if self.aux_detach else out_md, labels, ) L_total = L_total + aux_terms self.optimizer.zero_grad() L_total.backward() self.optimizer.step() self.ema_fused_loss = self.ema_alpha * self.ema_fused_loss + (1.0 - self.ema_alpha) * loss_f_val # Confusion tables for logging cm = { "fused": self._confusion(preds_f, labels), "img": self._confusion(out_img.argmax(dim=1), labels), "meta": self._confusion(out_md.argmax(dim=1), labels), } acc_i = (out_img.argmax(dim=1) == labels).float().mean().item() acc_m = (out_md.argmax(dim=1) == labels).float().mean().item() with torch.no_grad(): loss_img_val = self._classification_loss(out_img, labels).item() loss_md_val = self._classification_loss(out_md, labels).item() return loss_f_val, loss_img_val, loss_md_val, acc_f, acc_i, acc_m, cm else: img_feats = self.img_tower(imgs, geometry) meta_feats = self.md_tower(metas) out_fused, out_img, out_md = self.bridge(img_feats, meta_feats) preds_f = out_fused.argmax(dim=1) acc_f = (preds_f == labels).float().mean().item() loss_f = self._classification_loss(out_fused, labels) loss_f_val = float(loss_f.detach().item()) cm = {"fused": self._confusion(preds_f, labels)} loss_i = acc_i = loss_m = acc_m = None if out_img is not None: with torch.no_grad(): preds_i = out_img.argmax(dim=1) acc_i = (preds_i == labels).float().mean().item() loss_i = self._classification_loss(out_img, labels).item() cm["img"] = self._confusion(preds_i, labels) if out_md is not None: with torch.no_grad(): preds_m = out_md.argmax(dim=1) acc_m = (preds_m == labels).float().mean().item() loss_m = self._classification_loss(out_md, labels).item() cm["meta"] = self._confusion(preds_m, labels) # EMA-blended fused loss if self.ema_fused_loss is None: self.ema_fused_loss = loss_f_val L_total = (1.0 - self.ema_alpha) * loss_f + self.ema_alpha * loss_f.new_tensor(self.ema_fused_loss) # Auxiliary tower losses (small weights). Use detached features to calibrate heads only. aux_terms = 0.0 if self.aux_img > 0 and out_img is not None: logits_img_for_aux = self.bridge.classifier_img(img_feats.detach()) if self.aux_detach else out_img aux_terms = aux_terms + self.aux_img * self._classification_loss(logits_img_for_aux, labels) if self.aux_md > 0 and out_md is not None: logits_md_for_aux = self.bridge.classifier_md(meta_feats.detach()) if self.aux_detach else out_md aux_terms = aux_terms + self.aux_md * self._classification_loss(logits_md_for_aux, labels) L_total = L_total + aux_terms self.optimizer.zero_grad() L_total.backward() self.optimizer.step() # update EMA after the step self.ema_fused_loss = self.ema_alpha * self.ema_fused_loss + (1.0 - self.ema_alpha) * loss_f_val return loss_f_val, loss_i, loss_m, acc_f, acc_i, acc_m, cm def _make_loader_for_df(self, df, is_train: bool = False): # Rebuild a DataLoader for the current epoch's split view = _ClinicalView(self.clinical, df) geometry_provider = self.image_preprocessor if self.geometry_dim > 0 else None ds = ClinicalDataset( view, self.img_tower.transform, image_preprocessor=self.image_preprocessor, geometry_provider=geometry_provider, geometry_dim=self.geometry_dim, ) # Optional: class-balanced bootstrapped sampling for TRAIN only if is_train and getattr(self.args, "balanced_sampler", False): import numpy as _np y = _np.asarray(df[self.clinical.label_col].values) # inverse-frequency weights per class uniq, counts = _np.unique(y, return_counts=True) inv = {c: (1.0 / cnt if cnt > 0 else 0.0) for c, cnt in zip(uniq, counts)} w = _np.array([inv[c] for c in y], dtype=_np.float32) sampler = WeightedRandomSampler(weights=w, num_samples=len(y), replacement=True) return DataLoader(ds, batch_size=self.batch_size, sampler=sampler, shuffle=False) return DataLoader(ds, batch_size=self.batch_size, shuffle=not is_train) # ---- Dynamic BCD helpers ---- def _metric_value(self, name: str) -> float: # Pull either AUC or ACC from last_eval, falling back if NaN/zero if self.bcd_cfg["metric"] == "auc": v = self.last_eval.get(f"auc_{name}", 0.0) if v == v: # not NaN return float(v) # fallback to accuracy return float(self.last_eval.get(f"acc_{name}", 0.0)) return float(self.last_eval.get(f"acc_{name}", 0.0)) def _epoch_bcd_baseline(self): # Compute epoch-level baselines: p_bcd_epoch and tower weights w_i, w_m p0 = self.bcd_cfg["p0"]; k = self.bcd_cfg["k"] pmin = self.bcd_cfg["pmin"]; pmax = self.bcd_cfg["pmax"] eps = 1e-6 Af = self._metric_value("fused"); Ai = self._metric_value("img"); Am = self._metric_value("md") di = max(0.0, Af - Ai); dm = max(0.0, Af - Am) p_bcd_epoch = max(pmin, min(pmax, p0 + k * (di + dm) / 2.0)) wi = (di + eps) / (di + dm + 2 * eps) wm = 1.0 - wi return p_bcd_epoch, wi, wm, {"di": di, "dm": dm, "Af": Af, "Ai": Ai, "Am": Am} @staticmethod def _entropy_from_logits(logits: torch.Tensor, num_classes: int) -> float: # Returns entropy normalized to [0,1] using log(K) denominator with torch.no_grad(): probs = F.softmax(logits, dim=1) ent = -(probs * (probs.clamp_min(1e-12)).log()).sum(dim=1) ent = ent / np.log(num_classes) return float(ent.mean().item()) def _batch_uncertainty(self, imgs: torch.Tensor, metas: torch.Tensor, geometry: Optional[torch.Tensor] = None) -> dict: self.img_tower.eval(); self.md_tower.eval() if self.mode == "vote": self.head_img.eval(); self.head_md.eval(); self.vote.eval() with torch.no_grad(): if geometry is None and self.geometry_dim > 0: geometry = torch.zeros(imgs.size(0), self.geometry_dim, device=imgs.device, dtype=imgs.dtype) img_feats = self.img_tower(imgs, geometry) md_feats = self.md_tower(metas) out_img = self.head_img(img_feats) out_md = self.head_md(md_feats) out_fused = self.vote(out_img, out_md) K = out_fused.shape[1] e_f = self._entropy_from_logits(out_fused, K) e_i = self._entropy_from_logits(out_img, K) e_m = self._entropy_from_logits(out_md, K) # restore train() self.img_tower.train(); self.md_tower.train() self.head_img.train(); self.head_md.train(); self.vote.train() else: self.bridge.eval() with torch.no_grad(): if geometry is None and self.geometry_dim > 0: geometry = torch.zeros(imgs.size(0), self.geometry_dim, device=imgs.device, dtype=imgs.dtype) img_feats = self.img_tower(imgs, geometry) meta_feats = self.md_tower(metas) out_fused, out_img, out_md = self.bridge(img_feats, meta_feats) K = out_fused.shape[1] e_f = self._entropy_from_logits(out_fused, K) e_i = self._entropy_from_logits(out_img, K) if out_img is not None else 0.5 e_m = self._entropy_from_logits(out_md, K) if out_md is not None else 0.5 self.img_tower.train(); self.md_tower.train(); self.bridge.train() a = self.bcd_cfg["entropy_ema"] self.entropy_ema["fused"] = a * self.entropy_ema["fused"] + (1 - a) * e_f self.entropy_ema["img"] = a * self.entropy_ema["img"] + (1 - a) * e_i self.entropy_ema["md"] = a * self.entropy_ema["md"] + (1 - a) * e_m return {"fused": self.entropy_ema["fused"], "img": self.entropy_ema["img"], "md": self.entropy_ema["md"]} def _write_epoch_log(self, row: dict, path: str | Path | None = None): self.ht_logger.write_epoch_row(row=row, path=path) def _snapshot(self): state = { "img_tower": self.img_tower.state_dict(), "md_tower": self.md_tower.state_dict(), "optimizer": self.optimizer.state_dict(), "args": vars(self.args), } if hasattr(self, "bridge"): state["bridge"] = self.bridge.state_dict() if hasattr(self, "head_img"): state["head_img"] = self.head_img.state_dict() if hasattr(self, "head_md"): state["head_md"] = self.head_md.state_dict() return state def train(self): best_path = None if getattr(self.args, "checkpoint_best", False): best_path = str((self.models_dir / "best.pth").resolve()) # Always track the best model even if we don't stop early monitor = getattr(self.args, "early_metric", None) if not monitor: # Prefer AUC-based monitoring by default; fall back to loss only if specified if self.mode == "image_only": monitor = "auc_img" elif self.mode == "metadata_only": monitor = "auc_md" else: # fused or vote monitor = "auc_fused" mode = getattr(self.args, "early_mode", "auto") # Optional switch: monitor holdout metrics instead of validation use_holdout_monitor = bool(getattr(self.args, "early_monitor_holdout", False)) if use_holdout_monitor and self.holdout_loader is None: print("[early] Requested holdout monitoring but no holdout set is configured; falling back to validation metrics.") use_holdout_monitor = False if use_holdout_monitor: # If the monitor isn't already a holdout metric, prepend it if not monitor.startswith("holdout_"): monitor = f"holdout_{monitor}" # Force sensible default if no monitor was provided if monitor == "holdout_auc_fused" and mode == "auto": mode = "max" if mode == "auto": mode = ("min" if ("loss" in monitor.lower()) else "max") self._best_info = { "monitor": monitor, "mode": mode, "min_delta": float(getattr(self.args, "early_min_delta", 0.0)), "value": (-math.inf if mode == "max" else math.inf), "epoch": -1, "state": None, "save_path": best_path, } # keep args in sync so EarlyStopper uses the same monitor/mode self.args.early_metric = monitor self.args.early_mode = mode # Separate tracker for best holdout performance (if a holdout set is provided) self._holdout_best = None if self.holdout_loader is not None: self._holdout_best = { "monitor": "holdout_auc_fused", "mode": "max", "min_delta": 0.0, "value": -math.inf, "epoch": -1, "state": None, "save_path": str((self.models_dir / "model_holdout_best.pt").resolve()), } self._early = None if getattr(self.args, "early_stop", False): self._early = EarlyStopper( monitor=monitor, mode=mode, patience=self.args.early_patience, min_delta=self.args.early_min_delta, save_path=best_path, restore_best=True, ) # Gradual thaw config gradual = bool(getattr(self.args, "gradual_thaw", False)) thaw_ratio = float(getattr(self.args, "thaw_ratio", 0.33)) thaw_phase = int(getattr(self.args, "thaw_phase_duration", 5)) thaw_target = str(getattr(self.args, "thaw_target", "image")) thaw_start_epoch = int(getattr(self.args, "thaw_start_epoch", -1)) thaw_initial_freeze = bool(getattr(self.args, "thaw_initial_freeze", False)) if thaw_start_epoch < 0: thaw_start_epoch = int(getattr(self, "warmup_tower_epochs", 0)) for epoch in range(self.total_epochs): # Apply thaw schedule at epoch start: start fully frozen and unfreeze by ratios if gradual and thaw_phase > 0: try: if epoch < thaw_start_epoch: if thaw_initial_freeze: freeze_r = 1.0 if thaw_target in ("image", "both") and hasattr(self, "img_tower"): self.img_tower.set_freeze_ratio(freeze_r) if thaw_target in ("metadata", "both") and hasattr(self, "md_tower"): self.md_tower.set_freeze_ratio(freeze_r) # else: leave current requires_grad as constructed (no forced freeze) else: phase_idx = (epoch - thaw_start_epoch) // thaw_phase freeze_r = max(0.0, 1.0 - phase_idx * thaw_ratio) if thaw_target in ("image", "both") and hasattr(self, "img_tower"): self.img_tower.set_freeze_ratio(freeze_r) if thaw_target in ("metadata", "both") and hasattr(self, "md_tower"): self.md_tower.set_freeze_ratio(freeze_r) except Exception: pass # build splits/loaders per-epoch train_df, test_df = self.clinical.get_split_dfs(self.fold) if hasattr(self.bridge, "reset_se_stats"): self.bridge.reset_se_stats() if getattr(self, "args", None) is None: self.args = argparse.Namespace() # if not stashed already # Expect user to pass --eval_mode and --num-classes==2 for binary if getattr(self.args, "eval_mode", "multiclass") == "binary": # Drop Suspect (assumed label==2 in PAPILA spreadsheets) train_df = train_df[train_df[self.clinical.label_col].isin([0,1])].copy() test_df = test_df [test_df [self.clinical.label_col].isin([0,1])].copy() self.train_loader = self._make_loader_for_df(train_df, is_train=True) self.test_loader = self._make_loader_for_df(test_df, is_train=False) self.img_tower.train(); self.md_tower.train(); self.bridge.train() total_losses = {"fused":0.0,"image_only":0.0,"metadata_only":0.0} total_accs = {"fused":0.0,"image_only":0.0,"metadata_only":0.0} counts = {"fused":0, "image_only":0, "metadata_only":0} step_counts = {"fused":0, "image_only":0, "metadata_only":0} # phase selection in_tower_warmup = epoch < self.warmup_tower_epochs in_fused_warmup = (self.warmup_tower_epochs <= epoch < (self.warmup_tower_epochs + self.warmup_fused_epochs)) # epoch baselines (needed for splits and logging) p_bcd_epoch, wi_epoch, wm_epoch, diag = self._epoch_bcd_baseline() if in_tower_warmup: # tower-only warmup: alternate or use epoch split; default to 0.5 if empty wi = wi_epoch if (diag["Ai"] or diag["Am"]) else 0.5 for batch_idx, batch in enumerate(self.train_loader): imgs, metas, geometry, labels = self._unpack_batch(batch) imgs = imgs.to(self.device) metas = metas.to(self.device) labels = labels.to(self.device) geometry = geometry.to(self.device) if geometry is not None else None if random() < wi: step_type = "image" loss_f, loss_i, loss_m, acc_f, acc_i, acc_m, cm = self._step_image_only(imgs, metas, geometry, labels) step_counts["image_only"] += 1 else: step_type = "meta" loss_f, loss_i, loss_m, acc_f, acc_i, acc_m, cm = self._step_meta_only(imgs, metas, geometry, labels) step_counts["metadata_only"] += 1 self.logger.info(f"epoch={epoch} batch={batch_idx} phase=tower_warmup step={step_type} " f"loss_f={loss_f} loss_i={loss_i} loss_m={loss_m} acc_f={acc_f} acc_i={acc_i} acc_m={acc_m} cm={cm}") if loss_i is not None: total_losses["image_only"] += loss_i; total_accs["image_only"] += acc_i; counts["image_only"] += 1 if loss_m is not None: total_losses["metadata_only"] += loss_m; total_accs["metadata_only"] += acc_m; counts["metadata_only"] += 1 elif in_fused_warmup: # fused-only warmup for batch_idx, batch in enumerate(self.train_loader): imgs, metas, geometry, labels = self._unpack_batch(batch) imgs = imgs.to(self.device) metas = metas.to(self.device) labels = labels.to(self.device) geometry = geometry.to(self.device) if geometry is not None else None step_type = "fused" loss_f, loss_i, loss_m, acc_f, acc_i, acc_m, cm = self._step_fused(imgs, metas, geometry, labels) step_counts["fused"] += 1 self.logger.info(f"epoch={epoch} batch={batch_idx} phase=fused_warmup step={step_type} " f"loss_f={loss_f} acc_f={acc_f} cm={cm}") total_losses["fused"] += loss_f; total_accs["fused"] += acc_f; counts["fused"] += 1 else: # dynamic BCD (fusion-centric) self.logger.info( f"epoch={epoch} bcd_epoch={p_bcd_epoch:.3f} wi_epoch={wi_epoch:.3f} wm_epoch={wm_epoch:.3f} " f"deficits={{img:{diag['di']:.3f}, md:{diag['dm']:.3f}}} metrics={{Af:{diag['Af']:.3f}, Ai:{diag['Ai']:.3f}, Am:{diag['Am']:.3f}}}" ) for batch_idx, batch in enumerate(self.train_loader): imgs, metas, geometry, labels = self._unpack_batch(batch) imgs = imgs.to(self.device) metas = metas.to(self.device) labels = labels.to(self.device) geometry = geometry.to(self.device) if geometry is not None else None ents = self._batch_uncertainty(imgs, metas, geometry) # flip rule: higher fused entropy => fewer tower-only batches => more fused training alpha_b = self.bcd_cfg["alpha_batch"]; pmin = self.bcd_cfg["pmin"]; pmax = self.bcd_cfg["pmax"] p_bcd_batch = (1 - alpha_b) * p_bcd_epoch + alpha_b * (1.0 - ents["fused"]) p_bcd_batch = max(pmin, min(pmax, p_bcd_batch)) # image vs meta split alpha_t = self.bcd_cfg["alpha_tower"]; floor = self.bcd_cfg["explore_floor"]; eps = 1e-6 s_img = ents["img"] / (ents["img"] + ents["md"] + eps) p_img_batch = (1 - alpha_t) * wi_epoch + alpha_t * s_img p_img_batch = max(floor, min(1.0 - floor, p_img_batch)) u = random(); v = random() if u < p_bcd_batch: if v < p_img_batch: step_type = "image" loss_f, loss_i, loss_m, acc_f, acc_i, acc_m, cm = self._step_image_only(imgs, metas, geometry, labels) step_counts["image_only"] += 1 else: step_type = "meta" loss_f, loss_i, loss_m, acc_f, acc_i, acc_m, cm = self._step_meta_only(imgs, metas, geometry, labels) step_counts["metadata_only"] += 1 else: step_type = "fused" loss_f, loss_i, loss_m, acc_f, acc_i, acc_m, cm = self._step_fused(imgs, metas, geometry, labels) step_counts["fused"] += 1 self.logger.info( f"epoch={epoch} batch={batch_idx} phase=dynamic step={step_type} " f"p_bcd_batch={p_bcd_batch:.3f} p_img_batch={p_img_batch:.3f} " f"ents={{f:{ents['fused']:.3f}, i:{ents['img']:.3f}, m:{ents['md']:.3f}}} " f"loss_f={loss_f} loss_i={loss_i} loss_m={loss_m} acc_f={acc_f} acc_i={acc_i} acc_m={acc_m} cm={cm}" ) if loss_f is not None: total_losses["fused"] += loss_f; total_accs["fused"] += acc_f; counts["fused"] += 1 if loss_i is not None: total_losses["image_only"] += loss_i; total_accs["image_only"] += acc_i; counts["image_only"] += 1 if loss_m is not None: total_losses["metadata_only"] += loss_m; total_accs["metadata_only"] += acc_m; counts["metadata_only"] += 1 # epoch prints avg_loss = total_losses["fused"] / counts["fused"] if counts["fused"] else 0.0 avg_fused = total_accs["fused"] / counts["fused"] if counts["fused"] else 0.0 avg_img = total_accs["image_only"] / counts["image_only"] if counts["image_only"] else 0.0 avg_md = total_accs["metadata_only"] / counts["metadata_only"] if counts["metadata_only"] else 0.0 # realized mix for logging tot_steps = sum(step_counts.values()) or 1 self._last_step_mix = { "pct_fused": step_counts["fused"] / tot_steps, "pct_img": step_counts["image_only"] / tot_steps, "pct_md": step_counts["metadata_only"] / tot_steps, "phase": "tower_warmup" if in_tower_warmup else ("fused_warmup" if in_fused_warmup else "dynamic") } print( f"[Epoch {epoch+1}/{self.total_epochs}] Train Loss: {avg_loss:.4f} | " f"Fused Acc: {100*avg_fused:.2f}% | Img Acc: {100*avg_img:.2f}% | Md Acc: {100*avg_md:.2f}%" ) row, stop_now = self.evaluate(epoch) if stop_now: print(f"[early] stopping on '{self.args.early_metric}' " f"with best={self._early.best:.5f} at epoch {epoch+1 - self._early.bad_epochs}") break torch.save(self._snapshot(), str(self.models_dir / "model_last.pt")) # Ensure holdout-best checkpoint is written if getattr(self, "_holdout_best", None): hb = self._holdout_best if hb.get("state") is not None and hb.get("save_path"): try: torch.save(hb["state"], hb["save_path"]) except Exception: pass def _copy_best_roc(epoch_idx: int, tag: str): if epoch_idx is None or epoch_idx < 0: return src_dir = self.run_dir / "roc_curves" if not src_dir.exists(): return prefix = f"epoch{epoch_idx + 1}_" dest_dir = self.run_dir / f"roc_curves_{tag}" dest_dir.mkdir(parents=True, exist_ok=True) for path in src_dir.glob(f"{prefix}*.json"): try: shutil.copy2(path, dest_dir / path.name) except Exception: pass if getattr(self, "_early", None): print(f"[early] restoring best model at epoch {self._early.best_epoch+1}: {self._early.best:.5f}") self._early.restore(self) elif getattr(self, "_best_info", None) and self._best_info.get("state") is not None: be = self._best_info print(f"[best] restoring best model at epoch {be['epoch']+1}: {be['value']:.5f} (monitor={be['monitor']})") self._restore_from_state(be["state"]) torch.save(self._snapshot(), str(self.models_dir / "model_best.pt")) # Snapshot ROC curves for both main-best and holdout-best epochs if getattr(self, "_early", None) and getattr(self._early, "best_epoch", None) is not None: _copy_best_roc(self._early.best_epoch, "best") elif getattr(self, "_best_info", None): _copy_best_roc(self._best_info.get("epoch", -1), "best") if getattr(self, "_holdout_best", None): _copy_best_roc(self._holdout_best.get("epoch", -1), "holdout_best") self.ht_logger.close() def _evaluate_split(self, loader, epoch: int, split: str): """ Shared evaluation helper that computes core metrics (loss/accuracy/AUC/etc.) for a given dataloader. Returns a dict with scalar metrics or None if the loader is empty. """ if loader is None: return None total_loss = 0.0 total_counts = {"fused": 0, "image_only": 0, "metadata_only": 0} total_accs = {"fused": 0.0, "image_only": 0.0, "metadata_only": 0.0} all_y = [] probs_f_list, probs_i_list, probs_m_list = [], [], [] agree_f_img = 0 agree_f_md = 0 n_img_comp = 0 n_md_comp = 0 with torch.no_grad(): for batch in loader: imgs, metas, geometry, labels = self._unpack_batch(batch) imgs = imgs.to(self.device) metas = metas.to(self.device) labels = labels.to(self.device) geometry = geometry.to(self.device) if geometry is not None else None if self.mode == "vote": img_feats = self.img_tower(imgs, geometry) md_feats = self.md_tower(metas) out_img = self.head_img(img_feats) out_md = self.head_md(md_feats) out_fused = self.vote(out_img, out_md) else: img_feats = self.img_tower(imgs, geometry) md_feats = self.md_tower(metas) outputs = self.bridge(img_feats, md_feats) if isinstance(outputs, tuple): out_fused, out_img, out_md = outputs else: out_fused, out_img, out_md = outputs, None, None all_y.append(labels.cpu().numpy()) probs_f_list.append(F.softmax(out_fused, dim=1).cpu().numpy()) loss = self._classification_loss(out_fused, labels) total_loss += loss.item() preds_f = out_fused.argmax(dim=1) total_accs["fused"] += (preds_f == labels).float().mean().item() total_counts["fused"] += 1 if out_img is not None: preds_i = out_img.argmax(dim=1) total_accs["image_only"] += (preds_i == labels).float().mean().item() total_counts["image_only"] += 1 probs_i_list.append(F.softmax(out_img, dim=1).cpu().numpy()) agree_f_img += (preds_f == preds_i).sum().item() n_img_comp += preds_f.numel() if out_md is not None: preds_m = out_md.argmax(dim=1) total_accs["metadata_only"] += (preds_m == labels).float().mean().item() total_counts["metadata_only"] += 1 probs_m_list.append(F.softmax(out_md, dim=1).cpu().numpy()) agree_f_md += (preds_f == preds_m).sum().item() n_md_comp += preds_f.numel() if total_counts["fused"] == 0: return None avg_loss = total_loss / total_counts["fused"] avg_fused = total_accs["fused"] / total_counts["fused"] avg_img = (total_accs["image_only"] / total_counts["image_only"]) if total_counts["image_only"] else 0.0 avg_md = (total_accs["metadata_only"] / total_counts["metadata_only"]) if total_counts["metadata_only"] else 0.0 y_true = np.concatenate(all_y, axis=0) if all_y else np.array([]) y_prob_f = np.concatenate(probs_f_list, axis=0) if probs_f_list else None y_prob_i = np.concatenate(probs_i_list, axis=0) if probs_i_list else None y_prob_m = np.concatenate(probs_m_list, axis=0) if probs_m_list else None def compute_multiclass_roc(y_true_np, prob_np): if prob_np is None or prob_np.size == 0: return None C = int(prob_np.shape[1]) classes = list(range(C)) y_bin = label_binarize(y_true_np, classes=classes) per_class = {} aucs = [] for c in range(C): try: fpr, tpr, _ = roc_curve(y_bin[:, c], prob_np[:, c]) auc_val = float(auc(fpr, tpr)) if len(fpr) > 1 else float("nan") per_class[c] = {"fpr": fpr.tolist(), "tpr": tpr.tolist(), "auc": auc_val} aucs.append(auc_val) except Exception: per_class[c] = {"fpr": [0.0, 1.0], "tpr": [0.0, 1.0], "auc": float("nan")} try: fpr_micro, tpr_micro, _ = roc_curve(y_bin.ravel(), prob_np[:, :C].ravel()) auc_micro = float(auc(fpr_micro, tpr_micro)) micro = {"fpr": fpr_micro.tolist(), "tpr": tpr_micro.tolist(), "auc": auc_micro} except Exception: micro = None auc_vals = [v["auc"] for v in per_class.values() if v["auc"] == v["auc"]] macro_auc = float(np.mean(auc_vals)) if auc_vals else float("nan") return {"per_class": per_class, "micro": micro, "macro_auc": macro_auc} def macro_ovr_auc(y, p): try: y = np.asarray(y) if p is None: return float("nan") p = np.asarray(p) if p.ndim == 2 and p.shape[1] == 2: return roc_auc_score(y, p[:, 1]) if p.ndim == 1 or p.shape[1] == 1: return roc_auc_score(y, p.ravel()) return roc_auc_score(y, p, multi_class="ovr", average="macro") except Exception: return float("nan") def top2_acc(y, p): if p is None: return float("nan") if p.ndim != 2: return float("nan") k = 2 if p.shape[1] >= 2 else 1 topk = np.argpartition(-p, kth=k - 1, axis=1)[:, :k] return np.mean((topk == y[:, None]).any(axis=1).astype(np.float32)) def avg_margin(p): if p is None or p.ndim != 2: return float("nan") s = np.sort(p, axis=1)[:, ::-1] if s.shape[1] == 1: return float("nan") return float(np.mean(s[:, 0] - s[:, 1])) roc_f = compute_multiclass_roc(y_true, y_prob_f) roc_i = compute_multiclass_roc(y_true, y_prob_i) roc_m = compute_multiclass_roc(y_true, y_prob_m) roc_dir = self.run_dir / "roc_curves" roc_dir.mkdir(parents=True, exist_ok=True) def dump_roc(blob, head: str): if blob is None: return suffix = head if split == "val" else f"{split}_{head}" path = roc_dir / f"epoch{epoch + 1}_{suffix}.json" with open(path, "w") as f: json.dump(blob, f) dump_roc(roc_f, "fused") dump_roc(roc_i, "image") dump_roc(roc_m, "metadata") auc_f = macro_ovr_auc(y_true, y_prob_f) if y_prob_f is not None else float("nan") auc_i = macro_ovr_auc(y_true, y_prob_i) if y_prob_i is not None else float("nan") auc_m = macro_ovr_auc(y_true, y_prob_m) if y_prob_m is not None else float("nan") top2_f = top2_acc(y_true, y_prob_f) top2_i = top2_acc(y_true, y_prob_i) top2_m = top2_acc(y_true, y_prob_m) mar_f = avg_margin(y_prob_f) mar_i = avg_margin(y_prob_i) mar_m = avg_margin(y_prob_m) agree_rate_f_img = (agree_f_img / n_img_comp) if n_img_comp else float("nan") agree_rate_f_md = (agree_f_md / n_md_comp) if n_md_comp else float("nan") return { "loss": float(avg_loss), "acc_fused": float(avg_fused), "acc_img": float(avg_img), "acc_md": float(avg_md), "auc_fused": float(auc_f) if auc_f == auc_f else float("nan"), "auc_img": float(auc_i) if auc_i == auc_i else float("nan"), "auc_md": float(auc_m) if auc_m == auc_m else float("nan"), "top2_fused": float(top2_f) if top2_f == top2_f else float("nan"), "top2_img": float(top2_i) if top2_i == top2_i else float("nan"), "top2_md": float(top2_m) if top2_m == top2_m else float("nan"), "margin_fused": float(mar_f) if mar_f == mar_f else float("nan"), "margin_img": float(mar_i) if mar_i == mar_i else float("nan"), "margin_md": float(mar_m) if mar_m == mar_m else float("nan"), "agree_fused_img": float(agree_rate_f_img) if agree_rate_f_img == agree_rate_f_img else float("nan"), "agree_fused_md": float(agree_rate_f_md) if agree_rate_f_md == agree_rate_f_md else float("nan"), } def evaluate(self, epoch: int): # Evaluation additionally collects per-head probabilities to report macro AUC self.img_tower.eval() self.md_tower.eval() if self.mode == "vote": self.head_img.eval(); self.head_md.eval(); self.vote.eval() else: self.bridge.eval() metrics_val = self._evaluate_split(self.test_loader, epoch, "val") metrics_holdout = self._evaluate_split(self.holdout_loader, epoch, "holdout") if self.holdout_loader else None if metrics_val is None: raise RuntimeError("Validation loader produced no batches; cannot compute metrics.") def ffmt(x): return "NA" if (x != x) else f"{x:.3f}" print( f"[Epoch {epoch+1}/{self.total_epochs}] Eval Loss: {metrics_val['loss']:.4f} | " f"Fused Acc: {100*metrics_val['acc_fused']:.2f}% | " f"Img Acc: {100*metrics_val['acc_img']:.2f}% | " f"Md Acc: {100*metrics_val['acc_md']:.2f}% | " f"Fused AUC: {ffmt(metrics_val['auc_fused'])} | " f"Img AUC: {ffmt(metrics_val['auc_img'])} | " f"Md AUC: {ffmt(metrics_val['auc_md'])}" ) if metrics_holdout: print( f"[Epoch {epoch+1}/{self.total_epochs}] Holdout Loss: {metrics_holdout['loss']:.4f} | " f"Fused Acc: {100*metrics_holdout['acc_fused']:.2f}% | " f"Img Acc: {100*metrics_holdout['acc_img']:.2f}% | " f"Md Acc: {100*metrics_holdout['acc_md']:.2f}% | " f"Fused AUC: {ffmt(metrics_holdout['auc_fused'])} | " f"Img AUC: {ffmt(metrics_holdout['auc_img'])} | " f"Md AUC: {ffmt(metrics_holdout['auc_md'])}" ) self.last_eval = { "acc_fused": float(metrics_val["acc_fused"]), "acc_img": float(metrics_val["acc_img"]), "acc_md": float(metrics_val["acc_md"]), "auc_fused": float(metrics_val["auc_fused"]) if metrics_val["auc_fused"] == metrics_val["auc_fused"] else 0.0, "auc_img": float(metrics_val["auc_img"]) if metrics_val["auc_img"] == metrics_val["auc_img"] else 0.0, "auc_md": float(metrics_val["auc_md"]) if metrics_val["auc_md"] == metrics_val["auc_md"] else 0.0, } row = { "epoch": int(epoch + 1), "eval_loss": _num(metrics_val["loss"]), "acc_fused": _num(metrics_val["acc_fused"]), "acc_img": _num(metrics_val["acc_img"]), "acc_md": _num(metrics_val["acc_md"]), "auc_fused": _num(metrics_val["auc_fused"]), "auc_img": _num(metrics_val["auc_img"]), "auc_md": _num(metrics_val["auc_md"]), "top2_fused": _num(metrics_val["top2_fused"]), "top2_img": _num(metrics_val["top2_img"]), "top2_md": _num(metrics_val["top2_md"]), "margin_fused": _num(metrics_val["margin_fused"]), "margin_img": _num(metrics_val["margin_img"]), "margin_md": _num(metrics_val["margin_md"]), "agree_fused_img": _num(metrics_val["agree_fused_img"]), "agree_fused_md": _num(metrics_val["agree_fused_md"]), } if metrics_holdout: row.update({ "holdout_loss": _num(metrics_holdout["loss"]), "holdout_acc_fused": _num(metrics_holdout["acc_fused"]), "holdout_acc_img": _num(metrics_holdout["acc_img"]), "holdout_acc_md": _num(metrics_holdout["acc_md"]), "holdout_auc_fused": _num(metrics_holdout["auc_fused"]), "holdout_auc_img": _num(metrics_holdout["auc_img"]), "holdout_auc_md": _num(metrics_holdout["auc_md"]), }) else: row.setdefault("holdout_loss", None) row.setdefault("holdout_acc_fused", None) row.setdefault("holdout_acc_img", None) row.setdefault("holdout_acc_md", None) row.setdefault("holdout_auc_fused", None) row.setdefault("holdout_auc_img", None) row.setdefault("holdout_auc_md", None) # realized step mix from training (if present) mix = getattr(self, "_last_step_mix", None) if mix: row.update({ "pct_fused": _num(mix.get("pct_fused")), "pct_img": _num(mix.get("pct_img")), "pct_md": _num(mix.get("pct_md")), "phase": mix.get("phase"), }) # SE gate stats (if the bridge exposes them) # fetch without clearing bridge_module = getattr(self, "bridge", None) get_se_stats = getattr(bridge_module, "get_se_stats", None) if bridge_module is not None else None if callable(get_se_stats): se_stats = get_se_stats(reset=False) if se_stats: row.update({ "se_mean": _num(se_stats.get("mean")), "se_std": _num(se_stats.get("std")), "se_pct_lt_0.2": _num(se_stats.get("pct_lt_0.2")), "se_pct_gt_0.8": _num(se_stats.get("pct_gt_0.8")), }) # debug print BEFORE reset so you can see the true count try: print(f"SE gate samples seen: {getattr(getattr(bridge_module, 'se_log', None), '_n', 0)}") except Exception: pass # now clear for the next epoch try: get_se_stats(reset=True) except Exception: pass row.setdefault(self.args.early_metric, None) # early-stop decision (no break here) should_stop = False if getattr(self, "_early", None): should_stop = self._early.step(row, trainer=self, epoch=epoch) row.update({ "early_best_so_far" : float(self._early.best), "early_bad_epochs" : int(self._early.bad_epochs), "early_improved" : int(self._early.last_improved), "early_monitor" : self.args.early_metric, }) # concise console status each epoch mon = self.args.early_metric cur = row.get(mon, None) if cur is not None and cur == cur: # not NaN msg = ( f"[early] epoch {epoch+1}: {mon}={cur:.5f} | best={self._early.best:.5f} | " f"bad_epochs={self._early.bad_epochs}/{self._early.patience} | improved={'yes' if self._early.last_improved else 'no'}" ) print(msg) try: self.logger.info(msg) except Exception: pass # Holdout best-tracker (if a holdout set is present) if getattr(self, "_holdout_best", None): mon_h = self._holdout_best["monitor"] mode_h = self._holdout_best["mode"] min_delta_h = self._holdout_best["min_delta"] val_h = row.get(mon_h, None) improved_h = False if val_h is not None and val_h == val_h: best_val_h = self._holdout_best["value"] if mode_h == "max": improved_h = (val_h > best_val_h + min_delta_h) else: improved_h = (val_h < best_val_h - min_delta_h) if improved_h: self._holdout_best["value"] = float(val_h) self._holdout_best["epoch"] = int(epoch) st_h = self._snapshot() self._holdout_best["state"] = st_h if self._holdout_best.get("save_path"): try: torch.save(st_h, self._holdout_best["save_path"]) except Exception: pass print(f"[holdout_best] ↑ new best {mon_h}={val_h:.5f} at epoch {epoch+1}") row.update({ "holdout_best_monitor": mon_h, "holdout_best_so_far": (None if self._holdout_best["value"] in (math.inf, -math.inf) else float(self._holdout_best["value"])), "holdout_best_epoch": (None if self._holdout_best["epoch"] < 0 else int(self._holdout_best["epoch"] + 1)), }) # Best-tracker (always-on): save best model snapshot and optional checkpoint if getattr(self, "_best_info", None): # If early stopper is active, mirror its best info and avoid duplicate prints if getattr(self, "_early", None): row.update({ "best_monitor": self._best_info["monitor"], "best_so_far": float(self._early.best) if self._early.best == self._early.best else None, "best_epoch": (int(self._early.best_epoch) + 1) if (self._early.best_epoch is not None) else None, }) # skip independent tracking/prints to avoid duplication self._write_epoch_log(row) return row, should_stop mon = self._best_info["monitor"] mode = self._best_info["mode"] min_delta = self._best_info["min_delta"] val = row.get(mon, None) improved = False if val is not None and val == val: # not NaN best_val = self._best_info["value"] if mode == "max": improved = (val > best_val + min_delta) else: improved = (val < best_val - min_delta) if improved: self._best_info["value"] = float(val) self._best_info["epoch"] = int(epoch) st = self._snapshot() self._best_info["state"] = st if self._best_info.get("save_path"): try: torch.save(st, self._best_info["save_path"]) except Exception: pass print(f"[best] ↑ new best {mon}={val:.5f} at epoch {epoch+1}") # Add best-so-far to row row.update({ "best_monitor": mon, "best_so_far": (None if self._best_info["value"] in (math.inf, -math.inf) else float(self._best_info["value"])), "best_epoch": (None if self._best_info["epoch"] < 0 else int(self._best_info["epoch"] + 1)), }) # finally: write once, after all updates self._write_epoch_log(row) return row, should_stop def _restore_from_state(self, st: dict): if not st: return self.img_tower.load_state_dict(st["img_tower"]) self.md_tower.load_state_dict(st["md_tower"]) if "bridge" in st and hasattr(self, "bridge"): self.bridge.load_state_dict(st["bridge"]) if "head_img" in st and hasattr(self, "head_img"): self.head_img.load_state_dict(st["head_img"]) if "head_md" in st and hasattr(self, "head_md"): self.head_md.load_state_dict(st["head_md"]) if "optimizer" in st: self.optimizer.load_state_dict(st["optimizer"]) # --- V2 loader integration (no v1 hypertower import) --- from classes.v2 import build_papila_profile from classes.v2.slot_dataset import SlotDataset, slot_collate __all__ = [ "HyperTower", "V2HyperTower", "V2ModeComparisonOps", "V2ModeComparator", ] def _tuple_collate(batch: list[dict[str, object]]): """ Convert a slot-collated batch dict into (imgs, metas, labels) tuples so the legacy HyperTower training loop can consume it. """ data = slot_collate(batch) imgs = data.get("image_1") metas = data.get("matrix_1") labels = data.get("label_1") if labels is not None and not torch.is_tensor(labels): labels = torch.as_tensor(labels, dtype=torch.long) return imgs, metas, labels class V2HyperTower(HyperTower): """ HyperTower that swaps the loader pipeline for V2 slot datasets. This preserves the original training loop, logging, and metrics. """ def __init__(self, clinical, args): sample_mode = getattr(args, "sample_mode", "eye") self._v2_profile = build_papila_profile( patient_col="Patient ID", label_col=getattr(clinical, "label_col", "Diagnosis"), sample_mode=sample_mode, ) super().__init__(clinical, args) def _make_loader_for_df(self, df, is_train: bool = False): samples = self._v2_profile.build_samples(df=df, clinical=self.clinical) dataset = SlotDataset( samples, self._v2_profile.slot_descriptors(), image_transform=self.img_tower.transform, image_preprocessor=self.image_preprocessor, ) # Optional: class-balanced bootstrapped sampling for TRAIN only if is_train and getattr(self.args, "balanced_sampler", False): y = np.asarray(df[self.clinical.label_col].values) uniq, counts = np.unique(y, return_counts=True) inv = {c: (1.0 / cnt if cnt > 0 else 0.0) for c, cnt in zip(uniq, counts)} w = np.array([inv[c] for c in y], dtype=np.float32) sampler = WeightedRandomSampler(weights=w, num_samples=len(y), replacement=True) return DataLoader( dataset, batch_size=self.batch_size, sampler=sampler, shuffle=False, collate_fn=_tuple_collate, ) return DataLoader( dataset, batch_size=self.batch_size, shuffle=not is_train, collate_fn=_tuple_collate, ) class V2ModeComparisonOps: """ Shared mode-comparison training/eval ops hosted in V2 code so scripts remain orchestration-only. """ @staticmethod def _to_label_tensor(labels, device: torch.device) -> torch.Tensor: if torch.is_tensor(labels): return labels.to(device=device, dtype=torch.long) return torch.as_tensor(labels, dtype=torch.long, device=device) @staticmethod def _set_requires_grad(module: nn.Module, enabled: bool) -> None: for p in module.parameters(): p.requires_grad = enabled @staticmethod def _set_single_phase(model: nn.Module, phase: str) -> None: if phase == "tower_warmup": V2ModeComparisonOps._set_requires_grad(model.img_tower, True) V2ModeComparisonOps._set_requires_grad(model.md_tower, True) V2ModeComparisonOps._set_requires_grad(model.bridge.classifier_img, True) V2ModeComparisonOps._set_requires_grad(model.bridge.classifier_md, True) V2ModeComparisonOps._set_requires_grad(model.bridge.W_img, False) V2ModeComparisonOps._set_requires_grad(model.bridge.W_md, False) V2ModeComparisonOps._set_requires_grad(model.bridge.classifier_fused, False) return if phase == "fused_warmup": V2ModeComparisonOps._set_requires_grad(model.img_tower, False) V2ModeComparisonOps._set_requires_grad(model.md_tower, False) V2ModeComparisonOps._set_requires_grad(model.bridge.classifier_img, False) V2ModeComparisonOps._set_requires_grad(model.bridge.classifier_md, False) V2ModeComparisonOps._set_requires_grad(model.bridge.W_img, True) V2ModeComparisonOps._set_requires_grad(model.bridge.W_md, True) V2ModeComparisonOps._set_requires_grad(model.bridge.classifier_fused, True) return V2ModeComparisonOps._set_requires_grad(model, True) @staticmethod def _set_bilateral_phase(model: nn.Module, phase: str) -> None: if phase == "tower_warmup": V2ModeComparisonOps._set_requires_grad(model.eye_img_tower, True) V2ModeComparisonOps._set_requires_grad(model.eye_md_tower, True) V2ModeComparisonOps._set_requires_grad(model.joint_img, True) V2ModeComparisonOps._set_requires_grad(model.joint_md, True) V2ModeComparisonOps._set_requires_grad(model.aux_img, True) V2ModeComparisonOps._set_requires_grad(model.aux_md, True) V2ModeComparisonOps._set_requires_grad(model.bridge, False) return if phase == "fused_warmup": V2ModeComparisonOps._set_requires_grad(model.eye_img_tower, False) V2ModeComparisonOps._set_requires_grad(model.eye_md_tower, False) V2ModeComparisonOps._set_requires_grad(model.joint_img, False) V2ModeComparisonOps._set_requires_grad(model.joint_md, False) V2ModeComparisonOps._set_requires_grad(model.aux_img, False) V2ModeComparisonOps._set_requires_grad(model.aux_md, False) V2ModeComparisonOps._set_requires_grad(model.bridge, True) return V2ModeComparisonOps._set_requires_grad(model, True) @staticmethod def train_single_epoch(model: nn.Module, loader: DataLoader, opt, device: torch.device, *, phase: str, bcd_prob: float = 0.5): model.train() V2ModeComparisonOps._set_single_phase(model, phase) total_loss = total_correct = total_n = 0 for batch in loader: x = batch.get("image_1") m = batch.get("matrix_1") y = batch.get("label_1") if not torch.is_tensor(x) or not torch.is_tensor(m): continue x = x.to(device) m = m.to(device) y = V2ModeComparisonOps._to_label_tensor(y, device) img_feats = model.img_tower(x) md_feats = model.md_tower(m) if phase == "tower_warmup": logits_i = model.bridge.classifier_img(img_feats) logits_m = model.bridge.classifier_md(md_feats) loss = 0.5 * (F.cross_entropy(logits_i, y) + F.cross_entropy(logits_m, y)) logits = 0.5 * (F.softmax(logits_i, dim=1) + F.softmax(logits_m, dim=1)) elif phase == "fused_warmup": logits, _, _ = model.bridge(img_feats, md_feats) loss = F.cross_entropy(logits, y) else: if random() < bcd_prob: if random() < 0.5: logits = model.bridge.classifier_img(img_feats) else: logits = model.bridge.classifier_md(md_feats) loss = F.cross_entropy(logits, y) else: logits, _, _ = model.bridge(img_feats, md_feats) loss = F.cross_entropy(logits, y) opt.zero_grad() loss.backward() opt.step() bs = y.shape[0] total_loss += float(loss.item()) * bs total_correct += int((logits.argmax(1) == y).sum()) total_n += bs return ( total_loss / total_n if total_n else float("nan"), total_correct / total_n if total_n else float("nan"), ) @staticmethod def train_bilateral_epoch(model: nn.Module, loader: DataLoader, opt, device: torch.device, *, phase: str, bcd_prob: float = 0.5): model.train() V2ModeComparisonOps._set_bilateral_phase(model, phase) total_loss = total_correct = total_n = 0 for batch in loader: x1 = batch.get("image_1") m1 = batch.get("matrix_1") x2 = batch.get("image_2") m2 = batch.get("matrix_2") y = batch.get("label_1") if not (torch.is_tensor(x1) and torch.is_tensor(m1) and torch.is_tensor(x2) and torch.is_tensor(m2)): continue x1 = x1.to(device); m1 = m1.to(device) x2 = x2.to(device); m2 = m2.to(device) y = V2ModeComparisonOps._to_label_tensor(y, device) joint_img, joint_md = model.encode_joint(x1, m1, x2, m2) if phase == "tower_warmup": logits_i = model.aux_img(joint_img) logits_m = model.aux_md(joint_md) loss = 0.5 * (F.cross_entropy(logits_i, y) + F.cross_entropy(logits_m, y)) logits = 0.5 * (F.softmax(logits_i, dim=1) + F.softmax(logits_m, dim=1)) elif phase == "fused_warmup": logits, _, _ = model.bridge(joint_img, joint_md) loss = F.cross_entropy(logits, y) else: if random() < bcd_prob: if random() < 0.5: logits = model.aux_img(joint_img) else: logits = model.aux_md(joint_md) loss = F.cross_entropy(logits, y) else: logits, _, _ = model.bridge(joint_img, joint_md) loss = F.cross_entropy(logits, y) opt.zero_grad() loss.backward() opt.step() bs = y.shape[0] total_loss += float(loss.item()) * bs total_correct += int((logits.argmax(1) == y).sum()) total_n += bs return ( total_loss / total_n if total_n else float("nan"), total_correct / total_n if total_n else float("nan"), ) @staticmethod def collect_probs_classic(model: nn.Module, loader: DataLoader, device: torch.device): model.eval() y_chunks, p_chunks = [], [] with torch.no_grad(): for batch in loader: x1 = batch.get("image_1"); m1 = batch.get("matrix_1") x2 = batch.get("image_2"); m2 = batch.get("matrix_2") y = batch.get("label_1") if not (torch.is_tensor(x1) and torch.is_tensor(m1) and torch.is_tensor(x2) and torch.is_tensor(m2)): continue y_t = V2ModeComparisonOps._to_label_tensor(y, device) p_od = F.softmax(model(x1.to(device), m1.to(device)), dim=1) p_os = F.softmax(model(x2.to(device), m2.to(device)), dim=1) y_np = y_t.cpu().numpy() y_chunks += [y_np, y_np] p_chunks += [p_od.cpu().numpy(), p_os.cpu().numpy()] if not y_chunks: return np.array([], dtype=np.int64), np.zeros((0, 0), dtype=np.float32) return np.concatenate(y_chunks), np.concatenate(p_chunks, axis=0) @staticmethod def collect_probs_ensemble(model: nn.Module, loader: DataLoader, device: torch.device): model.eval() y_chunks, p_chunks = [], [] with torch.no_grad(): for batch in loader: x1 = batch.get("image_1"); m1 = batch.get("matrix_1") x2 = batch.get("image_2"); m2 = batch.get("matrix_2") y = batch.get("label_1") if not (torch.is_tensor(x1) and torch.is_tensor(m1) and torch.is_tensor(x2) and torch.is_tensor(m2)): continue y_t = V2ModeComparisonOps._to_label_tensor(y, device) p_od = F.softmax(model(x1.to(device), m1.to(device)), dim=1) p_os = F.softmax(model(x2.to(device), m2.to(device)), dim=1) p = 0.5 * (p_od + p_os) y_chunks.append(y_t.cpu().numpy()) p_chunks.append(p.cpu().numpy()) if not y_chunks: return np.array([], dtype=np.int64), np.zeros((0, 0), dtype=np.float32) return np.concatenate(y_chunks), np.concatenate(p_chunks, axis=0) @staticmethod def collect_probs_bilateral(model: nn.Module, loader: DataLoader, device: torch.device): model.eval() y_chunks, p_chunks = [], [] with torch.no_grad(): for batch in loader: x1 = batch.get("image_1"); m1 = batch.get("matrix_1") x2 = batch.get("image_2"); m2 = batch.get("matrix_2") y = batch.get("label_1") if not (torch.is_tensor(x1) and torch.is_tensor(m1) and torch.is_tensor(x2) and torch.is_tensor(m2)): continue y_t = V2ModeComparisonOps._to_label_tensor(y, device) p = F.softmax(model(x1.to(device), m1.to(device), x2.to(device), m2.to(device)), dim=1) y_chunks.append(y_t.cpu().numpy()) p_chunks.append(p.cpu().numpy()) if not y_chunks: return np.array([], dtype=np.int64), np.zeros((0, 0), dtype=np.float32) return np.concatenate(y_chunks), np.concatenate(p_chunks, axis=0) class V2ModeComparator: """ V2 entrypoint for the three-mode comparison workflow: classic, ensemble, bilateral. """ @staticmethod def build_parser(): return build_parser() @staticmethod def run(cli_args=None): parser = V2ModeComparator.build_parser() args = parser.parse_args(cli_args) return run_mode(args) # --- Mode comparison orchestration (migrated from mode_compare_engine) --- 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) # --------------------------------------------------------------------------- # Models # --------------------------------------------------------------------------- class SingleEyeHT(nn.Module): """ ImageTower + MDTower + Bridge, trained on eye-level samples. Supports both Classic (eye-level) and Ensemble (patient-level averaging) eval. """ def __init__( self, *, backbone: str, freeze_ratio: float, augment: bool, clinical_data, num_classes: int, md_hidden_dim: int = 128, fusion_dim: int = 256, ): super().__init__() self.img_tower = ImageTower( backbone=backbone, freeze_ratio=freeze_ratio, augment=augment, use_se=False, ) self.md_tower = MDTower( clinical_data=clinical_data, hidden_dim=md_hidden_dim, use_se=False, ) self.bridge = Bridge( img_dim=self.img_tower.out_dim, meta_dim=self.md_tower.out_dim, num_classes=num_classes, fusion_dim=fusion_dim, mode="fused", use_se=False, ) @property def transform(self): return self.img_tower.transform def forward(self, x: torch.Tensor, meta: torch.Tensor) -> torch.Tensor: img_feats = self.img_tower(x) md_feats = self.md_tower(meta) out_f, _, _ = self.bridge(img_feats, md_feats) return out_f class BilateralHT(nn.Module): """ Bilateral mode with joint towers: - shared eye-level towers encode OD/OS independently - joint image and metadata towers combine OD/OS embeddings - standard Bridge fuses joint image + joint metadata embeddings """ def __init__( self, *, backbone: str, freeze_ratio: float, augment: bool, clinical_data, num_classes: int, md_hidden_dim: int = 128, fusion_dim: int = 256, ): super().__init__() self.eye_img_tower = ImageTower( backbone=backbone, freeze_ratio=freeze_ratio, augment=augment, use_se=False, ) self.eye_md_tower = MDTower( clinical_data=clinical_data, hidden_dim=md_hidden_dim, use_se=False, ) img_dim = self.eye_img_tower.out_dim md_dim = self.eye_md_tower.out_dim self.joint_img = nn.Sequential( nn.Linear(2 * img_dim, fusion_dim), nn.LayerNorm(fusion_dim), nn.ReLU(), nn.Dropout(0.3), nn.Linear(fusion_dim, img_dim), ) self.joint_md = nn.Sequential( nn.Linear(2 * md_dim, fusion_dim), nn.LayerNorm(fusion_dim), nn.ReLU(), nn.Dropout(0.3), nn.Linear(fusion_dim, md_dim), ) self.bridge = Bridge( img_dim=img_dim, meta_dim=md_dim, num_classes=num_classes, fusion_dim=fusion_dim, mode="fused", use_se=False, ) # Auxiliary heads for tower warmup / BCD tower steps. self.aux_img = nn.Linear(img_dim, num_classes) self.aux_md = nn.Linear(md_dim, num_classes) @property def transform(self): return self.eye_img_tower.transform def encode_joint( self, x_od: torch.Tensor, meta_od: torch.Tensor, x_os: torch.Tensor, meta_os: torch.Tensor, ) -> tuple[torch.Tensor, torch.Tensor]: img_od = self.eye_img_tower(x_od) md_od = self.eye_md_tower(meta_od) img_os = self.eye_img_tower(x_os) md_os = self.eye_md_tower(meta_os) joint_img = self.joint_img(torch.cat([img_od, img_os], dim=1)) joint_md = self.joint_md(torch.cat([md_od, md_os], dim=1)) return joint_img, joint_md def forward( self, x_od: torch.Tensor, meta_od: torch.Tensor, x_os: torch.Tensor, meta_os: torch.Tensor, ) -> torch.Tensor: joint_img, joint_md = self.encode_joint(x_od, meta_od, x_os, meta_os) out_f, _, _ = self.bridge(joint_img, joint_md) return out_f def _set_requires_grad(module: nn.Module, enabled: bool) -> None: V2ModeComparisonOps._set_requires_grad(module, enabled) def _set_single_phase(model: SingleEyeHT, phase: str) -> None: V2ModeComparisonOps._set_single_phase(model, phase) def _set_bilateral_phase(model: BilateralHT, phase: str) -> None: V2ModeComparisonOps._set_bilateral_phase(model, phase) # --------------------------------------------------------------------------- # Data helpers # --------------------------------------------------------------------------- def filter_eye_samples(samples: list[dict]) -> list[dict]: """Keep any single-eye sample with a valid image, matrix, and label.""" return [ s for s in samples if s.get("image_1") is not None and s.get("matrix_1") is not None and s.get("label_1") is not None ] def filter_bilateral_samples(samples: list[dict]) -> list[dict]: """Keep only patient-level samples where both eyes are fully present.""" return [ s for s in samples if s.get("image_1") is not None and s.get("matrix_1") is not None and s.get("image_2") is not None and s.get("matrix_2") is not None and s.get("label_1") is not None ] def make_loader( samples: list[dict], slots: dict, *, image_transform, image_preprocessor=None, batch_size: int, shuffle: bool, num_workers: int, ) -> DataLoader: ds = SlotDataset( samples, slots, image_transform=image_transform, image_preprocessor=image_preprocessor, ) return DataLoader( ds, batch_size=batch_size, shuffle=shuffle, num_workers=num_workers, collate_fn=slot_collate, ) def to_label_tensor(labels, device: torch.device) -> torch.Tensor: if torch.is_tensor(labels): return labels.to(device=device, dtype=torch.long) return torch.as_tensor(labels, dtype=torch.long, device=device) def _drop_mixed_label_patients(df, *, patient_col: str, label_col: str): 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, *, patient_col: str, label_col: str): """Set all rows for each patient to that patient's max observed label.""" 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) per_patient_unique = out.groupby(patient_col)[label_col].nunique(dropna=True) still_mixed = per_patient_unique[per_patient_unique > 1].index.tolist() return out.reset_index(drop=True), changed_rows, still_mixed def build_image_preprocessor_from_args(args): crop_manifest = getattr(args, "img_crop_manifest", None) crop_weights = getattr(args, "img_crop_weights", None) use_gt = bool(getattr(args, "img_crop_gt", False)) if not crop_manifest: return None crop_cache = Path(getattr(args, "img_crop_cache", Path("analysis_data/hypertower_crops"))) if use_gt: pre = ManifestImageCropper( manifest_path=Path(crop_manifest), scale=getattr(args, "img_crop_scale", 2.5), target_size=getattr(args, "img_crop_size", 224), cache_dir=crop_cache, ) print(f"[V2 modes] GT disc cropper enabled -> cache at {crop_cache}", flush=True) return pre if crop_weights: pre = UNetImageCropper( manifest_path=Path(crop_manifest), weights_path=Path(crop_weights), normalize=getattr(args, "img_crop_normalize", "per_image"), threshold=getattr(args, "img_crop_threshold", 0.5), tta=getattr(args, "img_crop_tta", False), scale=getattr(args, "img_crop_scale", 2.5), target_size=getattr(args, "img_crop_size", 224), cache_dir=crop_cache, ) print(f"[V2 modes] UNet disc cropper enabled -> cache at {crop_cache}", flush=True) return pre print( "[V2 modes] img_crop_manifest provided but no --img-crop-gt or --img-crop-weights; cropping disabled.", flush=True, ) return None # --------------------------------------------------------------------------- # Metrics helpers # --------------------------------------------------------------------------- def compute_ece(y_true: np.ndarray, probs: np.ndarray, n_bins: int = 10) -> float: """Expected Calibration Error: weighted mean of |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 bin_acc = float((predictions[mask] == y_true[mask]).mean()) bin_conf = float(confidences[mask].mean()) ece += float(mask.sum()) / n * abs(bin_conf - bin_acc) 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 Exception: kappa = nan try: mcc = float(matthews_corrcoef(y_true, preds)) except Exception: mcc = nan try: macro_f1 = float(f1_score(y_true, preds, average="macro", zero_division=0)) except Exception: macro_f1 = nan try: pcr = recall_score( y_true, preds, average=None, labels=list(range(num_classes)), zero_division=0, ).astype(float) except Exception: pcr = np.full(num_classes, nan) ece = compute_ece(y_true, probs, n_bins=n_bins) return dict(kappa=kappa, mcc=mcc, macro_f1=macro_f1, per_class_recall=pcr, ece=ece) def tune_binary_threshold(y_true: np.ndarray, p1: np.ndarray) -> float: if y_true.size == 0: return 0.5 grid = np.linspace(0.0, 1.0, 1001) best_t, best_acc = 0.5, -1.0 for t in grid: pred = (p1 >= t).astype(int) acc = float((pred == y_true).mean()) if acc > best_acc or (acc == best_acc and abs(t - 0.5) < abs(best_t - 0.5)): best_acc, best_t = acc, float(t) return best_t def multiclass_acc_with_bias(y_true: np.ndarray, probs: np.ndarray, bias: np.ndarray) -> float: if y_true.size == 0: return float("nan") logits = np.log(np.clip(probs, 1e-8, 1.0)) + bias.reshape(1, -1) return float((np.argmax(logits, axis=1) == y_true).mean()) def tune_multiclass_bias(y_true: np.ndarray, probs: np.ndarray, *, iters: int = 2) -> np.ndarray: 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 def _svf(vec) -> Optional[str]: if vec is None: return None arr = np.asarray(vec, dtype=float) if arr.size == 0: return None return "|".join(f"{float(v):.4f}" for v in arr.tolist()) def _score_arrays(y_true: np.ndarray, probs: np.ndarray, num_classes: int): """Returns (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)) def build_eval_transform(backbone: str): """Deterministic eval transform matching backbone normalization.""" key = (backbone or "").lower() if key not in BACKBONES: raise ValueError(f"Unsupported backbone '{backbone}'") spec = BACKBONES[key] mean = getattr(spec.weights_default, "meta", {}).get("mean", (0.485, 0.456, 0.406)) std = getattr(spec.weights_default, "meta", {}).get("std", (0.229, 0.224, 0.225)) crop = 299 if key == "inception_v3" else 224 return transforms.Compose( [ transforms.Resize(256), transforms.CenterCrop(crop), transforms.ToTensor(), transforms.Normalize(mean=mean, std=std), ] ) # --------------------------------------------------------------------------- # Train / collect # --------------------------------------------------------------------------- def train_single_epoch( model: SingleEyeHT, loader: DataLoader, opt, device, *, phase: str, bcd_prob: float = 0.5, ): return V2ModeComparisonOps.train_single_epoch( model, loader, opt, device, phase=phase, bcd_prob=bcd_prob ) def train_bilateral_epoch( model: BilateralHT, loader: DataLoader, opt, device, *, phase: str, bcd_prob: float = 0.5, ): return V2ModeComparisonOps.train_bilateral_epoch( model, loader, opt, device, phase=phase, bcd_prob=bcd_prob ) def collect_probs_classic( model: SingleEyeHT, loader: DataLoader, device: torch.device, ) -> tuple[np.ndarray, np.ndarray]: """ Classic eye-level eval using the bilateral val loader. OD and OS are treated as independent samples (both contribute to the arrays with the same patient label). Returns (y_true [2N], probs [2N, C]). """ return V2ModeComparisonOps.collect_probs_classic(model, loader, device) def collect_probs_ensemble( model: SingleEyeHT, loader: DataLoader, device: torch.device, ) -> tuple[np.ndarray, np.ndarray]: """ Patient-level ensemble eval: average OD and OS softmax probabilities. Returns (y_true [N], probs [N, C]). """ return V2ModeComparisonOps.collect_probs_ensemble(model, loader, device) def collect_probs_single_components( model: SingleEyeHT, loader: DataLoader, device: torch.device, *, aggregate_patient: bool, ) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]: """ Collect fused/img/md probabilities for SingleEyeHT. - aggregate_patient=False: eye-level (OD/OS as independent samples) - aggregate_patient=True : patient-level (average OD/OS per head) """ model.eval() y_chunks = [] pf_chunks, pi_chunks, pm_chunks = [], [], [] with torch.no_grad(): for batch in loader: x1 = batch.get("image_1"); m1 = batch.get("matrix_1") x2 = batch.get("image_2"); m2 = batch.get("matrix_2") y = batch.get("label_1") if not (torch.is_tensor(x1) and torch.is_tensor(m1) and torch.is_tensor(x2) and torch.is_tensor(m2)): continue y_t = V2ModeComparisonOps._to_label_tensor(y, device) def _per_eye_probs(x, m): img_feats = model.img_tower(x.to(device)) md_feats = model.md_tower(m.to(device)) out_f, out_i, out_m = model.bridge(img_feats, md_feats) return ( F.softmax(out_f, dim=1), F.softmax(out_i, dim=1), F.softmax(out_m, dim=1), ) pf_od, pi_od, pm_od = _per_eye_probs(x1, m1) pf_os, pi_os, pm_os = _per_eye_probs(x2, m2) if aggregate_patient: y_chunks.append(y_t.cpu().numpy()) pf_chunks.append((0.5 * (pf_od + pf_os)).cpu().numpy()) pi_chunks.append((0.5 * (pi_od + pi_os)).cpu().numpy()) pm_chunks.append((0.5 * (pm_od + pm_os)).cpu().numpy()) else: y_np = y_t.cpu().numpy() y_chunks += [y_np, y_np] pf_chunks += [pf_od.cpu().numpy(), pf_os.cpu().numpy()] pi_chunks += [pi_od.cpu().numpy(), pi_os.cpu().numpy()] pm_chunks += [pm_od.cpu().numpy(), pm_os.cpu().numpy()] if not y_chunks: z = np.zeros((0, 0), dtype=np.float32) return np.array([], dtype=np.int64), z, z, z return ( np.concatenate(y_chunks), np.concatenate(pf_chunks, axis=0), np.concatenate(pi_chunks, axis=0), np.concatenate(pm_chunks, axis=0), ) def collect_probs_bilateral( model: BilateralHT, loader: DataLoader, device: torch.device, ) -> tuple[np.ndarray, np.ndarray]: """Patient-level bilateral eval. Returns (y_true [N], probs [N, C]).""" return V2ModeComparisonOps.collect_probs_bilateral(model, loader, device) def collect_probs_bilateral_components( model: BilateralHT, loader: DataLoader, device: torch.device, ) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]: """Collect fused/img/md probabilities for bilateral joint-tower model.""" model.eval() y_chunks = [] pf_chunks, pi_chunks, pm_chunks = [], [], [] with torch.no_grad(): for batch in loader: x1 = batch.get("image_1"); m1 = batch.get("matrix_1") x2 = batch.get("image_2"); m2 = batch.get("matrix_2") y = batch.get("label_1") if not (torch.is_tensor(x1) and torch.is_tensor(m1) and torch.is_tensor(x2) and torch.is_tensor(m2)): continue y_t = V2ModeComparisonOps._to_label_tensor(y, device) joint_img, joint_md = model.encode_joint( x1.to(device), m1.to(device), x2.to(device), m2.to(device) ) out_f, _, _ = model.bridge(joint_img, joint_md) out_i = model.aux_img(joint_img) out_m = model.aux_md(joint_md) y_chunks.append(y_t.cpu().numpy()) pf_chunks.append(F.softmax(out_f, dim=1).cpu().numpy()) pi_chunks.append(F.softmax(out_i, dim=1).cpu().numpy()) pm_chunks.append(F.softmax(out_m, dim=1).cpu().numpy()) if not y_chunks: z = np.zeros((0, 0), dtype=np.float32) return np.array([], dtype=np.int64), z, z, z return ( np.concatenate(y_chunks), np.concatenate(pf_chunks, axis=0), np.concatenate(pi_chunks, axis=0), np.concatenate(pm_chunks, axis=0), ) # --------------------------------------------------------------------------- # Tuning helpers # --------------------------------------------------------------------------- def _tune_and_snap( y: np.ndarray, p: np.ndarray, acc: float, num_classes: int, args, n_bins: int, ) -> tuple[dict, float, Optional[np.ndarray], Optional[np.ndarray]]: """ Apply threshold/bias tuning and compute extended metrics. Returns (snap_dict, tuned_auc, threshold, bias). """ thr = 0.5 if num_classes == 2 else float("nan") bias = None ext_preds = None if args.tune_binary_threshold and num_classes == 2 and y.size > 0: thr = tune_binary_threshold(y, p[:, 1]) ext_preds = (p[:, 1] >= thr).astype(int) acc = float((ext_preds == y).mean()) elif args.tune_multiclass_bias and num_classes > 2 and y.size > 0: bias = tune_multiclass_bias(y, p) logits = np.log(np.clip(p, 1e-8, 1.0)) + bias.reshape(1, -1) ext_preds = np.argmax(logits, axis=1) acc = float((ext_preds == y).mean()) ext = compute_extended_metrics(y, p, num_classes, n_bins=n_bins, preds_override=ext_preds) _, auc, n = _score_arrays(y, p, num_classes) snap = dict( auc=auc, acc=acc, n=n, kappa=ext["kappa"], mcc=ext["mcc"], macro_f1=ext["macro_f1"], per_class_recall=ext["per_class_recall"], ece=ext["ece"], threshold=thr, bias=bias, ) return snap, auc, thr, bias # --------------------------------------------------------------------------- # Result dataclass # --------------------------------------------------------------------------- def _nan() -> float: return float("nan") @dataclass class FoldResult: mode: str fold: int # Epoch where each model hit its peak val AUC best_epoch_single: int # SingleEyeHT — selected by ensemble val AUC best_epoch_bilat: int # BilateralHT — selected by bilateral val AUC # Classic (eye-level eval of SingleEyeHT; n = 2 * ensemble_val_n) classic_val_auc: float classic_val_acc: float classic_val_kappa: float classic_val_mcc: float classic_val_f1: float classic_val_recall: Optional[str] classic_val_ece: float classic_val_threshold: float classic_val_bias: Optional[str] classic_val_n: int # Ensemble (patient-level eval of same SingleEyeHT) ensemble_val_auc: float ensemble_val_acc: float ensemble_val_kappa: float ensemble_val_mcc: float ensemble_val_f1: float ensemble_val_recall: Optional[str] ensemble_val_ece: float ensemble_val_threshold: float ensemble_val_bias: Optional[str] ensemble_val_n: int # Bilateral (BilateralHT patient-level) bilat_val_auc: float bilat_val_acc: float bilat_val_kappa: float bilat_val_mcc: float bilat_val_f1: float bilat_val_recall: Optional[str] bilat_val_ece: float bilat_val_threshold: float bilat_val_bias: Optional[str] bilat_val_n: int # Training sample counts single_train_n: int bilat_train_n: int @dataclass class FoldArtifacts: y_true_classic: Optional[np.ndarray] probs_classic: Optional[np.ndarray] y_true_ensemble: Optional[np.ndarray] probs_ensemble: Optional[np.ndarray] y_true_bilat: Optional[np.ndarray] probs_bilat: Optional[np.ndarray] # --------------------------------------------------------------------------- # Output helpers # --------------------------------------------------------------------------- def _f(v) -> Optional[float]: if v is None or (isinstance(v, float) and np.isnan(v)): return None return round(float(v), 6) def _sv(vec) -> Optional[str]: if vec is None: return None return "|".join(f"{float(v):.4f}" for v in vec) # --------------------------------------------------------------------------- # Main fold runner # --------------------------------------------------------------------------- def run_fold( fold: int, split, mode: str, args, device: torch.device, data, num_classes: int, profile_eye, profile_patient, image_preprocessor, fold_dir: Path, tower_mode: str, ) -> tuple[FoldResult, FoldArtifacts]: nan = _nan() tower_mode = "single" if tower_mode == "classic" else tower_mode run_single = tower_mode in ("single", "ensemble") run_bilat = tower_mode == "bilateral" global_warmup_tower = getattr(args, "warmup_tower_epochs", None) global_warmup_fused = getattr(args, "warmup_fused_epochs", None) single_warmup_tower = ( int(args.single_warmup_tower_epochs) if getattr(args, "single_warmup_tower_epochs", None) is not None else int(global_warmup_tower) if global_warmup_tower is not None else 2 ) single_warmup_fused = ( int(args.single_warmup_fused_epochs) if getattr(args, "single_warmup_fused_epochs", None) is not None else int(global_warmup_fused) if global_warmup_fused is not None else 2 ) bilat_warmup_tower = ( int(args.bilat_warmup_tower_epochs) if getattr(args, "bilat_warmup_tower_epochs", None) is not None else int(global_warmup_tower) if global_warmup_tower is not None else 4 ) bilat_warmup_fused = ( int(args.bilat_warmup_fused_epochs) if getattr(args, "bilat_warmup_fused_epochs", None) is not None else int(global_warmup_fused) if global_warmup_fused is not None else 3 ) if not run_single: single_warmup_tower = 0 single_warmup_fused = 0 if not run_bilat: bilat_warmup_tower = 0 bilat_warmup_fused = 0 main_epochs = int(args.epochs) total_single_epochs = (single_warmup_tower + single_warmup_fused + main_epochs) if run_single else 0 total_bilat_epochs = (bilat_warmup_tower + bilat_warmup_fused + main_epochs) if run_bilat else 0 total_epochs = max(total_single_epochs, total_bilat_epochs) # ---- samples ------------------------------------------------------------- eye_train = filter_eye_samples( profile_eye.build_samples(df=split.train, clinical=data) ) bilat_train = filter_bilateral_samples( profile_patient.build_samples(df=split.train, clinical=data) ) bilat_val = filter_bilateral_samples( profile_patient.build_samples(df=split.val, clinical=data) ) if len(bilat_val) == 0: print(f" [fold {fold+1}] WARNING: no bilateral val samples; skipping fold.", flush=True) empty = FoldResult( mode=mode, fold=fold, best_epoch_single=0, best_epoch_bilat=0, classic_val_auc=nan, classic_val_acc=nan, classic_val_kappa=nan, classic_val_mcc=nan, classic_val_f1=nan, classic_val_recall=None, classic_val_ece=nan, classic_val_threshold=nan, classic_val_bias=None, classic_val_n=0, ensemble_val_auc=nan, ensemble_val_acc=nan, ensemble_val_kappa=nan, ensemble_val_mcc=nan, ensemble_val_f1=nan, ensemble_val_recall=None, ensemble_val_ece=nan, ensemble_val_threshold=nan, ensemble_val_bias=None, ensemble_val_n=0, bilat_val_auc=nan, bilat_val_acc=nan, bilat_val_kappa=nan, bilat_val_mcc=nan, bilat_val_f1=nan, bilat_val_recall=None, bilat_val_ece=nan, bilat_val_threshold=nan, bilat_val_bias=None, bilat_val_n=0, single_train_n=len(eye_train), bilat_train_n=len(bilat_train), ) return empty, FoldArtifacts( y_true_classic=None, probs_classic=None, y_true_ensemble=None, probs_ensemble=None, y_true_bilat=None, probs_bilat=None, ) # ---- models -------------------------------------------------------------- single = None bilateral = None if run_single: single = SingleEyeHT( backbone=args.backbone, freeze_ratio=args.freeze_ratio, augment=args.augment, clinical_data=data, num_classes=num_classes, md_hidden_dim=args.md_hidden_dim, fusion_dim=args.fusion_dim, ).to(device) if run_bilat: bilateral = BilateralHT( backbone=args.backbone, freeze_ratio=args.freeze_ratio, augment=args.augment, clinical_data=data, num_classes=num_classes, md_hidden_dim=args.md_hidden_dim, fusion_dim=args.fusion_dim, ).to(device) slots_eye = profile_eye.slot_descriptors() slots_patient = profile_patient.slot_descriptors() loader_kw = dict(batch_size=args.batch_size, num_workers=args.num_workers) # ---- loaders ------------------------------------------------------------- train_single_loader = None train_bilat_loader = None if run_single: train_single_loader = make_loader( eye_train, slots_eye, image_transform=single.transform, image_preprocessor=image_preprocessor, shuffle=True, **loader_kw, ) if run_bilat: train_bilat_loader = make_loader( bilat_train, slots_patient, image_transform=bilateral.transform, image_preprocessor=image_preprocessor, shuffle=True, **loader_kw, ) eval_transform = build_eval_transform(args.backbone) val_loader = make_loader( bilat_val, slots_patient, image_transform=eval_transform, image_preprocessor=image_preprocessor, shuffle=False, **loader_kw, ) opt_single = torch.optim.Adam(single.parameters(), lr=args.lr) if run_single else None opt_bilateral = torch.optim.Adam(bilateral.parameters(), lr=args.lr) if run_bilat else None # ---- epoch log ----------------------------------------------------------- epoch_fields = [ "fold", "epoch", "phase_single", "phase_bilat", "main_epoch_single", "main_epoch_bilat", "single_active", "bilat_active", "single_train_loss", "single_train_acc", "classic_val_auc", "classic_val_acc", "classic_val_n", "ensemble_val_auc", "ensemble_val_acc", "ensemble_val_n", "bilat_train_loss", "bilat_train_acc", "bilat_val_auc", "bilat_val_acc", "bilat_val_n", "is_best_single", "is_best_bilat", ] fold_logger = HypertowerLogger(run_dir=fold_dir) # ---- best-epoch trackers ------------------------------------------------- best_single_auc = -1.0 # tracked by ensemble AUC best_bilat_auc = -1.0 best_epoch_single = 0 best_epoch_bilat = 0 best_single_state: Optional[dict] = None best_bilat_state: Optional[dict] = None snap_classic: dict = {} snap_ensemble: dict = {} snap_bilat: dict = {} if run_single: print( f" [fold {fold+1}] single_train_n={len(eye_train)} (eye-level) " f"val_n={len(bilat_val)} " f"single_warmup={single_warmup_tower}+{single_warmup_fused} total={total_single_epochs}", flush=True, ) else: print( f" [fold {fold+1}] bilat_train_n={len(bilat_train)} (bilateral) " f"val_n={len(bilat_val)} " f"bilat_warmup={bilat_warmup_tower}+{bilat_warmup_fused} total={total_bilat_epochs}", flush=True, ) # ---- epoch loop ---------------------------------------------------------- for epoch in range(total_epochs): if not run_single: phase_single, main_epoch_single, single_active = "inactive", 0, False elif epoch < single_warmup_tower: phase_single, main_epoch_single, single_active = "tower_warmup", 0, True elif epoch < (single_warmup_tower + single_warmup_fused): phase_single, main_epoch_single, single_active = "fused_warmup", 0, True elif epoch < total_single_epochs: phase_single, main_epoch_single, single_active = ( "main", epoch - single_warmup_tower - single_warmup_fused + 1, True, ) else: phase_single, main_epoch_single, single_active = "done", main_epochs, False if not run_bilat: phase_bilat, main_epoch_bilat, bilat_active = "inactive", 0, False elif epoch < bilat_warmup_tower: phase_bilat, main_epoch_bilat, bilat_active = "tower_warmup", 0, True elif epoch < (bilat_warmup_tower + bilat_warmup_fused): phase_bilat, main_epoch_bilat, bilat_active = "fused_warmup", 0, True elif epoch < total_bilat_epochs: phase_bilat, main_epoch_bilat, bilat_active = ( "main", epoch - bilat_warmup_tower - bilat_warmup_fused + 1, True, ) else: phase_bilat, main_epoch_bilat, bilat_active = "done", main_epochs, False if run_single and single_active: sl_loss, sl_acc = train_single_epoch( single, train_single_loader, opt_single, device, phase=phase_single, bcd_prob=float(args.bcd_prob), ) else: sl_loss, sl_acc = nan, nan if run_bilat and bilat_active: bl_loss, bl_acc = train_bilateral_epoch( bilateral, train_bilat_loader, opt_bilateral, device, phase=phase_bilat, bcd_prob=float(args.bcd_prob), ) else: bl_loss, bl_acc = nan, nan if run_single and tower_mode == "single": y_cl, p_cl, p_cl_img, p_cl_md = collect_probs_single_components( single, val_loader, device, aggregate_patient=False ) cl_acc, cl_auc, cl_n = _score_arrays(y_cl, p_cl, num_classes) cl_acc_img = float((p_cl_img.argmax(1) == y_cl).mean()) if y_cl.size else nan cl_acc_md = float((p_cl_md.argmax(1) == y_cl).mean()) if y_cl.size else nan _, cl_auc_img, _ = _score_arrays(y_cl, p_cl_img, num_classes) _, cl_auc_md, _ = _score_arrays(y_cl, p_cl_md, num_classes) y_en = np.array([], dtype=np.int64) p_en = np.zeros((0, 0), dtype=np.float32) en_acc = en_auc = nan en_n = 0 en_acc_img = en_acc_md = en_auc_img = en_auc_md = nan elif run_single and tower_mode == "ensemble": y_en, p_en, p_en_img, p_en_md = collect_probs_single_components( single, val_loader, device, aggregate_patient=True ) en_acc, en_auc, en_n = _score_arrays(y_en, p_en, num_classes) en_acc_img = float((p_en_img.argmax(1) == y_en).mean()) if y_en.size else nan en_acc_md = float((p_en_md.argmax(1) == y_en).mean()) if y_en.size else nan _, en_auc_img, _ = _score_arrays(y_en, p_en_img, num_classes) _, en_auc_md, _ = _score_arrays(y_en, p_en_md, num_classes) y_cl = np.array([], dtype=np.int64) p_cl = np.zeros((0, 0), dtype=np.float32) cl_acc = cl_auc = nan cl_n = 0 cl_acc_img = cl_acc_md = cl_auc_img = cl_auc_md = nan else: y_cl = y_en = np.array([], dtype=np.int64) p_cl = p_en = np.zeros((0, 0), dtype=np.float32) cl_acc = cl_auc = en_acc = en_auc = nan cl_n = en_n = 0 cl_acc_img = cl_acc_md = en_acc_img = en_acc_md = nan cl_auc_img = cl_auc_md = en_auc_img = en_auc_md = nan if run_bilat: y_bi, p_bi, p_bi_img, p_bi_md = collect_probs_bilateral_components(bilateral, val_loader, device) bi_acc, bi_auc, bi_n = _score_arrays(y_bi, p_bi, num_classes) bi_acc_img = float((p_bi_img.argmax(1) == y_bi).mean()) if y_bi.size else nan bi_acc_md = float((p_bi_md.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_md, _ = _score_arrays(y_bi, p_bi_md, num_classes) else: y_bi = np.array([], dtype=np.int64) p_bi = np.zeros((0, 0), dtype=np.float32) bi_acc = bi_auc = nan bi_n = 0 bi_acc_img = bi_acc_md = bi_auc_img = bi_auc_md = nan # Best-epoch checks: checkpointing is restricted to the main phase only. target_single_auc = cl_auc if tower_mode == "single" else en_auc single_ckpt_eligible = run_single and (phase_single == "main") is_best_single = ( single_ckpt_eligible and (not np.isnan(target_single_auc)) and (target_single_auc > best_single_auc) ) if is_best_single: best_single_auc = target_single_auc best_epoch_single = epoch + 1 best_single_state = copy.deepcopy(single.state_dict()) if tower_mode == "single": snap_cl, _, _, _ = _tune_and_snap(y_cl, p_cl, cl_acc, num_classes, args, args.ece_bins) snap_classic = snap_cl else: snap_en, _, _, _ = _tune_and_snap(y_en, p_en, en_acc, num_classes, args, args.ece_bins) snap_ensemble = snap_en bilat_ckpt_eligible = run_bilat and (phase_bilat == "main") is_best_bilat = ( bilat_ckpt_eligible and (not np.isnan(bi_auc)) and (bi_auc > best_bilat_auc) ) if is_best_bilat: best_bilat_auc = bi_auc best_epoch_bilat = epoch + 1 best_bilat_state = copy.deepcopy(bilateral.state_dict()) snap_bi, _, _, _ = _tune_and_snap(y_bi, p_bi, bi_acc, num_classes, args, args.ece_bins) snap_bilat = snap_bi fold_logger.write_epoch_row({ "fold": fold, "epoch": epoch + 1, "phase_single": phase_single, "phase_bilat": phase_bilat, "main_epoch_single": main_epoch_single, "main_epoch_bilat": main_epoch_bilat, "single_active": int(single_active), "bilat_active": int(bilat_active), "single_train_loss": _f(sl_loss), "single_train_acc": _f(sl_acc), "classic_val_auc": _f(cl_auc), "classic_val_acc": _f(cl_acc), "classic_val_n": cl_n, "ensemble_val_auc": _f(en_auc), "ensemble_val_acc": _f(en_acc), "ensemble_val_n": en_n, "bilat_train_loss": _f(bl_loss), "bilat_train_acc": _f(bl_acc), "bilat_val_auc": _f(bi_auc), "bilat_val_acc": _f(bi_acc), "bilat_val_n": bi_n, "is_best_single": int(is_best_single), "is_best_bilat": int(is_best_bilat), }, optional_cols=epoch_fields) if args.log_every > 0 and (epoch + 1) % args.log_every == 0: if run_single: if tower_mode == "single": msg = ( f" ep {epoch+1:>3}/{total_epochs} " f"[single:{phase_single} {main_epoch_single}/{main_epochs}] " f"fused(acc={cl_acc:.4f},auc={cl_auc:.4f}) " f"img(acc={cl_acc_img:.4f},auc={cl_auc_img:.4f}) " f"md(acc={cl_acc_md:.4f},auc={cl_auc_md:.4f}) " f"(best_fused={best_single_auc:.4f} @ep{best_epoch_single})" ) else: msg = ( f" ep {epoch+1:>3}/{total_epochs} " f"[single:{phase_single} {main_epoch_single}/{main_epochs}] " f"fused(acc={en_acc:.4f},auc={en_auc:.4f}) " f"img(acc={en_acc_img:.4f},auc={en_auc_img:.4f}) " f"md(acc={en_acc_md:.4f},auc={en_auc_md:.4f}) " f"(best_fused={best_single_auc:.4f} @ep{best_epoch_single})" ) else: msg = ( f" ep {epoch+1:>3}/{total_epochs} " f"[bilat:{phase_bilat} {main_epoch_bilat}/{main_epochs}] " f"fused(acc={bi_acc:.4f},auc={bi_auc:.4f}) " f"img(acc={bi_acc_img:.4f},auc={bi_auc_img:.4f}) " f"md(acc={bi_acc_md:.4f},auc={bi_auc_md:.4f}) " f"(best_bilat={best_bilat_auc:.4f} @ep{best_epoch_bilat})" ) print(msg, flush=True) fold_logger.info(msg) fold_logger.close() if args.save_checkpoints: if best_single_state is not None: torch.save(best_single_state, fold_dir / "best_single.pt") if best_bilat_state is not None: torch.save(best_bilat_state, fold_dir / "best_bilateral.pt") if run_single: if tower_mode == "single": print( f" [fold {fold+1}] BEST " f"fused(acc={snap_classic.get('acc', nan):.4f},auc={snap_classic.get('auc', nan):.4f}) " f"kappa={snap_classic.get('kappa', nan):.4f} " f"F1={snap_classic.get('macro_f1', nan):.4f} " f"ECE={snap_classic.get('ece', nan):.4f} @ep{best_epoch_single}", flush=True, ) else: print( f" [fold {fold+1}] BEST " f"ensemble(acc={snap_ensemble.get('acc', nan):.4f},auc={snap_ensemble.get('auc', nan):.4f}) " f"kappa={snap_ensemble.get('kappa', nan):.4f} " f"F1={snap_ensemble.get('macro_f1', nan):.4f} " f"ECE={snap_ensemble.get('ece', nan):.4f} @ep{best_epoch_single}", flush=True, ) else: print( f" [fold {fold+1}] BEST " f"fused(acc={snap_bilat.get('acc', nan):.4f},auc={snap_bilat.get('auc', nan):.4f}) " f"kappa={snap_bilat.get('kappa', nan):.4f} " f"F1={snap_bilat.get('macro_f1', nan):.4f} " f"ECE={snap_bilat.get('ece', nan):.4f} @ep{best_epoch_bilat}", flush=True, ) # Export best-epoch prediction artifacts for easier side-by-side analysis. if run_single and best_single_state is not None: single.load_state_dict(best_single_state) if run_bilat and best_bilat_state is not None: bilateral.load_state_dict(best_bilat_state) if run_single and tower_mode == "single": y_cl_best, p_cl_best = collect_probs_classic(single, val_loader, device) y_en_best = p_en_best = None elif run_single and tower_mode == "ensemble": y_en_best, p_en_best = collect_probs_ensemble(single, val_loader, device) y_cl_best = p_cl_best = None else: y_cl_best = y_en_best = None p_cl_best = p_en_best = None if run_bilat: y_bi_best, p_bi_best = collect_probs_bilateral(bilateral, val_loader, device) else: y_bi_best = p_bi_best = None return FoldResult( mode=mode, fold=fold, best_epoch_single=best_epoch_single, best_epoch_bilat=best_epoch_bilat, classic_val_auc=snap_classic.get("auc", nan), classic_val_acc=snap_classic.get("acc", nan), classic_val_kappa=snap_classic.get("kappa", nan), classic_val_mcc=snap_classic.get("mcc", nan), classic_val_f1=snap_classic.get("macro_f1", nan), classic_val_recall=_sv(snap_classic.get("per_class_recall")), classic_val_ece=snap_classic.get("ece", nan), classic_val_threshold=snap_classic.get("threshold", nan), classic_val_bias=_svf(snap_classic.get("bias")), classic_val_n=snap_classic.get("n", 0), ensemble_val_auc=snap_ensemble.get("auc", nan), ensemble_val_acc=snap_ensemble.get("acc", nan), ensemble_val_kappa=snap_ensemble.get("kappa", nan), ensemble_val_mcc=snap_ensemble.get("mcc", nan), ensemble_val_f1=snap_ensemble.get("macro_f1", nan), ensemble_val_recall=_sv(snap_ensemble.get("per_class_recall")), ensemble_val_ece=snap_ensemble.get("ece", nan), ensemble_val_threshold=snap_ensemble.get("threshold", nan), ensemble_val_bias=_svf(snap_ensemble.get("bias")), ensemble_val_n=snap_ensemble.get("n", 0), bilat_val_auc=snap_bilat.get("auc", nan), bilat_val_acc=snap_bilat.get("acc", nan), bilat_val_kappa=snap_bilat.get("kappa", nan), bilat_val_mcc=snap_bilat.get("mcc", nan), bilat_val_f1=snap_bilat.get("macro_f1", nan), bilat_val_recall=_sv(snap_bilat.get("per_class_recall")), bilat_val_ece=snap_bilat.get("ece", nan), bilat_val_threshold=snap_bilat.get("threshold", nan), bilat_val_bias=_svf(snap_bilat.get("bias")), bilat_val_n=snap_bilat.get("n", 0), single_train_n=len(eye_train), bilat_train_n=len(bilat_train), ), FoldArtifacts( y_true_classic=y_cl_best, probs_classic=p_cl_best, y_true_ensemble=y_en_best, probs_ensemble=p_en_best, y_true_bilat=y_bi_best, probs_bilat=p_bi_best, ) # --------------------------------------------------------------------------- # Summary helpers # --------------------------------------------------------------------------- def _summary(results: list[FoldResult]) -> dict: def _ms(vals): v = np.array([x for x in vals if x is not None and not np.isnan(float(x))], dtype=float) return (float(np.mean(v)) if v.size else None, float(np.std(v)) if v.size else None) out = {} for label, prefix in [ ("classic_best_val", "classic_val"), ("ensemble_best_val", "ensemble_val"), ("bilat_best_val", "bilat_val"), ]: sub = {} for m in ["auc", "acc", "kappa", "mcc", "f1", "ece", "threshold"]: vals = [getattr(r, f"{prefix}_{m}") for r in results] mean, std = _ms(vals) sub[f"{m}_mean"] = mean if m in ("auc", "f1", "kappa"): sub[f"{m}_std"] = std out[label] = sub # Deltas: ensemble − classic (eval strategy effect, same model) # bilateral − ensemble (bilateral training effect) for delta_label, prefix_a, prefix_b in [ ("delta_ensemble_vs_classic", "classic_val", "ensemble_val"), ("delta_bilat_vs_ensemble", "ensemble_val", "bilat_val"), ]: delta = {} for m in ["auc", "f1", "kappa"]: pairs = [ getattr(r, f"{prefix_b}_{m}") - getattr(r, f"{prefix_a}_{m}") for r in results if not np.isnan(float(getattr(r, f"{prefix_a}_{m}"))) and not np.isnan(float(getattr(r, f"{prefix_b}_{m}"))) ] delta[f"{m}_mean"] = float(np.mean(pairs)) if pairs else None delta[f"{m}_std"] = float(np.std(pairs)) if pairs else None out[delta_label] = delta out["n_folds_completed"] = len(results) out["single_train_mode"] = "eye-level (all OD+OS samples)" out["bilat_train_mode"] = "patient-level (bilateral only)" out["eval_note"] = ( "classic=eye-level SingleEyeHT; " "ensemble=patient-level SingleEyeHT (OD+OS averaged); " "bilateral=patient-level BilateralHT" ) return out def _print_summary(mode: str, s: dict, tower_mode: str | None = None) -> None: def f(v): return "nan" if v is None else f"{v:.4f}" cv = s["classic_best_val"] ev = s["ensemble_best_val"] bv = s["bilat_best_val"] d1 = s["delta_ensemble_vs_classic"] d2 = s["delta_bilat_vs_ensemble"] print(f"\n=== Summary [{mode}] — best-epoch val ===") print(f" {'':26s} {'AUC':>8} {'ACC':>8} {'Kappa':>8} {'F1-mac':>8} {'ECE':>8}") if tower_mode == "single": rows = [("single (eye-lvl eval)", cv)] elif tower_mode == "ensemble": rows = [("ensemble (pat-lvl eval)", ev)] elif tower_mode == "bilateral": rows = [("bilateral (bilat eval)", bv)] else: rows = [ ("classic (eye-lvl eval)", cv), ("ensemble (pat-lvl eval)", ev), ("bilateral (bilat eval)", bv), ] for label, d in rows: print( f" {label:26s} " f"{f(d['auc_mean']):>8} {f(d['acc_mean']):>8} " f"{f(d['kappa_mean']):>8} {f(d['f1_mean']):>8} {f(d['ece_mean']):>8}" ) if tower_mode is None: print( f" {'Δ ensemble−classic':26s} " f"{f(d1['auc_mean']):>8} {'':>8} " f"{f(d1['kappa_mean']):>8} {f(d1['f1_mean']):>8}" ) print( f" {'Δ bilateral−ensemble':26s} " f"{f(d2['auc_mean']):>8} {'':>8} " f"{f(d2['kappa_mean']):>8} {f(d2['f1_mean']):>8}" ) # --------------------------------------------------------------------------- # CLI # --------------------------------------------------------------------------- def build_parser() -> argparse.ArgumentParser: ap = argparse.ArgumentParser( description=( "Three HyperTower modes: Classic (eye-level), Ensemble (patient-level avg), " "Bilateral (BilateralBridge with shared towers). Pure k-fold CV." ) ) ap.add_argument("--image-dir", default="Papila/FundusImages") ap.add_argument("--clinical-dir", default="Papila/ClinicalData") ap.add_argument("--label-col", default="Diagnosis") ap.add_argument("--cat-cols", nargs="*", default=["Gender", "Phakic/Pseudophakic"]) ap.add_argument("--eval-mode", choices=["binary", "multiclass"], default="multiclass") ap.add_argument( "--tower-mode", choices=["single", "ensemble", "bilateral", "classic"], default="single", help="Train/evaluate a single tower mode.", ) ap.add_argument("--n-splits", type=int, default=5) ap.add_argument("--fold-seed", type=int, default=42) ap.add_argument( "--folds", type=int, default=None, help="Optional cap on how many folds to run (default: all --n-splits).", ) ap.add_argument("--epochs", type=int, default=40) ap.add_argument( "--warmup-tower-epochs", type=int, default=None, help="Extra tower warmup epochs (added before main epochs). Default: auto by mode.", ) ap.add_argument( "--warmup-fused-epochs", type=int, default=None, help="Extra fused warmup epochs (added before main epochs). Default: auto by mode.", ) ap.add_argument("--single-warmup-tower-epochs", type=int, default=None, help="Single-eye model tower warmup (overrides --warmup-tower-epochs).") ap.add_argument("--single-warmup-fused-epochs", type=int, default=None, help="Single-eye model fused warmup (overrides --warmup-fused-epochs).") ap.add_argument("--bilat-warmup-tower-epochs", type=int, default=None, help="Bilateral model tower warmup (overrides --warmup-tower-epochs).") ap.add_argument("--bilat-warmup-fused-epochs", type=int, default=None, help="Bilateral model fused warmup (overrides --warmup-fused-epochs).") ap.add_argument("--batch-size", type=int, default=8) ap.add_argument("--lr", type=float, default=1e-4) ap.add_argument("--bcd-prob", type=float, default=0.5, help="Tower-only step probability during main phase (per model).") ap.add_argument("--backbone", default="refugelike") ap.add_argument("--freeze-ratio", type=float, default=0.0) ap.add_argument("--augment", action="store_true") ap.add_argument("--num-workers", type=int, default=0) ap.add_argument("--device", choices=["auto", "cpu", "cuda"], default="auto") ap.add_argument("--seed", type=int, default=1234) ap.add_argument("--run-name", default=None) ap.add_argument("--output-root", default="analysis_data") # Optional ROI cropping (GT manifest or UNet-generated mask crop) ap.add_argument("--img-crop-manifest", type=str, default=None, help="Path to crop manifest CSV for ROI cropping.") ap.add_argument("--img-crop-gt", action="store_true", help="Use ground-truth masks/contours from manifest for ROI crop.") ap.add_argument("--img-crop-weights", type=str, default=None, help="UNet weights path for ROI cropping from predicted masks.") ap.add_argument("--img-crop-normalize", type=str, default="per_image", choices=["per_image", "imagenet"], help="UNet input normalization mode.") ap.add_argument("--img-crop-threshold", type=float, default=0.5, help="UNet mask threshold for ROI extraction.") ap.add_argument("--img-crop-tta", action="store_true", help="Enable flip-TTA during UNet mask inference.") ap.add_argument("--img-crop-scale", type=float, default=2.5, help="Disc-radius multiplier for square crop.") ap.add_argument("--img-crop-size", type=int, default=224, help="Output ROI size before tower transforms.") ap.add_argument("--img-crop-cache", type=str, default="analysis_data/hypertower_crops", help="Cache directory for cropped images and geometry sidecars.") # Architecture ap.add_argument("--md-hidden-dim", type=int, default=128, help="MDTower hidden dimension.") ap.add_argument("--fusion-dim", type=int, default=256, help="Bridge/BilateralBridge fusion dimension.") # Mixed patients ap.add_argument( "--exclude-mixed-patients", dest="exclude_mixed_patients", action="store_true", help="Drop patients whose two eyes have different labels before splitting.", ) ap.add_argument( "--include-mixed-patients", dest="exclude_mixed_patients", action="store_false", ) ap.add_argument( "--relabel-mixed-patients-to-max", dest="relabel_mixed_patients_to_max", action="store_true", help="When mixed patients are included, relabel both eyes to patient max severity.", ) # Backward-compatible alias; default behavior is now to keep raw labels. ap.add_argument( "--keep-mixed-raw-labels", dest="relabel_mixed_patients_to_max", action="store_false", help=argparse.SUPPRESS, ) ap.set_defaults(exclude_mixed_patients=False, relabel_mixed_patients_to_max=False) # Tuning ap.add_argument( "--tune-binary-threshold", action="store_true", help="Tune per-model binary threshold on validation each epoch.", ) ap.add_argument( "--tune-multiclass-bias", action="store_true", help="Tune per-model multiclass log-prob bias on validation each epoch.", ) ap.add_argument("--ece-bins", type=int, default=10) ap.add_argument("--log-every", type=int, default=1) ap.add_argument("--save-checkpoints", action="store_true") return ap def parse_args(): return build_parser().parse_args() def choose_device(name: str) -> torch.device: if name == "cuda": if not torch.cuda.is_available(): raise RuntimeError("--device cuda requested but CUDA is not available.") return torch.device("cuda") if name == "cpu": return torch.device("cpu") return torch.device("cuda" if torch.cuda.is_available() else "cpu") # --------------------------------------------------------------------------- # Entry point # --------------------------------------------------------------------------- def run_mode(args) -> Path: # Defensive local import for CLI/entrypoint execution paths. import csv device = choose_device(args.device) seed_everything(args.seed) print(f"Device: {device}", flush=True) print("Loading PAPILA data...", flush=True) data = build_papila_data( image_dir=args.image_dir, clinical_dir=args.clinical_dir, label_col=args.label_col, cat_cols=list(args.cat_cols), n_splits=args.n_splits, random_seed=args.fold_seed, ) print(f"Loaded: {len(data.df)} rows feature_dim={data.feature_dim}", flush=True) image_preprocessor = build_image_preprocessor_from_args(args) ts = time.strftime("%Y%m%d_%H%M%S") run_name = args.run_name or f"hypertower_modes_{ts}" out_dir = Path(args.output_root) / run_name out_dir.mkdir(parents=True, exist_ok=True) mode = args.eval_mode tower_mode = "single" if args.tower_mode == "classic" else args.tower_mode df_mode = data.df.copy() if args.exclude_mixed_patients: before = df_mode["Patient ID"].nunique() df_mode, mixed = _drop_mixed_label_patients( df_mode, patient_col="Patient ID", label_col=args.label_col ) print( f"[{mode}] dropped {len(mixed)} mixed-label patients " f"({before} → {df_mode['Patient ID'].nunique()})", flush=True, ) else: if args.relabel_mixed_patients_to_max: before_rows = len(df_mode) df_mode, changed_rows, still_mixed = _relabel_mixed_patients_to_max( df_mode, patient_col="Patient ID", label_col=args.label_col ) print( f"[{mode}] relabeled mixed patients to max severity " f"(changed={changed_rows}, rows={before_rows}→{len(df_mode)}, " f"remaining_mixed={len(still_mixed)}).", flush=True, ) else: print(f"[{mode}] keeping mixed-label patients with raw per-eye labels.", flush=True) if mode == "binary": df_mode = df_mode[df_mode[args.label_col].isin([0, 1])].reset_index(drop=True) num_classes = 2 if mode == "binary" else int(df_mode[args.label_col].nunique()) print( f"\n[{mode}] num_classes={num_classes} rows={len(df_mode)} " f"patients={df_mode['Patient ID'].nunique()}", flush=True, ) split_manager = PatientFirstSplitManager( patient_col="Patient ID", label_col=args.label_col ) split_args = SimpleNamespace( eval_mode=mode, holdout_per_class=0, holdout_seed=123, n_splits=args.n_splits, fold_seed=args.fold_seed, ) clinical_ns = SimpleNamespace(df=df_mode, label_col=args.label_col) plans = split_manager.build_plans(clinical=clinical_ns, args=split_args, profile=None) requested_folds = args.n_splits if args.folds is None else int(args.folds) n_folds = min(requested_folds, len(plans)) profile_eye = build_papila_profile( patient_col="Patient ID", label_col=args.label_col, sample_mode="eye" ) profile_patient = build_papila_profile( patient_col="Patient ID", label_col=args.label_col, sample_mode="patient" ) tm_dir = out_dir / mode / tower_mode tm_dir.mkdir(parents=True, exist_ok=True) fold_results: list[FoldResult] = [] for fold in range(n_folds): seed_everything(args.seed + fold * 100) fold_dir = tm_dir / f"fold{fold}" fold_dir.mkdir(exist_ok=True) print(f"\n[{mode}:{tower_mode}] fold {fold+1}/{n_folds}", flush=True) result, artifacts = run_fold( fold=fold, split=plans[fold], mode=mode, args=args, device=device, data=data, num_classes=num_classes, profile_eye=profile_eye, profile_patient=profile_patient, image_preprocessor=image_preprocessor, fold_dir=fold_dir, tower_mode=tower_mode, ) fold_results.append(result) if artifacts.y_true_ensemble is not None: np.save(fold_dir / "y_true.npy", artifacts.y_true_ensemble) if artifacts.probs_ensemble is not None: np.save(fold_dir / "probs_fused.npy", artifacts.probs_ensemble) if artifacts.probs_classic is not None: np.save(fold_dir / "probs_classic.npy", artifacts.probs_classic) if artifacts.probs_bilat is not None: np.save(fold_dir / "probs_bilat.npy", artifacts.probs_bilat) fold_csv = tm_dir / "fold_results.csv" csv_fields = list(FoldResult.__dataclass_fields__.keys()) with fold_csv.open("w", newline="", encoding="utf-8") as fh: w = csv.DictWriter(fh, fieldnames=csv_fields) w.writeheader() for r in fold_results: w.writerow({k: getattr(r, k) for k in csv_fields}) summary = _summary(fold_results) _print_summary(f"{mode}:{tower_mode}", summary, tower_mode=tower_mode) metric_key = { "single": "classic_val_auc", "ensemble": "ensemble_val_auc", "bilateral": "bilat_val_auc", }[tower_mode] fold_metrics = [] best_vals = [] for r in fold_results: best_val = getattr(r, metric_key) fold_metrics.append( { "fold": r.fold, "best_metric_value": _f(best_val), "best_epoch": (r.best_epoch_bilat if tower_mode == "bilateral" else r.best_epoch_single), "monitor": metric_key, } ) if not np.isnan(float(best_val)): best_vals.append(float(best_val)) mode_summary = { "run_id": run_name, "backbone": args.backbone, "epochs": args.epochs, "warmup_tower_epochs": args.warmup_tower_epochs, "warmup_fused_epochs": args.warmup_fused_epochs, "single_warmup_tower_epochs": args.single_warmup_tower_epochs, "single_warmup_fused_epochs": args.single_warmup_fused_epochs, "bilat_warmup_tower_epochs": args.bilat_warmup_tower_epochs, "bilat_warmup_fused_epochs": args.bilat_warmup_fused_epochs, "batch_size": args.batch_size, "lr": args.lr, "eval_mode": mode, "tower_mode": tower_mode, "n_splits": n_folds, "best_metric": metric_key, "best_metric_mode": "max", "best_metric_mean": (float(np.mean(best_vals)) if best_vals else None), "best_metric_std": (float(np.std(best_vals)) if best_vals else None), "fold_metrics": fold_metrics, "mode_summary": summary, } (tm_dir / "summary.json").write_text(json.dumps(mode_summary, indent=2), encoding="utf-8") # Merge run-level summary across sequential invocations. root_summary_path = out_dir / "summary.json" if root_summary_path.exists(): try: payload = json.loads(root_summary_path.read_text(encoding="utf-8")) except Exception: payload = {} else: payload = {} payload.setdefault("run_name", run_name) payload.setdefault("timestamp", ts) payload["config"] = vars(args) payload.setdefault("summaries", {}) payload["summaries"][f"{mode}:{tower_mode}"] = summary root_summary_path.write_text(json.dumps(payload, indent=2), encoding="utf-8") print(f"\nOutputs written to: {out_dir}") return out_dir