tools/generate_beliefs.py

"""Generate per-frame <|BELIEF|> content for DoTA and DADA datasets.

Final belief type rules:
  Type 1 (GPT-4o):        Keep as-is (already in corpus)
  Type 2 (DADA acc_type):  Keep — accident_type text at accident_time frame
  Type 3 (DoTA acc_name):  Convert to natural language; normal → diverse safe phrases
  Type 4 (Template):       ❌ DELETE ALL
  Type 5 (DADA human):     Keep only for SILENT, 1 frame/video:
                            negative → random frame
                            positive → first frame (only if frame 0 < risky_time, else skip)

This script writes 'per_frame_beliefs' into each annotation.json.
"""
from __future__ import annotations
import json, glob, random, hashlib, logging
from pathlib import Path
from collections import Counter

ROOT = Path("PROJECT_ROOT")
DADA_ROOT = ROOT / "DADA-2000"
DOTA_ROOT = ROOT / "DoTA"

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

# ─── DoTA accident_name → natural language ───
ACCIDENT_NAME_MAP = {
    "normal":                     None,  # handled separately
    "turning":                    "turning",
    "lateral":                    "lateral collision",
    "moving_ahead_or_waiting":    "moving ahead or waiting",
    "leave_to_left":              "leaving lane to the left",
    "leave_to_right":             "leaving lane to the right",
    "oncoming":                   "oncoming vehicle",
    "obstacle":                   "obstacle on road",
    "pedestrian":                 "pedestrian in path",
    "start_stop_or_stationary":   "start, stop, or stationary vehicle",
    "unknown":                    "unknown anomaly",
}

# ─── Diverse "normal driving" belief bank (50 phrases) ───
NORMAL_BELIEFS = [
    "clear road ahead, normal traffic flow, no hazards detected",
    "steady driving, lane markings visible, surroundings stable",
    "open road with no immediate threats, maintaining safe speed",
    "traffic moving smoothly, no sudden changes in surrounding vehicles",
    "routine driving conditions, road surface in good condition",
    "normal lane keeping, no vehicles encroaching from adjacent lanes",
    "safe following distance maintained, lead vehicle steady",
    "no pedestrians or cyclists in the immediate vicinity",
    "driving straight ahead, visibility is clear, no obstructions",
    "surrounding traffic is predictable, no erratic behavior observed",
    "road is clear, weather conditions appear normal for driving",
    "no signs of developing hazard, all lanes flowing freely",
    "ego vehicle maintaining course, no steering correction needed",
    "intersection clear, no conflicting traffic approaching",
    "highway driving, vehicles spaced evenly, no sudden braking ahead",
    "urban road with normal density, traffic signals functioning",
    "residential area, low traffic volume, no unexpected obstacles",
    "gentle curve ahead, road conditions suitable, maintaining speed",
    "parked vehicles on roadside, no doors opening, path clear",
    "green traffic light, proceeding normally through intersection",
    "overpass approach, structural clearance adequate, no concerns",
    "multilane road, adjacent vehicles maintaining their lanes",
    "slight uphill grade, engine load normal, visibility unaffected",
    "road markings intact, lane boundaries well defined",
    "bridge crossing, road surface stable, wind conditions manageable",
    "traffic circle ahead, yielding as required, flow is orderly",
    "school zone but outside active hours, speed limit noted",
    "construction zone ended, resuming normal driving speed",
    "ramp merging area, checking mirrors, gap available",
    "tunnel exit, adjusting to ambient light, road ahead visible",
    "no emergency vehicles detected, audio environment calm",
    "fuel station visible on right, no vehicles entering from driveway",
    "median barrier present, oncoming traffic fully separated",
    "crosswalk ahead but no pedestrians waiting to cross",
    "bus stop area, no bus currently stopped, lane unobstructed",
    "speed bump traversed, resuming normal speed smoothly",
    "rail crossing clear, no signals active, proceeding safely",
    "driveway entrance on left, no vehicles emerging",
    "road gradient flattening, coasting at target speed",
    "passing a slower vehicle in the adjacent lane, safe clearance",
    "street lighting adequate, nighttime visibility acceptable",
    "wet road surface but no standing water, traction appears normal",
    "slight fog in distance, current visibility still sufficient",
    "delivery truck parked with hazards on, passing with clearance",
    "motorcycle in adjacent lane, maintaining steady position",
    "roundabout exit taken, straightening into destination lane",
    "shopping area with moderate pedestrian activity on sidewalk",
    "cyclist on bike lane to the right, separated by marking",
    "ambulance parked at curb with lights off, no obstruction",
    "dust or debris visible on road shoulder, driving lane clear",
]


def _pick_normal_belief(video_name: str, frame_id: int) -> str:
    """Deterministic diverse pick based on hash."""
    h = int(hashlib.md5(f"{video_name}_{frame_id}".encode()).hexdigest(), 16)
    return NORMAL_BELIEFS[h % len(NORMAL_BELIEFS)]


def _anomaly_belief(accident_name: str) -> str:
    """Convert DoTA accident_name to natural-language belief."""
    natural = ACCIDENT_NAME_MAP.get(accident_name, accident_name.replace("_", " "))
    return f"{natural} — Loss of control"


# ═══════════════════════════════════════════════════════════════
#  DoTA: generate per-frame beliefs from per-frame accident_name
# ═══════════════════════════════════════════════════════════════
def process_dota():
    stats = Counter()
    ann_dir = DOTA_ROOT / "annotations"
    for ann_path in sorted(ann_dir.glob("*.json")):
        d = json.load(open(ann_path))
        vname = d.get("video_name", ann_path.stem)
        labels = d.get("labels", [])
        if not labels:
            stats["skip_no_labels"] += 1
            continue

        beliefs = []
        for L in labels:
            fid = L.get("frame_id", 0)
            aname = L.get("accident_name", "normal")
            if aname == "normal":
                beliefs.append(_pick_normal_belief(vname, fid))
                stats["dota_normal"] += 1
            else:
                beliefs.append(_anomaly_belief(aname))
                stats["dota_anomaly"] += 1

        d["per_frame_beliefs"] = beliefs
        ann_path.write_text(json.dumps(d, indent=2, ensure_ascii=False))
        stats["dota_clips"] += 1

    return stats


# ═══════════════════════════════════════════════════════════════
#  DADA: generate beliefs from accident_type + Type 5 rules
# ═══════════════════════════════════════════════════════════════
def _dada_type5_belief(ann: dict) -> str:
    """DADA human annotation belief from metadata fields."""
    weather = ann.get("weather", "normal")
    road = ann.get("road_type", "road")
    speed = ann.get("car_speed", "normal")
    tod = ann.get("time_of_day", "day")
    return f"Normal driving on {road}, {weather} weather, {speed} speed, {tod}"


def process_dada():
    stats = 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.load(open(ann_path))
            is_positive = str(ann.get("accident", "False")).lower() == "true"
            accident_time = int(ann.get("accident_time", -1))
            risky_time = int(ann.get("risky_time", -1))
            accident_type = ann.get("accident_type", "")
            n_frames = len(ann.get("per_frame_labels", []))
            if n_frames == 0:
                # Fallback: count images
                n_frames = len(list(clip_dir.glob("*.jpg"))) + len(list(clip_dir.glob("*.png")))
                if (clip_dir / "images").is_dir():
                    n_frames = max(n_frames,
                                   len(list((clip_dir / "images").glob("*.jpg"))) +
                                   len(list((clip_dir / "images").glob("*.png"))))

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

            beliefs = [None] * n_frames  # None = no belief for this frame

            # Type 2: accident_type at accident_time frame
            if is_positive and accident_time >= 0 and accident_type:
                if accident_time < n_frames:
                    beliefs[accident_time] = accident_type
                    stats["dada_type2"] += 1

            # Type 5: DADA human annotation, 1 SILENT frame per video
            if cat == "negative":
                # Random frame
                rng = random.Random(hash(str(clip_dir)))
                idx = rng.randint(0, n_frames - 1)
                beliefs[idx] = _dada_type5_belief(ann)
                stats["dada_type5_neg"] += 1
            elif is_positive:
                # First frame, only if frame 0 < risky_time
                if risky_time > 0:  # frame 0 is before risky_time
                    beliefs[0] = _dada_type5_belief(ann)
                    stats["dada_type5_pos"] += 1
                else:
                    stats["dada_type5_pos_skip"] += 1

            ann["per_frame_beliefs"] = beliefs
            ann_path.write_text(json.dumps(ann, indent=2, ensure_ascii=False))
            stats[f"dada_{cat}"] += 1

    return stats


def main():
    logger.info("=== Generating DoTA beliefs ===")
    dota_stats = process_dota()
    for k, v in sorted(dota_stats.items()):
        logger.info(f"  {k}: {v}")

    logger.info("\n=== Generating DADA beliefs ===")
    dada_stats = process_dada()
    for k, v in sorted(dada_stats.items()):
        logger.info(f"  {k}: {v}")

    # ═══ Summary with examples ═══
    print("\n" + "=" * 80)
    print("  BELIEF GENERATION COMPLETE")
    print("=" * 80)

    # DoTA examples
    print("\n── DoTA Examples ──")
    ann = json.load(open(next((DOTA_ROOT / "annotations").glob("*.json"))))
    vname = ann["video_name"]
    labels = ann["labels"]
    beliefs = ann["per_frame_beliefs"]
    a_start = ann.get("anomaly_start", -1)
    print(f"  Clip: {vname}  anomaly_start={a_start}")
    # Show 2 normal + 2 anomaly
    shown_n = shown_a = 0
    for i, (L, b) in enumerate(zip(labels, beliefs)):
        aname = L["accident_name"]
        if aname == "normal" and shown_n < 2:
            print(f"    frame {L['frame_id']:>3d} [normal]:  <|BELIEF|> {b} </|BELIEF|>")
            shown_n += 1
        elif aname != "normal" and shown_a < 2:
            print(f"    frame {L['frame_id']:>3d} [{aname}]:  <|BELIEF|> {b} </|BELIEF|>")
            shown_a += 1
        if shown_n >= 2 and shown_a >= 2:
            break

    # DADA examples
    print("\n── DADA Examples ──")
    for cat in ["positive", "negative"]:
        cat_dir = DADA_ROOT / cat
        for clip_dir in sorted(cat_dir.iterdir())[:20]:
            ann_path = clip_dir / "annotation.json"
            if not ann_path.exists():
                continue
            ann = json.load(open(ann_path))
            beliefs = ann.get("per_frame_beliefs", [])
            non_none = [(i, b) for i, b in enumerate(beliefs) if b is not None]
            if non_none:
                print(f"  {cat}/{clip_dir.name}:")
                for idx, b in non_none[:2]:
                    label = ann.get("per_frame_labels", ["?"] * len(beliefs))[idx] if idx < len(ann.get("per_frame_labels", [])) else "?"
                    print(f"    frame {idx:>3d} [{label}]:  <|BELIEF|> {b} </|BELIEF|>")
                break

    # Final count
    print(f"\n  DoTA: {dota_stats.get('dota_clips', 0)} clips, "
          f"{dota_stats.get('dota_normal', 0)} normal beliefs + "
          f"{dota_stats.get('dota_anomaly', 0)} anomaly beliefs")
    print(f"  DADA: Type2 (accident_type) = {dada_stats.get('dada_type2', 0)}, "
          f"Type5 (human) = {dada_stats.get('dada_type5_neg', 0) + dada_stats.get('dada_type5_pos', 0)} "
          f"(neg={dada_stats.get('dada_type5_neg', 0)}, pos={dada_stats.get('dada_type5_pos', 0)}, "
          f"skip={dada_stats.get('dada_type5_pos_skip', 0)})")


if __name__ == "__main__":
    main()