training/Policy/make_policy_labels.py

#!/usr/bin/env python3
"""
Generate per-window action labels for Stage 1 supervised policy warm-start.

Reuses SFTDataset for window generation — no duplication of frame-sampling or
window-stride logic.  Only the label assignment is new.

Label rules (conservative: only high-confidence assignments enter Stage 1 CE)
───────────────────────────────────────────────────────────────────────────────
  ego_positive, TTA ∈ [1.5, 5.0)   → ALERT   (2), ce_weight = 1.0
  ego_positive, TTA ∈ [5.5, 8.0]   → OBSERVE (1), ce_weight = 1.0
  ego_positive, TTA > 8.0           → SILENT  (0), ce_weight = 0.8
    (includes censored windows with tta_raw > MAX_TTA = 10.0)

  ego_positive, TTA ∈ [5.0, 5.5)   → EXCLUDE (boundary zone)
  ego_positive, TTA < 1.5           → EXCLUDE (too late, semantically complex)

  non_ego                           → OBSERVE (1), ce_weight = 0.4
    (gentle push only; semantics ambiguous; treated separately in metrics)

  safe_neg                          → SILENT  (0), ce_weight = 1.0
  safe_neg with neg_tag="pre_risky" → SILENT  (0), ce_weight = 0.8
    (pre_risky: early window from a crash video, before risk onset)

Usage:
  cd PROJECT_ROOT
  python -m training.Policy.make_policy_labels \
      --manifest_dir data/sft_manifests \
      --out_dir      data/policy_labels
"""

from __future__ import annotations

import argparse
import json
import logging
import sys
from pathlib import Path
from typing import Dict, List, Optional, Tuple

sys.path.insert(0, str(Path(__file__).resolve().parent.parent.parent))

from training.SFT.dataset import SFTDataset, TTASample

logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
logger = logging.getLogger("Policy.make_labels")

# ── action space ──────────────────────────────────────────────────────────────
SILENT  = 0
OBSERVE = 1
ALERT   = 2
ACTION_NAMES = {SILENT: "SILENT", OBSERVE: "OBSERVE", ALERT: "ALERT"}

# ── TTA boundaries (seconds) ──────────────────────────────────────────────────
ALERT_TTA_MIN   = 1.5    # below this: too late, exclude
ALERT_TTA_MAX   = 5.0    # [ALERT_TTA_MIN, ALERT_TTA_MAX) → ALERT
BOUNDARY_LO     = 5.0    # [BOUNDARY_LO, BOUNDARY_HI)    → exclude
BOUNDARY_HI     = 5.5
OBSERVE_TTA_MAX = 8.0    # [BOUNDARY_HI, OBSERVE_TTA_MAX] → OBSERVE
                          # > OBSERVE_TTA_MAX               → SILENT


# ── label derivation ──────────────────────────────────────────────────────────

def _derive_label(s: TTASample) -> Optional[Tuple[int, float]]:
    """
    Returns (action_label, ce_weight) or None to exclude from Stage 1.

    Uses tta_raw (not the capped tta_label) for ego_positive decisions so that
    censored windows (tta_raw > 10.0) fall into the TTA > 8.0 → SILENT bucket
    rather than being ambiguous.
    """
    if s.is_ego_positive:
        tta = s.tta_raw
        if tta < ALERT_TTA_MIN:
            return None                            # too late
        if BOUNDARY_LO <= tta < BOUNDARY_HI:
            return None                            # boundary zone
        if tta < ALERT_TTA_MAX:
            return (ALERT, 1.0)                    # [1.5, 5.0)
        if tta <= OBSERVE_TTA_MAX:
            return (OBSERVE, 1.0)                  # [5.5, 8.0]
        return (SILENT, 0.8)                       # > 8.0 (incl. censored > 10.0)

    if s.is_non_ego:
        return (OBSERVE, 0.4)                      # gentle push, not dogmatic

    # safe_neg (includes pre_risky windows from crash videos)
    weight = 0.8 if s.metadata.get("neg_tag") == "pre_risky" else 1.0
    return (SILENT, weight)


# ── per-split processing ──────────────────────────────────────────────────────

def process_split(
    manifests:     List[Path],
    split_name:    str,
    sft_split:     str,         # "train" or "val" for SFTDataset (affects frame sampling)
    debug:         bool = False,
    debug_samples: int  = 200,
) -> dict:
    """Build policy label manifest for one split from SFT video manifests."""
    logger.info(f"\n{'='*60}")
    logger.info(f"Processing split: {split_name}")

    # Instantiate SFTDataset with a huge neg_pos_ratio so no samples are capped.
    # We only use dataset.samples (TTASample objects) — no frame I/O here.
    ds = SFTDataset(
        manifests       = manifests,
        split           = sft_split,
        seed            = 42,
        debug           = False,
        neg_pos_ratio   = 10_000,   # effectively disable sample capping
        multi_window    = True,
    )

    samples_out = []
    excluded = {"tta_too_late": 0, "tta_boundary": 0}

    for s in ds.samples:
        result = _derive_label(s)
        if result is None:
            if s.is_ego_positive:
                if s.tta_raw < ALERT_TTA_MIN:
                    excluded["tta_too_late"] += 1
                else:
                    excluded["tta_boundary"] += 1
            continue

        action_label, ce_weight = result
        samples_out.append({
            "video_id":      s.video_id,
            "source":        s.source,
            "category":      s.category,
            "source_dir":    s.source_dir,
            "frame_indices": s.frame_indices,
            # tta_raw: store -1.0 for non_ego / safe_neg (inf is not JSON-serialisable)
            "tta_raw":       float(s.tta_raw) if s.tta_raw != float("inf") else -1.0,
            "action_label":  action_label,
            "ce_weight":     ce_weight,
            "metadata":      s.metadata,
        })

    if debug:
        import random
        rng = random.Random(42)
        rng.shuffle(samples_out)
        samples_out = samples_out[:debug_samples]

    # ── statistics ────────────────────────────────────────────────────────────
    label_counts: Dict[str, int] = {v: 0 for v in ACTION_NAMES.values()}
    cat_action:   Dict[str, Dict[str, int]] = {}
    for s in samples_out:
        lname = ACTION_NAMES[s["action_label"]]
        label_counts[lname] += 1
        cat = s["category"]
        cat_action.setdefault(cat, {})
        cat_action[cat][lname] = cat_action[cat].get(lname, 0) + 1

    logger.info(f"  Kept: {len(samples_out)}  |  Excluded: {excluded}")
    logger.info(f"  Label counts: {label_counts}")
    for cat, dist in sorted(cat_action.items()):
        logger.info(f"  {cat}: {dict(sorted(dist.items()))}")

    return {
        "name":          split_name,
        "split":         sft_split,
        "total_samples": len(samples_out),
        "label_counts":  label_counts,
        "excluded":      excluded,
        "samples":       samples_out,
    }


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

def main():
    parser = argparse.ArgumentParser("make_policy_labels")
    parser.add_argument("--manifest_dir",  default="data/sft_manifests")
    parser.add_argument("--out_dir",       default="data/policy_labels")
    parser.add_argument("--debug",         action="store_true")
    parser.add_argument("--debug_samples", type=int, default=200)
    args = parser.parse_args()

    mdir = Path(args.manifest_dir)
    odir = Path(args.out_dir)
    odir.mkdir(parents=True, exist_ok=True)

    splits = {
        "train": {
            "manifests": [
                mdir / "nexar_train.json",
                mdir / "dada_pos_train.json",
                mdir / "dada_noneego_train.json",
                mdir / "dada_neg_train.json",
            ],
            "sft_split": "train",
        },
        "val": {
            "manifests": [
                mdir / "nexar_val.json",
                mdir / "dada_pos_val.json",
                mdir / "dada_noneego_val.json",
            ],
            "sft_split": "val",
        },
    }

    for split_name, cfg in splits.items():
        existing = [p for p in cfg["manifests"] if p.exists()]
        if not existing:
            logger.warning(f"  No manifests found for {split_name}, skipping.")
            continue

        data = process_split(
            manifests     = existing,
            split_name    = split_name,
            sft_split     = cfg["sft_split"],
            debug         = args.debug,
            debug_samples = args.debug_samples,
        )
        out = odir / f"{split_name}.json"
        with open(out, "w") as f:
            json.dump(data, f)
        logger.info(f"  Saved {data['total_samples']} samples → {out}")

    logger.info("\n✅ Policy label manifests generated.")


if __name__ == "__main__":
    main()