"""Full v9b OOD eval: JEPA + DDPM + SDF + two-tower combine. Compares 4 v9b score variants against the v8 segmentation baseline: - v9b_app: JEPA appearance anomaly only (p95 of prediction_error_map) - v9b_geo: SDF geometric anomaly only (p95 of SDF deviation) - v9b_combo: two-tower weighted_sum (lambda_app=0.6, lambda_geo=0.4) - v9b_residual: DDPM healthy-counterfactual residual (|x - x_healthy|) (much slower because of DDIM sampling) For each variant, per-image p95 score → tumor/no_tumor verdict at the threshold sweep optimum. AUC computed against ground truth. """ from __future__ import annotations import csv import json import sys import time from pathlib import Path import numpy as np import torch from PIL import Image ROOT = Path(__file__).resolve().parent.parent sys.path.insert(0, str(ROOT)) from src.research.v9b_model import V9BModel from scripts.eval_ood_cascade import ( SEG_ONNX, MIN_TUMOR_AREA, _sess, _preprocess_seg, seg_tta, GT, ) JEPA_CKPT = ROOT / 'v9b_artifacts' / 'v9b_jepa' / 'last.pt' STAGE2_CKPT = ROOT / 'v9b_artifacts' / 'v9b_stage2' / 'last.pt' CONFORMAL_JSON = ROOT / 'v9b_artifacts' / 'v9b_conformal.json' SAMPLES_DIR = ROOT / 'samples' / 'ood' IMAGE_SIZE = 256 def preprocess(img_pil: Image.Image, device: str) -> torch.Tensor: img = img_pil.convert('RGB').resize((IMAGE_SIZE, IMAGE_SIZE), Image.BILINEAR) arr = np.asarray(img, dtype=np.float32) / 255.0 return torch.from_numpy(arr.transpose(2, 0, 1)).unsqueeze(0).to(device) def _stats(rows, score_key, t): TP = sum(1 for r in rows if r['gt']=='tumor' and r[score_key] > t) FN = sum(1 for r in rows if r['gt']=='tumor' and r[score_key] <= t) FP = sum(1 for r in rows if r['gt']=='no_tumor' and r[score_key] > t) TN = sum(1 for r in rows if r['gt']=='no_tumor' and r[score_key] <= t) recall = TP/(TP+FN) if TP+FN else 0 fpr = FP/(FP+TN) if FP+TN else 0 acc = (TP+TN)/len(rows) f1 = 2*TP/(2*TP+FP+FN) if 2*TP+FP+FN else 0 return TP, FN, FP, TN, recall, fpr, acc, f1 def _auc(rows, score_key): pos = [r[score_key] for r in rows if r['gt']=='tumor'] neg = [r[score_key] for r in rows if r['gt']=='no_tumor'] if not pos or not neg: return float('nan') wins = ties = total = 0 for sp in pos: for sn in neg: if sp > sn: wins += 1 elif sp == sn: ties += 1 total += 1 return (wins + 0.5*ties) / total def _best_threshold(rows, score_key, metric='f1'): candidates = sorted(set(round(r[score_key], 4) for r in rows)) best = (None, -1.0) for t in candidates: TP, FN, FP, TN, re, fp, acc, f1 = _stats(rows, score_key, t) m = f1 if metric == 'f1' else acc if m > best[1]: best = (t, m, re, fp, acc, f1) return best def main(): if not all(p.exists() for p in (JEPA_CKPT, STAGE2_CKPT, CONFORMAL_JSON)): sys.exit('missing one of the v9b artefacts') device = 'cuda' if torch.cuda.is_available() else 'cpu' print(f'[init] device={device}' + (f' ({torch.cuda.get_device_name(0)})' if device == 'cuda' else '')) print('[init] loading V9BModel (JEPA + DDPM + SDF + conformal) ...') t0 = time.perf_counter() model = V9BModel.from_checkpoints( str(JEPA_CKPT), str(STAGE2_CKPT), str(CONFORMAL_JSON), image_size=IMAGE_SIZE, device=device, ) print(f' loaded in {time.perf_counter()-t0:.1f}s' f' (JEPA={model.jepa is not None}, DDPM={model.ddpm is not None}, ' f'SDF={model.sdf_tower is not None}, conformal_q={model.conformal.q:.4f})') seg = _sess(SEG_ONNX) samples = sorted(p for p in SAMPLES_DIR.rglob('*') if p.suffix.lower() in ('.png','.jpg','.jpeg') and p.parent.name in GT) print(f'[init] {len(samples)} OOD samples') # Decide if we run the DDPM residual (50 DDIM steps per image, slow). # Default: skip to save time; turn on with V9B_RUN_DDPM=1 import os run_ddpm = os.environ.get('V9B_RUN_DDPM', '0').strip() == '1' print(f'[init] DDPM residual: {"ON (slow, ~30s/sample)" if run_ddpm else "OFF (set V9B_RUN_DDPM=1 to enable)"}') print() rows = [] t0 = time.perf_counter() last = t0 for i, p in enumerate(samples): img = Image.open(p) x = preprocess(img, device) out = model.infer(x, combine_mode='weighted_sum', lambda_app=0.6, lambda_geo=0.4, ddpm_num_steps=50 if run_ddpm else 0) # Per-image scores: 95th percentile of each anomaly map app = out['appearance_anomaly'].squeeze().cpu().numpy() geo = out['geometry_anomaly'].squeeze().cpu().numpy() if out['geometry_anomaly'] is not None else None combo = out['combined_anomaly'].squeeze().cpu().numpy() rec = { 'source': p.parent.name, 'file': p.name, 'gt': GT[p.parent.name], 'v9b_app_p95': float(np.percentile(app, 95)), 'v9b_app_max': float(app.max()), 'v9b_geo_p95': float(np.percentile(geo, 95)) if geo is not None else 0.0, 'v9b_combo_p95': float(np.percentile(combo, 95)), 'v9b_combo_max': float(combo.max()), } if run_ddpm and out['residual'] is not None: res = out['residual'].squeeze().cpu().numpy() rec['v9b_residual_p95'] = float(np.percentile(res, 95)) rec['v9b_residual_mean'] = float(res.mean()) # v8 segmentation baseline prob = seg_tta(seg, _preprocess_seg(img)) rec['v8_area_020'] = int((prob >= 0.20).sum()) rec['v8_verdict_020'] = 'TUMOR' if rec['v8_area_020'] >= MIN_TUMOR_AREA else 'no_tumor' rows.append(rec) if time.perf_counter() - last > 20: last = time.perf_counter() print(f' [{i+1}/{len(samples)}] elapsed={time.perf_counter()-t0:.0f}s') print(f'\n[done] {len(rows)} samples in {(time.perf_counter()-t0)/60:.1f} min\n') # ===================== AUC for each score variant ===================== print('='*82) print('AUC for tumor-vs-healthy on OOD per scoring variant') print('='*82) score_keys = ['v9b_app_p95', 'v9b_geo_p95', 'v9b_combo_p95'] if run_ddpm: score_keys.append('v9b_residual_p95') for k in score_keys: print(f' AUC({k}) = {_auc(rows, k):.4f}') # ===================== best operating point per variant ================= print('\n' + '='*82) print('BEST F1 OPERATING POINT per scoring variant') print('='*82) print(f' {"variant":18s} thr recall FPR acc F1') for k in score_keys: t, _, re, fp, acc, f1 = _best_threshold(rows, k, 'f1') print(f' {k:18s} {t:.3f} {re:.0%} {fp:.0%} {acc:.0%} {f1:.2f}') # v8 baseline for reference TP, FN, FP, TN, re_v8, fp_v8, acc_v8, f1_v8 = _stats(rows, 'v8_area_020', MIN_TUMOR_AREA - 1) print(f' {"v8 seg @ 0.20":18s} ---- {re_v8:.0%} {fp_v8:.0%} {acc_v8:.0%} {f1_v8:.2f}') # ===================== Pareto curve for combined score ================= print('\n' + '='*82) print('THRESHOLD SWEEP — v9b_combo_p95 (best variant)') print('='*82) print(f' {"thr":>6s} recall FPR accuracy F1') candidates = sorted(set(round(r['v9b_combo_p95'], 3) for r in rows)) for t in candidates: TP, FN, FP, TN, re, fp, acc, f1 = _stats(rows, 'v9b_combo_p95', t) print(f' {t:>6.3f} {re:>5.0%} {fp:>5.0%} {acc:>6.0%} {f1:.2f}') # ===================== per-source on best combo ===================== print('\n' + '='*82) print('PER-SOURCE on v9b_combo @ best-F1 threshold') print('='*82) best_t = _best_threshold(rows, 'v9b_combo_p95', 'f1')[0] by_src = {} for r in rows: by_src.setdefault(r['source'], []).append(r) print(f'\nthreshold = {best_t:.3f}') print(f'{"source":48s} GT n v9b_combo_recall/FPR v8_recall/FPR') for src in sorted(by_src): rs = by_src[src]; gt = rs[0]['gt']; n = len(rs) v9b_hits = sum(1 for r in rs if r['v9b_combo_p95'] > best_t) v8_hits = sum(1 for r in rs if r['v8_verdict_020'] == 'TUMOR') kind = 'recall' if gt=='tumor' else 'FPR' print(f' {src:46s} {gt[:6]:6s} {n:3d} v9b_{kind}={v9b_hits/n:.0%}'.ljust(80) + f' v8_{kind}={v8_hits/n:.0%}') # ===================== final scoreboard ===================== print('\n' + '='*82) print('FINAL SCOREBOARD — every policy on this OOD test bench') print('='*82) rows_final = [] rows_final.append(('OLD 3 classifiers (Kaggle-only)', '28-42%', '8-58%', '31-54%')) rows_final.append(('v8-MVMM 3 classifiers (multi-view)', '25-47%', '0%', '44-60%')) rows_final.append((f'v8 segmentation only @ 0.20', f'{re_v8:.0%}', f'{fp_v8:.0%}', f'{acc_v8:.0%}')) for k, label in (('v9b_app_p95', 'v9b JEPA appearance (Stage 1 only)'), ('v9b_geo_p95', 'v9b SDF geometry tower (Stage 2 only)'), ('v9b_combo_p95', 'v9b two-tower combo (full stack)')): t, _, re, fp, acc, f1 = _best_threshold(rows, k, 'f1') rows_final.append((label, f'{re:.0%}', f'{fp:.0%}', f'{acc:.0%}')) print(f'\n {"policy":42s} recall FPR accuracy') for label, r, f, a in rows_final: print(f' {label:42s} {r:>9s} {f:>9s} {a:>9s}') # ===================== persist ===================== out_csv = SAMPLES_DIR / 'eval_v9b_full_results.csv' fields = list(rows[0].keys()) with out_csv.open('w', newline='', encoding='utf-8') as f: w = csv.DictWriter(f, fieldnames=fields) w.writeheader() for r in rows: w.writerow(r) print(f'\n[csv] {out_csv}') if __name__ == '__main__': main()