yovision/Brain/yovision_brain/runtime.py

from __future__ import annotations

import math
import threading
import time
from collections import deque
from datetime import datetime, timezone
from typing import Any

from .domain import Detection, Point, Zone, ZoneEntryEvaluator, zone_from_payload
from .source import CentroidTracker, HOGPersonDetector, require_opencv

try:
    import cv2  # type: ignore
except ImportError:  # pragma: no cover
    cv2 = None


DEFAULT_ZONE = Zone(
    "zone-demo-01",
    "楼梯口危险区",
    1,
    (Point(0.55, 0.30), Point(0.90, 0.30), Point(0.90, 0.88), Point(0.55, 0.88)),
)


class DemoEngine:
    def __init__(self, source: Any, fps: float = 2.0, event_limit: int = 100) -> None:
        if not math.isfinite(fps) or fps <= 0.0 or fps > 30.0:
            raise ValueError("fps must be within (0, 30]")
        if not 1 <= event_limit <= 100:
            raise ValueError("event_limit must be within [1, 100]")
        require_opencv()
        self._source = source
        self._fps = fps
        self._detector = None if source.fixture else HOGPersonDetector()
        self._tracker = CentroidTracker()
        self._evaluator = ZoneEntryEvaluator(source.source_ref, source.fixture)
        self._zone = DEFAULT_ZONE
        self._events: deque[dict[str, object]] = deque(maxlen=event_limit)
        self._lock = threading.RLock()
        self._stop = threading.Event()
        self._thread: threading.Thread | None = None
        self._sequence = 0
        self._frame_jpeg: bytes | None = None
        self._frame_width = 0
        self._frame_height = 0
        self._captured_at: str | None = None
        self._detections: list[dict[str, object]] = []
        self._latency_ms: float | None = None
        self._connected = False
        self._last_error_code: str | None = None

    def start(self) -> None:
        if self._thread is not None:
            return
        self._thread = threading.Thread(target=self._run, name="brain-demo", daemon=True)
        self._thread.start()

    def stop(self) -> None:
        self._stop.set()
        if self._thread is not None:
            self._thread.join(timeout=3.0)
            self._thread = None
        self._source.close()

    def _run(self) -> None:
        interval = 1.0 / self._fps
        while not self._stop.is_set():
            started = time.perf_counter()
            try:
                self.step()
            except RuntimeError:
                with self._lock:
                    self._connected = False
                    self._last_error_code = "source_unavailable"
            elapsed = time.perf_counter() - started
            self._stop.wait(max(0.0, interval - elapsed))

    def step(self) -> None:
        started = time.perf_counter()
        packet = self._source.read()
        self._sequence += 1
        if packet.scripted_detections is not None:
            detections = list(packet.scripted_detections)
        else:
            raw_boxes = self._detector.detect(packet.frame) if self._detector is not None else []
            detections = self._tracker.update(raw_boxes, self._sequence)
        with self._lock:
            zone = self._zone
        new_events, inside_by_track = self._evaluator.evaluate(self._sequence, packet.captured_at, zone, detections)
        ok, encoded = cv2.imencode(".jpg", packet.frame, [int(cv2.IMWRITE_JPEG_QUALITY), 82])
        if not ok:
            raise RuntimeError("frame encoding failed")
        height, width = packet.frame.shape[:2]
        serialized_detections = [self._serialize_detection(item, inside_by_track.get(item.track_id, False)) for item in detections]
        with self._lock:
            for event in new_events:
                self._events.appendleft(event.as_dict())
            self._frame_jpeg = encoded.tobytes()
            self._frame_width = int(width)
            self._frame_height = int(height)
            self._captured_at = packet.captured_at.astimezone(timezone.utc).isoformat().replace("+00:00", "Z")
            self._detections = serialized_detections
            self._latency_ms = round((time.perf_counter() - started) * 1000.0, 1)
            self._connected = True
            self._last_error_code = None

    @staticmethod
    def _serialize_detection(detection: Detection, inside: bool) -> dict[str, object]:
        return {
            "track_id": detection.track_id,
            "class": detection.class_name,
            "bbox": [detection.box.x1, detection.box.y1, detection.box.x2, detection.box.y2],
            "detector_score": detection.detector_score,
            "inside_zone": inside,
        }

    def update_zone(self, payload: object) -> Zone:
        with self._lock:
            updated = zone_from_payload(payload, self._zone)
            self._zone = updated
            self._evaluator.reset()
            return updated

    def frame_jpeg(self) -> bytes | None:
        with self._lock:
            return self._frame_jpeg

    def state(self) -> dict[str, object]:
        with self._lock:
            zone = self._zone
            return {
                "prototype": True,
                "notice": "工程原型；合成回放不是模型输出，HOG 适配器不是生产检测模型。",
                "source": {
                    "mode": self._source.mode,
                    "label": self._source.label,
                    "fixture": self._source.fixture,
                    "connected": self._connected,
                    "ref": self._source.source_ref,
                },
                "detector": {
                    "name": "scripted_fixture" if self._source.fixture else self._detector.name,
                    "production_ready": False,
                },
                "frame": {
                    "sequence": self._sequence,
                    "width": self._frame_width,
                    "height": self._frame_height,
                    "captured_at": self._captured_at,
                },
                "inference": {"target_fps": self._fps, "latency_ms": self._latency_ms},
                "zone": {
                    "id": zone.zone_id,
                    "name": zone.name,
                    "version": zone.version,
                    "points": [{"x": point.x, "y": point.y} for point in zone.points],
                },
                "detections": list(self._detections),
                "events": list(self._events),
                "last_error_code": self._last_error_code,
                "generated_at": datetime.now(timezone.utc).isoformat().replace("+00:00", "Z"),
            }