294 lines
11 KiB
Python
294 lines
11 KiB
Python
"""Recognition worker process for Motion Agent.
|
|
|
|
The WebSocket server stays in the parent process. This worker owns OpenCV and
|
|
MediaPipe so native camera/model work cannot block or crash the control plane.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import json
|
|
import sys
|
|
import threading
|
|
import time
|
|
from dataclasses import replace
|
|
from typing import Any
|
|
|
|
from .cameras import NullCameraInput, UrlCameraInput, UrlCameraSpec, UsbCameraInput, UsbCameraSpec
|
|
from .config import MotionAgentConfig
|
|
from .recognizer import MediaPipeGestureRecognizer, NullGestureRecognizer
|
|
|
|
|
|
class OutputClosed(Exception):
|
|
"""Signal that the parent process stopped consuming worker events."""
|
|
|
|
|
|
def _parse_indexes(raw: str | None, fallback: tuple[int, ...]) -> tuple[int, ...]:
|
|
if raw is None:
|
|
return fallback
|
|
values = [item.strip() for item in raw.split(",") if item.strip()]
|
|
return tuple(int(item) for item in values) or fallback
|
|
|
|
|
|
def _parse_urls(raw: str | None, fallback: tuple[str, ...]) -> tuple[str, ...]:
|
|
if raw is None:
|
|
return fallback
|
|
return tuple(item.strip() for item in raw.split(",") if item.strip())
|
|
|
|
|
|
def _build_parser() -> argparse.ArgumentParser:
|
|
parser = argparse.ArgumentParser(description="Run the Planet motion recognition worker.")
|
|
parser.add_argument("--camera-indexes", default=None)
|
|
parser.add_argument("--camera-urls", default=None)
|
|
parser.add_argument(
|
|
"--mode",
|
|
choices=["auto", "single", "dual", "dual_redundant", "single_fallback", "calibrated_3d"],
|
|
default=None,
|
|
)
|
|
parser.add_argument("--width", type=int, default=None)
|
|
parser.add_argument("--height", type=int, default=None)
|
|
parser.add_argument("--fps", type=int, default=None)
|
|
parser.add_argument("--max-event-hz", type=int, default=None)
|
|
parser.add_argument("--max-skeleton-hz", type=int, default=None)
|
|
parser.add_argument("--dry-run", action="store_true")
|
|
return parser
|
|
|
|
|
|
def _config_from_args(argv: list[str] | None = None) -> MotionAgentConfig:
|
|
args = _build_parser().parse_args(argv)
|
|
config = MotionAgentConfig.from_env()
|
|
return replace(
|
|
config,
|
|
camera_indexes=_parse_indexes(args.camera_indexes, config.camera_indexes),
|
|
camera_urls=_parse_urls(args.camera_urls, config.camera_urls),
|
|
camera_width=args.width or config.camera_width,
|
|
camera_height=args.height or config.camera_height,
|
|
camera_fps=args.fps or config.camera_fps,
|
|
max_event_hz=args.max_event_hz or config.max_event_hz,
|
|
max_skeleton_hz=args.max_skeleton_hz if args.max_skeleton_hz is not None else config.max_skeleton_hz,
|
|
mode=args.mode or config.mode,
|
|
dry_run=args.dry_run or config.dry_run,
|
|
)
|
|
|
|
|
|
def _build_cameras(config: MotionAgentConfig) -> list[Any]:
|
|
if config.dry_run:
|
|
return [NullCameraInput()]
|
|
if config.camera_urls:
|
|
return [
|
|
UrlCameraInput(UrlCameraSpec(url=url, camera_id=f"url:{index}"))
|
|
for index, url in enumerate(config.camera_urls)
|
|
]
|
|
return [
|
|
UsbCameraInput(
|
|
UsbCameraSpec(
|
|
index=index,
|
|
width=config.camera_width,
|
|
height=config.camera_height,
|
|
fps=config.camera_fps,
|
|
)
|
|
)
|
|
for index in config.camera_indexes
|
|
]
|
|
|
|
|
|
def _resolve_mode(config: MotionAgentConfig, camera_count: int) -> str:
|
|
if config.mode == "dual":
|
|
return "dual_redundant"
|
|
if config.mode in {"single", "dual_redundant", "single_fallback", "calibrated_3d"}:
|
|
return config.mode
|
|
return "dual_redundant" if camera_count >= 2 else "single"
|
|
|
|
|
|
def _emit(kind: str, payload: Any) -> None:
|
|
try:
|
|
print(
|
|
json.dumps({"kind": kind, "payload": payload}, ensure_ascii=False, separators=(",", ":")),
|
|
flush=True,
|
|
)
|
|
except BrokenPipeError as exc:
|
|
# Prevent Python's interpreter shutdown flush from printing another
|
|
# BrokenPipeError after the parent process closes the JSONL pipe.
|
|
sys.stdout = None
|
|
raise OutputClosed from exc
|
|
|
|
|
|
class LatestFrameReader:
|
|
"""Continuously read one camera and expose only the newest frame.
|
|
|
|
OpenCV and RTSP sources can buffer frames when recognition is slower than
|
|
capture. The recognizer should never drain that backlog; it should always
|
|
work on the latest frame so gesture latency stays bounded.
|
|
"""
|
|
|
|
def __init__(self, camera: Any, *, fallback_fps: int) -> None:
|
|
self.camera = camera
|
|
self.camera_id = str(getattr(camera, "camera_id", "unknown"))
|
|
self._fallback_interval = 1 / max(1, fallback_fps)
|
|
self._lock = threading.Lock()
|
|
self._stop = threading.Event()
|
|
self._thread: threading.Thread | None = None
|
|
self._frame: Any = None
|
|
self._frame_seq = 0
|
|
self._frame_at = 0.0
|
|
self._error: str | None = None
|
|
|
|
def start(self) -> None:
|
|
if self._thread and self._thread.is_alive():
|
|
return
|
|
self._stop.clear()
|
|
self._thread = threading.Thread(target=self._run, name=f"motion-frame-reader-{self.camera_id}", daemon=True)
|
|
self._thread.start()
|
|
|
|
def stop(self) -> None:
|
|
self._stop.set()
|
|
if self._thread and self._thread.is_alive():
|
|
self._thread.join(timeout=1)
|
|
self._thread = None
|
|
|
|
@property
|
|
def frame_count(self) -> int:
|
|
with self._lock:
|
|
return self._frame_seq
|
|
|
|
@property
|
|
def error(self) -> str | None:
|
|
with self._lock:
|
|
return self._error
|
|
|
|
def latest_after(self, last_seq: int) -> tuple[bool, int, Any, float]:
|
|
with self._lock:
|
|
if self._frame_seq <= last_seq:
|
|
return False, last_seq, None, self._frame_at
|
|
return True, self._frame_seq, self._frame, self._frame_at
|
|
|
|
def _run(self) -> None:
|
|
while not self._stop.is_set():
|
|
try:
|
|
frame = self.camera.read()
|
|
except Exception as exc:
|
|
with self._lock:
|
|
self._error = str(exc)
|
|
return
|
|
with self._lock:
|
|
self._frame = frame
|
|
self._frame_seq += 1
|
|
self._frame_at = time.monotonic()
|
|
self._error = None
|
|
if frame is None:
|
|
time.sleep(self._fallback_interval)
|
|
|
|
|
|
def main(argv: list[str] | None = None) -> int:
|
|
config = _config_from_args(argv)
|
|
cameras = _build_cameras(config)
|
|
recognizer = NullGestureRecognizer() if config.dry_run else MediaPipeGestureRecognizer()
|
|
mode = _resolve_mode(config, len(cameras))
|
|
min_interval = 1 / max(1, config.max_event_hz)
|
|
skeleton_interval = 1 / config.max_skeleton_hz if config.max_skeleton_hz > 0 else None
|
|
last_skeleton_at = 0.0
|
|
opened: list[Any] = []
|
|
readers: list[LatestFrameReader] = []
|
|
last_frame_seq: dict[str, int] = {}
|
|
last_stats_at = time.monotonic()
|
|
last_frame_total = 0
|
|
recognition_total = 0
|
|
try:
|
|
for camera in cameras:
|
|
camera.open()
|
|
opened.append(camera)
|
|
reader = LatestFrameReader(camera, fallback_fps=config.camera_fps)
|
|
reader.start()
|
|
readers.append(reader)
|
|
last_frame_seq[reader.camera_id] = 0
|
|
_emit(
|
|
"status",
|
|
{
|
|
"devices_open": True,
|
|
"recognizer": recognizer.name,
|
|
"error": None,
|
|
"fps": 0.0,
|
|
"recognition_fps": 0.0,
|
|
},
|
|
)
|
|
while True:
|
|
loop_started_at = time.monotonic()
|
|
observations = []
|
|
for reader in readers:
|
|
error = reader.error
|
|
if error:
|
|
raise RuntimeError(error)
|
|
has_frame, seq, frame, _frame_at = reader.latest_after(last_frame_seq.get(reader.camera_id, 0))
|
|
if not has_frame:
|
|
continue
|
|
last_frame_seq[reader.camera_id] = seq
|
|
observation = recognizer.recognize(frame)
|
|
recognition_total += 1
|
|
matched_gesture = None
|
|
matched_confidence = 0.0
|
|
if observation is not None:
|
|
matched_gesture = observation.gesture
|
|
matched_confidence = observation.confidence
|
|
observations.append(
|
|
{
|
|
"gesture": observation.gesture,
|
|
"confidence": observation.confidence,
|
|
"intensity": observation.intensity,
|
|
"timestamp_ms": observation.timestamp_ms,
|
|
"camera_id": reader.camera_id,
|
|
}
|
|
)
|
|
now = time.monotonic()
|
|
if skeleton_interval is not None and now - last_skeleton_at >= skeleton_interval:
|
|
skeleton = recognizer.debug_skeleton(
|
|
frame,
|
|
camera_id=reader.camera_id,
|
|
mode=mode,
|
|
matched_gesture=matched_gesture,
|
|
confidence=matched_confidence,
|
|
)
|
|
if skeleton is not None:
|
|
_emit("skeleton", skeleton.to_dict())
|
|
last_skeleton_at = now
|
|
if observations:
|
|
_emit("observations", observations)
|
|
now = time.monotonic()
|
|
if now - last_stats_at >= 1:
|
|
frame_total = sum(reader.frame_count for reader in readers)
|
|
elapsed = max(0.001, now - last_stats_at)
|
|
_emit(
|
|
"status",
|
|
{
|
|
"devices_open": True,
|
|
"recognizer": recognizer.name,
|
|
"error": None,
|
|
"fps": round((frame_total - last_frame_total) / elapsed, 2),
|
|
"recognition_fps": round(recognition_total / elapsed, 2),
|
|
},
|
|
)
|
|
last_stats_at = now
|
|
last_frame_total = frame_total
|
|
recognition_total = 0
|
|
time.sleep(max(0.001, min_interval - (time.monotonic() - loop_started_at)))
|
|
except (KeyboardInterrupt, OutputClosed):
|
|
return 0
|
|
except Exception as exc:
|
|
_emit("status", {"devices_open": False, "recognizer": recognizer.name, "error": str(exc)})
|
|
return 2
|
|
finally:
|
|
for reader in readers:
|
|
reader.stop()
|
|
for camera in opened:
|
|
try:
|
|
camera.close()
|
|
except Exception:
|
|
pass
|
|
try:
|
|
_emit("status", {"devices_open": False, "recognizer": recognizer.name, "error": None})
|
|
except OutputClosed:
|
|
pass
|
|
|
|
|
|
if __name__ == "__main__":
|
|
raise SystemExit(main(sys.argv[1:]))
|