assistant.py

"""Driving perception assistant — dashcam replay entry point.

Run:  python assistant.py --source dashcam.mp4
Keys: q quits the replay window.
"""

from __future__ import annotations

import argparse
import dataclasses
import json
import os
import tempfile
import time
from pathlib import Path

import cv2

from perception.engine import PerceptionEngine
from perception.hud import render
from perception.risk import RiskEngine
from perception.sources import VideoFileSource
from perception.types import PolicyConfig
from pothole.video import encode_h264

VALID_DISABLE = {"potholes", "scene", "depth"}
_STRIDE_FLAGS = {
    "stride_objects": "--stride-objects",
    "stride_potholes": "--stride-potholes",
    "stride_scene": "--stride-scene",
    "stride_depth": "--stride-depth",
}


def _depth_disabled_warning(disable: set[str]) -> str | None:
    """Critical-1 guard: with depth disabled, no object ever gets a
    distance_m, and risk.py's `consider()` bails out on distance_m is None
    for every single candidate — so disabling depth silently disables ALL
    alerting, not just distance readouts. Returns the stdout warning to
    print, or None when depth is enabled."""
    if "depth" in disable:
        return ("WARNING: depth disabled -- object distances are "
                 "unavailable, so NO alerts will be raised.")
    return None


def parse_args(argv=None) -> argparse.Namespace:
    p = argparse.ArgumentParser(description=__doc__)
    p.add_argument("--source", required=True, help="dashcam video file")
    p.add_argument("--output", default=None, help="output directory")
    p.add_argument("--device", default="auto", choices=["auto", "cpu", "cuda"])
    p.add_argument("--conf", type=float, default=0.35)
    p.add_argument("--disable", default="", help="comma list: depth,scene,potholes")
    p.add_argument("--stride-objects", type=int, default=1)
    p.add_argument("--stride-potholes", type=int, default=2)
    p.add_argument("--stride-scene", type=int, default=2)
    p.add_argument("--stride-depth", type=int, default=3)
    p.add_argument("--max-frames", type=int, default=0, help="0 = all")
    p.add_argument("--no-window", action="store_true")
    args = p.parse_args(argv)
    args.disable = {s.strip() for s in args.disable.split(",") if s.strip()}

    unknown = args.disable - VALID_DISABLE
    if unknown:
        p.error(
            f"unknown --disable value(s): {', '.join(sorted(unknown))} "
            f"(valid: {', '.join(sorted(VALID_DISABLE))})"
        )

    for attr, flag in _STRIDE_FLAGS.items():
        value = getattr(args, attr)
        if value < 1:
            p.error(f"{flag} must be >= 1, got {value}")

    return args


def _build_engine(args) -> PerceptionEngine:
    import torch

    device = args.device
    if device == "auto":
        device = "cuda" if torch.cuda.is_available() else "cpu"

    from perception.objects import CocoObjectDetector

    objects_fn = CocoObjectDetector(device=device, confidence=args.conf)

    potholes_fn = scene_fn = depth_fn = None
    if "potholes" not in args.disable:
        from perception.potholes import PotholeLayer

        potholes_fn = PotholeLayer(device=device)
    if "scene" not in args.disable:
        from perception.scene import Yolopv2SceneModel

        scene_fn = Yolopv2SceneModel(device=device)
    if "depth" not in args.disable:
        from perception.depth import MetricDepthEstimator

        depth_fn = MetricDepthEstimator(device=device)

    return PerceptionEngine(
        objects_fn=objects_fn, potholes_fn=potholes_fn,
        scene_fn=scene_fn, depth_fn=depth_fn,
        strides={"objects": args.stride_objects,
                 "potholes": args.stride_potholes,
                 "scene": args.stride_scene,
                 "depth": args.stride_depth},
    )


def run(args) -> dict:
    out_dir = Path(args.output or Path("runs/assistant") / Path(args.source).stem)
    out_dir.mkdir(parents=True, exist_ok=True)

    depth_warning = _depth_disabled_warning(args.disable)
    if depth_warning:
        print(depth_warning)

    engine = _build_engine(args)
    risk = RiskEngine(PolicyConfig())
    config = risk.config

    fd, raw_name = tempfile.mkstemp(suffix=".mp4")
    os.close(fd)
    raw_path = Path(raw_name)
    writer = None
    frames = 0
    last_time_s = 0.0
    t_wall = time.perf_counter()

    try:
        # Inner try/finally: on ANY exit from the frame loop — normal
        # completion, `q` pressed, the writer-open RuntimeError, or an
        # exception raised by engine.update/risk.update/render/etc. — the
        # writer must be released and the display window torn down. Mirrors
        # the cap/writer cleanup pattern in pothole/video.py's process_video.
        try:
            with VideoFileSource(args.source) as src:
                for frame_idx, time_s, frame in src:
                    if args.max_frames and frames >= args.max_frames:
                        break
                    state = engine.update(frame_idx, time_s, frame, config)
                    alerts = risk.update(state)
                    frames += 1
                    last_time_s = time_s
                    fps_now = frames / max(time.perf_counter() - t_wall, 1e-6)
                    annotated = render(frame, state, alerts, engine.last_timings,
                                       fps_now)

                    if writer is None:
                        h, w = annotated.shape[:2]
                        writer = cv2.VideoWriter(str(raw_path),
                                                 cv2.VideoWriter_fourcc(*"mp4v"),
                                                 src.fps, (w, h))
                        if not writer.isOpened():
                            raise RuntimeError("Could not open video writer (mp4v).")
                    writer.write(annotated)

                    if not args.no_window:
                        cv2.imshow("Driving Perception Assistant", annotated)
                        if cv2.waitKey(1) & 0xFF == ord("q"):
                            break
        finally:
            if writer is not None:
                writer.release()
            if not args.no_window:
                cv2.destroyAllWindows()

        if frames == 0:
            raise ValueError(f"No frames decoded from {args.source}")

        # Only encode to the output dir once the loop above completed
        # without raising — a mid-loop failure must not produce a partial
        # annotated.mp4.
        encode_h264(raw_path, out_dir / "annotated.mp4")
    finally:
        # Outer finally: the temp raw file is always removed, whether the
        # run succeeded, failed with zero frames, or raised mid-loop.
        raw_path.unlink(missing_ok=True)

    # An alert still active on the last processed frame would otherwise keep
    # t_end=None forever — force-end everything at the clip's last known
    # timestamp so alerts.json lets consumers compute a duration for it.
    # last_time_s (set only after a frame is actually processed) rather than
    # the loop variable directly: on a --max-frames cutoff, the for-loop has
    # already fetched (but not processed) one more frame by the time the
    # break executes, and that frame's timestamp would overshoot the true
    # last-processed time.
    risk.close_all(last_time_s)

    history = [dataclasses.asdict(a) for a in risk.history]
    (out_dir / "alerts.json").write_text(json.dumps(history, indent=2),
                                         encoding="utf-8")

    counts: dict[str, int] = {}
    for a in risk.history:
        counts[a.type] = counts.get(a.type, 0) + 1
    wall = time.perf_counter() - t_wall
    print(f"frames: {frames}  wall: {wall:.1f}s  fps: {frames / max(wall, 1e-6):.2f}")
    print(f"alerts: {counts or 'none'}")
    print(f"artifacts: {out_dir}")
    return {"frames": frames, "alerts": risk.history, "output_dir": out_dir}


if __name__ == "__main__":
    run(parse_args())