scripts/eval_ood_v9b_jepa.py

"""Thorough OOD evaluation of v9b normative JEPA + conformal calibration.

Pipeline per image:
  1. Resize to 256x256, RGB, [0,1]
  2. JEPA prediction_error_map(x) — for each of the 256 patches, mask it
     out, predict its latent from the other 255, measure smooth-L1
     residual. Returns (256, 256) anomaly map.
  3. Per-image score = 95th percentile of all pixel residuals (matches
     what conformal calibration used).
  4. Conformal verdict: per_image_score > q  ->  anomaly (tumor).
     q=0.308 loaded from v9b_conformal.json (alpha=0.10, empirical
     coverage 0.90 on 17,487 healthy calibration samples).

Compares v9b verdict against:
  - v8 segmentation alone @ threshold 0.20 (the current production baseline)
  - the 4 historical classifier sets (OLD, v8-RAW, v8-BAL, v8-MVMM)
    via per-classifier accuracy already in our earlier audit CSVs.

Outputs:
  - per-image table (csv)
  - per-source aggregate
  - 4-policy comparison table for the executive summary
"""
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.jepa import IJEPAModel
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'
CONFORMAL_JSON = ROOT / 'v9b_artifacts' / 'v9b_conformal.json'
SAMPLES_DIR = ROOT / 'samples' / 'ood'
IMAGE_SIZE = 256


def load_jepa(device: str = 'cuda') -> IJEPAModel:
    ckpt = torch.load(str(JEPA_CKPT), map_location=device, weights_only=False)
    a = ckpt.get('args') or {}
    model = IJEPAModel(
        image_size=a.get('image_size', 256),
        patch_size=a.get('patch_size', 16),
        in_chans=a.get('in_chans', 3),
        embed_dim=a.get('embed_dim', 384),
        depth=a.get('depth', 12),
        heads=a.get('heads', 6),
    )
    sd = ckpt.get('model_state_dict') or ckpt
    miss, unexp = model.load_state_dict(sd, strict=False)
    if miss or unexp:
        print(f'  [warn] checkpoint load: missing={len(miss)} unexpected={len(unexp)}')
    model = model.to(device).eval()
    return model


def jepa_anomaly_map(model: IJEPAModel, img_pil: Image.Image, device: str = 'cuda') -> np.ndarray:
    img = img_pil.convert('RGB').resize((IMAGE_SIZE, IMAGE_SIZE), Image.BILINEAR)
    arr = np.asarray(img, dtype=np.float32) / 255.0
    x = torch.from_numpy(arr.transpose(2, 0, 1)).unsqueeze(0).to(device)
    with torch.no_grad():
        emap = model.prediction_error_map(x)  # (1, 1, IMAGE_SIZE, IMAGE_SIZE)
    return emap.squeeze().cpu().numpy()


def main():
    if not JEPA_CKPT.exists():
        sys.exit(f'missing {JEPA_CKPT}')
    if not CONFORMAL_JSON.exists():
        sys.exit(f'missing {CONFORMAL_JSON}')
    conformal = json.loads(CONFORMAL_JSON.read_text())
    q = float(conformal['q'])
    alpha = float(conformal['alpha'])
    print(f'[init] conformal q={q:.4f} alpha={alpha} (target coverage = {1-alpha:.0%})')
    print(f'[init] calibrated on {conformal["report"]["n_calib"]} healthy samples, '
          f'empirical_coverage={conformal["report"]["empirical_coverage"]:.4f}')

    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 ''))
    model = load_jepa(device)
    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\n')

    rows = []
    t0 = time.perf_counter()
    last_print = t0
    for i, p in enumerate(samples):
        img = Image.open(p)
        gt = GT[p.parent.name]

        # v9b JEPA + conformal
        emap = jepa_anomaly_map(model, img, device)
        p95 = float(np.percentile(emap, 95))
        max_err = float(emap.max())
        mean_err = float(emap.mean())
        # Anomaly mask (pixels with error > q)
        ano_mask = (emap > q).astype(np.uint8)
        ano_area = int(ano_mask.sum())
        v9b_verdict = 'TUMOR' if p95 > q else 'no_tumor'

        # v8 segmentation comparison
        prob = seg_tta(seg, _preprocess_seg(img))
        v8_area = int((prob >= 0.20).sum())
        v8_verdict = 'TUMOR' if v8_area >= MIN_TUMOR_AREA else 'no_tumor'

        rows.append({
            'source': p.parent.name,
            'file': p.name,
            'gt': gt,
            'v9b_p95': round(p95, 4),
            'v9b_max': round(max_err, 4),
            'v9b_mean': round(mean_err, 4),
            'v9b_ano_area': ano_area,
            'v9b_verdict': v9b_verdict,
            'v8_area_020': v8_area,
            'v8_verdict': v8_verdict,
        })
        if time.perf_counter() - last_print > 20:
            last_print = time.perf_counter()
            print(f'  [{i+1}/{len(samples)}]  elapsed={time.perf_counter()-t0:.0f}s')

    elapsed = time.perf_counter() - t0
    print(f'\n[done] {len(rows)} samples in {elapsed:.0f}s ({elapsed/len(rows):.2f}/sample)\n')

    # ======================== aggregate ============================
    def _stats(rs, col):
        TP = sum(1 for r in rs if r['gt']=='tumor' and r[col]=='TUMOR')
        FN = sum(1 for r in rs if r['gt']=='tumor' and r[col]=='no_tumor')
        FP = sum(1 for r in rs if r['gt']=='no_tumor' and r[col]=='TUMOR')
        TN = sum(1 for r in rs if r['gt']=='no_tumor' and r[col]=='no_tumor')
        recall = TP/(TP+FN) if TP+FN else 0
        fpr = FP/(FP+TN) if FP+TN else 0
        acc = (TP+TN)/(TP+FN+FP+TN)
        f1 = 2*TP/(2*TP+FP+FN) if 2*TP+FP+FN else 0
        return TP, FN, FP, TN, recall, fpr, acc, f1

    print('='*82)
    print('v9b NORMATIVE JEPA + CONFORMAL — OOD SCORECARD')
    print('='*82)
    for col, label in (('v9b_verdict', 'v9b JEPA + conformal (q=0.308)'),
                        ('v8_verdict', 'v8 segmentation @ 0.20 (baseline)')):
        TP, FN, FP, TN, re, fp, acc, f1 = _stats(rows, col)
        print(f'\n  {label}')
        print(f'    TP={TP:2d}  FN={FN:2d}  FP={FP:2d}  TN={TN:2d}')
        print(f'    recall={re:.0%}   FPR={fp:.0%}   accuracy={acc:.0%}   F1={f1:.2f}')

    # ======================== per-source ===========================
    print('\n' + '='*82)
    print('PER-SOURCE BREAKDOWN')
    print('='*82)
    by_src = {}
    for r in rows:
        by_src.setdefault(r['source'], []).append(r)
    print(f'\n{"source":48s} GT     n   v9b_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_verdict']=='TUMOR')
        v8_hits  = sum(1 for r in rs if r['v8_verdict']=='TUMOR')
        if gt == 'tumor':
            print(f'  {src:46s} pos    {n:3d}   recall={v9b_hits/n:.0%}'.ljust(76)
                  + f'     recall={v8_hits/n:.0%}')
        else:
            print(f'  {src:46s} neg    {n:3d}   FPR={v9b_hits/n:.0%}'.ljust(76)
                  + f'        FPR={v8_hits/n:.0%}')

    # ======================== 4-policy comparison =================
    print('\n' + '='*82)
    print('EXECUTIVE SUMMARY — every policy we have measured to date')
    print('='*82)
    # Historical numbers from the audit scripts we already ran
    OLD = {'recall_range': '28-42%', 'fpr_range': '8-58%'}
    BAL = {'recall_range': '31-44%', 'fpr_range': '0%'}
    MVMM = {'recall_range': '25-47%', 'fpr_range': '0%'}
    TP, FN, FP, TN, re_v9b, fp_v9b, acc_v9b, f1_v9b = _stats(rows, 'v9b_verdict')
    _, _, _, _, re_v8, fp_v8, acc_v8, _ = _stats(rows, 'v8_verdict')
    print(f'\n  {"policy":40s}  recall      FPR        accuracy')
    print(f'  {"OLD 3 classifiers (Kaggle-only)":40s}  {OLD["recall_range"]:>10s}  {OLD["fpr_range"]:>9s}  31-54%')
    print(f'  {"v8-BAL 3 classifiers (+OpenNeuro)":40s}  {BAL["recall_range"]:>10s}  {BAL["fpr_range"]:>9s}  46-58%')
    print(f'  {"v8-MVMM 3 classifiers (+multi-view)":40s}  {MVMM["recall_range"]:>10s}  {MVMM["fpr_range"]:>9s}  44-60%')
    print(f'  {"v8 segmentation only @ 0.20":40s}  {re_v8:>10.0%}  {fp_v8:>9.0%}  {acc_v8:>8.0%}')
    print(f'  {"v9b JEPA + conformal":40s}  {re_v9b:>10.0%}  {fp_v9b:>9.0%}  {acc_v9b:>8.0%}   <-- NEW')

    # ======================== persist ==============================
    out_csv = SAMPLES_DIR / 'eval_v9b_jepa_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()