""" Leave-One-Image-Out Cross-Validation (LOOCV) evaluation runner. For each fold: test: held-out image val: next image (for threshold tuning) train: remaining images CRITICAL: Image-level splits ONLY. Patch-level splits inflate F1 by 5-15%. Usage: python evaluate_loocv.py --config config/config.yaml python evaluate_loocv.py --config config/config.yaml --ensemble-dir checkpoints/ """ import argparse import json from pathlib import Path import numpy as np import pandas as pd import torch import yaml from src.evaluate import match_detections_to_gt from src.heatmap import extract_peaks from src.model import ImmunogoldCenterNet from src.postprocess import ( apply_structural_mask_filter, cross_class_nms, sweep_confidence_threshold, ) from src.preprocessing import discover_synapse_data, load_synapse from src.ensemble import ensemble_predict, sliding_window_inference from src.visualize import overlay_annotations def parse_args(): parser = argparse.ArgumentParser(description="LOOCV evaluation") parser.add_argument("--config", type=str, default="config/config.yaml") parser.add_argument("--ensemble-dir", type=str, default="checkpoints", help="Directory containing fold_*/phase3_*.pth") parser.add_argument("--device", type=str, default="auto") parser.add_argument("--use-tta", action="store_true") parser.add_argument("--fold", type=str, default=None, help="Evaluate a single fold (e.g., S1). If omitted, runs all folds.") parser.add_argument("--output", type=str, default="results/loocv_metrics.csv") return parser.parse_args() def load_fold_models(ensemble_dir: Path, fold_id: str, cfg: dict, device: torch.device): """Load all models for a fold (5 seeds × 3 snapshots = 15 models).""" models = [] n_seeds = cfg["training"]["n_seeds"] snapshot_epochs = cfg["training"]["n_snapshot_epochs"] for seed_idx in range(n_seeds): seed = seed_idx + 42 # seeds start at 42 fold_dir = ensemble_dir / f"fold_{fold_id}_seed{seed}" for epoch in snapshot_epochs: ckpt_path = fold_dir / f"phase3_{epoch}.pth" if not ckpt_path.exists(): # Try best checkpoint instead ckpt_path = fold_dir / "phase3_best.pth" if not ckpt_path.exists(): continue model = ImmunogoldCenterNet( bifpn_channels=cfg["model"]["bifpn_channels"], bifpn_rounds=cfg["model"]["bifpn_rounds"], num_classes=cfg["model"]["num_classes"], ) ckpt = torch.load(ckpt_path, map_location="cpu", weights_only=False) model.load_state_dict(ckpt["model_state_dict"]) model.to(device) model.eval() models.append(model) return models def main(): args = parse_args() with open(args.config) as f: cfg = yaml.safe_load(f) device = torch.device( "cuda" if args.device == "auto" and torch.cuda.is_available() else args.device if args.device != "auto" else "cpu" ) records = discover_synapse_data(cfg["data"]["root"], cfg["data"]["synapse_ids"]) synapse_ids = cfg["data"]["synapse_ids"] incomplete_6nm = set(cfg["data"].get("incomplete_6nm", [])) ensemble_dir = Path(args.ensemble_dir) all_results = [] match_radii = {k: float(v) for k, v in cfg["evaluation"]["match_radii_px"].items()} val_offset = cfg["evaluation"]["loocv_val_offset"] # Support single-fold mode for SLURM array jobs if args.fold: eval_folds = [(synapse_ids.index(args.fold), args.fold)] else: eval_folds = list(enumerate(synapse_ids)) for test_idx, test_sid in eval_folds: print(f"\n{'='*60}") print(f"Fold {test_idx}: test={test_sid}") # Val image for threshold tuning val_idx = (test_idx + val_offset) % len(synapse_ids) val_sid = synapse_ids[val_idx] # Load test and val data test_record = [r for r in records if r.synapse_id == test_sid][0] val_record = [r for r in records if r.synapse_id == val_sid][0] test_data = load_synapse(test_record) val_data = load_synapse(val_record) has_6nm = test_sid not in incomplete_6nm # Load ensemble models models = load_fold_models(ensemble_dir, test_sid, cfg, device) if not models: print(f" No models found for fold {test_sid}, skipping") all_results.append({ "fold": test_sid, "n_models": 0, "6nm_f1": float("nan"), "12nm_f1": float("nan"), "mean_f1": float("nan"), }) continue print(f" Loaded {len(models)} ensemble members") # Tune threshold on validation image val_hm, val_off = ensemble_predict( models, val_data["image"], device, use_tta=args.use_tta, ) val_hm_t = torch.from_numpy(val_hm) val_off_t = torch.from_numpy(val_off) # Get all detections at low threshold for sweep val_dets = extract_peaks( val_hm_t, val_off_t, stride=cfg["data"]["stride"], conf_threshold=0.05, nms_kernel_sizes=cfg["postprocessing"]["nms_kernel_size"], ) best_thresholds = sweep_confidence_threshold( val_dets, val_data["annotations"], match_radii, ) print(f" Best thresholds: {best_thresholds}") # Test inference test_hm, test_off = ensemble_predict( models, test_data["image"], device, use_tta=args.use_tta, ) test_hm_t = torch.from_numpy(test_hm) test_off_t = torch.from_numpy(test_off) # Use per-class thresholds all_detections = [] for cls in ["6nm", "12nm"]: thr = best_thresholds.get(cls, 0.3) cls_dets = extract_peaks( test_hm_t, test_off_t, stride=cfg["data"]["stride"], conf_threshold=thr, nms_kernel_sizes=cfg["postprocessing"]["nms_kernel_size"], ) all_detections.extend([d for d in cls_dets if d["class"] == cls]) # Post-processing if test_data["mask"] is not None: all_detections = apply_structural_mask_filter( all_detections, test_data["mask"], margin_px=cfg["postprocessing"]["mask_filter_margin_px"], ) all_detections = cross_class_nms( all_detections, cfg["postprocessing"]["cross_class_nms_distance_px"], ) # Evaluate results = match_detections_to_gt( all_detections, test_data["annotations"].get("6nm", np.empty((0, 2))), test_data["annotations"].get("12nm", np.empty((0, 2))), match_radii, ) fold_result = { "fold": test_sid, "n_models": len(models), "6nm_f1": results["6nm"]["f1"] if has_6nm else float("nan"), "6nm_precision": results["6nm"]["precision"] if has_6nm else float("nan"), "6nm_recall": results["6nm"]["recall"] if has_6nm else float("nan"), "12nm_f1": results["12nm"]["f1"], "12nm_precision": results["12nm"]["precision"], "12nm_recall": results["12nm"]["recall"], "mean_f1": results["mean_f1"], } all_results.append(fold_result) for cls in ["6nm", "12nm"]: r = results[cls] note = " (N/A)" if cls == "6nm" and not has_6nm else "" print(f" {cls}: F1={r['f1']:.3f}, P={r['precision']:.3f}, " f"R={r['recall']:.3f}{note}") print(f" Mean F1: {results['mean_f1']:.3f}") # Save per-fold visualization overlay_annotations( test_data["image"], test_data["annotations"], title=f"Fold {test_sid} — F1={results['mean_f1']:.3f}", save_path=Path("results/per_fold_predictions") / f"{test_sid}.png", predictions=all_detections, ) # Summary df = pd.DataFrame(all_results) output_path = Path(args.output) output_path.parent.mkdir(parents=True, exist_ok=True) df.to_csv(output_path, index=False) print(f"\n{'='*60}") print("LOOCV Results:") f1_6nm = df["6nm_f1"].dropna() f1_12nm = df["12nm_f1"].dropna() mean_f1 = df["mean_f1"].dropna() print(f" 6nm F1: {f1_6nm.mean():.3f} ± {f1_6nm.std():.3f} (n={len(f1_6nm)})") print(f" 12nm F1: {f1_12nm.mean():.3f} ± {f1_12nm.std():.3f} (n={len(f1_12nm)})") print(f" Mean F1: {mean_f1.mean():.3f} ± {mean_f1.std():.3f}") print(f"\nResults saved to {output_path}") if __name__ == "__main__": main()