tools/relabel_dada_nexar.py

"""Relabel DADA-2000 and Nexar per-frame actions using accident_time + risky_time.

Rule (at 20Hz, L = 2.0s = 40 frames):
  Case A (accident_time - 40 >= risky_time):
      [risky_time, accident_time - 40)  → OBSERVE
      [accident_time - 40, accident_time] → ALERT
  Case B (accident_time - 40 < risky_time):
      [risky_time, accident_time] → ALL ALERT (no OBSERVE room)
  Everything else → SILENT
  Negative clips (no accident) → ALL SILENT

Updates annotation.json in-place: adds "per_frame_labels" list.

Usage:
    python tools/relabel_dada_nexar.py
"""
from __future__ import annotations
import json
import logging
from collections import Counter
from pathlib import Path

ROOT = Path("PROJECT_ROOT")
DADA_ROOT = ROOT / "DADA-2000"
NEXAR_ROOT = ROOT / "NEXAR_COLLISION" / "dataset"

FPS = 20
L_SEC = 2.0
L_FRAMES = int(L_SEC * FPS)  # 40

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


def label_one_clip(n_frames: int, accident_time: int, risky_time: int) -> list[str]:
    """Generate per-frame label for one clip."""
    labels = ["SILENT"] * n_frames

    if accident_time is None or accident_time <= 0:
        return labels  # negative clip

    alert_start = max(accident_time - L_FRAMES, risky_time)

    for f in range(n_frames):
        if alert_start <= f <= accident_time:
            labels[f] = "ALERT"
        elif risky_time is not None and risky_time <= f < alert_start:
            labels[f] = "OBSERVE"
        # else SILENT (already default)

    return labels


def count_images(folder: Path) -> int:
    """Count .jpg or .png images in a folder."""
    n = len(list(folder.glob("*.jpg"))) + len(list(folder.glob("*.png")))
    return n


def process_dada():
    """Process all DADA-2000 clips."""
    stats = Counter()
    label_dist = Counter()

    for cat in ["positive", "non-ego", "negative"]:
        cat_dir = DADA_ROOT / cat
        if not cat_dir.exists():
            continue
        for clip_dir in sorted(cat_dir.iterdir()):
            ann_path = clip_dir / "annotation.json"
            if not ann_path.exists():
                continue

            ann = json.loads(ann_path.read_text())
            accident = ann.get("accident", "False")
            is_positive = str(accident).lower() == "true"
            accident_time = int(ann.get("accident_time", -1))
            risky_time = int(ann.get("risky_time", -1))

            # Count frames in folder
            n_frames = count_images(clip_dir)
            if n_frames == 0:
                # Try images/ subfolder
                if (clip_dir / "images").is_dir():
                    n_frames = count_images(clip_dir / "images")

            if n_frames == 0:
                stats["dada_skip_no_frames"] += 1
                continue

            if not is_positive or accident_time <= 0:
                labels = ["SILENT"] * n_frames
                risky_time = -1
            else:
                if risky_time < 0:
                    risky_time = max(0, accident_time - L_FRAMES)
                labels = label_one_clip(n_frames, accident_time, risky_time)

            # Save back
            ann["per_frame_labels"] = labels
            ann["label_rule"] = f"L={L_SEC}s, fps={FPS}, L_frames={L_FRAMES}"
            ann_path.write_text(json.dumps(ann, indent=2, ensure_ascii=False))

            for la in labels:
                label_dist[f"dada_{cat}_{la}"] += 1
            stats[f"dada_{cat}"] += 1

            # Log case type
            if is_positive and accident_time > 0:
                case = "A" if (accident_time - L_FRAMES >= risky_time) else "B"
                stats[f"dada_{cat}_case_{case}"] += 1

    return stats, label_dist


def process_nexar():
    """Process all Nexar clips."""
    stats = Counter()
    label_dist = Counter()

    for split in ["train", "test-public", "test-private"]:
        for polarity in ["positive", "negative"]:
            parent = NEXAR_ROOT / split / polarity
            if not parent.exists():
                continue
            for clip_dir in sorted(parent.iterdir()):
                if not clip_dir.is_dir():
                    continue
                ann_path = clip_dir / "annotation.json"
                if not ann_path.exists():
                    stats[f"nexar_{split}_{polarity}_no_ann"] += 1
                    continue

                ann = json.loads(ann_path.read_text())
                is_positive = bool(ann.get("accident", False))

                # Use LOCAL frame indices (20fps extracted); handle None
                at_raw = ann.get("accident_time_local") or ann.get("accident_time")
                rt_raw = ann.get("risky_time_local") or ann.get("risky_time")
                accident_time = int(at_raw) if at_raw is not None else -1
                risky_time = int(rt_raw) if rt_raw is not None else -1

                # Count frames
                n_frames = count_images(clip_dir)
                if n_frames == 0:
                    stats[f"nexar_{split}_skip_no_frames"] += 1
                    continue

                if not is_positive or accident_time <= 0:
                    labels = ["SILENT"] * n_frames
                else:
                    if risky_time < 0:
                        risky_time = max(0, accident_time - L_FRAMES)
                    labels = label_one_clip(n_frames, accident_time, risky_time)

                # Save back
                ann["per_frame_labels"] = labels
                ann["label_rule"] = f"L={L_SEC}s, fps={FPS}, L_frames={L_FRAMES}"
                ann_path.write_text(json.dumps(ann, indent=2, ensure_ascii=False))

                for la in labels:
                    label_dist[f"nexar_{split}_{polarity}_{la}"] += 1
                stats[f"nexar_{split}_{polarity}"] += 1

                if is_positive and accident_time > 0:
                    case = "A" if (accident_time - L_FRAMES >= risky_time) else "B"
                    stats[f"nexar_{split}_{polarity}_case_{case}"] += 1

    return stats, label_dist


def main():
    logger.info("=== Processing DADA-2000 ===")
    dada_stats, dada_dist = process_dada()
    for k, v in sorted(dada_stats.items()):
        logger.info(f"  {k}: {v}")
    logger.info("  label distribution:")
    for k, v in sorted(dada_dist.items()):
        logger.info(f"    {k}: {v}")

    logger.info("\n=== Processing Nexar ===")
    nexar_stats, nexar_dist = process_nexar()
    for k, v in sorted(nexar_stats.items()):
        logger.info(f"  {k}: {v}")
    logger.info("  label distribution:")
    for k, v in sorted(nexar_dist.items()):
        logger.info(f"    {k}: {v}")

    # Summary
    print("\n" + "=" * 70)
    print("  DADA + Nexar Relabeling Summary")
    print("=" * 70)
    total_clips = sum(v for k, v in {**dada_stats, **nexar_stats}.items()
                       if not k.endswith(("_A", "_B", "_no_ann", "_no_frames")))
    total_A = sum(v for k, v in {**dada_stats, **nexar_stats}.items() if k.endswith("case_A"))
    total_B = sum(v for k, v in {**dada_stats, **nexar_stats}.items() if k.endswith("case_B"))
    print(f"  Total clips processed:  {total_clips}")
    print(f"  Case A (OBSERVE+ALERT): {total_A}  (risky_time > 2s before accident)")
    print(f"  Case B (ALL ALERT):     {total_B}  (risky_time within 2s of accident)")
    print(f"\n  Label distribution (frames):")
    all_dist = {**dada_dist, **nexar_dist}
    for la in ["SILENT", "OBSERVE", "ALERT"]:
        n = sum(v for k, v in all_dist.items() if k.endswith(f"_{la}"))
        print(f"    {la}: {n:>8d}")
    print()


if __name__ == "__main__":
    main()