training/Nexar/mvit_submit.py

#!/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()