tools/build_v6_dataset.py

#!/usr/bin/env python
"""Generate v6 jsonl from v5 with corrected post-accident labels + discard.

Policy:
  DADA / Nexar (both at 20 fps annotation convention):
    frame_indices[-1] < accident_frame                              → keep original label
    frame_indices[-1] in [accident_frame, accident_frame + 100)     → ALERT (5s window)
    frame_indices[-1] >= accident_frame + 100                       → DISCARD tick
  DoTA (unchanged from prior fix):
    frame in [anomaly_start, anomaly_end) → ALERT
    frame >= anomaly_end                  → SILENT
    no discard
  DAD: untouched

Outputs:
  data/cot_corpus_v3/v5_sft_train_v6.jsonl
  data/cot_corpus_v3/v5_sft_val_v6.jsonl
  data/cot_corpus_v3/v6_changelog.json

Also propagates the new tick_action to actions_per_frame[-1] (the last of the 8
frames in the tick), so downstream "use last frame as GT" stays consistent.
"""
import json, csv, logging
from pathlib import Path
from collections import Counter, defaultdict

ROOT = Path("PROJECT_ROOT")

logging.basicConfig(level=logging.INFO, format="%(asctime)s %(message)s")
log = logging.getLogger("v6")

WINDOW_FRAMES_DADA_NEXAR = 100   # 5s @ 20 fps


def build_accident_lookup():
    ACC = {}; END = {}
    # DADA — accident_time is JPG index at 20 fps
    for cat in ["positive", "non-ego", "negative"]:
        for d in (ROOT / f"DADA-2000/{cat}").glob("images_*"):
            ann = d / "annotation.json"
            if ann.exists():
                a = json.load(open(ann))
                if a.get("accident_time") is not None:
                    ACC[f"dada_{d.name}"] = a["accident_time"]
    # DoTA — anomaly_start at native (10 fps for DoTA)
    for f in (ROOT / "DoTA/annotations").glob("*.json"):
        a = json.load(open(f))
        s = a.get("anomaly_start"); e = a.get("anomaly_end")
        if s is not None:
            ACC[f"dota_{f.stem}"] = s
            if e is not None: END[f"dota_{f.stem}"] = e
    # Nexar — time_of_event(sec) × 20 fps (per user convention)
    for split in ["train", "test-public", "test-private"]:
        for po in ["positive", "negative"]:
            mp = ROOT / f"NEXAR_COLLISION/{split}/{po}/metadata.csv"
            if not mp.exists(): continue
            for row in csv.DictReader(open(mp)):
                fn = row["file_name"].replace(".mp4", "")
                toe = row.get("time_of_event", "").strip()
                if toe:
                    ACC[f"nexar_{fn}"] = round(float(toe) * 20)
    return ACC, END


def process_split(in_path, out_path, ACC, END):
    stats = {"total": 0, "discarded": 0, "no_meta_kept": 0,
             "flips": Counter(), "by_src_kept": Counter(),
             "by_src_discarded": Counter(),
             "old_dist": Counter(), "new_dist": Counter()}
    kept_records = []

    with open(in_path) as f:
        for ln in f:
            d = json.loads(ln)
            stats["total"] += 1
            src = d["source"]; vid = d["video_id"]
            cur = d["frame_indices"][-1]
            ta = d.get("tick_action", "SILENT")
            stats["old_dist"][ta] += 1

            acc = ACC.get(vid)
            new_action = None  # None = keep original; "DISCARD" = drop

            if acc is None:
                # No metadata → keep as-is (DAD + half of nexar)
                new_action = ta
                stats["no_meta_kept"] += 1
            elif src in ("dada", "nexar"):
                if cur < acc:
                    new_action = ta
                elif cur < acc + WINDOW_FRAMES_DADA_NEXAR:
                    new_action = "ALERT"
                else:
                    new_action = "DISCARD"
            elif src == "dota":
                end = END.get(vid)
                if cur < acc:
                    new_action = ta
                elif end is None or cur < end:
                    new_action = "ALERT"
                else:
                    new_action = "SILENT"
            else:
                new_action = ta

            if new_action == "DISCARD":
                stats["discarded"] += 1
                stats["by_src_discarded"][src] += 1
                continue

            # Apply
            if new_action != ta:
                stats["flips"][f"{src}:{ta}→{new_action}"] += 1
                d["tick_action"] = new_action
                # Also patch actions_per_frame[-1] so downstream consumers see it
                if d.get("actions_per_frame"):
                    d["actions_per_frame"] = list(d["actions_per_frame"])
                    d["actions_per_frame"][-1] = new_action

            stats["new_dist"][new_action] += 1
            stats["by_src_kept"][src] += 1
            kept_records.append(d)

    # Write
    out_path.parent.mkdir(parents=True, exist_ok=True)
    with open(out_path, "w") as f:
        for d in kept_records:
            f.write(json.dumps(d) + "\n")

    return stats


def main():
    ACC, END = build_accident_lookup()
    log.info(f"Lookup built: {len(ACC)} videos, {len(END)} with anomaly_end")
    log.info(f"5s window for DADA/Nexar = {WINDOW_FRAMES_DADA_NEXAR} frames (20 fps)")

    out_stats = {}
    for split in ["train", "val"]:
        in_p  = ROOT / f"data/cot_corpus_v3/v5_sft_{split}.jsonl"
        out_p = ROOT / f"data/cot_corpus_v3/v5_sft_{split}_v6.jsonl"
        log.info(f"\nProcessing {in_p.name} → {out_p.name}")
        st = process_split(in_p, out_p, ACC, END)
        out_stats[split] = st
        kept = st["total"] - st["discarded"]
        log.info(f"  total={st['total']:,}  discarded={st['discarded']:,}  kept={kept:,}")
        log.info(f"  no_meta_kept={st['no_meta_kept']:,}")
        log.info(f"  flips: {sum(st['flips'].values()):,}")
        log.info(f"  OLD dist: {dict(st['old_dist'])}")
        log.info(f"  NEW dist: {dict(st['new_dist'])}")
        log.info(f"  discarded by src: {dict(st['by_src_discarded'])}")

    # Changelog
    changelog = {
        "policy": {
            "DADA_Nexar": "frame in [acc, acc+5s] → ALERT; frame > acc+5s → DISCARD. fps=20.",
            "DoTA": "frame in [anom_start, anom_end) → ALERT; >= anom_end → SILENT.",
            "DAD": "untouched (no per-video accident metadata)",
            "window_frames": WINDOW_FRAMES_DADA_NEXAR,
        },
        "splits": {
            split: {
                "total": s["total"],
                "discarded": s["discarded"],
                "kept": s["total"] - s["discarded"],
                "no_meta_kept": s["no_meta_kept"],
                "flips": dict(s["flips"]),
                "old_dist": dict(s["old_dist"]),
                "new_dist": dict(s["new_dist"]),
                "discarded_by_src": dict(s["by_src_discarded"]),
            }
            for split, s in out_stats.items()
        },
    }
    cl_path = ROOT / "data/cot_corpus_v3/v6_changelog.json"
    json.dump(changelog, open(cl_path, "w"), indent=2)
    log.info(f"\nChangelog → {cl_path}")


if __name__ == "__main__":
    main()