rpotter6298
1026 lines
38 KiB
Python
1026 lines
38 KiB
Python
#!/usr/bin/env python3
|
||
"""
|
||
Compare all-eye baseline (ImageTower) vs SiameseImageTower bilateral model.
|
||
|
||
Key differences from compare_siamese_tower.py (v1):
|
||
- Baseline trains on ALL eye samples (both OD + OS rows) — matches the
|
||
original grid-search training regime, not just OD-from-bilateral.
|
||
- No holdout set: pure k-fold CV is sufficient for architecture comparison.
|
||
- Both models evaluated at patient level on the same bilateral val set:
|
||
- Baseline: runs on OD and OS separately, mean-pools probabilities.
|
||
- Siamese: runs on both eyes simultaneously.
|
||
- Extended metrics at best-epoch snapshots: kappa, MCC, macro-F1,
|
||
per-class recall, and ECE (Expected Calibration Error).
|
||
"""
|
||
from __future__ import annotations
|
||
|
||
import argparse
|
||
import copy
|
||
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 (
|
||
cohen_kappa_score,
|
||
f1_score,
|
||
matthews_corrcoef,
|
||
recall_score,
|
||
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 (
|
||
ImageTower,
|
||
PatientFirstSplitManager,
|
||
SiameseImageTower,
|
||
SlotDataset,
|
||
build_papila_data,
|
||
build_papila_profile,
|
||
slot_collate,
|
||
)
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Reproducibility
|
||
# ---------------------------------------------------------------------------
|
||
|
||
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)
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Models
|
||
# ---------------------------------------------------------------------------
|
||
|
||
class BaselineCNN(nn.Module):
|
||
"""Single-eye image tower with a linear head."""
|
||
|
||
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:
|
||
return self.head(self.tower(x))
|
||
|
||
|
||
class SiameseCNN(nn.Module):
|
||
"""Bilateral image model using SiameseImageTower (shared backbone)."""
|
||
|
||
def __init__(self, *, backbone: str, freeze_ratio: float, num_classes: int, augment: bool):
|
||
super().__init__()
|
||
self.tower = SiameseImageTower(
|
||
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_od: torch.Tensor, x_os: torch.Tensor) -> torch.Tensor:
|
||
return self.head(self.tower(x_od, x_os))
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Data helpers
|
||
# ---------------------------------------------------------------------------
|
||
|
||
def filter_eye_samples(samples: list[dict]) -> list[dict]:
|
||
"""Keep any single-eye sample with a valid image and label (OD or OS)."""
|
||
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]:
|
||
"""Keep only patient-level samples where both eyes are present."""
|
||
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 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 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)
|
||
)
|
||
mixed_patients = per_patient_unique[per_patient_unique > 1].index.tolist()
|
||
return out.reset_index(drop=True), changed_rows, mixed_patients
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# 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:
|
||
"""
|
||
Returns kappa, mcc, macro_f1, per_class_recall (np.ndarray), ece.
|
||
All float('nan') on empty input or single-class edge cases.
|
||
"""
|
||
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 = 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 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
|
||
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):
|
||
"""Score pre-collected arrays. 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))
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Train / collect
|
||
# ---------------------------------------------------------------------------
|
||
|
||
def train_baseline_epoch(model: BaselineCNN, loader: DataLoader, opt, device):
|
||
"""Train on single-eye batches (image_1)."""
|
||
model.train()
|
||
total_loss = total_correct = 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)
|
||
opt.zero_grad()
|
||
loss.backward()
|
||
opt.step()
|
||
bs = y.shape[0]
|
||
total_loss += float(loss.item()) * bs
|
||
total_correct += int((logits.argmax(1) == y).sum())
|
||
total_n += bs
|
||
return (
|
||
total_loss / total_n if total_n else float("nan"),
|
||
total_correct / total_n if total_n else float("nan"),
|
||
)
|
||
|
||
|
||
def train_siamese_epoch(model: SiameseCNN, loader: DataLoader, opt, device):
|
||
"""Train on bilateral patient batches (image_1 = OD, image_2 = OS)."""
|
||
model.train()
|
||
total_loss = total_correct = total_n = 0
|
||
for batch in loader:
|
||
x1 = batch.get("image_1")
|
||
x2 = batch.get("image_2")
|
||
y = batch.get("label_1")
|
||
if not torch.is_tensor(x1) or not torch.is_tensor(x2):
|
||
continue
|
||
y = to_label_tensor(y, device)
|
||
logits = model(x1.to(device), x2.to(device))
|
||
loss = F.cross_entropy(logits, y)
|
||
opt.zero_grad()
|
||
loss.backward()
|
||
opt.step()
|
||
bs = y.shape[0]
|
||
total_loss += float(loss.item()) * bs
|
||
total_correct += int((logits.argmax(1) == y).sum())
|
||
total_n += bs
|
||
return (
|
||
total_loss / total_n if total_n else float("nan"),
|
||
total_correct / total_n if total_n else float("nan"),
|
||
)
|
||
|
||
|
||
def collect_probs_baseline_bilateral(
|
||
model: BaselineCNN,
|
||
loader: DataLoader,
|
||
device: torch.device,
|
||
) -> tuple[np.ndarray, np.ndarray]:
|
||
"""
|
||
Evaluate baseline on bilateral (patient-level) samples.
|
||
|
||
For each patient, runs the single-eye model on OD (image_1) and OS
|
||
(image_2) separately, then mean-pools the softmax probabilities.
|
||
Returns (y_true [N], probs [N, C]) at patient level.
|
||
"""
|
||
model.eval()
|
||
y_chunks, p_chunks = [], []
|
||
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(x1.to(device)), dim=1)
|
||
p_os = F.softmax(model(x2.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)
|
||
|
||
|
||
def collect_probs_siamese(
|
||
model: SiameseCNN,
|
||
loader: DataLoader,
|
||
device: torch.device,
|
||
) -> tuple[np.ndarray, np.ndarray]:
|
||
"""Evaluate siamese on bilateral (patient-level) samples."""
|
||
model.eval()
|
||
y_chunks, p_chunks = [], []
|
||
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 = F.softmax(model(x1.to(device), x2.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)
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# 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_base: int
|
||
best_epoch_siam: int
|
||
# Baseline metrics at best-val epoch (patient-level, prob-aggregated)
|
||
base_val_auc: float
|
||
base_val_acc: float
|
||
base_val_kappa: float
|
||
base_val_mcc: float
|
||
base_val_f1: float
|
||
base_val_recall: Optional[str] # pipe-delimited per-class recall
|
||
base_val_ece: float
|
||
base_val_threshold: float
|
||
base_val_bias: Optional[str]
|
||
base_val_n: int
|
||
# Siamese metrics at best-val epoch
|
||
siam_val_auc: float
|
||
siam_val_acc: float
|
||
siam_val_kappa: float
|
||
siam_val_mcc: float
|
||
siam_val_f1: float
|
||
siam_val_recall: Optional[str]
|
||
siam_val_ece: float
|
||
siam_val_threshold: float
|
||
siam_val_bias: Optional[str]
|
||
siam_val_n: int
|
||
# Train sample sizes (informational)
|
||
base_train_n: int
|
||
siam_train_n: int
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Output helpers
|
||
# ---------------------------------------------------------------------------
|
||
|
||
def _f(v) -> Optional[float]:
|
||
"""Nan-safe float serialiser."""
|
||
if v is None or (isinstance(v, float) and np.isnan(v)):
|
||
return None
|
||
return round(float(v), 6)
|
||
|
||
|
||
def _sv(vec) -> Optional[str]:
|
||
"""Serialise a numeric vector to a pipe-delimited string."""
|
||
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,
|
||
fold_dir: Path,
|
||
) -> FoldResult:
|
||
# ---- build samples -------------------------------------------------------
|
||
# Baseline trains on ALL eye-level samples (OD + OS as separate rows)
|
||
eye_train = filter_eye_samples(
|
||
profile_eye.build_samples(df=split.train, clinical=data)
|
||
)
|
||
# Siamese trains on bilateral patient-level samples
|
||
bilat_train = filter_bilateral_samples(
|
||
profile_patient.build_samples(df=split.train, clinical=data)
|
||
)
|
||
# Val: bilateral patients only — shared between both model evaluations
|
||
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)
|
||
nan = _nan()
|
||
return FoldResult(
|
||
mode=mode, fold=fold,
|
||
best_epoch_base=0, best_epoch_siam=0,
|
||
base_val_auc=nan, base_val_acc=nan, base_val_kappa=nan,
|
||
base_val_mcc=nan, base_val_f1=nan, base_val_recall=None,
|
||
base_val_ece=nan, base_val_threshold=nan, base_val_bias=None, base_val_n=0,
|
||
siam_val_auc=nan, siam_val_acc=nan, siam_val_kappa=nan,
|
||
siam_val_mcc=nan, siam_val_f1=nan, siam_val_recall=None,
|
||
siam_val_ece=nan, siam_val_threshold=nan, siam_val_bias=None, siam_val_n=0,
|
||
base_train_n=len(eye_train), siam_train_n=len(bilat_train),
|
||
)
|
||
|
||
# ---- models --------------------------------------------------------------
|
||
baseline = BaselineCNN(
|
||
backbone=args.backbone, freeze_ratio=args.freeze_ratio,
|
||
num_classes=num_classes, augment=args.augment,
|
||
).to(device)
|
||
siamese = SiameseCNN(
|
||
backbone=args.backbone, freeze_ratio=args.freeze_ratio,
|
||
num_classes=num_classes, augment=args.augment,
|
||
).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 -------------------------------------------------------------
|
||
# Baseline uses eye-level transform; siamese shares the same backbone
|
||
# transform so both models see the same normalisation at eval time.
|
||
train_base = make_loader(
|
||
eye_train, slots_eye,
|
||
image_transform=baseline.tower.transform, shuffle=True, **loader_kw,
|
||
)
|
||
train_siam = make_loader(
|
||
bilat_train, slots_patient,
|
||
image_transform=siamese.tower.transform, shuffle=True, **loader_kw,
|
||
)
|
||
# Shared val loader — both models read from this
|
||
val_loader = make_loader(
|
||
bilat_val, slots_patient,
|
||
image_transform=baseline.tower.transform, shuffle=False, **loader_kw,
|
||
)
|
||
|
||
opt_base = torch.optim.Adam(baseline.parameters(), lr=args.lr)
|
||
opt_siam = torch.optim.Adam(siamese.parameters(), lr=args.lr)
|
||
|
||
# ---- epoch log -----------------------------------------------------------
|
||
epoch_fields = [
|
||
"fold", "epoch",
|
||
"base_train_loss", "base_train_acc",
|
||
"base_val_auc", "base_val_acc", "base_val_n",
|
||
"base_val_threshold",
|
||
"base_val_bias",
|
||
"siam_train_loss", "siam_train_acc",
|
||
"siam_val_auc", "siam_val_acc", "siam_val_n",
|
||
"siam_val_threshold",
|
||
"siam_val_bias",
|
||
"is_best_base", "is_best_siam",
|
||
]
|
||
epoch_fp = (fold_dir / "epoch_log.csv").open("w", newline="", encoding="utf-8")
|
||
epoch_writer = csv.DictWriter(epoch_fp, fieldnames=epoch_fields)
|
||
epoch_writer.writeheader()
|
||
|
||
# ---- best-epoch trackers -------------------------------------------------
|
||
best_base_auc = -1.0
|
||
best_siam_auc = -1.0
|
||
best_base_state: Optional[dict] = None
|
||
best_siam_state: Optional[dict] = None
|
||
best_epoch_base = 0
|
||
best_epoch_siam = 0
|
||
snap_base: dict = {}
|
||
snap_siam: dict = {}
|
||
|
||
print(
|
||
f" [fold {fold+1}] base_train_n={len(eye_train)} (eye-level) "
|
||
f"siam_train_n={len(bilat_train)} (bilateral) val_n={len(bilat_val)}",
|
||
flush=True,
|
||
)
|
||
|
||
# ---- epoch loop ----------------------------------------------------------
|
||
for epoch in range(args.epochs):
|
||
bl_loss, bl_acc = train_baseline_epoch(baseline, train_base, opt_base, device)
|
||
si_loss, si_acc = train_siamese_epoch(siamese, train_siam, opt_siam, device)
|
||
|
||
y_b, p_b = collect_probs_baseline_bilateral(baseline, val_loader, device)
|
||
y_s, p_s = collect_probs_siamese(siamese, val_loader, device)
|
||
|
||
b_acc, b_auc, b_n = _score_arrays(y_b, p_b, num_classes)
|
||
s_acc, s_auc, s_n = _score_arrays(y_s, p_s, num_classes)
|
||
b_thr = 0.5
|
||
s_thr = 0.5
|
||
b_bias = None
|
||
s_bias = None
|
||
b_ext_preds = None
|
||
s_ext_preds = None
|
||
if args.tune_binary_threshold and num_classes == 2 and b_n > 0 and s_n > 0:
|
||
b_thr = tune_binary_threshold(y_b, p_b[:, 1])
|
||
s_thr = tune_binary_threshold(y_s, p_s[:, 1])
|
||
b_ext_preds = (p_b[:, 1] >= b_thr).astype(int)
|
||
s_ext_preds = (p_s[:, 1] >= s_thr).astype(int)
|
||
b_acc = float((b_ext_preds == y_b).mean())
|
||
s_acc = float((s_ext_preds == y_s).mean())
|
||
elif args.tune_multiclass_bias and num_classes > 2 and b_n > 0 and s_n > 0:
|
||
b_bias = tune_multiclass_bias(y_b, p_b)
|
||
s_bias = tune_multiclass_bias(y_s, p_s)
|
||
b_logits = np.log(np.clip(p_b, 1e-8, 1.0)) + b_bias.reshape(1, -1)
|
||
s_logits = np.log(np.clip(p_s, 1e-8, 1.0)) + s_bias.reshape(1, -1)
|
||
b_ext_preds = np.argmax(b_logits, axis=1)
|
||
s_ext_preds = np.argmax(s_logits, axis=1)
|
||
b_acc = float((b_ext_preds == y_b).mean())
|
||
s_acc = float((s_ext_preds == y_s).mean())
|
||
|
||
# Independent best-epoch update per model
|
||
is_best_base = (not np.isnan(b_auc)) and (b_auc > best_base_auc)
|
||
if is_best_base:
|
||
best_base_auc = b_auc
|
||
best_base_state = copy.deepcopy(baseline.state_dict())
|
||
best_epoch_base = epoch + 1
|
||
ext = compute_extended_metrics(
|
||
y_b, p_b, num_classes, n_bins=args.ece_bins, preds_override=b_ext_preds
|
||
)
|
||
snap_base = dict(
|
||
auc=b_auc, acc=b_acc, n=b_n,
|
||
kappa=ext["kappa"], mcc=ext["mcc"], macro_f1=ext["macro_f1"],
|
||
per_class_recall=ext["per_class_recall"], ece=ext["ece"], threshold=b_thr, bias=b_bias,
|
||
)
|
||
|
||
is_best_siam = (not np.isnan(s_auc)) and (s_auc > best_siam_auc)
|
||
if is_best_siam:
|
||
best_siam_auc = s_auc
|
||
best_siam_state = copy.deepcopy(siamese.state_dict())
|
||
best_epoch_siam = epoch + 1
|
||
ext = compute_extended_metrics(
|
||
y_s, p_s, num_classes, n_bins=args.ece_bins, preds_override=s_ext_preds
|
||
)
|
||
snap_siam = dict(
|
||
auc=s_auc, acc=s_acc, n=s_n,
|
||
kappa=ext["kappa"], mcc=ext["mcc"], macro_f1=ext["macro_f1"],
|
||
per_class_recall=ext["per_class_recall"], ece=ext["ece"], threshold=s_thr, bias=s_bias,
|
||
)
|
||
|
||
epoch_writer.writerow({
|
||
"fold": fold, "epoch": epoch + 1,
|
||
"base_train_loss": _f(bl_loss), "base_train_acc": _f(bl_acc),
|
||
"base_val_auc": _f(b_auc), "base_val_acc": _f(b_acc), "base_val_n": b_n,
|
||
"base_val_threshold": _f(b_thr if num_classes == 2 else float("nan")),
|
||
"base_val_bias": _svf(b_bias if num_classes > 2 else None),
|
||
"siam_train_loss": _f(si_loss), "siam_train_acc": _f(si_acc),
|
||
"siam_val_auc": _f(s_auc), "siam_val_acc": _f(s_acc), "siam_val_n": s_n,
|
||
"siam_val_threshold": _f(s_thr if num_classes == 2 else float("nan")),
|
||
"siam_val_bias": _svf(s_bias if num_classes > 2 else None),
|
||
"is_best_base": int(is_best_base),
|
||
"is_best_siam": int(is_best_siam),
|
||
})
|
||
epoch_fp.flush()
|
||
|
||
if args.log_every > 0 and (epoch + 1) % args.log_every == 0:
|
||
print(
|
||
f" ep {epoch+1:>3}/{args.epochs} "
|
||
f"base val AUC={b_auc:.4f} siam val AUC={s_auc:.4f} "
|
||
f"(best base={best_base_auc:.4f} @ep{best_epoch_base} "
|
||
f"best siam={best_siam_auc:.4f} @ep{best_epoch_siam})",
|
||
flush=True,
|
||
)
|
||
|
||
epoch_fp.close()
|
||
|
||
if args.save_checkpoints:
|
||
if best_base_state is not None:
|
||
torch.save(best_base_state, fold_dir / "best_baseline.pt")
|
||
if best_siam_state is not None:
|
||
torch.save(best_siam_state, fold_dir / "best_siamese.pt")
|
||
|
||
nan = _nan()
|
||
|
||
print(
|
||
f" [fold {fold+1}] BEST "
|
||
f"base AUC={snap_base.get('auc', nan):.4f} "
|
||
f"kappa={snap_base.get('kappa', nan):.4f} "
|
||
f"F1={snap_base.get('macro_f1', nan):.4f} "
|
||
f"ECE={snap_base.get('ece', nan):.4f} @ep{best_epoch_base} | "
|
||
f"siam AUC={snap_siam.get('auc', nan):.4f} "
|
||
f"kappa={snap_siam.get('kappa', nan):.4f} "
|
||
f"F1={snap_siam.get('macro_f1', nan):.4f} "
|
||
f"ECE={snap_siam.get('ece', nan):.4f} @ep{best_epoch_siam}",
|
||
flush=True,
|
||
)
|
||
|
||
return FoldResult(
|
||
mode=mode, fold=fold,
|
||
best_epoch_base=best_epoch_base, best_epoch_siam=best_epoch_siam,
|
||
base_val_auc=snap_base.get("auc", nan),
|
||
base_val_acc=snap_base.get("acc", nan),
|
||
base_val_kappa=snap_base.get("kappa", nan),
|
||
base_val_mcc=snap_base.get("mcc", nan),
|
||
base_val_f1=snap_base.get("macro_f1", nan),
|
||
base_val_recall=_sv(snap_base.get("per_class_recall")),
|
||
base_val_ece=snap_base.get("ece", nan),
|
||
base_val_threshold=snap_base.get("threshold", nan),
|
||
base_val_bias=_svf(snap_base.get("bias")),
|
||
base_val_n=snap_base.get("n", 0),
|
||
siam_val_auc=snap_siam.get("auc", nan),
|
||
siam_val_acc=snap_siam.get("acc", nan),
|
||
siam_val_kappa=snap_siam.get("kappa", nan),
|
||
siam_val_mcc=snap_siam.get("mcc", nan),
|
||
siam_val_f1=snap_siam.get("macro_f1", nan),
|
||
siam_val_recall=_sv(snap_siam.get("per_class_recall")),
|
||
siam_val_ece=snap_siam.get("ece", nan),
|
||
siam_val_threshold=snap_siam.get("threshold", nan),
|
||
siam_val_bias=_svf(snap_siam.get("bias")),
|
||
siam_val_n=snap_siam.get("n", 0),
|
||
base_train_n=len(eye_train),
|
||
siam_train_n=len(bilat_train),
|
||
)
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Summary helpers
|
||
# ---------------------------------------------------------------------------
|
||
|
||
def _summary(results: list[FoldResult]) -> dict:
|
||
def _ms(vals):
|
||
v = np.array([x for x in vals if not np.isnan(float(x)) if x is not None], dtype=float)
|
||
return (
|
||
float(np.mean(v)) if v.size else None,
|
||
float(np.std(v)) if v.size else None,
|
||
)
|
||
|
||
metrics = ["auc", "acc", "kappa", "mcc", "f1", "ece", "threshold"]
|
||
out = {}
|
||
for label, prefix in [("baseline_best_val", "base_val"), ("siamese_best_val", "siam_val")]:
|
||
sub = {}
|
||
for m in metrics:
|
||
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
|
||
|
||
# Per-fold deltas (siamese − baseline)
|
||
delta = {}
|
||
for m in ["auc", "f1", "kappa"]:
|
||
pairs = [
|
||
getattr(r, f"siam_val_{m}") - getattr(r, f"base_val_{m}")
|
||
for r in results
|
||
if not np.isnan(float(getattr(r, f"base_val_{m}")))
|
||
and not np.isnan(float(getattr(r, f"siam_val_{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_val"] = delta
|
||
|
||
out["n_folds_completed"] = len(results)
|
||
out["base_train_mode"] = "eye-level (all OD+OS samples)"
|
||
out["siam_train_mode"] = "patient-level (bilateral only)"
|
||
out["eval_mode"] = "patient-level bilateral (both models, same val set)"
|
||
return out
|
||
|
||
|
||
def _print_summary(mode: str, s: dict) -> None:
|
||
def f(v):
|
||
return "nan" if v is None else f"{v:.4f}"
|
||
|
||
bv = s["baseline_best_val"]
|
||
sv = s["siamese_best_val"]
|
||
dv = s["delta_val"]
|
||
|
||
print(f"\n=== Summary [{mode}] — best-epoch, patient-level bilateral val ===")
|
||
print(f" {'':22s} {'AUC':>8} {'ACC':>8} {'Kappa':>8} {'F1-mac':>8} {'ECE':>8} {'Thr':>8}")
|
||
print(
|
||
f" {'baseline (eye-lvl tr)':22s} "
|
||
f"{f(bv['auc_mean']):>8} {f(bv['acc_mean']):>8} "
|
||
f"{f(bv['kappa_mean']):>8} {f(bv['f1_mean']):>8} {f(bv['ece_mean']):>8} {f(bv['threshold_mean']):>8}"
|
||
)
|
||
print(
|
||
f" {'siamese (bilateral tr)':22s} "
|
||
f"{f(sv['auc_mean']):>8} {f(sv['acc_mean']):>8} "
|
||
f"{f(sv['kappa_mean']):>8} {f(sv['f1_mean']):>8} {f(sv['ece_mean']):>8} {f(sv['threshold_mean']):>8}"
|
||
)
|
||
print(
|
||
f" {'delta (siam − base)':22s} "
|
||
f"{f(dv['auc_mean']):>8} {'':>8} "
|
||
f"{f(dv['kappa_mean']):>8} {f(dv['f1_mean']):>8}"
|
||
)
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# CLI
|
||
# ---------------------------------------------------------------------------
|
||
|
||
def parse_args():
|
||
ap = argparse.ArgumentParser(
|
||
description=(
|
||
"All-eye baseline (ImageTower) vs SiameseImageTower bilateral model. "
|
||
"Pure k-fold CV — no holdout set."
|
||
)
|
||
)
|
||
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(
|
||
"--eval-modes", nargs="+", choices=["binary", "multiclass"], default=None,
|
||
help="Run multiple eval modes in one pass.",
|
||
)
|
||
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=5)
|
||
ap.add_argument("--epochs", type=int, default=40)
|
||
ap.add_argument("--batch-size", type=int, default=8)
|
||
ap.add_argument("--lr", type=float, default=1e-4)
|
||
ap.add_argument("--backbone", default="refugelike")
|
||
ap.add_argument("--freeze-ratio", type=float, default=0.0)
|
||
ap.add_argument("--augment", action="store_true")
|
||
ap.add_argument("--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-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",
|
||
help="Keep mixed-label patients (default behavior).",
|
||
)
|
||
ap.add_argument(
|
||
"--keep-mixed-raw-labels",
|
||
action="store_true",
|
||
help="When mixed patients are included, keep original per-eye labels (default is relabel to patient max label).",
|
||
)
|
||
ap.set_defaults(exclude_mixed_patients=False)
|
||
ap.add_argument("--log-every", type=int, default=5)
|
||
ap.add_argument(
|
||
"--tune-binary-threshold",
|
||
action="store_true",
|
||
help="Tune per-model binary threshold on validation each epoch and use it for ACC/F1/Kappa/MCC/recall.",
|
||
)
|
||
ap.add_argument(
|
||
"--tune-multiclass-bias",
|
||
action="store_true",
|
||
help="Tune per-model multiclass log-prob bias on validation each epoch and use it for ACC/F1/Kappa/MCC/recall.",
|
||
)
|
||
ap.add_argument("--ece-bins", type=int, default=10,
|
||
help="Number of bins for ECE calibration calculation.")
|
||
ap.add_argument("--save-checkpoints", action="store_true",
|
||
help="Save best model state dicts (disabled by default to save disk).")
|
||
return ap.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 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: {len(data.df)} rows", flush=True)
|
||
|
||
ts = time.strftime("%Y%m%d_%H%M%S")
|
||
run_name = args.run_name or f"siamese_v2_{ts}"
|
||
out_dir = Path(args.output_root) / run_name
|
||
out_dir.mkdir(parents=True, exist_ok=True)
|
||
|
||
eval_modes = args.eval_modes if args.eval_modes else [args.eval_mode]
|
||
|
||
all_results: dict[str, list[FoldResult]] = {}
|
||
summaries: dict[str, dict] = {}
|
||
|
||
for mode in eval_modes:
|
||
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.keep_mixed_raw_labels:
|
||
print(f"[{mode}] keeping mixed-label patients with raw per-eye labels.", flush=True)
|
||
else:
|
||
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}] included mixed-label patients; relabeled to patient max severity "
|
||
f"(changed_rows={changed_rows}, rows={before_rows}->{len(df_mode)}, remaining_mixed={len(still_mixed)}).",
|
||
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, # No holdout by design
|
||
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)
|
||
n_folds = min(args.folds, len(plans))
|
||
|
||
# Two profiles: eye-level for baseline training, patient-level for
|
||
# siamese training and shared bilateral val evaluation.
|
||
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"
|
||
)
|
||
|
||
mode_dir = out_dir / mode
|
||
mode_dir.mkdir(exist_ok=True)
|
||
|
||
fold_results: list[FoldResult] = []
|
||
for fold in range(n_folds):
|
||
fold_seed = args.seed + fold * 100
|
||
seed_everything(fold_seed)
|
||
fold_dir = mode_dir / f"fold{fold}"
|
||
fold_dir.mkdir(exist_ok=True)
|
||
|
||
print(f"\n[{mode}] fold {fold+1}/{n_folds}", flush=True)
|
||
result = 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,
|
||
fold_dir=fold_dir,
|
||
)
|
||
fold_results.append(result)
|
||
|
||
# Write per-mode fold CSV
|
||
fold_csv = out_dir / f"{mode}_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(mode, summary)
|
||
|
||
all_results[mode] = fold_results
|
||
summaries[mode] = summary
|
||
|
||
payload = {
|
||
"run_name": run_name,
|
||
"timestamp": ts,
|
||
"config": vars(args),
|
||
"summaries": summaries,
|
||
}
|
||
(out_dir / "summary.json").write_text(json.dumps(payload, indent=2), encoding="utf-8")
|
||
print(f"\nOutputs written to: {out_dir}")
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main()
|