#!/usr/bin/env python3 from __future__ import annotations import argparse import csv import json import random import sys import time from dataclasses import dataclass from pathlib import Path from types import SimpleNamespace from typing import Optional import numpy as np import pandas as pd import torch import torch.nn.functional as F from sklearn.metrics import roc_auc_score from torch import nn from torch.utils.data import DataLoader REPO_ROOT = Path(__file__).resolve().parents[2] if str(REPO_ROOT) not in sys.path: sys.path.insert(0, str(REPO_ROOT)) from classes.v2 import ( Bridge, ImageTower, PatientFirstSplitManager, SlotDataset, build_papila_data, build_papila_profile, slot_collate, ) def seed_everything(seed: int) -> None: random.seed(seed) np.random.seed(seed) torch.manual_seed(seed) if torch.cuda.is_available(): torch.cuda.manual_seed_all(seed) @dataclass class FoldMetrics: mode: str fold: int baseline_acc: float baseline_auc: float bilateral_acc: float bilateral_auc: float baseline_n: int bilateral_n: int os_acc: float os_auc: float os_n: int holdout_baseline_acc: float holdout_baseline_auc: float holdout_bilateral_acc: float holdout_bilateral_auc: float holdout_baseline_n: int holdout_bilateral_n: int holdout_os_acc: float holdout_os_auc: float holdout_os_n: int class EyeLevelCNN(nn.Module): def __init__(self, *, backbone: str, freeze_ratio: float, num_classes: int, augment: bool): super().__init__() self.tower = ImageTower( backbone=backbone, freeze_ratio=freeze_ratio, augment=augment, use_se=False, ) self.head = nn.Linear(self.tower.out_dim, num_classes) def forward(self, x: torch.Tensor) -> torch.Tensor: feats = self.tower(x) return self.head(feats) class BilateralFusionCNN(nn.Module): def __init__( self, *, backbone: str, freeze_ratio: float, num_classes: int, augment: bool, use_se: bool, fusion_dim: int, ): super().__init__() self.tower_od = ImageTower( backbone=backbone, freeze_ratio=freeze_ratio, augment=augment, use_se=False, ) self.tower_os = ImageTower( backbone=backbone, freeze_ratio=freeze_ratio, augment=augment, use_se=False, ) self.bridge = Bridge( img_dim=self.tower_od.out_dim, meta_dim=self.tower_os.out_dim, num_classes=num_classes, fusion_dim=fusion_dim, mode="fused", use_se=use_se, ) def forward(self, od: torch.Tensor, os: torch.Tensor) -> torch.Tensor: f_od = self.tower_od(od) f_os = self.tower_os(os) out_fused, _, _ = self.bridge(f_od, f_os) return out_fused def patient_to_single_eye_samples(patient_samples: list[dict], eye_key: str) -> list[dict]: out = [] for s in patient_samples: img = s.get(eye_key) lbl = s.get("label_1") if img is None or lbl is None: continue out.append({"id_1": s.get("id_1"), "image_1": img, "label_1": lbl}) return out def make_loader( samples: list[dict], slots: dict, *, image_transform, batch_size: int, shuffle: bool, num_workers: int, ) -> DataLoader: ds = SlotDataset(samples, slots, image_transform=image_transform) return DataLoader( ds, batch_size=batch_size, shuffle=shuffle, num_workers=num_workers, collate_fn=slot_collate, ) def filter_eye_samples(samples: list[dict]) -> list[dict]: return [s for s in samples if s.get("image_1") is not None and s.get("label_1") is not None] def filter_bilateral_samples(samples: list[dict]) -> list[dict]: return [ s for s in samples if s.get("image_1") is not None and s.get("image_2") is not None and s.get("label_1") is not None ] 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 train_eye_epoch( model: EyeLevelCNN, loader: DataLoader, optimizer: torch.optim.Optimizer, device: torch.device, ): model.train() total_loss = 0.0 total_correct = 0 total_n = 0 for batch in loader: x = batch.get("image_1") y = batch.get("label_1") if not torch.is_tensor(x): continue y = to_label_tensor(y, device) x = x.to(device) logits = model(x) loss = F.cross_entropy(logits, y) optimizer.zero_grad() loss.backward() optimizer.step() bs = int(y.shape[0]) total_loss += float(loss.item()) * bs total_correct += int((logits.argmax(dim=1) == y).sum().item()) total_n += bs avg_loss = float(total_loss / total_n) if total_n > 0 else float("nan") acc = float(total_correct / total_n) if total_n > 0 else float("nan") return avg_loss, acc, total_n def train_bilateral_epoch( model: BilateralFusionCNN, loader: DataLoader, optimizer: torch.optim.Optimizer, device: torch.device, ): model.train() total_loss = 0.0 total_correct = 0 total_n = 0 for batch in loader: x1 = batch.get("image_1") x2 = batch.get("image_2") y = batch.get("label_1") if not torch.is_tensor(x1) or not torch.is_tensor(x2): continue y = to_label_tensor(y, device) x1 = x1.to(device) x2 = x2.to(device) logits = model(x1, x2) loss = F.cross_entropy(logits, y) optimizer.zero_grad() loss.backward() optimizer.step() bs = int(y.shape[0]) total_loss += float(loss.item()) * bs total_correct += int((logits.argmax(dim=1) == y).sum().item()) total_n += bs avg_loss = float(total_loss / total_n) if total_n > 0 else float("nan") acc = float(total_correct / total_n) if total_n > 0 else float("nan") return avg_loss, acc, total_n def evaluate_eye(model: EyeLevelCNN, loader: DataLoader, device: torch.device, num_classes: int): model.eval() y_true = [] y_prob = [] total_loss = 0.0 total_n = 0 with torch.no_grad(): for batch in loader: x = batch.get("image_1") y = batch.get("label_1") if not torch.is_tensor(x): continue y_t = to_label_tensor(y, device) logits = model(x.to(device)) bs = int(y_t.shape[0]) total_loss += float(F.cross_entropy(logits, y_t).item()) * bs total_n += bs probs = F.softmax(logits, dim=1).cpu().numpy() y_true.append(y_t.cpu().numpy()) y_prob.append(probs) acc, auc, n = _score_arrays(y_true, y_prob, num_classes) avg_loss = float(total_loss / total_n) if total_n > 0 else float("nan") return avg_loss, acc, auc, n def evaluate_bilateral( model: BilateralFusionCNN, loader: DataLoader, device: torch.device, num_classes: int, ): model.eval() y_true = [] y_prob = [] total_loss = 0.0 total_n = 0 with torch.no_grad(): for batch in loader: x1 = batch.get("image_1") x2 = batch.get("image_2") y = batch.get("label_1") if not torch.is_tensor(x1) or not torch.is_tensor(x2): continue y_t = to_label_tensor(y, device) logits = model(x1.to(device), x2.to(device)) bs = int(y_t.shape[0]) total_loss += float(F.cross_entropy(logits, y_t).item()) * bs total_n += bs probs = F.softmax(logits, dim=1).cpu().numpy() y_true.append(y_t.cpu().numpy()) y_prob.append(probs) acc, auc, n = _score_arrays(y_true, y_prob, num_classes) avg_loss = float(total_loss / total_n) if total_n > 0 else float("nan") return avg_loss, acc, auc, n def evaluate_two_single_merge( model_od: EyeLevelCNN, model_os: EyeLevelCNN, loader: DataLoader, device: torch.device, num_classes: int, ): model_od.eval() model_os.eval() y_true = [] y_prob = [] total_loss = 0.0 total_n = 0 with torch.no_grad(): for batch in loader: x1 = batch.get("image_1") x2 = batch.get("image_2") y = batch.get("label_1") if not torch.is_tensor(x1) or not torch.is_tensor(x2): continue y_t = to_label_tensor(y, device) p1 = F.softmax(model_od(x1.to(device)), dim=1) p2 = F.softmax(model_os(x2.to(device)), dim=1) p = 0.5 * (p1 + p2) bs = int(y_t.shape[0]) total_loss += float(F.nll_loss(torch.log(p.clamp_min(1e-8)), y_t).item()) * bs total_n += bs y_true.append(y_t.cpu().numpy()) y_prob.append(p.cpu().numpy()) acc, auc, n = _score_arrays(y_true, y_prob, num_classes) avg_loss = float(total_loss / total_n) if total_n > 0 else float("nan") return avg_loss, acc, auc, n def collect_binary_probs_eye(model: EyeLevelCNN, loader: DataLoader, device: torch.device): model.eval() y_true = [] p1 = [] with torch.no_grad(): for batch in loader: x = batch.get("image_1") y = batch.get("label_1") if not torch.is_tensor(x): continue y_t = to_label_tensor(y, device) probs = F.softmax(model(x.to(device)), dim=1)[:, 1] y_true.append(y_t.cpu().numpy()) p1.append(probs.cpu().numpy()) if not y_true: return np.array([]), np.array([]) return np.concatenate(y_true), np.concatenate(p1) def collect_binary_probs_bilateral(model: BilateralFusionCNN, loader: DataLoader, device: torch.device): model.eval() y_true = [] p1 = [] with torch.no_grad(): for batch in loader: x1 = batch.get("image_1") x2 = batch.get("image_2") y = batch.get("label_1") if not torch.is_tensor(x1) or not torch.is_tensor(x2): continue y_t = to_label_tensor(y, device) probs = F.softmax(model(x1.to(device), x2.to(device)), dim=1)[:, 1] y_true.append(y_t.cpu().numpy()) p1.append(probs.cpu().numpy()) if not y_true: return np.array([]), np.array([]) return np.concatenate(y_true), np.concatenate(p1) def collect_binary_probs_merge(model_od: EyeLevelCNN, model_os: EyeLevelCNN, loader: DataLoader, device: torch.device): model_od.eval() model_os.eval() y_true = [] p1 = [] with torch.no_grad(): for batch in loader: x1 = batch.get("image_1") x2 = batch.get("image_2") y = batch.get("label_1") if not torch.is_tensor(x1) or not torch.is_tensor(x2): continue y_t = to_label_tensor(y, device) p_od = F.softmax(model_od(x1.to(device)), dim=1)[:, 1] p_os = F.softmax(model_os(x2.to(device)), dim=1)[:, 1] p = 0.5 * (p_od + p_os) y_true.append(y_t.cpu().numpy()) p1.append(p.cpu().numpy()) if not y_true: return np.array([]), np.array([]) return np.concatenate(y_true), np.concatenate(p1) 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 = 0.5 best_acc = -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 = acc best_t = float(t) return best_t def binary_acc_at_threshold(y_true: np.ndarray, p1: np.ndarray, t: float) -> float: if y_true.size == 0: return float("nan") pred = (p1 >= t).astype(int) return float((pred == y_true).mean()) def collect_probs_eye(model: EyeLevelCNN, loader: DataLoader, device: torch.device): model.eval() y_true = [] probs_all = [] with torch.no_grad(): for batch in loader: x = batch.get("image_1") y = batch.get("label_1") if not torch.is_tensor(x): continue y_t = to_label_tensor(y, device) probs = F.softmax(model(x.to(device)), dim=1).cpu().numpy() y_true.append(y_t.cpu().numpy()) probs_all.append(probs) if not y_true: return np.array([]), np.zeros((0, 0), dtype=float) return np.concatenate(y_true), np.concatenate(probs_all, axis=0) def collect_probs_bilateral(model: BilateralFusionCNN, loader: DataLoader, device: torch.device): model.eval() y_true = [] probs_all = [] with torch.no_grad(): for batch in loader: x1 = batch.get("image_1") x2 = batch.get("image_2") y = batch.get("label_1") if not torch.is_tensor(x1) or not torch.is_tensor(x2): continue y_t = to_label_tensor(y, device) probs = F.softmax(model(x1.to(device), x2.to(device)), dim=1).cpu().numpy() y_true.append(y_t.cpu().numpy()) probs_all.append(probs) if not y_true: return np.array([]), np.zeros((0, 0), dtype=float) return np.concatenate(y_true), np.concatenate(probs_all, axis=0) def collect_probs_merge(model_od: EyeLevelCNN, model_os: EyeLevelCNN, loader: DataLoader, device: torch.device): model_od.eval() model_os.eval() y_true = [] probs_all = [] with torch.no_grad(): for batch in loader: x1 = batch.get("image_1") x2 = batch.get("image_2") y = batch.get("label_1") if not torch.is_tensor(x1) or not torch.is_tensor(x2): continue y_t = to_label_tensor(y, device) p_od = F.softmax(model_od(x1.to(device)), dim=1) p_os = F.softmax(model_os(x2.to(device)), dim=1) probs = (0.5 * (p_od + p_os)).cpu().numpy() y_true.append(y_t.cpu().numpy()) probs_all.append(probs) if not y_true: return np.array([]), np.zeros((0, 0), dtype=float) return np.concatenate(y_true), np.concatenate(probs_all, axis=0) 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) pred = np.argmax(logits, axis=1) return float((pred == 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 = acc best_v = float(v) bias[k] = best_v # small stabilization around baseline if np.isnan(best_acc): bias[k] = old return bias def _serialize_vec(vec: np.ndarray | None) -> str | None: if vec is None: return None if vec.size == 0: return None return "|".join(f"{float(v):.4f}" for v in vec.tolist()) def _score_arrays(y_true_chunks, y_prob_chunks, num_classes: int): if not y_true_chunks: return float("nan"), float("nan"), 0 y = np.concatenate(y_true_chunks, axis=0) p = np.concatenate(y_prob_chunks, axis=0) acc = float((p.argmax(axis=1) == y).mean()) try: if num_classes == 2: auc = float(roc_auc_score(y, p[:, 1])) else: auc = float(roc_auc_score(y, p, multi_class="ovr", average="macro")) except Exception: auc = float("nan") return acc, auc, int(y.shape[0]) def _drop_mixed_label_patients(df: pd.DataFrame, *, 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_ids = [pid for pid, labels in per_patient.items() if len(labels) > 1] if not mixed_ids: return df, [] keep = ~df[patient_col].isin(mixed_ids) return df[keep].reset_index(drop=True), mixed_ids def choose_device(name: str) -> torch.device: if name == "cuda": if not torch.cuda.is_available(): raise RuntimeError("Requested --device cuda 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") def parse_args(): ap = argparse.ArgumentParser( description="Compare eye-level single-tower CNN vs patient-level bilateral dual-tower fusion model." ) 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=["multiclass", "binary"], default="binary") ap.add_argument( "--eval-modes", nargs="+", choices=["multiclass", "binary"], default=None, help="Optional list of modes to run in one pass (e.g. --eval-modes binary multiclass).", ) ap.add_argument("--n-splits", type=int, default=5) ap.add_argument("--fold-seed", type=int, default=42) ap.add_argument("--holdout-per-class", type=int, default=0) ap.add_argument("--holdout-seed", type=int, default=123) ap.add_argument("--folds", type=int, default=5, help="How many folds to run (<= n-splits).") ap.add_argument("--epochs", type=int, default=8) ap.add_argument("--batch-size", type=int, default=8) ap.add_argument("--lr", type=float, default=1e-4) ap.add_argument("--backbone", default="resnet50") ap.add_argument("--freeze-ratio", type=float, default=0.0) ap.add_argument("--augment", action="store_true", help="Enable image augmentation during training.") ap.add_argument("--fusion-dim", type=int, default=256) ap.add_argument("--bridge-se", action="store_true", help="Enable SE in bilateral bridge.") ap.add_argument( "--bilateral-method", choices=["bridge", "two-single-merge"], default="bridge", help="Bilateral comparator: learned dual-tower bridge or two separate single-eye models merged by prob average.", ) 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/basic_analysis") ap.add_argument( "--exclude-binary-mixed-patients", action="store_true", help="Drop patients with mixed eye labels (any disagreement across eyes) before splitting.", ) ap.add_argument( "--log-every", type=int, default=0, help="If > 0, print epoch progress every N epochs within each fold.", ) ap.add_argument( "--tune-binary-threshold", action="store_true", help="In binary mode, tune decision thresholds on validation probs and apply them to val/holdout accuracy.", ) ap.add_argument( "--tune-multiclass-bias", action="store_true", help="In multiclass mode, tune per-class log-prob biases on validation and apply to val/holdout accuracy.", ) return ap.parse_args() def _serialize_float(v: float) -> Optional[float]: return None if np.isnan(v) else float(v) def _rows_from_metrics(metrics: list[FoldMetrics]) -> list[dict]: rows = [] for m in metrics: rows.append( { "mode": m.mode, "fold": m.fold, "baseline_acc": _serialize_float(m.baseline_acc), "baseline_auc": _serialize_float(m.baseline_auc), "bilateral_acc": _serialize_float(m.bilateral_acc), "bilateral_auc": _serialize_float(m.bilateral_auc), "baseline_n": m.baseline_n, "bilateral_n": m.bilateral_n, "os_acc": _serialize_float(m.os_acc), "os_auc": _serialize_float(m.os_auc), "os_n": m.os_n, "holdout_baseline_acc": _serialize_float(m.holdout_baseline_acc), "holdout_baseline_auc": _serialize_float(m.holdout_baseline_auc), "holdout_bilateral_acc": _serialize_float(m.holdout_bilateral_acc), "holdout_bilateral_auc": _serialize_float(m.holdout_bilateral_auc), "holdout_baseline_n": m.holdout_baseline_n, "holdout_bilateral_n": m.holdout_bilateral_n, "holdout_os_acc": _serialize_float(m.holdout_os_acc), "holdout_os_auc": _serialize_float(m.holdout_os_auc), "holdout_os_n": m.holdout_os_n, } ) return rows def _summary_for_mode(mode: str, metrics: list[FoldMetrics]) -> dict: b_accs = np.array([m.baseline_acc for m in metrics], dtype=float) b_aucs = np.array([m.baseline_auc for m in metrics], dtype=float) d_accs = np.array([m.bilateral_acc for m in metrics], dtype=float) d_aucs = np.array([m.bilateral_auc for m in metrics], dtype=float) hb_accs = np.array([m.holdout_baseline_acc for m in metrics], dtype=float) hb_aucs = np.array([m.holdout_baseline_auc for m in metrics], dtype=float) hd_accs = np.array([m.holdout_bilateral_acc for m in metrics], dtype=float) hd_aucs = np.array([m.holdout_bilateral_auc for m in metrics], dtype=float) return { "mode": mode, "baseline": { "acc_mean": _serialize_float(float(np.nanmean(b_accs))), "acc_std": _serialize_float(float(np.nanstd(b_accs))), "auc_mean": _serialize_float(float(np.nanmean(b_aucs))), "auc_std": _serialize_float(float(np.nanstd(b_aucs))), }, "bilateral": { "acc_mean": _serialize_float(float(np.nanmean(d_accs))), "acc_std": _serialize_float(float(np.nanstd(d_accs))), "auc_mean": _serialize_float(float(np.nanmean(d_aucs))), "auc_std": _serialize_float(float(np.nanstd(d_aucs))), }, "delta_bilateral_minus_baseline": { "acc_mean": _serialize_float(float(np.nanmean(d_accs - b_accs))), "auc_mean": _serialize_float(float(np.nanmean(d_aucs - b_aucs))), }, "holdout_baseline": { "acc_mean": _serialize_float(float(np.nanmean(hb_accs))), "acc_std": _serialize_float(float(np.nanstd(hb_accs))), "auc_mean": _serialize_float(float(np.nanmean(hb_aucs))), "auc_std": _serialize_float(float(np.nanstd(hb_aucs))), }, "holdout_bilateral": { "acc_mean": _serialize_float(float(np.nanmean(hd_accs))), "acc_std": _serialize_float(float(np.nanstd(hd_accs))), "auc_mean": _serialize_float(float(np.nanmean(hd_aucs))), "auc_std": _serialize_float(float(np.nanstd(hd_aucs))), }, "holdout_delta_bilateral_minus_baseline": { "acc_mean": _serialize_float(float(np.nanmean(hd_accs - hb_accs))), "auc_mean": _serialize_float(float(np.nanmean(hd_aucs - hb_aucs))), }, } def _print_summary(mode: str, summary: dict) -> None: def fmt(v): return "nan" if v is None else f"{v:.4f}" print(f"\n=== Summary ({mode}) ===") print( "baseline_eye_cnn " f"acc={fmt(summary['baseline']['acc_mean'])}±{fmt(summary['baseline']['acc_std'])} " f"auc={fmt(summary['baseline']['auc_mean'])}±{fmt(summary['baseline']['auc_std'])}" ) print( "bilateral_dual_img " f"acc={fmt(summary['bilateral']['acc_mean'])}±{fmt(summary['bilateral']['acc_std'])} " f"auc={fmt(summary['bilateral']['auc_mean'])}±{fmt(summary['bilateral']['auc_std'])}" ) print( "delta(bilateral-baseline) " f"acc={fmt(summary['delta_bilateral_minus_baseline']['acc_mean'])} " f"auc={fmt(summary['delta_bilateral_minus_baseline']['auc_mean'])}" ) if summary["holdout_baseline"]["acc_mean"] is not None: print( "holdout baseline_eye_cnn " f"acc={fmt(summary['holdout_baseline']['acc_mean'])}±{fmt(summary['holdout_baseline']['acc_std'])} " f"auc={fmt(summary['holdout_baseline']['auc_mean'])}±{fmt(summary['holdout_baseline']['auc_std'])}" ) print( "holdout bilateral_dual_img " f"acc={fmt(summary['holdout_bilateral']['acc_mean'])}±{fmt(summary['holdout_bilateral']['acc_std'])} " f"auc={fmt(summary['holdout_bilateral']['auc_mean'])}±{fmt(summary['holdout_bilateral']['auc_std'])}" ) print( "holdout delta(bilateral-baseline) " f"acc={fmt(summary['holdout_delta_bilateral_minus_baseline']['acc_mean'])} " f"auc={fmt(summary['holdout_delta_bilateral_minus_baseline']['auc_mean'])}" ) def run_mode(args, mode: str, device: torch.device, data, out_dir: Path) -> tuple[list[FoldMetrics], dict]: df_mode = data.df.copy() if args.exclude_binary_mixed_patients: before_rows = len(df_mode) before_patients = int(df_mode["Patient ID"].nunique()) df_mode, mixed_ids = _drop_mixed_label_patients( df_mode, patient_col="Patient ID", label_col=args.label_col ) print( f"[mode={mode}] excluded {len(mixed_ids)} mixed-label patients " f"(rows {before_rows}->{len(df_mode)}, patients {before_patients}->{df_mode['Patient ID'].nunique()})", 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={mode}] preparing splits (num_classes={num_classes}, rows={len(df_mode)}, patients={df_mode['Patient ID'].nunique()})...", flush=True, ) split_manager = PatientFirstSplitManager(patient_col="Patient ID", label_col=args.label_col) split_args = SimpleNamespace( eval_mode=mode, holdout_per_class=args.holdout_per_class, holdout_seed=args.holdout_seed, n_splits=args.n_splits, fold_seed=args.fold_seed, ) clinical_for_split = SimpleNamespace(df=df_mode, label_col=args.label_col) plans = split_manager.build_plans(clinical=clinical_for_split, args=split_args, profile=None) n_folds = min(args.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") fold_metrics: list[FoldMetrics] = [] mode_dir = out_dir / mode mode_dir.mkdir(parents=True, exist_ok=True) for fold in range(n_folds): split = plans[fold] holdout_df = split.holdout fold_seed = args.seed + fold * 100 seed_everything(fold_seed) print( f"[mode={mode}] fold {fold+1}/{n_folds}: building models/loaders...", flush=True, ) fold_dir = mode_dir / f"fold{fold}" fold_dir.mkdir(parents=True, exist_ok=True) train_log_path = fold_dir / "train.log" epoch_log_path = fold_dir / "epoch_log.csv" baseline_model = None bilateral_model = None model_od = None model_os = None if args.bilateral_method == "bridge": baseline_model = EyeLevelCNN( backbone=args.backbone, freeze_ratio=args.freeze_ratio, num_classes=num_classes, augment=args.augment, ).to(device) bilateral_model = BilateralFusionCNN( backbone=args.backbone, freeze_ratio=args.freeze_ratio, num_classes=num_classes, augment=args.augment, use_se=args.bridge_se, fusion_dim=args.fusion_dim, ).to(device) else: model_od = EyeLevelCNN( backbone=args.backbone, freeze_ratio=args.freeze_ratio, num_classes=num_classes, augment=args.augment, ).to(device) model_os = EyeLevelCNN( backbone=args.backbone, freeze_ratio=args.freeze_ratio, num_classes=num_classes, augment=args.augment, ).to(device) baseline_model = model_od eye_train_samples = filter_eye_samples(profile_eye.build_samples(df=split.train, clinical=data)) eye_val_samples = filter_eye_samples(profile_eye.build_samples(df=split.val, clinical=data)) bilat_train_samples = filter_bilateral_samples(profile_patient.build_samples(df=split.train, clinical=data)) bilat_val_samples = filter_bilateral_samples(profile_patient.build_samples(df=split.val, clinical=data)) od_train_samples = patient_to_single_eye_samples(bilat_train_samples, "image_1") od_val_samples = patient_to_single_eye_samples(bilat_val_samples, "image_1") os_train_samples = patient_to_single_eye_samples(bilat_train_samples, "image_2") os_val_samples = patient_to_single_eye_samples(bilat_val_samples, "image_2") eye_holdout_samples = [] bilat_holdout_samples = [] od_holdout_samples = [] os_holdout_samples = [] if holdout_df is not None and not holdout_df.empty: eye_holdout_samples = filter_eye_samples(profile_eye.build_samples(df=holdout_df, clinical=data)) bilat_holdout_samples = filter_bilateral_samples( profile_patient.build_samples(df=holdout_df, clinical=data) ) od_holdout_samples = patient_to_single_eye_samples(bilat_holdout_samples, "image_1") os_holdout_samples = patient_to_single_eye_samples(bilat_holdout_samples, "image_2") baseline_train = None baseline_val = None if args.bilateral_method == "bridge": baseline_train = make_loader( eye_train_samples, profile_eye.slot_descriptors(), image_transform=baseline_model.tower.transform, batch_size=args.batch_size, shuffle=True, num_workers=args.num_workers, ) baseline_val = make_loader( eye_val_samples, profile_eye.slot_descriptors(), image_transform=baseline_model.tower.transform, batch_size=args.batch_size, shuffle=False, num_workers=args.num_workers, ) else: baseline_train = make_loader( od_train_samples, profile_eye.slot_descriptors(), image_transform=model_od.tower.transform, batch_size=args.batch_size, shuffle=True, num_workers=args.num_workers, ) baseline_val = make_loader( od_val_samples, profile_eye.slot_descriptors(), image_transform=model_od.tower.transform, batch_size=args.batch_size, shuffle=False, num_workers=args.num_workers, ) bilateral_train = make_loader( bilat_train_samples, profile_patient.slot_descriptors(), image_transform=(bilateral_model.tower_od.transform if bilateral_model is not None else model_od.tower.transform), batch_size=args.batch_size, shuffle=True, num_workers=args.num_workers, ) bilateral_val = make_loader( bilat_val_samples, profile_patient.slot_descriptors(), image_transform=(bilateral_model.tower_od.transform if bilateral_model is not None else model_od.tower.transform), batch_size=args.batch_size, shuffle=False, num_workers=args.num_workers, ) od_train = None od_val = None os_train = None os_val = None if args.bilateral_method == "two-single-merge": od_train = make_loader( od_train_samples, profile_eye.slot_descriptors(), image_transform=model_od.tower.transform, batch_size=args.batch_size, shuffle=True, num_workers=args.num_workers, ) os_train = make_loader( os_train_samples, profile_eye.slot_descriptors(), image_transform=model_os.tower.transform, batch_size=args.batch_size, shuffle=True, num_workers=args.num_workers, ) od_val = make_loader( od_val_samples, profile_eye.slot_descriptors(), image_transform=model_od.tower.transform, batch_size=args.batch_size, shuffle=False, num_workers=args.num_workers, ) os_val = make_loader( os_val_samples, profile_eye.slot_descriptors(), image_transform=model_os.tower.transform, batch_size=args.batch_size, shuffle=False, num_workers=args.num_workers, ) baseline_holdout = None bilateral_holdout = None os_holdout = None if eye_holdout_samples: if args.bilateral_method == "bridge": baseline_holdout = make_loader( eye_holdout_samples, profile_eye.slot_descriptors(), image_transform=baseline_model.tower.transform, batch_size=args.batch_size, shuffle=False, num_workers=args.num_workers, ) else: baseline_holdout = make_loader( od_holdout_samples, profile_eye.slot_descriptors(), image_transform=model_od.tower.transform, batch_size=args.batch_size, shuffle=False, num_workers=args.num_workers, ) os_holdout = make_loader( os_holdout_samples, profile_eye.slot_descriptors(), image_transform=model_os.tower.transform, batch_size=args.batch_size, shuffle=False, num_workers=args.num_workers, ) if bilat_holdout_samples: bilateral_holdout = make_loader( bilat_holdout_samples, profile_patient.slot_descriptors(), image_transform=(bilateral_model.tower_od.transform if bilateral_model is not None else model_od.tower.transform), batch_size=args.batch_size, shuffle=False, num_workers=args.num_workers, ) opt_base = None opt_bilat = None opt_od = None opt_os = None if args.bilateral_method == "bridge": opt_base = torch.optim.Adam(baseline_model.parameters(), lr=args.lr) opt_bilat = torch.optim.Adam(bilateral_model.parameters(), lr=args.lr) else: opt_od = torch.optim.Adam(model_od.parameters(), lr=args.lr) opt_os = torch.optim.Adam(model_os.parameters(), lr=args.lr) epoch_fields = [ "mode", "fold", "epoch", "baseline_train_loss", "baseline_train_acc", "baseline_train_n", "baseline_val_loss", "baseline_val_acc", "baseline_val_auc", "baseline_val_n", "baseline_threshold", "baseline_bias", "bilateral_train_loss", "bilateral_train_acc", "bilateral_train_n", "bilateral_val_loss", "bilateral_val_acc", "bilateral_val_auc", "bilateral_val_n", "bilateral_threshold", "bilateral_bias", "os_val_loss", "os_val_acc", "os_val_auc", "os_val_n", "os_threshold", "os_bias", "holdout_baseline_loss", "holdout_baseline_acc", "holdout_baseline_auc", "holdout_baseline_n", "holdout_bilateral_loss", "holdout_bilateral_acc", "holdout_bilateral_auc", "holdout_bilateral_n", "holdout_os_loss", "holdout_os_acc", "holdout_os_auc", "holdout_os_n", ] epoch_fp = epoch_log_path.open("w", newline="", encoding="utf-8") epoch_writer = csv.DictWriter(epoch_fp, fieldnames=epoch_fields) epoch_writer.writeheader() print( f"[mode={mode}] fold {fold+1}/{n_folds}: training " f"(epochs={args.epochs}, baseline_train_n={len(eye_train_samples)}, " f"bilateral_train_n={len(bilat_train_samples)}, method={args.bilateral_method})", flush=True, ) b_loss = b_acc = b_auc = float("nan") b_n = 0 b_thr = 0.5 b_bias = None os_loss = os_acc = os_auc = float("nan") os_n = 0 os_thr = 0.5 os_bias = None d_loss = d_acc = d_auc = float("nan") d_n = 0 d_thr = 0.5 d_bias = None hb_loss = hb_acc = hb_auc = float("nan") hb_n = 0 hos_loss = hos_acc = hos_auc = float("nan") hos_n = 0 hd_loss = hd_acc = hd_auc = float("nan") hd_n = 0 with train_log_path.open("w", encoding="utf-8") as train_log: for epoch in range(args.epochs): bt_loss = bt_acc = float("nan") bt_n = 0 ot_loss = ot_acc = float("nan") ot_n = 0 if args.bilateral_method == "bridge": bt_loss, bt_acc, bt_n = train_eye_epoch(baseline_model, baseline_train, opt_base, device) else: bt_loss, bt_acc, bt_n = train_eye_epoch(model_od, od_train, opt_od, device) ot_loss, ot_acc, ot_n = train_eye_epoch(model_os, os_train, opt_os, device) if args.bilateral_method == "bridge": dt_loss, dt_acc, dt_n = train_bilateral_epoch(bilateral_model, bilateral_train, opt_bilat, device) else: dt_loss = float(np.nanmean([bt_loss, ot_loss])) dt_acc = float(np.nanmean([bt_acc, ot_acc])) dt_n = int(min(bt_n, ot_n)) b_loss, b_acc, b_auc, b_n = evaluate_eye(baseline_model, baseline_val, device, num_classes) if args.bilateral_method == "two-single-merge": os_loss, os_acc, os_auc, os_n = evaluate_eye(model_os, os_val, device, num_classes) if args.bilateral_method == "bridge": d_loss, d_acc, d_auc, d_n = evaluate_bilateral(bilateral_model, bilateral_val, device, num_classes) else: d_loss, d_acc, d_auc, d_n = evaluate_two_single_merge( model_od, model_os, bilateral_val, device, num_classes ) if args.tune_binary_threshold and num_classes == 2: yb, pb = collect_binary_probs_eye(baseline_model, baseline_val, device) b_thr = tune_binary_threshold(yb, pb) b_acc = binary_acc_at_threshold(yb, pb, b_thr) if args.bilateral_method == "bridge": yd, pd = collect_binary_probs_bilateral(bilateral_model, bilateral_val, device) else: yd, pd = collect_binary_probs_merge(model_od, model_os, bilateral_val, device) d_thr = tune_binary_threshold(yd, pd) d_acc = binary_acc_at_threshold(yd, pd, d_thr) if args.bilateral_method == "two-single-merge": yo, po = collect_binary_probs_eye(model_os, os_val, device) os_thr = tune_binary_threshold(yo, po) os_acc = binary_acc_at_threshold(yo, po, os_thr) elif args.tune_multiclass_bias and num_classes > 2: yb, pb = collect_probs_eye(baseline_model, baseline_val, device) b_bias = tune_multiclass_bias(yb, pb) b_acc = multiclass_acc_with_bias(yb, pb, b_bias) if args.bilateral_method == "bridge": yd, pd = collect_probs_bilateral(bilateral_model, bilateral_val, device) else: yd, pd = collect_probs_merge(model_od, model_os, bilateral_val, device) d_bias = tune_multiclass_bias(yd, pd) d_acc = multiclass_acc_with_bias(yd, pd, d_bias) if args.bilateral_method == "two-single-merge": yo, po = collect_probs_eye(model_os, os_val, device) os_bias = tune_multiclass_bias(yo, po) os_acc = multiclass_acc_with_bias(yo, po, os_bias) hb_loss = hb_acc = hb_auc = float("nan") hb_n = 0 hos_loss = hos_acc = hos_auc = float("nan") hos_n = 0 hd_loss = hd_acc = hd_auc = float("nan") hd_n = 0 if baseline_holdout is not None: hb_loss, hb_acc, hb_auc, hb_n = evaluate_eye(baseline_model, baseline_holdout, device, num_classes) if os_holdout is not None: hos_loss, hos_acc, hos_auc, hos_n = evaluate_eye(model_os, os_holdout, device, num_classes) if bilateral_holdout is not None: if args.bilateral_method == "bridge": hd_loss, hd_acc, hd_auc, hd_n = evaluate_bilateral( bilateral_model, bilateral_holdout, device, num_classes ) else: hd_loss, hd_acc, hd_auc, hd_n = evaluate_two_single_merge( model_od, model_os, bilateral_holdout, device, num_classes ) if args.tune_binary_threshold and num_classes == 2: if baseline_holdout is not None: yhb, phb = collect_binary_probs_eye(baseline_model, baseline_holdout, device) hb_acc = binary_acc_at_threshold(yhb, phb, b_thr) if os_holdout is not None: yho, pho = collect_binary_probs_eye(model_os, os_holdout, device) hos_acc = binary_acc_at_threshold(yho, pho, os_thr) if bilateral_holdout is not None: if args.bilateral_method == "bridge": yhd, phd = collect_binary_probs_bilateral(bilateral_model, bilateral_holdout, device) else: yhd, phd = collect_binary_probs_merge(model_od, model_os, bilateral_holdout, device) hd_acc = binary_acc_at_threshold(yhd, phd, d_thr) elif args.tune_multiclass_bias and num_classes > 2: if baseline_holdout is not None and b_bias is not None: yhb, phb = collect_probs_eye(baseline_model, baseline_holdout, device) hb_acc = multiclass_acc_with_bias(yhb, phb, b_bias) if os_holdout is not None and os_bias is not None: yho, pho = collect_probs_eye(model_os, os_holdout, device) hos_acc = multiclass_acc_with_bias(yho, pho, os_bias) if bilateral_holdout is not None and d_bias is not None: if args.bilateral_method == "bridge": yhd, phd = collect_probs_bilateral(bilateral_model, bilateral_holdout, device) else: yhd, phd = collect_probs_merge(model_od, model_os, bilateral_holdout, device) hd_acc = multiclass_acc_with_bias(yhd, phd, d_bias) row = { "mode": mode, "fold": fold, "epoch": epoch + 1, "baseline_train_loss": _serialize_float(bt_loss), "baseline_train_acc": _serialize_float(bt_acc), "baseline_train_n": bt_n, "baseline_val_loss": _serialize_float(b_loss), "baseline_val_acc": _serialize_float(b_acc), "baseline_val_auc": _serialize_float(b_auc), "baseline_val_n": b_n, "baseline_threshold": _serialize_float(b_thr if num_classes == 2 else float("nan")), "baseline_bias": _serialize_vec(b_bias if num_classes > 2 else None), "bilateral_train_loss": _serialize_float(dt_loss), "bilateral_train_acc": _serialize_float(dt_acc), "bilateral_train_n": dt_n, "bilateral_val_loss": _serialize_float(d_loss), "bilateral_val_acc": _serialize_float(d_acc), "bilateral_val_auc": _serialize_float(d_auc), "bilateral_val_n": d_n, "bilateral_threshold": _serialize_float(d_thr if num_classes == 2 else float("nan")), "bilateral_bias": _serialize_vec(d_bias if num_classes > 2 else None), "os_val_loss": _serialize_float(os_loss), "os_val_acc": _serialize_float(os_acc), "os_val_auc": _serialize_float(os_auc), "os_val_n": os_n, "os_threshold": _serialize_float(os_thr if (num_classes == 2 and args.bilateral_method == "two-single-merge") else float("nan")), "os_bias": _serialize_vec(os_bias if (num_classes > 2 and args.bilateral_method == "two-single-merge") else None), "holdout_baseline_loss": _serialize_float(hb_loss), "holdout_baseline_acc": _serialize_float(hb_acc), "holdout_baseline_auc": _serialize_float(hb_auc), "holdout_baseline_n": hb_n, "holdout_bilateral_loss": _serialize_float(hd_loss), "holdout_bilateral_acc": _serialize_float(hd_acc), "holdout_bilateral_auc": _serialize_float(hd_auc), "holdout_bilateral_n": hd_n, "holdout_os_loss": _serialize_float(hos_loss), "holdout_os_acc": _serialize_float(hos_acc), "holdout_os_auc": _serialize_float(hos_auc), "holdout_os_n": hos_n, } epoch_writer.writerow(row) epoch_fp.flush() line_train = ( f"[mode={mode} fold={fold+1}/{n_folds} epoch={epoch+1}/{args.epochs}] " f"train baseline(loss={bt_loss:.4f}, acc={bt_acc:.4f}, n={bt_n}) " f"bilateral(loss={dt_loss:.4f}, acc={dt_acc:.4f}, n={dt_n})" ) line_val = ( f"[mode={mode} fold={fold+1}/{n_folds} epoch={epoch+1}/{args.epochs}] " f"val baseline(loss={b_loss:.4f}, acc={b_acc:.4f}, auc={b_auc:.4f}, n={b_n}) " f"bilateral(loss={d_loss:.4f}, acc={d_acc:.4f}, auc={d_auc:.4f}, n={d_n})" ) if args.bilateral_method == "two-single-merge": line_val += f" os(loss={os_loss:.4f}, acc={os_acc:.4f}, auc={os_auc:.4f}, n={os_n})" print(line_train, flush=True) print(line_val, flush=True) train_log.write(line_train + "\n") train_log.write(line_val + "\n") if hb_n > 0 or hd_n > 0: line_holdout = ( f"[mode={mode} fold={fold+1}/{n_folds} epoch={epoch+1}/{args.epochs}] " f"holdout baseline(loss={hb_loss:.4f}, acc={hb_acc:.4f}, auc={hb_auc:.4f}, n={hb_n}) " f"bilateral(loss={hd_loss:.4f}, acc={hd_acc:.4f}, auc={hd_auc:.4f}, n={hd_n})" ) if args.bilateral_method == "two-single-merge": line_holdout += ( f" os(loss={hos_loss:.4f}, acc={hos_acc:.4f}, auc={hos_auc:.4f}, n={hos_n})" ) print(line_holdout, flush=True) train_log.write(line_holdout + "\n") if args.log_every > 0 and ((epoch + 1) % args.log_every == 0 or (epoch + 1) == args.epochs): print( f"[mode={mode}] fold {fold+1}/{n_folds}: epoch {epoch+1}/{args.epochs} checkpoint", flush=True, ) epoch_fp.close() fold_metrics.append( FoldMetrics( mode=mode, fold=fold, baseline_acc=b_acc, baseline_auc=b_auc, bilateral_acc=d_acc, bilateral_auc=d_auc, baseline_n=b_n, bilateral_n=d_n, os_acc=os_acc, os_auc=os_auc, os_n=os_n, holdout_baseline_acc=hb_acc, holdout_baseline_auc=hb_auc, holdout_bilateral_acc=hd_acc, holdout_bilateral_auc=hd_auc, holdout_baseline_n=hb_n, holdout_bilateral_n=hd_n, holdout_os_acc=hos_acc, holdout_os_auc=hos_auc, holdout_os_n=hos_n, ) ) msg = ( f"mode={mode} fold={fold} baseline(val loss={b_loss:.4f}, acc={b_acc:.4f}, auc={b_auc:.4f}, n={b_n}) " f"bilateral(val loss={d_loss:.4f}, acc={d_acc:.4f}, auc={d_auc:.4f}, n={d_n})" ) if args.bilateral_method == "two-single-merge": msg += f" | os(val loss={os_loss:.4f}, acc={os_acc:.4f}, auc={os_auc:.4f}, n={os_n})" if hb_n > 0 or hd_n > 0: msg += ( f" | holdout baseline(loss={hb_loss:.4f}, acc={hb_acc:.4f}, auc={hb_auc:.4f}, n={hb_n}) " f"bilateral(loss={hd_loss:.4f}, acc={hd_acc:.4f}, auc={hd_auc:.4f}, n={hd_n})" ) if args.bilateral_method == "two-single-merge": msg += ( f" os(loss={hos_loss:.4f}, acc={hos_acc:.4f}, auc={hos_auc:.4f}, n={hos_n})" ) print(msg) summary = _summary_for_mode(mode, fold_metrics) _print_summary(mode, summary) return fold_metrics, summary def main(): args = parse_args() 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 PAPILA data: rows={len(data.df)}", flush=True) eval_modes = args.eval_modes if args.eval_modes else [args.eval_mode] print(f"Eval modes: {eval_modes}", flush=True) ts = time.strftime("%Y%m%d_%H%M%S") run_name = args.run_name or f"dual_eye_compare_{ts}" out_dir = Path(args.output_root) / run_name out_dir.mkdir(parents=True, exist_ok=True) all_rows = [] summaries = {} for mode in eval_modes: fold_metrics, summary = run_mode(args, mode, device, data, out_dir) rows = _rows_from_metrics(fold_metrics) all_rows.extend(rows) summaries[mode] = summary mode_csv = out_dir / f"{mode}_fold_metrics.csv" if rows: with mode_csv.open("w", newline="", encoding="utf-8") as fh: writer = csv.DictWriter(fh, fieldnames=list(rows[0].keys())) writer.writeheader() writer.writerows(rows) all_csv = out_dir / "all_fold_metrics.csv" if all_rows: with all_csv.open("w", newline="", encoding="utf-8") as fh: writer = csv.DictWriter(fh, fieldnames=list(all_rows[0].keys())) writer.writeheader() writer.writerows(all_rows) payload = { "run_name": run_name, "timestamp": ts, "config": vars(args), "summaries": summaries, } summary_json = out_dir / "summary.json" summary_json.write_text(json.dumps(payload, indent=2), encoding="utf-8") print(f"\nOutputs written to: {out_dir}") if __name__ == "__main__": main()