#!/usr/bin/env python3 """ Generate Nexar submission using fine-tuned MViT-v2-s. Modes: mvit_only — MViT scores only mvit_ensemble — blend MViT scores with sample_submission baseline Usage: python -m training.Nexar.mvit_submit \ --model_dir checkpoints/Nexar/mvit_v1 \ --test_dir nexar-collision-prediction/test \ --test_csv nexar-collision-prediction/test.csv \ --out_csv submissions/nexar_mvit_v1.csv \ --evaluate NEXAR_COLLISION/solution.csv # Ensemble: python -m training.Nexar.mvit_submit \ --model_dir checkpoints/Nexar/mvit_v1 \ --test_dir nexar-collision-prediction/test \ --test_csv nexar-collision-prediction/test.csv \ --baseline_csv NEXAR_COLLISION/sample_submission.csv \ --ensemble_alpha 0.6 \ --out_csv submissions/nexar_mvit_ensemble_0.6.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 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.mvit_dataset import NexarMViTDataset, N_FRAMES, IMG_SIZE logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s") logger = logging.getLogger("Nexar.mvit_submit") def build_test_csv(test_dir: str, test_csv: str) -> str: """Create a temporary CSV for the test set (id column only, target=-1).""" df = pd.read_csv(test_csv) # Add dummy columns needed by NexarMViTDataset df["target"] = 0 df["time_of_event"] = None df["time_of_alert"] = None tmp = Path(test_csv).parent / "_test_with_dummy.csv" df.to_csv(tmp, index=False) return str(tmp) @torch.no_grad() def score_test_clips( model_dir: str, test_dir: str, test_csv: str, batch_size: int = 16, ) -> Dict[str, float]: 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) n_frames = meta.get("n_frames", N_FRAMES) img_size = meta.get("img_size", IMG_SIZE) # Load model from torchvision.models.video import mvit_v2_s import torch.nn as nn model = mvit_v2_s(weights=None) in_features = model.head[1].in_features model.head[1] = nn.Linear(in_features, 1) model.load_state_dict(torch.load(Path(model_dir) / "best_model.pt", map_location=device)) model = model.to(device) model.eval() logger.info(f"Loaded MViT-v2-s from {model_dir}") # Build test dataset tmp_csv = build_test_csv(test_dir, test_csv) ds = NexarMViTDataset( tmp_csv, test_dir, train_mode=False, min_warning_s=0.0, is_test=True, n_frames=n_frames, img_size=img_size, ) loader = DataLoader(ds, batch_size=batch_size, shuffle=False, num_workers=4, pin_memory=True) scores = {} for batch in tqdm(loader, desc="Scoring test clips"): videos = batch["video"].to(device) with torch.amp.autocast(device_type="cuda", dtype=torch.bfloat16): logits = model(videos).squeeze(-1) probs = torch.sigmoid(logits) for vid_id, score in zip(batch["vid_id"], probs.cpu().tolist()): scores[vid_id] = float(np.clip(score, 0, 1)) return scores def evaluate_submission(submission: Dict[str, float], solution_csv: str): 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) aps = [] for g in sorted(merged["group"].unique()): g_df = merged[merged["group"] == g] if g_df["target"].nunique() < 2: continue aps.append(float(average_precision_score(g_df["target"], g_df["score"]))) print(f"mAP ({usage}): {np.mean(aps):.6f}" if aps else f"mAP ({usage}): nan") def main(): parser = argparse.ArgumentParser("mvit_submit") parser.add_argument("--model_dir", required=True) parser.add_argument("--test_dir", default="nexar-collision-prediction/test") parser.add_argument("--test_csv", default="nexar-collision-prediction/test.csv") parser.add_argument("--batch_size", type=int, default=16) parser.add_argument("--baseline_csv", default=None, help="NEXAR_COLLISION/sample_submission.csv for ensemble") parser.add_argument("--ensemble_alpha", type=float, default=0.6, help="Weight for MViT (1-alpha = baseline weight)") parser.add_argument("--out_csv", required=True) parser.add_argument("--evaluate", default=None, help="Path to solution.csv for local evaluation") args = parser.parse_args() scores = score_test_clips(args.model_dir, args.test_dir, args.test_csv, args.batch_size) if args.baseline_csv and Path(args.baseline_csv).exists(): b_df = pd.read_csv(args.baseline_csv) baseline = {str(row["id"]).zfill(5): float(row["score"]) for _, row in b_df.iterrows()} blended = {} for vid_id in set(scores) | set(baseline): m = scores.get(vid_id, 0.5) b = baseline.get(vid_id, 0.5) blended[vid_id] = float(np.clip(args.ensemble_alpha * m + (1 - args.ensemble_alpha) * b, 0, 1)) scores = blended logger.info(f"Ensemble: {args.ensemble_alpha:.2f}×MViT + {1-args.ensemble_alpha:.2f}×baseline") out = Path(args.out_csv) out.parent.mkdir(parents=True, exist_ok=True) ids = pd.read_csv(args.test_csv)["id"].astype(str).str.zfill(5).tolist() rows = [{"id": vid_id, "score": scores.get(vid_id, 0.5)} for vid_id in ids] pd.DataFrame(rows).to_csv(args.out_csv, index=False) vals = [r["score"] for r in rows] logger.info( f"Saved → {args.out_csv} mean={np.mean(vals):.3f} std={np.std(vals):.3f} " f"min={np.min(vals):.3f} max={np.max(vals):.3f}" ) if args.evaluate: logger.info(f"\nLocal evaluation:") evaluate_submission(scores, args.evaluate) if __name__ == "__main__": main()