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