Upload evaluate_loocv.py with huggingface_hub

evaluate_loocv.py

@@ -0,0 +1,243 @@
+"""
+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()