run_phase2.py

#!/usr/bin/env python3
"""
Phase 2 — Spatial Synthesis: Real-time 3D Object Detection.

Runs YOLO-World-S and Depth Anything V2 in parallel threads, fuses their
outputs via 3D back-projection, and displays per-object spatial positions.

Usage:
    python run_phase2.py
    python run_phase2.py --show-depth --confidence 0.4
    python run_phase2.py --no-display
"""

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.phase2")

_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]:
    """Load the class vocabulary from a YAML configuration file."""
    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)

    # Create pipeline (workers start loading models in background)
    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,
    )
    pipeline.start()

    annotator = FrameAnnotator(class_names=class_names)
    show_display = not args.no_display
    show_depth = args.show_depth
    window_name = "Spatial AI — Phase 2 Spatial Synthesis"

    # Allow models time to load in threads before flooding with frames
    logger.info("Waiting for models to initialise ...")
    time.sleep(2.0)

    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

            # Process through parallel pipeline
            detections_3d, depth_map, _, _, timings = pipeline.process_frame(
                frame_id, frame
            )

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

            # Stdout output
            if detections_3d:
                for d in detections_3d:
                    pos = d.get("position_3d", [0, 0, 0])
                    dist = d.get("distance", 0)
                    print(
                        f"[Frame {frame_id:>5d}] "
                        f"{d['label']:>15s} | "
                        f"bbox=({d['x1']:.0f},{d['y1']:.0f},"
                        f"{d['x2']:.0f},{d['y2']:.0f}) | "
                        f"pos=({pos[0]:+.2f},{pos[1]:+.2f},{pos[2]:+.2f})m | "
                        f"dist={dist:.2f}m | "
                        f"conf={d['confidence']:.2f} | "
                        f"YOLO={timings['yolo_ms']:.0f}ms "
                        f"Depth={timings['depth_ms']:.0f}ms"
                    )
            else:
                if frame_id % 30 == 0:
                    print(f"[Frame {frame_id:>5d}] No detections | "
                          f"YOLO={timings['yolo_ms']:.0f}ms "
                          f"Depth={timings['depth_ms']:.0f}ms")

            # Display
            if show_display:
                annotator.draw(
                    frame, detections_3d, frame_id,
                    timings["yolo_ms"], loop_ms,
                    timings=timings,
                    depth_map=depth_map,
                    show_depth_overlay=show_depth,
                )
                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 2 pipeline stopped.")


# ──────────────────────────────────────────────────────────────────────
# CLI
# ──────────────────────────────────────────────────────────────────────
def parse_args() -> argparse.Namespace:
    p = argparse.ArgumentParser(
        description="Phase 2 — Spatial Synthesis: 3D Object Detection",
        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
    )
    p.add_argument("--camera-id", type=int, default=0,
                    help="OpenCV VideoCapture device index")
    p.add_argument("--config", type=str, default="config/detection_classes.yaml",
                    help="Detection class vocabulary YAML")
    p.add_argument("--model", type=str, default="yolov8s-worldv2.pt",
                    help="YOLO-World model weights path")
    p.add_argument("--confidence", type=float, default=0.35,
                    help="Min detection confidence threshold")
    p.add_argument("--intrinsics", type=str,
                    default="config/camera_intrinsics.json",
                    help="Camera calibration JSON file")
    p.add_argument("--depth-model", type=str, default=None,
                    help="HuggingFace depth model ID (default: DA2-Small)")
    p.add_argument("--show-depth", action="store_true", default=False,
                    help="Overlay colourised depth map on display")
    p.add_argument("--no-display", action="store_true", default=False,
                    help="Suppress OpenCV display window")
    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)