cmvr_edge_ai/detect_server/show_detections.py

#!/usr/bin/env python3
"""Display live cmvr-es detection results in a local OpenCV window.

This is an isolated demo entrypoint.  It registers a temporary viewer Sink and
does not change the production detection pipeline or its HTTP alert behavior.
"""

from __future__ import annotations

import argparse
import asyncio
import logging
import math
import os
import signal
import sys
from collections.abc import Callable, Mapping
from concurrent.futures import Executor
from importlib import import_module
from pathlib import Path
from time import monotonic
from typing import Any

from pydantic import ValidationError

from cmvr_edge_ai.application import (
    EdgeAIApplication,
    create_default_model_registry,
    create_default_registry,
    validate_application,
)
from cmvr_edge_ai.compiler import PipelineCompileError
from cmvr_edge_ai.config import AppConfig, ConfigLoadError, load_config
from cmvr_edge_ai.contracts import Detection, DetectionResult, ImageFrame
from cmvr_edge_ai.core import ComponentContext, Envelope, Sink
from cmvr_edge_ai.observability import configure_logging
from cmvr_edge_ai.plugins import PluginKind, PluginRegistry, PluginSpec
from cmvr_edge_ai.workers import run_blocking


_LOGGER = logging.getLogger("cmvr_edge_ai.demo.detection_viewer")
_DEFAULT_CONFIG = (
    Path(__file__).resolve().parents[1]
    / "configs"
    / "debug"
    / "detection_viewer.yaml"
)
_VIEWER_PLUGIN_ID = "demo.opencv_detection_viewer@1"
_VIEWER_INPUT_PORTS = ("input", "ppe", "phone")


class DetectionViewerError(RuntimeError):
    """Raised when a detection frame cannot be displayed safely."""


class OpenCvDetectionViewerSink(Sink):
    """Pair same-frame model results and render their boxes in one window."""

    _PARAM_KEYS = frozenset(
        {
            "window_name",
            "window_width",
            "window_height",
            "wait_key_ms",
            "box_thickness",
            "font_scale",
            "show_stats",
            "expected_inputs",
            "max_pending_frames",
        }
    )

    def __init__(
        self,
        node_id: str,
        params: Mapping[str, Any],
        *,
        request_stop: Callable[[], None],
        module_loader: Callable[[str], Any] = import_module,
    ) -> None:
        unknown = set(params) - self._PARAM_KEYS
        if unknown:
            raise ValueError(
                "unknown detection viewer parameter(s): "
                + ", ".join(sorted(str(value) for value in unknown))
            )
        if not callable(request_stop):
            raise TypeError("request_stop must be callable")

        self._node_id = node_id
        self._window_name = _non_empty_string(
            params.get("window_name", "CMVR PPE Detection"),
            "window_name",
        )
        self._window_width = _bounded_int(
            params.get("window_width", 1280),
            "window_width",
            minimum=1,
            maximum=16384,
        )
        self._window_height = _bounded_int(
            params.get("window_height", 720),
            "window_height",
            minimum=1,
            maximum=16384,
        )
        self._wait_key_ms = _bounded_int(
            params.get("wait_key_ms", 1),
            "wait_key_ms",
            minimum=1,
            maximum=20,
        )
        self._box_thickness = _bounded_int(
            params.get("box_thickness", 2),
            "box_thickness",
            minimum=1,
            maximum=20,
        )
        self._font_scale = _bounded_float(
            params.get("font_scale", 0.6),
            "font_scale",
            minimum=0.1,
            maximum=5.0,
        )
        self._show_stats = _strict_bool(
            params.get("show_stats", True),
            "show_stats",
        )
        self._expected_inputs = _input_ports(
            params.get("expected_inputs", ["input"])
        )
        self._expected_input_set = frozenset(self._expected_inputs)
        self._max_pending_frames = _bounded_int(
            params.get("max_pending_frames", 8),
            "max_pending_frames",
            minimum=1,
            maximum=256,
        )
        self._request_stop = request_stop
        self._module_loader = module_loader
        self._executor: Executor | None = None
        self._cv2: Any = None
        self._numpy: Any = None
        self._window_open = False
        self._stop_requested = False
        self._event_pump_task: asyncio.Task[None] | None = None
        self._window_seen_visible = False
        self._frames_shown = 0
        self._last_frame_at: float | None = None
        self._display_fps: float | None = None
        self._pending_results: dict[
            tuple[str, str, str | None, int, int],
            dict[str, DetectionResult],
        ] = {}
        self._pending_evicted = 0

    async def setup(self, context: ComponentContext) -> None:
        executor = context.metadata.get("thread_executor")
        if executor is not None and not isinstance(executor, Executor):
            raise TypeError("component context thread_executor must be an Executor")
        self._executor = executor

        if sys.platform.startswith("linux") and not _linux_display_available(os.environ):
            raise DetectionViewerError(
                "OpenCV viewer needs a graphical session; DISPLAY and "
                "WAYLAND_DISPLAY are both unset. Run it on the desktop or use "
                "SSH X11 forwarding."
            )
        try:
            cv2 = self._module_loader("cv2")
            numpy = self._module_loader("numpy")
        except ImportError as exc:
            raise DetectionViewerError(
                "OpenCV viewer dependencies are missing; run "
                "'bash scripts/bootstrap.sh' first"
            ) from exc

        gui_backend = _opencv_gui_backend(cv2.getBuildInformation())
        if gui_backend.upper() in {"NONE", "NO"}:
            raise DetectionViewerError(
                "installed OpenCV has no GUI backend; install opencv-python "
                "instead of opencv-python-headless"
            )

        self._cv2 = cv2
        self._numpy = numpy
        try:
            cv2.namedWindow(self._window_name, cv2.WINDOW_NORMAL)
            self._window_open = True
            cv2.resizeWindow(
                self._window_name,
                self._window_width,
                self._window_height,
            )
        except Exception as exc:
            if self._window_open:
                try:
                    cv2.destroyWindow(self._window_name)
                except Exception:
                    pass
                self._window_open = False
            raise DetectionViewerError(
                "OpenCV could not create the viewer window; verify the local "
                "desktop session and DISPLAY configuration"
            ) from exc
        _LOGGER.info(
            "detection viewer opened node=%s window=%s gui_backend=%s "
            "quit_keys=q,esc",
            self._node_id,
            self._window_name,
            gui_backend,
        )

    async def start(self) -> None:
        if not self._window_open or self._cv2 is None:
            raise RuntimeError("detection viewer has not been set up")
        if self._event_pump_task is not None:
            raise RuntimeError("detection viewer has already been started")
        self._event_pump_task = asyncio.create_task(
            self._pump_window_events(),
            name=f"detection-viewer-events:{self._node_id}",
        )

    async def consume(
        self,
        envelope: Envelope[Any],
        input_port: str = "input",
    ) -> None:
        if self._stop_requested:
            return
        if not self._window_open or self._cv2 is None or self._numpy is None:
            raise RuntimeError("detection viewer has not been set up")
        if input_port not in self._expected_input_set:
            expected = ", ".join(self._expected_inputs)
            raise DetectionViewerError(
                f"{self._node_id} received unexpected input port "
                f"{input_port!r}; configured expected_inputs: {expected}"
            )
        result = envelope.payload
        if not isinstance(result, DetectionResult):
            raise TypeError(
                f"{self._node_id} expected DetectionResult, "
                f"got {type(result).__name__}"
            )
        frame = result.source_frame
        if frame is None:
            raise DetectionViewerError(
                "DetectionResult has no source_frame; set detector "
                "attach_frame: true in the demo config"
            )

        if len(self._expected_inputs) == 1:
            await self._display_results(((input_port, result),))
            return

        frame_key = (
            envelope.trace_id,
            envelope.source_id,
            envelope.session_id,
            envelope.sequence,
            id(frame),
        )
        pending = self._pending_results.setdefault(frame_key, {})
        if input_port in pending:
            raise DetectionViewerError(
                f"duplicate result for input port {input_port!r} and frame "
                f"source={envelope.source_id!r} sequence={envelope.sequence}"
            )
        pending[input_port] = result
        if self._expected_input_set.issubset(pending):
            self._pending_results.pop(frame_key, None)
            await self._display_results(
                tuple((port, pending[port]) for port in self._expected_inputs)
            )
            return

        while len(self._pending_results) > self._max_pending_frames:
            oldest_key = next(iter(self._pending_results))
            dropped = self._pending_results.pop(oldest_key)
            self._pending_evicted += 1
            _LOGGER.warning(
                "detection viewer evicted unmatched frame node=%s "
                "source=%s sequence=%s received_inputs=%s evicted_total=%s",
                self._node_id,
                oldest_key[1],
                oldest_key[3],
                ",".join(sorted(dropped)),
                self._pending_evicted,
            )

    async def _display_results(
        self,
        results: tuple[tuple[str, DetectionResult], ...],
    ) -> None:
        frame = results[0][1].source_frame
        if frame is None:  # guarded in consume; retained for type narrowing
            raise DetectionViewerError("DetectionResult has no source_frame")
        if any(result.source_frame is not frame for _, result in results[1:]):
            raise DetectionViewerError(
                "paired detection results do not reference the same source "
                "frame; keep the shared decoded-frame fan-out in the demo graph"
            )

        now = monotonic()
        if self._last_frame_at is not None and now > self._last_frame_at:
            instantaneous_fps = 1.0 / (now - self._last_frame_at)
            self._display_fps = (
                instantaneous_fps
                if self._display_fps is None
                else self._display_fps * 0.85 + instantaneous_fps * 0.15
            )
        self._last_frame_at = now
        header = None
        if self._show_stats:
            display_fps = self._display_fps or 0.0
            if len(results) == 1:
                result = results[0][1]
                model_name = result.model_name or result.model_id
                header = (
                    f"{model_name} | boxes={len(result.detections)} | "
                    f"inference={result.inference_ms:.1f} ms | "
                    f"display={display_fps:.1f} FPS"
                )
            else:
                model_stats = " | ".join(
                    f"{port.upper()} boxes={len(result.detections)} "
                    f"{result.inference_ms:.1f}ms"
                    for port, result in results
                )
                header = f"{model_stats} | display={display_fps:.1f} FPS"

        detections: list[Detection] = []
        for port, result in results:
            for detection in result.detections:
                if not isinstance(detection, Detection):
                    raise DetectionViewerError(
                        "DetectionResult.detections must contain Detection instances"
                    )
                if len(results) == 1:
                    detections.append(detection)
                else:
                    detections.append(
                        Detection(
                            label=f"{port.upper()}:{detection.label}",
                            confidence=detection.confidence,
                            box=detection.box,
                            track_id=detection.track_id,
                        )
                    )

        image = await run_blocking(
            _render_detection_frame,
            frame,
            tuple(detections),
            cv2_module=self._cv2,
            numpy_module=self._numpy,
            box_thickness=self._box_thickness,
            font_scale=self._font_scale,
            header=header,
            executor=self._executor,
        )
        self._cv2.imshow(self._window_name, image)
        self._frames_shown += 1

    async def stop(self) -> None:
        self._pending_results.clear()
        if not self._window_open:
            return
        self._window_open = False
        if self._event_pump_task is not None:
            self._event_pump_task.cancel()
            await asyncio.gather(self._event_pump_task, return_exceptions=True)
            self._event_pump_task = None
        try:
            self._cv2.destroyWindow(self._window_name)
        except Exception:
            _LOGGER.warning(
                "detection viewer window cleanup failed node=%s window=%s",
                self._node_id,
                self._window_name,
                exc_info=True,
            )
        _LOGGER.info(
            "detection viewer stopped node=%s frames_shown=%s "
            "unmatched_evicted=%s",
            self._node_id,
            self._frames_shown,
            self._pending_evicted,
        )

    async def _pump_window_events(self) -> None:
        """Keep the GUI responsive even while the camera produces no frames."""

        try:
            while self._window_open:
                try:
                    key = int(self._cv2.waitKey(self._wait_key_ms)) & 0xFF
                    if key in {27, ord("q"), ord("Q")}:
                        self._request_demo_stop("keyboard")
                        return
                    visible = float(
                        self._cv2.getWindowProperty(
                            self._window_name,
                            self._cv2.WND_PROP_VISIBLE,
                        )
                    )
                except Exception:
                    # Visibility queries are optional in some GUI backends.
                    visible = -1.0

                if visible >= 1.0:
                    self._window_seen_visible = True
                elif visible == 0.0 or (
                    visible < 0.0 and self._window_seen_visible
                ):
                    self._request_demo_stop("window_closed")
                    return

                # waitKey processes native events; this small cooperative pause
                # prevents an idle, frame-less stream from spinning one CPU core.
                await asyncio.sleep(max(0.01, self._wait_key_ms / 1000.0))
        except asyncio.CancelledError:
            raise

    def _request_demo_stop(self, reason: str) -> None:
        if self._stop_requested:
            return
        self._stop_requested = True
        _LOGGER.info(
            "detection viewer stop requested node=%s reason=%s frames_shown=%s",
            self._node_id,
            reason,
            self._frames_shown,
        )
        self._request_stop()


def _render_detection_frame(
    frame: ImageFrame,
    detections: tuple[Detection, ...],
    *,
    cv2_module: Any,
    numpy_module: Any,
    box_thickness: int,
    font_scale: float,
    header: str | None,
) -> Any:
    """Copy one packed frame and draw validated boxes into a BGR ndarray."""

    _validate_display_frame(frame)
    image = (
        numpy_module.frombuffer(frame.data, dtype=numpy_module.uint8)
        .reshape((frame.height, frame.width, 3))
        .copy()
    )
    if frame.pixel_format.strip().upper() == "RGB8":
        image = cv2_module.cvtColor(image, cv2_module.COLOR_RGB2BGR)

    if header:
        cv2_module.rectangle(
            image,
            (0, 0),
            (frame.width - 1, min(frame.height - 1, 30)),
            (24, 24, 24),
            -1,
        )
        cv2_module.putText(
            image,
            header,
            (8, min(frame.height - 1, 21)),
            cv2_module.FONT_HERSHEY_SIMPLEX,
            0.55,
            (255, 255, 255),
            1,
            cv2_module.LINE_AA,
        )

    # Draw evidence after the status bar so boxes at the top of the image are
    # never hidden behind presentation-only statistics.
    for detection in detections:
        if not isinstance(detection, Detection):
            raise DetectionViewerError(
                "DetectionResult.detections must contain Detection instances"
            )
        box = _clipped_box(detection, frame.width, frame.height)
        if box is None:
            continue
        color = _label_color_bgr(detection.label)
        x_min, y_min, x_max, y_max = box
        cv2_module.rectangle(
            image,
            (x_min, y_min),
            (x_max, y_max),
            color,
            box_thickness,
        )
        label = f"{detection.label} {detection.confidence:.2f}"
        (text_width, text_height), baseline = cv2_module.getTextSize(
            label,
            cv2_module.FONT_HERSHEY_SIMPLEX,
            font_scale,
            1,
        )
        text_bottom = max(text_height + baseline + 4, y_min)
        background_top = max(0, text_bottom - text_height - baseline - 6)
        background_right = min(frame.width - 1, x_min + text_width + 6)
        cv2_module.rectangle(
            image,
            (x_min, background_top),
            (background_right, text_bottom),
            color,
            -1,
        )
        cv2_module.putText(
            image,
            label,
            (x_min + 3, max(text_height + 1, text_bottom - baseline - 3)),
            cv2_module.FONT_HERSHEY_SIMPLEX,
            font_scale,
            (255, 255, 255),
            1,
            cv2_module.LINE_AA,
        )

    return image


def _validate_display_frame(frame: ImageFrame) -> None:
    if not isinstance(frame, ImageFrame):
        raise DetectionViewerError(
            f"expected ImageFrame, got {type(frame).__name__}"
        )
    if not isinstance(frame.codec, str):
        raise DetectionViewerError("frame codec must be a string")
    if frame.is_encoded:
        raise DetectionViewerError("viewer requires a decoded ImageFrame")
    if not isinstance(frame.pixel_format, str):
        raise DetectionViewerError("frame pixel_format must be a string")
    pixel_format = frame.pixel_format.strip().upper()
    if pixel_format not in {"BGR8", "RGB8"}:
        raise DetectionViewerError(
            f"viewer supports BGR8 or RGB8, got {frame.pixel_format!r}"
        )
    if (
        isinstance(frame.width, bool)
        or not isinstance(frame.width, int)
        or frame.width < 1
        or isinstance(frame.height, bool)
        or not isinstance(frame.height, int)
        or frame.height < 1
    ):
        raise DetectionViewerError("frame dimensions must be positive integers")
    if not isinstance(frame.data, bytes):
        raise DetectionViewerError("packed frame buffer must be bytes")
    expected_size = frame.width * frame.height * 3
    if len(frame.data) != expected_size:
        raise DetectionViewerError(
            f"packed frame has {len(frame.data)} bytes; expected {expected_size}"
        )


def _clipped_box(
    detection: Detection,
    width: int,
    height: int,
) -> tuple[int, int, int, int] | None:
    values = (
        detection.box.x_min,
        detection.box.y_min,
        detection.box.x_max,
        detection.box.y_max,
    )
    if any(
        isinstance(value, bool)
        or not isinstance(value, (int, float))
        or not math.isfinite(float(value))
        for value in values
    ):
        return None
    x_min, y_min, x_max, y_max = (float(value) for value in values)
    if x_max <= x_min or y_max <= y_min:
        return None
    left = min(max(int(math.floor(x_min)), 0), width - 1)
    top = min(max(int(math.floor(y_min)), 0), height - 1)
    right = min(max(int(math.ceil(x_max)), 0), width - 1)
    bottom = min(max(int(math.ceil(y_max)), 0), height - 1)
    if right <= left or bottom <= top:
        return None
    return left, top, right, bottom


def _label_color_bgr(label: str) -> tuple[int, int, int]:
    seed = sum((index + 1) * byte for index, byte in enumerate(label.encode("utf-8")))
    return (
        64 + (seed * 11) % 192,
        64 + (seed * 5) % 192,
        64 + seed % 192,
    )


def _linux_display_available(environment: Mapping[str, str]) -> bool:
    return bool(environment.get("DISPLAY") or environment.get("WAYLAND_DISPLAY"))


def _opencv_gui_backend(build_information: str) -> str:
    for line in str(build_information).splitlines():
        stripped = line.strip()
        if stripped.upper().startswith("GUI:"):
            return stripped.split(":", maxsplit=1)[1].strip() or "unknown"
    return "unknown"


def _non_empty_string(value: Any, name: str) -> str:
    if not isinstance(value, str) or not value.strip():
        raise ValueError(f"{name} must be a non-empty string")
    return value.strip()


def _bounded_int(value: Any, name: str, *, minimum: int, maximum: int) -> int:
    if isinstance(value, bool) or not isinstance(value, int):
        raise ValueError(f"{name} must be an integer")
    if not minimum <= value <= maximum:
        raise ValueError(f"{name} must be between {minimum} and {maximum}")
    return value


def _bounded_float(
    value: Any,
    name: str,
    *,
    minimum: float,
    maximum: float,
) -> float:
    if isinstance(value, bool):
        raise ValueError(f"{name} must be a number")
    try:
        parsed = float(value)
    except (TypeError, ValueError) as exc:
        raise ValueError(f"{name} must be a number") from exc
    if not math.isfinite(parsed) or not minimum <= parsed <= maximum:
        raise ValueError(f"{name} must be between {minimum} and {maximum}")
    return parsed


def _strict_bool(value: Any, name: str) -> bool:
    if type(value) is not bool:
        raise ValueError(f"{name} must be a boolean")
    return value


def _input_ports(value: Any) -> tuple[str, ...]:
    if not isinstance(value, (list, tuple)):
        raise ValueError("expected_inputs must be a list of input port names")
    ports = tuple(_non_empty_string(item, "expected_inputs item") for item in value)
    if not ports:
        raise ValueError("expected_inputs must contain at least one input port")
    if len(set(ports)) != len(ports):
        raise ValueError("expected_inputs must not contain duplicate ports")
    unknown = sorted(set(ports) - set(_VIEWER_INPUT_PORTS))
    if unknown:
        raise ValueError(
            "expected_inputs contains unsupported viewer port(s): "
            + ", ".join(unknown)
        )
    return ports


def _build_registry(request_stop: Callable[[], None]) -> PluginRegistry:
    model_registry = create_default_model_registry()
    registry = create_default_registry(model_registry=model_registry)
    registry.register(
        PluginSpec(
            plugin_id=_VIEWER_PLUGIN_ID,
            kind=PluginKind.SINK,
            factory=lambda node_id, params: OpenCvDetectionViewerSink(
                node_id,
                params,
                request_stop=request_stop,
            ),
            inputs={
                "input": "DetectionResult/v1",
                "ppe": "DetectionResult/v1",
                "phone": "DetectionResult/v1",
            },
            description=(
                "Pair same-frame model results and show their boxes in OpenCV"
            ),
            tags=frozenset({"demo", "visualization"}),
        )
    )
    return registry


def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(
        description="Show live cmvr-es YOLO detections in an OpenCV window"
    )
    parser.add_argument("--config", "-c", type=Path, default=_DEFAULT_CONFIG)
    parser.add_argument("--pipeline", default="detection_show")
    parser.add_argument("--log-level", default="INFO")
    parser.add_argument("--log-format", choices=("text", "json"), default="text")
    parser.add_argument(
        "--validate-only",
        action="store_true",
        help="validate the demo graph without opening a camera or GUI window",
    )
    return parser


def main(argv: list[str] | None = None) -> int:
    args = build_parser().parse_args(argv)
    configure_logging(args.log_level, args.log_format)
    try:
        config = load_config(args.config)
        if args.validate_only:
            validate_application(config, _build_registry(lambda: None), (args.pipeline,))
            print(f"demo configuration is valid; pipeline: {args.pipeline}")
            return 0
        return asyncio.run(_run_demo(config, args.pipeline))
    except (ConfigLoadError, ValidationError, PipelineCompileError, ValueError) as exc:
        print(f"configuration error: {exc}", file=sys.stderr)
        return 2
    except KeyboardInterrupt:
        return 130
    except Exception as exc:
        print(f"runtime error: {exc}", file=sys.stderr)
        return 1


async def _run_demo(config: AppConfig, pipeline_id: str) -> int:
    stop_event = asyncio.Event()
    application = EdgeAIApplication(config, _build_registry(stop_event.set))
    loop = asyncio.get_running_loop()
    for signum in (signal.SIGINT, signal.SIGTERM):
        try:
            loop.add_signal_handler(signum, stop_event.set)
        except NotImplementedError:
            pass

    await application.start((pipeline_id,))
    _LOGGER.info(
        "detection viewer demo running pipeline=%s; press q or Esc in the window to stop",
        pipeline_id,
    )
    wait_task = asyncio.create_task(application.wait(), name="demo-application-wait")
    stop_task = asyncio.create_task(stop_event.wait(), name="demo-stop-request")
    try:
        done, _ = await asyncio.wait(
            (wait_task, stop_task),
            return_when=asyncio.FIRST_COMPLETED,
        )
        if wait_task in done:
            await wait_task
        else:
            await application.stop(graceful=True)
            wait_task.cancel()
            await asyncio.gather(wait_task, return_exceptions=True)
    finally:
        stop_task.cancel()
        await asyncio.gather(stop_task, return_exceptions=True)
        await application.stop(graceful=True)
    return 0


if __name__ == "__main__":
    raise SystemExit(main())