training/Nexar/train_maskflow.py

#!/usr/bin/env python3
"""MaskFlow Track-A trainer (3rd-place style, visible-clip-only).

Protocol (CORRECTED 2026-04-28):
  - TRAIN  = all 1500 nexar train clips (data/maskflow_cache/train.pt)
  - VAL    = test-public subset (~672 clips) of data/maskflow_cache/test.pt,
              filtered by NEXAR_COLLISION/solution.csv `Usage == "Public"`.
              We use Kaggle public-LB labels for early stopping / model
              selection — they are publicly visible and never leak to
              test-private.
  - TEST   = test-private subset (~672 clips, `Usage == "Private"`).
              Reported only at final eval time; never used for tuning.

This replaces the earlier 1280/220 internal-val split, which wasted
training data and validated on a non-Kaggle distribution.

Architecture:
  - RGB branch    : ResNet18 on each of the last 3 frames → [3, 512]
  - Mask*Flow br. : ResNet18 on `concat(mask, flow_x, flow_y)` (3-ch) for
                    each of the last 3 frames → [3, 512]
  - Per-frame fusion : MLP([rgb, motion]) → score logit per frame
  - Clip score    : weighted average of last-3 frame logits, weights
                    [0.2, 0.3, 0.5]

Loss:
  - BCEWithLogitsLoss on the clip score
  - Optional aux: `--badas_soft_label` enables BCE distillation against
    BADAS deployable soft labels with mixing α

Usage:
    python -m training.Nexar.train_maskflow --seed 0 \\
      --train_cache data/maskflow_cache/train.pt \\
      --test_cache  data/maskflow_cache/test.pt \\
      --solution_csv NEXAR_COLLISION/solution.csv \\
      --label_csv    nexar-collision-prediction/train.csv \\
      --output       checkpoints/Nexar/maskflow_seed0
"""
from __future__ import annotations

import argparse
import csv
import json
import random
from pathlib import Path

import numpy as np
import torch
import torch.nn as nn
import torchvision.models as tvm
from sklearn.metrics import average_precision_score, roc_auc_score
from torch.utils.data import DataLoader, Dataset


def set_seed(s: int):
    random.seed(s); np.random.seed(s); torch.manual_seed(s)
    torch.cuda.manual_seed_all(s)


# ─── dataset ─────────────────────────────────────────────────────────────

class MaskFlowDataset(Dataset):
    """Wraps a maskflow cache + binary labels."""
    def __init__(self, cache_path: Path, labels: dict[str, int],
                 last_n: int = 3, badas_soft: dict | None = None,
                 keep_ids: set[str] | None = None):
        c = torch.load(cache_path, weights_only=False, map_location="cpu")
        self.ids:  list[str] = c["ids"]
        self.rgb  = c["rgb"]
        self.flow = c["flow"]
        self.mask = c["mask"]
        self.last_n = last_n
        self.labels = labels
        self.badas_soft = badas_soft or {}
        kept = []
        for i, vid in enumerate(self.ids):
            if vid not in labels:
                continue
            if keep_ids is not None and vid not in keep_ids:
                continue
            kept.append(i)
        self.kept = kept
        self.mean = torch.tensor([0.485, 0.456, 0.406]).view(1, 3, 1, 1)
        self.std  = torch.tensor([0.229, 0.224, 0.225]).view(1, 3, 1, 1)

    def __len__(self):
        return len(self.kept)

    def __getitem__(self, i):
        idx = self.kept[i]
        vid = self.ids[idx]
        n = self.last_n
        rgb = self.rgb[idx, -n:].float()                 # [n, 3, H, W]
        rgb = (rgb - self.mean) / self.std
        flow_full = self.flow[idx].float()                # [K-1, 2, H, W]
        flow = flow_full[-n:] if flow_full.shape[0] >= n else \
               torch.cat([flow_full[:1].repeat(n - flow_full.shape[0], 1, 1, 1),
                          flow_full], dim=0)
        flow = flow / (flow.flatten(1).std(dim=1, unbiased=False)
                         .clamp(min=1e-2)[:, None, None, None] + 1e-6)
        mask = self.mask[idx, -n:].float()                # [n, 1, H, W]
        motion = torch.cat([mask, flow], dim=1)            # [n, 3, H, W]
        y = float(self.labels[vid])
        soft = float(self.badas_soft.get(vid, y))
        return rgb, motion, torch.tensor(y), torch.tensor(soft), vid


# ─── model ───────────────────────────────────────────────────────────────

class MaskFlowHead(nn.Module):
    def __init__(self, last_n: int = 3, weights: list[float] | None = None,
                 dropout: float = 0.3):
        super().__init__()
        self.last_n = last_n
        if weights is None:
            weights = [0.2, 0.3, 0.5]
        assert len(weights) == last_n
        w = torch.tensor(weights); w = w / w.sum()
        self.register_buffer("frame_w", w)
        self.rgb_backbone    = tvm.resnet18(weights=tvm.ResNet18_Weights.IMAGENET1K_V1)
        self.motion_backbone = tvm.resnet18(weights=tvm.ResNet18_Weights.IMAGENET1K_V1)
        self.rgb_backbone.fc    = nn.Identity()  # → 512
        self.motion_backbone.fc = nn.Identity()
        self.fusion = nn.Sequential(
            nn.Linear(1024, 256), nn.ReLU(inplace=True),
            nn.Dropout(dropout),  nn.Linear(256, 1),
        )

    def forward(self, rgb: torch.Tensor, motion: torch.Tensor) -> torch.Tensor:
        # rgb, motion: [B, n, 3, H, W]
        B, n, C, H, W = rgb.shape
        rgb_f    = self.rgb_backbone(rgb.reshape(B * n, C, H, W))      # [B*n, 512]
        motion_f = self.motion_backbone(motion.reshape(B * n, 3, H, W))  # [B*n, 512]
        feats    = torch.cat([rgb_f, motion_f], dim=1)                  # [B*n, 1024]
        per_frame = self.fusion(feats).view(B, n)                        # [B, n]
        clip = (per_frame * self.frame_w[None, :]).sum(dim=1)            # [B]
        return clip


# ─── label loaders ───────────────────────────────────────────────────────

def load_labels(csv_path: Path) -> dict[str, int]:
    """Load nexar-collision-prediction/train.csv labels."""
    rows = list(csv.DictReader(open(csv_path)))
    return {r["id"]: int(r["target"] or 0) for r in rows
            if r.get("target") is not None}


def load_solution_split(csv_path: Path) -> tuple[set[str], set[str], dict[str, int]]:
    """Read solution.csv → (public_ids, private_ids, {id: target}).

    The Kaggle test split: 50% Public (visible LB), 50% Private (hidden).
    We use Public as our val set; Private is held out for final eval.
    """
    rows = list(csv.DictReader(open(csv_path)))
    pub  = {r["id"] for r in rows if r["Usage"] == "Public"}
    priv = {r["id"] for r in rows if r["Usage"] == "Private"}
    targets = {r["id"]: int(r["target"]) for r in rows}
    return pub, priv, targets


def load_badas_soft(per_clip_path: Path) -> dict[str, float]:
    if not per_clip_path.exists():
        return {}
    j = json.loads(per_clip_path.read_text())
    out = {}
    for cid, rec in j.items():
        s = rec.get("score_last4s")
        if s is not None and not np.isnan(s):
            out[cid] = float(s)
    return out


# ─── eval ────────────────────────────────────────────────────────────────

@torch.no_grad()
def eval_split(model, loader, device) -> tuple[float, float]:
    model.eval()
    ys, ps = [], []
    for rgb, motion, y, _soft, _vid in loader:
        rgb = rgb.to(device); motion = motion.to(device)
        logit = model(rgb, motion).cpu().numpy()
        ys.extend(y.numpy().tolist()); ps.extend(logit.tolist())
    ys, ps = np.asarray(ys), np.asarray(ps)
    if len(np.unique(ys)) < 2:
        return float("nan"), float("nan")
    return float(average_precision_score(ys, ps)), float(roc_auc_score(ys, ps))


@torch.no_grad()
def eval_kaggle_mAP(model, ds: MaskFlowDataset, device, solution: Path,
                     batch: int = 16) -> tuple[float, float]:
    """Compute Kaggle bucket-mean AP_500/1000/1500 on the dataset, given
    the solution.csv that maps id → group ∈ {0,1,2}. Used for val tracking."""
    rows = list(csv.DictReader(open(solution)))
    group = {r["id"]: int(r["group"]) for r in rows}
    usage = {r["id"]: r["Usage"] for r in rows}
    model.eval()
    score: dict[str, float] = {}
    target: dict[str, int] = {}
    dl = DataLoader(ds, batch_size=batch, shuffle=False, num_workers=4)
    for rgb, motion, y, _soft, vids in dl:
        rgb = rgb.to(device); motion = motion.to(device)
        logit = model(rgb, motion).cpu().numpy()
        for v, p, t in zip(vids, logit.tolist(), y.numpy().tolist()):
            score[v] = float(p); target[v] = int(t)
    common = sorted(set(score) & set(group))
    if not common:
        return float("nan"), float("nan")
    pub_aps, priv_aps = [], []
    for g in (0, 1, 2):
        for u, sink in (("Public", pub_aps), ("Private", priv_aps)):
            ids = [v for v in common if usage.get(v) == u and group.get(v) == g]
            if len(ids) < 2: continue
            y = np.array([target[v] for v in ids])
            s = np.array([score[v] for v in ids])
            if len(np.unique(y)) < 2: continue
            sink.append(float(average_precision_score(y, s)))
    pub_mAP  = float(np.mean(pub_aps))  if pub_aps  else float("nan")
    priv_mAP = float(np.mean(priv_aps)) if priv_aps else float("nan")
    return pub_mAP, priv_mAP


# ─── main ────────────────────────────────────────────────────────────────

def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--seed",        type=int, default=0)
    ap.add_argument("--train_cache", default="data/maskflow_cache/train.pt")
    ap.add_argument("--test_cache",  default="data/maskflow_cache/test.pt")
    ap.add_argument("--solution_csv", default="NEXAR_COLLISION/solution.csv")
    ap.add_argument("--label_csv",   default="nexar-collision-prediction/train.csv")
    ap.add_argument("--output",      default="checkpoints/Nexar/maskflow_seed0")
    ap.add_argument("--epochs",   type=int, default=12)
    ap.add_argument("--batch",    type=int, default=16)
    ap.add_argument("--lr",       type=float, default=2e-4)
    ap.add_argument("--wd",       type=float, default=1e-4)
    ap.add_argument("--dropout",  type=float, default=0.3)
    ap.add_argument("--last_n",   type=int, default=3)
    ap.add_argument("--frame_weights", nargs=3, type=float,
                     default=[0.2, 0.3, 0.5])
    ap.add_argument("--badas_soft_label", default=None)
    ap.add_argument("--alpha_soft", type=float, default=0.3)
    ap.add_argument("--num_workers", type=int, default=4)
    ap.add_argument("--report_private", action="store_true",
                    help="ALSO log private mAP each epoch (visibility only; "
                         "still selects best by public)")
    args = ap.parse_args()

    set_seed(args.seed)
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    out = Path(args.output); out.mkdir(parents=True, exist_ok=True)

    badas_soft = (load_badas_soft(Path(args.badas_soft_label))
                   if args.badas_soft_label else {})
    if badas_soft:
        print(f"[init] loaded {len(badas_soft)} BADAS soft labels")

    # train labels — ALL nexar train clips
    train_labels = load_labels(Path(args.label_csv))
    print(f"[init] train labels: {len(train_labels)}")

    # val/test split from Kaggle solution.csv
    pub_ids, priv_ids, test_targets = load_solution_split(Path(args.solution_csv))
    print(f"[init] solution split: public={len(pub_ids)}  private={len(priv_ids)}")

    # datasets
    train_ds = MaskFlowDataset(Path(args.train_cache), train_labels,
                                last_n=args.last_n, badas_soft=badas_soft)
    val_ds   = MaskFlowDataset(Path(args.test_cache), test_targets,
                                last_n=args.last_n, keep_ids=pub_ids)
    priv_ds  = MaskFlowDataset(Path(args.test_cache), test_targets,
                                last_n=args.last_n, keep_ids=priv_ids)
    print(f"[init] datasets: train_n={len(train_ds)}  val_pub_n={len(val_ds)}  "
          f"priv_n={len(priv_ds)}")

    train_dl = DataLoader(train_ds, batch_size=args.batch, shuffle=True,
                           num_workers=args.num_workers, pin_memory=True)
    val_dl   = DataLoader(val_ds,   batch_size=args.batch, shuffle=False,
                           num_workers=args.num_workers, pin_memory=True)

    model = MaskFlowHead(last_n=args.last_n,
                          weights=args.frame_weights,
                          dropout=args.dropout).to(device)
    opt = torch.optim.AdamW(model.parameters(), lr=args.lr,
                              weight_decay=args.wd)
    sched = torch.optim.lr_scheduler.CosineAnnealingLR(opt, T_max=args.epochs)
    bce_logit = nn.BCEWithLogitsLoss()

    best_pub_mAP = -1.0; best_path = out / "best.pt"
    history = []
    for ep in range(args.epochs):
        model.train(); running = 0.0; n_batch = 0
        for rgb, motion, y, soft, _vid in train_dl:
            rgb = rgb.to(device); motion = motion.to(device)
            y   = y.to(device);   soft   = soft.to(device)
            logit = model(rgb, motion)
            loss_gt = bce_logit(logit, y)
            loss = ((1 - args.alpha_soft) * loss_gt
                    + args.alpha_soft * bce_logit(logit, soft)) if badas_soft else loss_gt
            opt.zero_grad(); loss.backward(); opt.step()
            running += float(loss.item()); n_batch += 1
        sched.step()

        # evaluate on PUBLIC val (model selection signal)
        pub_ap, pub_auc = eval_split(model, val_dl, device)
        # bucket-mean Kaggle mAP on public side
        pub_mAP, _ = eval_kaggle_mAP(model, val_ds, device,
                                        Path(args.solution_csv),
                                        batch=args.batch)
        priv_mAP_str = ""
        if args.report_private:
            priv_dl = DataLoader(priv_ds, batch_size=args.batch, shuffle=False,
                                   num_workers=args.num_workers, pin_memory=True)
            priv_ap, _ = eval_split(model, priv_dl, device)
            _, priv_mAP = eval_kaggle_mAP(model, priv_ds, device,
                                            Path(args.solution_csv),
                                            batch=args.batch)
            priv_mAP_str = f"  [info] priv_mAP={priv_mAP:.4f} priv_AP={priv_ap:.4f}"

        avg = running / max(n_batch, 1)
        line = (f"epoch {ep+1:02d}/{args.epochs}  loss={avg:.4f}  "
                f"pub_AP={pub_ap:.4f}  pub_AUC={pub_auc:.4f}  "
                f"pub_mAP={pub_mAP:.4f}{priv_mAP_str}")
        print(line, flush=True)
        history.append({"epoch": ep + 1, "loss": avg,
                          "pub_AP": pub_ap, "pub_AUC": pub_auc,
                          "pub_mAP": pub_mAP})
        # selection: bucket-mean public mAP (matches Kaggle scoring)
        if pub_mAP > best_pub_mAP:
            best_pub_mAP = pub_mAP
            torch.save({
                "head_state": model.state_dict(),
                "args":       vars(args),
                "epoch":      ep + 1,
                "pub_AP":     pub_ap,
                "pub_AUC":    pub_auc,
                "pub_mAP":    pub_mAP,
            }, best_path)
            print(f"  ↑ saved best to {best_path} (pub_mAP={pub_mAP:.4f})")

    (out / "history.json").write_text(json.dumps(history, indent=2))
    print(f"[done] best pub_mAP={best_pub_mAP:.4f}  ckpt={best_path}")


if __name__ == "__main__":
    main()