run_phase3.py

#!/usr/bin/env python3
"""
Phase 3 — Scene Graph Engine: Real-time 3D Scene Understanding.

Extends Phase 2 with object tracking, spatial predicate assignment,
dwell-time filtering, and scene graph serialisation.

Usage:
    python run_phase3.py
    python run_phase3.py --show-graph --show-depth
    python run_phase3.py --dwell-frames 6 --ema-alpha 0.25
"""

import argparse
import logging
import signal
import time
from pathlib import Path

import cv2
import yaml

from src.video_capture import VideoCaptureManager
from src.pipeline import SpatialPipeline
from src.display import FrameAnnotator

# ──────────────────────────────────────────────────────────────────────
# Logging
# ──────────────────────────────────────────────────────────────────────
logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s | %(levelname)-7s | %(name)s | %(message)s",
    datefmt="%H:%M:%S",
)
logger = logging.getLogger("spatial_ai.phase3")

_shutdown = False


def _signal_handler(signum, frame):
    global _shutdown
    _shutdown = True
    logger.info("Shutdown signal received.")


signal.signal(signal.SIGINT, _signal_handler)


# ──────────────────────────────────────────────────────────────────────
# Config
# ──────────────────────────────────────────────────────────────────────
def load_class_vocab(config_path: str) -> list[str]:
    path = Path(config_path)
    if not path.exists():
        logger.warning("Config '%s' not found — using defaults.", config_path)
        return ["person", "chair", "table", "door", "backpack",
                "bottle", "cup", "laptop", "couch", "bed"]
    with open(path, "r") as f:
        data = yaml.safe_load(f)
    classes = data.get("classes", [])
    if not classes:
        raise ValueError(f"No 'classes' key in {config_path}")
    logger.info("Loaded %d classes from '%s'", len(classes), config_path)
    return classes


# ──────────────────────────────────────────────────────────────────────
# Main
# ──────────────────────────────────────────────────────────────────────
def main(args: argparse.Namespace) -> None:
    class_names = load_class_vocab(args.config)

    pipeline = SpatialPipeline(
        class_names=class_names,
        conf_threshold=args.confidence,
        yolo_model_path=args.model,
        depth_model_id=args.depth_model,
        intrinsics_path=args.intrinsics,
        scene_graph_enabled=True,
        scene_config_path=args.scene_config,
        dwell_frames=args.dwell_frames,
        ema_alpha=args.ema_alpha,
    )
    pipeline.start()

    annotator = FrameAnnotator(class_names=class_names)
    show_display = not args.no_display
    show_depth = args.show_depth
    show_graph = args.show_graph
    window_name = "Spatial AI — Phase 3 Scene Graph"

    logger.info("Waiting for models to initialise ...")
    time.sleep(2.0)

    t_start = time.perf_counter()

    with VideoCaptureManager(
        camera_id=args.camera_id,
        frame_interval_s=args.frame_interval,
        flush_grabs=2 if args.frame_interval == 0 else 0,
        stream_flush_grabs=5 if args.frame_interval > 0 else 0,
    ) as cam:
        logger.info("Pipeline running — press 'q' to quit.")
        prev_time = time.perf_counter()

        while not _shutdown:
            frame_id, frame = cam.read()
            if frame is None:
                logger.warning("Empty frame — stopping.")
                break

            elapsed = time.perf_counter() - t_start

            dets, depth_map, graph_str, scene_edges, timings = pipeline.process_frame(
                frame_id, frame, timestamp=elapsed,
            )

            now = time.perf_counter()
            loop_ms = (now - prev_time) * 1000.0
            prev_time = now

            # Print scene graph
            if show_graph and graph_str:
                print(f"\n{graph_str}")
            elif dets:
                for d in dets:
                    pos = d.get("smoothed_position", d.get("position_3d", [0, 0, 0]))
                    dist = d.get("distance", 0)
                    tid = d.get("track_id", "?")
                    print(
                        f"[Frame {frame_id:>5d}] "
                        f"#{tid:<3d} {d['label']:>12s} | "
                        f"pos=({pos[0]:+.2f},{pos[1]:+.2f},{pos[2]:+.2f})m | "
                        f"dist={dist:.2f}m | conf={d['confidence']:.2f} | "
                        f"YOLO={timings['yolo_ms']:.0f}ms "
                        f"Depth={timings['depth_ms']:.0f}ms "
                        f"Graph={timings['graph_ms']:.0f}ms"
                    )
            else:
                if frame_id % 30 == 0:
                    print(f"[Frame {frame_id:>5d}] No detections")

            if show_display:
                annotator.draw(
                    frame, dets, frame_id,
                    timings["yolo_ms"], loop_ms,
                    timings=timings,
                    depth_map=depth_map,
                    show_depth_overlay=show_depth,
                    scene_edges=scene_edges,
                )
                cv2.imshow(window_name, frame)
                key = cv2.waitKey(1) & 0xFF
                if key == ord("q"):
                    logger.info("'q' pressed — shutting down.")
                    break

    pipeline.stop()
    if show_display:
        cv2.destroyAllWindows()
    logger.info("Phase 3 pipeline stopped.")


# ──────────────────────────────────────────────────────────────────────
# CLI
# ──────────────────────────────────────────────────────────────────────
def parse_args() -> argparse.Namespace:
    p = argparse.ArgumentParser(
        description="Phase 3 — Scene Graph Engine: Real-time 3D Scene Understanding",
        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
    )
    p.add_argument("--camera-id", type=int, default=0)
    p.add_argument("--config", type=str, default="config/detection_classes.yaml")
    p.add_argument("--model", type=str, default="yolov8s-worldv2.pt")
    p.add_argument("--confidence", type=float, default=0.35)
    p.add_argument("--intrinsics", type=str, default="config/camera_intrinsics.json")
    p.add_argument("--depth-model", type=str, default=None)
    p.add_argument("--show-depth", action="store_true", default=False,
                    help="Overlay colourised depth map")
    p.add_argument("--no-display", action="store_true", default=False)
    # Phase 3 specific
    p.add_argument("--scene-config", type=str, default="config/scene_graph.yaml",
                    help="Scene graph predicate config YAML")
    p.add_argument("--dwell-frames", type=int, default=4,
                    help="Predicate dwell filter frame count")
    p.add_argument("--ema-alpha", type=float, default=0.35,
                    help="EMA position smoothing coefficient")
    p.add_argument("--show-graph", action="store_true", default=False,
                    help="Print serialised scene graph to stdout")
    p.add_argument("--frame-interval", type=float, default=0.0,
                    help="Seconds between captured frames (0 = live)")
    return p.parse_args()


if __name__ == "__main__":
    args = parse_args()
    main(args)