#!/usr/bin/env python3 """ Generate Nexar submission CSV and optionally evaluate against solution.csv. Modes: 1. zero_shot — use P(ALERT) from cached PolicyHead directly (no extra training) 2. trained — use fine-tuned NexarTemporalHead/NexarSimpleHead 3. ensemble — weighted blend of trained scores + baseline sample_submission scores Usage: # Zero-shot (fastest, no training needed): python -m training.Nexar.nexar_submit \ --mode zero_shot \ --test_cache data/nexar_cache/test.pt \ --out_csv submissions/nexar_zero_shot.csv # Trained model: python -m training.Nexar.nexar_submit \ --mode trained \ --test_cache data/nexar_cache/test.pt \ --model_dir checkpoints/Nexar/nexar_v1 \ --out_csv submissions/nexar_trained.csv # Ensemble: python -m training.Nexar.nexar_submit \ --mode ensemble \ --test_cache data/nexar_cache/test.pt \ --model_dir checkpoints/Nexar/nexar_v1 \ --baseline_csv nexar-collision-prediction/sample_submission.csv \ --ensemble_alpha 0.5 \ --out_csv submissions/nexar_ensemble.csv \ --evaluate NEXAR_COLLISION/solution.csv """ from __future__ import annotations import argparse import json import logging from pathlib import Path from typing import Dict, List, Optional import numpy as np import pandas as pd import torch import torch.nn.functional as F from torch.utils.data import DataLoader from tqdm import tqdm import sys sys.path.insert(0, str(Path(__file__).resolve().parent.parent.parent)) from training.Nexar.nexar_dataset import NexarTestDataset, nexar_collate_test from training.Nexar.nexar_model import build_model logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s") logger = logging.getLogger("Nexar.submit") # ── scoring functions ───────────────────────────────────────────────────────── def scores_zero_shot(cache_file: str, n_windows: int = 3, agg: str = "max_last") -> Dict[str, float]: """ Zero-shot collision scores from cached P(ALERT) values. agg strategies: last — P(ALERT) of the last (most recent) window only max — max P(ALERT) over all windows max_last — 0.7 * last + 0.3 * max (emphasises recency) weighted — linearly increasing weights over windows (latest = highest) tta — 1 / (1 + tta_mean_last) combined with p_alert_last """ cache = torch.load(cache_file, map_location="cpu", weights_only=False) scores = {} for vid_id, feat in cache["features"].items(): p = feat["p_alert"].float() # [T] tta = feat["tta_means"].float() # [T] if agg == "last": s = p[-1].item() elif agg == "max": s = p.max().item() elif agg == "max_last": s = 0.6 * p[-1].item() + 0.4 * p.max().item() elif agg == "weighted": T = p.shape[0] w = torch.linspace(0.5, 1.0, T) w = w / w.sum() s = (p * w).sum().item() elif agg == "tta": # Combine P(ALERT) with recency-adjusted 1/(1+tta) tta_score = 1.0 / (1.0 + tta[-1].item()) s = 0.6 * p[-1].item() + 0.4 * tta_score else: s = p[-1].item() scores[vid_id] = float(np.clip(s, 0, 1)) return scores @torch.no_grad() def scores_trained( cache_file: str, model_dir: str, batch_size: int = 128, n_windows: int = 3, ) -> Dict[str, float]: """Run fine-tuned NexarHead on test features.""" device = torch.device("cuda" if torch.cuda.is_available() else "cpu") meta_path = Path(model_dir) / "best_meta.json" with open(meta_path) as f: meta = json.load(f) hidden_dim = meta["hidden_dim"] arch = meta["arch"] n_windows = meta.get("n_windows", n_windows) model = build_model(hidden_dim, arch).to(device) model.load_state_dict(torch.load(Path(model_dir) / "best_model.pt", map_location=device)) model.eval() ds = NexarTestDataset(cache_file, n_windows=n_windows) loader = DataLoader(ds, batch_size=batch_size, shuffle=False, num_workers=4, collate_fn=nexar_collate_test, pin_memory=True) scores = {} for batch in tqdm(loader, desc="Scoring"): beliefs = batch["beliefs"].to(device) tta_means = batch["tta_means"].to(device) tta_vars = batch["tta_vars"].to(device) p_alerts = batch["p_alerts"].to(device) if hasattr(model, "lstm"): s = model(beliefs, tta_means, tta_vars, p_alerts) else: s = model(beliefs[:, -1, :], tta_means[:, -1], tta_vars[:, -1], p_alerts[:, -1]) for vid_id, score in zip(batch["video_ids"], s.cpu().tolist()): scores[vid_id] = float(np.clip(score, 0, 1)) return scores def scores_ensemble( primary: Dict[str, float], baseline: Dict[str, float], alpha: float = 0.5, ) -> Dict[str, float]: """ Blend primary scores with baseline (sample_submission) scores. score = alpha * primary + (1 - alpha) * baseline """ all_ids = set(primary) | set(baseline) result = {} for vid_id in all_ids: p = primary.get(vid_id, 0.5) b = baseline.get(vid_id, 0.5) result[vid_id] = float(np.clip(alpha * p + (1 - alpha) * b, 0, 1)) return result def evaluate_submission(submission: Dict[str, float], solution_csv: str): """Compute mAP (Public) and mAP (Private) against solution.csv.""" from sklearn.metrics import average_precision_score sol = pd.read_csv(solution_csv) sol["id"] = sol["id"].astype(str).str.zfill(5) sub_df = pd.DataFrame(list(submission.items()), columns=["id", "score"]) sub_df["id"] = sub_df["id"].astype(str).str.zfill(5) for usage in ["Public", "Private"]: subset = sol[sol["Usage"] == usage].copy() merged = subset.merge(sub_df, on="id", how="left").fillna(0.5) ap_list = [] for g in sorted(merged["group"].unique()): g_df = merged[merged["group"] == g] if g_df["target"].nunique() < 2: continue ap = average_precision_score(g_df["target"], g_df["score"]) ap_list.append(ap) mean_ap = np.mean(ap_list) if ap_list else float("nan") print(f"mAP ({usage}): {mean_ap:.6f}") def write_submission(scores: Dict[str, float], out_csv: str, test_csv: Optional[str] = None): """Write submission CSV. Aligns with test.csv order if provided.""" out = Path(out_csv) out.parent.mkdir(parents=True, exist_ok=True) if test_csv and Path(test_csv).exists(): ids = pd.read_csv(test_csv)["id"].astype(str).str.zfill(5).tolist() else: ids = sorted(scores.keys()) rows = [{"id": vid_id, "score": scores.get(vid_id, 0.5)} for vid_id in ids] df = pd.DataFrame(rows) df.to_csv(out_csv, index=False) logger.info(f"Submission saved → {out_csv} ({len(df)} rows)") # Print score distribution vals = [r["score"] for r in rows] logger.info( f"Score stats: mean={np.mean(vals):.3f} std={np.std(vals):.3f} " f"min={np.min(vals):.3f} max={np.max(vals):.3f} " f"p50={np.median(vals):.3f} p90={np.percentile(vals,90):.3f}" ) def main(): parser = argparse.ArgumentParser("nexar_submit") parser.add_argument("--mode", required=True, choices=["zero_shot", "trained", "ensemble"]) parser.add_argument("--test_cache", required=True) parser.add_argument("--model_dir", default=None) parser.add_argument("--baseline_csv", default=None, help="sample_submission.csv for ensemble mode") parser.add_argument("--ensemble_alpha", type=float, default=0.5, help="Weight for primary model (0=baseline only, 1=model only)") parser.add_argument("--zero_shot_agg", default="max_last", choices=["last", "max", "max_last", "weighted", "tta"]) parser.add_argument("--n_windows", type=int, default=3) parser.add_argument("--out_csv", required=True) parser.add_argument("--test_csv", default="nexar-collision-prediction/test.csv") parser.add_argument("--evaluate", default=None, help="Path to solution.csv for local evaluation") args = parser.parse_args() # ── compute scores ──────────────────────────────────────────────────────── if args.mode == "zero_shot": logger.info(f"Zero-shot inference (agg={args.zero_shot_agg}) ...") primary_scores = scores_zero_shot(args.test_cache, args.n_windows, args.zero_shot_agg) elif args.mode == "trained": if not args.model_dir: parser.error("--model_dir required for mode=trained") logger.info("Running trained NexarHead ...") primary_scores = scores_trained(args.test_cache, args.model_dir, n_windows=args.n_windows) elif args.mode == "ensemble": if not args.model_dir: parser.error("--model_dir required for mode=ensemble") logger.info("Running trained NexarHead + ensemble ...") trained = scores_trained(args.test_cache, args.model_dir, n_windows=args.n_windows) baseline = {} if args.baseline_csv and Path(args.baseline_csv).exists(): b_df = pd.read_csv(args.baseline_csv) for _, row in b_df.iterrows(): vid_id = str(row["id"]).zfill(5) baseline[vid_id] = float(row["score"]) logger.info(f"Baseline scores loaded: {len(baseline)} entries") primary_scores = scores_ensemble(trained, baseline, args.ensemble_alpha) # ── write output ────────────────────────────────────────────────────────── write_submission(primary_scores, args.out_csv, args.test_csv) # ── optional local evaluation ───────────────────────────────────────────── if args.evaluate: logger.info(f"\nLocal evaluation against {args.evaluate}:") evaluate_submission(primary_scores, args.evaluate) if __name__ == "__main__": main()