rpotter6298
462 lines
16 KiB
Python
462 lines
16 KiB
Python
#!/usr/bin/env python3
|
|
"""
|
|
V2-parity metadata-only runner.
|
|
|
|
Goal:
|
|
- Match V2HyperTower single-model metadata-only behavior as closely as possible.
|
|
- Avoid image tower/image IO overhead in forward/training.
|
|
|
|
How parity is achieved:
|
|
- Uses PatientFirstSplitManager (same split policy).
|
|
- Uses PAPILA profile builders + V2 filters:
|
|
- eye_train = filter_eye_samples(...)
|
|
- bilat_val/test = filter_bilateral_samples(...)
|
|
- Uses V2 training/eval helpers directly:
|
|
- train_single_epoch(...)
|
|
- collect_probs_single_components(...)
|
|
- Uses bridge_mode="metadata_only".
|
|
|
|
Implementation detail:
|
|
- Batch dictionaries still include image slots to satisfy shared V2 helpers,
|
|
but these are tiny dummy tensors and are never consumed in metadata-only mode.
|
|
|
|
Outputs:
|
|
analysis_data/{run_name}/{eval_mode}/{tower_mode}/fold{N}/
|
|
y_true.npy
|
|
probs_classic.npy or probs_ensemble.npy
|
|
y_true_holdout.npy
|
|
probs_classic_holdout.npy or probs_ensemble_holdout.npy
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import copy
|
|
import json
|
|
import sys
|
|
import time
|
|
from pathlib import Path
|
|
from types import SimpleNamespace
|
|
|
|
# ensure repo root is on sys.path when run directly
|
|
_REPO_ROOT = Path(__file__).resolve().parents[3]
|
|
if str(_REPO_ROOT) not in sys.path:
|
|
sys.path.insert(0, str(_REPO_ROOT))
|
|
|
|
import numpy as np
|
|
import torch
|
|
from torch import nn
|
|
from torch.utils.data import DataLoader, Dataset
|
|
|
|
from classes.v2.bridges import Bridge
|
|
from classes.v2.loader_factory import (
|
|
build_balanced_sampler,
|
|
filter_bilateral_samples,
|
|
filter_eye_samples,
|
|
)
|
|
from classes.v2.metrics import _score_arrays
|
|
from classes.v2.models import collect_probs_single_components, train_single_epoch
|
|
from classes.v2.papila_builders import build_papila_data
|
|
from classes.v2.profiles import build_papila_profile
|
|
from classes.v2.split_manager import PatientFirstSplitManager
|
|
from classes.v2.towers import MDTower
|
|
from classes.v2.utils import choose_device, seed_everything
|
|
|
|
|
|
class MetadataOnlySingleHT(nn.Module):
|
|
"""SingleEyeHT-compatible shell without real image tower usage."""
|
|
|
|
def __init__(
|
|
self,
|
|
*,
|
|
clinical_data,
|
|
num_classes: int,
|
|
md_hidden_dim: int,
|
|
fusion_dim: int,
|
|
dropout: float,
|
|
use_se: bool,
|
|
se_reduction: int,
|
|
se_pre_norm: bool,
|
|
):
|
|
super().__init__()
|
|
# Placeholder module to satisfy phase toggling logic.
|
|
self.img_tower = nn.Identity()
|
|
self.md_tower = MDTower(
|
|
clinical_data=clinical_data,
|
|
hidden_dim=md_hidden_dim,
|
|
dropout=dropout,
|
|
use_se=use_se,
|
|
se_reduction=se_reduction,
|
|
se_pre_norm=se_pre_norm,
|
|
)
|
|
# img_dim is irrelevant in metadata_only mode, but Bridge defines img head params.
|
|
self.bridge = Bridge(
|
|
img_dim=1,
|
|
meta_dim=self.md_tower.out_dim,
|
|
num_classes=num_classes,
|
|
fusion_dim=fusion_dim,
|
|
mode="metadata_only",
|
|
use_se=False,
|
|
se_reduction=16,
|
|
se_pre_norm=True,
|
|
)
|
|
|
|
|
|
class EyeMetaDataset(Dataset):
|
|
"""Eye-level dataset for V2 train_single_epoch input contract."""
|
|
|
|
def __init__(self, samples: list[dict]):
|
|
self.samples = samples
|
|
|
|
def __len__(self) -> int:
|
|
return len(self.samples)
|
|
|
|
def __getitem__(self, idx: int) -> dict[str, torch.Tensor]:
|
|
s = self.samples[idx]
|
|
return {
|
|
"image_1": torch.zeros(1, dtype=torch.float32),
|
|
"matrix_1": torch.as_tensor(s["matrix_1"], dtype=torch.float32),
|
|
"label_1": torch.tensor(int(s["label_1"]), dtype=torch.long),
|
|
}
|
|
|
|
|
|
class BilatMetaDataset(Dataset):
|
|
"""Patient-level bilateral dataset for collect_probs_single_components."""
|
|
|
|
def __init__(self, samples: list[dict]):
|
|
self.samples = samples
|
|
|
|
def __len__(self) -> int:
|
|
return len(self.samples)
|
|
|
|
def __getitem__(self, idx: int) -> dict[str, torch.Tensor]:
|
|
s = self.samples[idx]
|
|
return {
|
|
"image_1": torch.zeros(1, dtype=torch.float32),
|
|
"image_2": torch.zeros(1, dtype=torch.float32),
|
|
"matrix_1": torch.as_tensor(s["matrix_1"], dtype=torch.float32),
|
|
"matrix_2": torch.as_tensor(s["matrix_2"], dtype=torch.float32),
|
|
"label_1": torch.tensor(int(s["label_1"]), dtype=torch.long),
|
|
}
|
|
|
|
|
|
def _phase_for_epoch(epoch_idx: int, warm_tower: int, warm_fused: int, main_epochs: int) -> tuple[str, int]:
|
|
if epoch_idx < warm_tower:
|
|
return "tower_warmup", 0
|
|
if epoch_idx < (warm_tower + warm_fused):
|
|
return "fused_warmup", 0
|
|
if epoch_idx < (warm_tower + warm_fused + main_epochs):
|
|
main_ep = epoch_idx - warm_tower - warm_fused + 1
|
|
return "main", main_ep
|
|
return "done", main_epochs
|
|
|
|
|
|
def _evaluate_single(
|
|
model: nn.Module,
|
|
loader: DataLoader,
|
|
device: torch.device,
|
|
num_classes: int,
|
|
aggregate_patient: bool,
|
|
) -> tuple[np.ndarray, np.ndarray, float, float]:
|
|
y, p_fused, _, p_md = collect_probs_single_components(
|
|
model, loader, device, aggregate_patient=aggregate_patient
|
|
)
|
|
# In metadata_only mode p_fused == p_md; keep md explicitly for clarity.
|
|
probs = p_md if p_md.size else p_fused
|
|
acc, auc, _ = _score_arrays(y, probs, num_classes)
|
|
return y, probs, float(auc), float(acc)
|
|
|
|
|
|
def main() -> None:
|
|
ap = argparse.ArgumentParser(description="V2-parity metadata-only runner.")
|
|
ap.add_argument("--eval-mode", required=True, choices=["binary", "multiclass"])
|
|
ap.add_argument("--tower-mode", default="single", choices=["single", "ensemble"])
|
|
ap.add_argument("--run-name", required=True)
|
|
|
|
ap.add_argument("--epochs", type=int, default=40, help="Main-phase epochs.")
|
|
ap.add_argument("--warmup-tower-epochs", type=int, default=None)
|
|
ap.add_argument("--warmup-fused-epochs", type=int, default=None)
|
|
ap.add_argument("--batch-size", type=int, default=8)
|
|
ap.add_argument("--lr", type=float, default=1e-4)
|
|
ap.add_argument("--weight-decay", type=float, default=0.0)
|
|
ap.add_argument("--bcd-prob", type=float, default=0.5)
|
|
|
|
ap.add_argument("--md-hidden-dim", type=int, default=128)
|
|
ap.add_argument("--fusion-dim", type=int, default=256)
|
|
ap.add_argument("--dropout", type=float, default=0.1)
|
|
ap.add_argument("--use-se", action="store_true")
|
|
ap.add_argument("--se-reduction", type=int, default=16)
|
|
ap.add_argument("--se-pre-norm", action="store_true")
|
|
|
|
ap.add_argument("--n-splits", type=int, default=5)
|
|
ap.add_argument("--holdout-per-class", type=int, default=5)
|
|
ap.add_argument("--holdout-seed", type=int, default=123)
|
|
ap.add_argument("--fold-seed", type=int, default=42)
|
|
ap.add_argument("--seed", type=int, default=1234)
|
|
ap.add_argument("--balanced-sampling", action=argparse.BooleanOptionalAction, default=False)
|
|
|
|
ap.add_argument("--analysis-dir", default="analysis_data")
|
|
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("--patient-col", default="Patient ID")
|
|
ap.add_argument("--cat-cols", nargs="*", default=["Gender", "Phakic/Pseudophakic"])
|
|
|
|
args = ap.parse_args()
|
|
|
|
seed_everything(args.seed)
|
|
device = choose_device(None)
|
|
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)
|
|
|
|
num_classes = 2 if args.eval_mode == "binary" else 3
|
|
df_mode = data.df.copy()
|
|
if args.eval_mode == "binary":
|
|
df_mode = df_mode[df_mode[args.label_col].isin([0, 1])].reset_index(drop=True)
|
|
print(f"[{args.eval_mode}] rows={len(df_mode)}", flush=True)
|
|
|
|
class _ClinicalShim:
|
|
label_col = args.label_col
|
|
|
|
def __init__(self, df):
|
|
self.df = df
|
|
|
|
split_args = SimpleNamespace(
|
|
eval_mode=args.eval_mode,
|
|
holdout_per_class=args.holdout_per_class,
|
|
holdout_seed=args.holdout_seed,
|
|
n_splits=args.n_splits,
|
|
fold_seed=args.fold_seed,
|
|
)
|
|
splitter = PatientFirstSplitManager(patient_col=args.patient_col, label_col=args.label_col)
|
|
plans = splitter.build_plans(clinical=_ClinicalShim(df_mode), args=split_args)
|
|
|
|
profile_eye = build_papila_profile(
|
|
patient_col=args.patient_col,
|
|
label_col=args.label_col,
|
|
sample_mode="eye",
|
|
)
|
|
profile_patient = build_papila_profile(
|
|
patient_col=args.patient_col,
|
|
label_col=args.label_col,
|
|
sample_mode="patient",
|
|
)
|
|
|
|
warm_tower = int(args.warmup_tower_epochs) if args.warmup_tower_epochs is not None else 2
|
|
warm_fused = int(args.warmup_fused_epochs) if args.warmup_fused_epochs is not None else 2
|
|
total_epochs = warm_tower + warm_fused + int(args.epochs)
|
|
|
|
out_root = Path(args.analysis_dir) / args.run_name / args.eval_mode / args.tower_mode
|
|
out_root.mkdir(parents=True, exist_ok=True)
|
|
|
|
fold_metrics = []
|
|
aggregate_patient = args.tower_mode == "ensemble"
|
|
|
|
for fold_idx, split in enumerate(plans[: args.n_splits]):
|
|
fold_dir = out_root / f"fold{fold_idx}"
|
|
fold_dir.mkdir(parents=True, exist_ok=True)
|
|
|
|
eye_train = filter_eye_samples(profile_eye.build_samples(df=split.train, clinical=data))
|
|
bilat_val = filter_bilateral_samples(profile_patient.build_samples(df=split.val, clinical=data))
|
|
|
|
holdout_bilat = []
|
|
if split.holdout is not None and not split.holdout.empty:
|
|
holdout_bilat = filter_bilateral_samples(
|
|
profile_patient.build_samples(df=split.holdout, clinical=data)
|
|
)
|
|
|
|
if not eye_train or not bilat_val:
|
|
print(f"[fold {fold_idx+1}] skipped (eye_train={len(eye_train)} bilat_val={len(bilat_val)})", flush=True)
|
|
fold_metrics.append(
|
|
{
|
|
"fold": fold_idx,
|
|
"best_epoch": None,
|
|
"best_phase": None,
|
|
"val_auc": float("nan"),
|
|
"val_acc": float("nan"),
|
|
"hld_auc": float("nan"),
|
|
"hld_acc": float("nan"),
|
|
"eye_train_n": len(eye_train),
|
|
"bilat_val_n": len(bilat_val),
|
|
"bilat_holdout_n": len(holdout_bilat),
|
|
}
|
|
)
|
|
continue
|
|
|
|
sampler = build_balanced_sampler(eye_train) if args.balanced_sampling else None
|
|
train_loader = DataLoader(
|
|
EyeMetaDataset(eye_train),
|
|
batch_size=args.batch_size,
|
|
shuffle=(sampler is None),
|
|
sampler=sampler,
|
|
)
|
|
val_loader = DataLoader(BilatMetaDataset(bilat_val), batch_size=args.batch_size, shuffle=False)
|
|
holdout_loader = (
|
|
DataLoader(BilatMetaDataset(holdout_bilat), batch_size=args.batch_size, shuffle=False)
|
|
if holdout_bilat
|
|
else None
|
|
)
|
|
|
|
model = MetadataOnlySingleHT(
|
|
clinical_data=data,
|
|
num_classes=num_classes,
|
|
md_hidden_dim=args.md_hidden_dim,
|
|
fusion_dim=args.fusion_dim,
|
|
dropout=args.dropout,
|
|
use_se=bool(args.use_se),
|
|
se_reduction=int(args.se_reduction),
|
|
se_pre_norm=bool(args.se_pre_norm),
|
|
).to(device)
|
|
opt = torch.optim.Adam(model.parameters(), lr=args.lr, weight_decay=args.weight_decay)
|
|
|
|
best_auc = -1.0
|
|
best_epoch = 0
|
|
best_phase = ""
|
|
best_state = None
|
|
|
|
print(
|
|
f"\n[fold {fold_idx+1}/{args.n_splits}] "
|
|
f"eye_train_n={len(eye_train)} bilat_val_n={len(bilat_val)} "
|
|
f"holdout_n={len(holdout_bilat)} warmup={warm_tower}+{warm_fused} total={total_epochs}",
|
|
flush=True,
|
|
)
|
|
|
|
for ep in range(total_epochs):
|
|
phase, main_ep = _phase_for_epoch(ep, warm_tower, warm_fused, int(args.epochs))
|
|
tr_loss, tr_acc = train_single_epoch(
|
|
model,
|
|
train_loader,
|
|
opt,
|
|
device,
|
|
phase=phase,
|
|
bcd_prob=float(args.bcd_prob),
|
|
)
|
|
|
|
_, p_val, val_auc, val_acc = _evaluate_single(
|
|
model,
|
|
val_loader,
|
|
device,
|
|
num_classes,
|
|
aggregate_patient=aggregate_patient,
|
|
)
|
|
|
|
is_main = phase == "main"
|
|
if is_main and (not np.isnan(val_auc)) and val_auc > best_auc:
|
|
best_auc = float(val_auc)
|
|
best_state = copy.deepcopy(model.state_dict())
|
|
best_epoch = ep + 1
|
|
best_phase = phase
|
|
|
|
if ep == 0 or (ep + 1) % 10 == 0 or (ep + 1) == total_epochs:
|
|
print(
|
|
f" ep {ep+1:>3}/{total_epochs} [{phase}:{main_ep}/{args.epochs}] "
|
|
f"loss={tr_loss:.4f} acc={tr_acc:.4f} "
|
|
f"val_auc={val_auc:.4f} val_acc={val_acc:.4f} "
|
|
f"best_auc={best_auc:.4f}",
|
|
flush=True,
|
|
)
|
|
|
|
if best_state is not None:
|
|
model.load_state_dict(best_state)
|
|
|
|
y_val, p_val, val_auc, val_acc = _evaluate_single(
|
|
model,
|
|
val_loader,
|
|
device,
|
|
num_classes,
|
|
aggregate_patient=aggregate_patient,
|
|
)
|
|
|
|
if args.tower_mode == "single":
|
|
probs_name = "probs_classic.npy"
|
|
probs_h_name = "probs_classic_holdout.npy"
|
|
else:
|
|
probs_name = "probs_ensemble.npy"
|
|
probs_h_name = "probs_ensemble_holdout.npy"
|
|
|
|
np.save(fold_dir / "y_true.npy", y_val)
|
|
np.save(fold_dir / probs_name, p_val)
|
|
|
|
hld_auc = float("nan")
|
|
hld_acc = float("nan")
|
|
if holdout_loader is not None:
|
|
y_h, p_h, hld_auc, hld_acc = _evaluate_single(
|
|
model,
|
|
holdout_loader,
|
|
device,
|
|
num_classes,
|
|
aggregate_patient=aggregate_patient,
|
|
)
|
|
np.save(fold_dir / "y_true_holdout.npy", y_h)
|
|
np.save(fold_dir / probs_h_name, p_h)
|
|
|
|
print(
|
|
f" [fold {fold_idx+1}] best_epoch={best_epoch} best_auc={best_auc:.4f} "
|
|
f"val_auc={val_auc:.4f} val_acc={val_acc:.4f} "
|
|
f"hld_auc={hld_auc:.4f} hld_acc={hld_acc:.4f}",
|
|
flush=True,
|
|
)
|
|
|
|
fold_metrics.append(
|
|
{
|
|
"fold": fold_idx,
|
|
"best_epoch": best_epoch,
|
|
"best_phase": best_phase,
|
|
"val_auc": float(val_auc),
|
|
"val_acc": float(val_acc),
|
|
"hld_auc": float(hld_auc),
|
|
"hld_acc": float(hld_acc),
|
|
"eye_train_n": len(eye_train),
|
|
"bilat_val_n": len(bilat_val),
|
|
"bilat_holdout_n": len(holdout_bilat),
|
|
}
|
|
)
|
|
|
|
val_aucs = [m["val_auc"] for m in fold_metrics if not np.isnan(m["val_auc"])]
|
|
hld_aucs = [m["hld_auc"] for m in fold_metrics if not np.isnan(m["hld_auc"])]
|
|
if val_aucs:
|
|
print(f"\nMean val AUC: {np.mean(val_aucs):.4f} ± {np.std(val_aucs):.4f}", flush=True)
|
|
if hld_aucs:
|
|
print(f"Mean hld AUC: {np.mean(hld_aucs):.4f} ± {np.std(hld_aucs):.4f}", flush=True)
|
|
|
|
summary = {
|
|
"run_name": args.run_name,
|
|
"eval_mode": args.eval_mode,
|
|
"tower_mode": args.tower_mode,
|
|
"bridge_mode": "metadata_only",
|
|
"model": "MetadataOnlySingleHT",
|
|
"epochs": int(args.epochs),
|
|
"warmup_tower_epochs": warm_tower,
|
|
"warmup_fused_epochs": warm_fused,
|
|
"md_hidden_dim": int(args.md_hidden_dim),
|
|
"fusion_dim": int(args.fusion_dim),
|
|
"dropout": float(args.dropout),
|
|
"lr": float(args.lr),
|
|
"weight_decay": float(args.weight_decay),
|
|
"bcd_prob": float(args.bcd_prob),
|
|
"balanced_sampling": bool(args.balanced_sampling),
|
|
"feature_dim": int(data.feature_dim),
|
|
"timestamp": time.strftime("%Y%m%d_%H%M%S"),
|
|
"fold_metrics": fold_metrics,
|
|
"val_auc_mean": float(np.mean(val_aucs)) if val_aucs else None,
|
|
"val_auc_std": float(np.std(val_aucs)) if val_aucs else None,
|
|
"hld_auc_mean": float(np.mean(hld_aucs)) if hld_aucs else None,
|
|
"hld_auc_std": float(np.std(hld_aucs)) if hld_aucs else None,
|
|
}
|
|
(out_root / "summary.json").write_text(json.dumps(summary, indent=2))
|
|
print(f"\nOutputs written to: {out_root}", flush=True)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|