135 lines
4.8 KiB
Python
135 lines
4.8 KiB
Python
"""Background perception thread.
|
|
|
|
Owns the camera->YOLO->road-mask pipeline on its own daemon thread so the main
|
|
control loop never blocks on inference. Each frame it publishes an immutable
|
|
:class:`gowelcome.types.PerceptionResult` snapshot that the control loop reads
|
|
via :meth:`PerceptionThread.latest`.
|
|
|
|
Heavy/optional deps (``ultralytics``, ``cv2``) are only ever imported by the
|
|
:class:`YoloDetector` / :class:`RoadDetector` helpers (lazily), never here at
|
|
module top-level.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import logging
|
|
import threading
|
|
import time
|
|
from typing import Optional
|
|
|
|
from config import GoWelcomeConfig
|
|
from gowelcome.perception.detector import YoloDetector
|
|
from gowelcome.perception.road_mask import RoadDetector
|
|
from gowelcome.robot.interface import RobotInterface
|
|
from gowelcome.types import PerceptionResult, RoadInfo
|
|
|
|
_LOG = logging.getLogger(__name__)
|
|
|
|
# How long to nap when no frame is available yet (seconds).
|
|
_NO_FRAME_SLEEP = 0.02
|
|
|
|
|
|
class PerceptionThread(threading.Thread):
|
|
"""Run perception in the background and publish the latest result."""
|
|
|
|
def __init__(self, robot: RobotInterface, cfg: GoWelcomeConfig) -> None:
|
|
"""Build the detector/road-mask pipeline and prepare the thread.
|
|
|
|
Args:
|
|
robot: Robot interface to pull frames from (``get_frame``).
|
|
cfg: Full GoWelcome configuration; ``cfg.perception`` configures
|
|
the detectors and ``cfg.camera.target_fps`` paces the loop.
|
|
"""
|
|
super().__init__(daemon=True, name="PerceptionThread")
|
|
self.robot = robot
|
|
self.cfg = cfg
|
|
|
|
self.detector = YoloDetector(cfg.perception)
|
|
self.road_detector = RoadDetector(cfg.perception)
|
|
|
|
self._stop_event = threading.Event()
|
|
self._lock = threading.Lock()
|
|
self._latest: Optional[PerceptionResult] = None
|
|
self._seq = 0
|
|
|
|
fps = max(1e-3, cfg.camera.target_fps)
|
|
self._period = 1.0 / fps
|
|
|
|
def run(self) -> None:
|
|
"""Main loop: grab a frame, detect + analyse, publish a result.
|
|
|
|
Runs until :meth:`stop` is called. The whole per-iteration body is
|
|
guarded so a single bad frame logs and continues rather than killing
|
|
the thread. The loop is throttled to roughly ``camera.target_fps``.
|
|
"""
|
|
pcfg = self.cfg.perception
|
|
while not self._stop_event.is_set():
|
|
start = time.monotonic()
|
|
try:
|
|
frame = self.robot.get_frame()
|
|
if frame is None:
|
|
# No frame yet: nap briefly so we don't busy-spin.
|
|
time.sleep(_NO_FRAME_SLEEP)
|
|
continue
|
|
|
|
dets = self.detector.detect(frame)
|
|
h, w = frame.shape[:2]
|
|
|
|
persons = [
|
|
d
|
|
for d in dets
|
|
if d.label in pcfg.person_classes
|
|
and d.conf >= pcfg.person_conf
|
|
]
|
|
dangers = [
|
|
d
|
|
for d in dets
|
|
if d.label in pcfg.danger_classes
|
|
and d.conf >= pcfg.danger_conf
|
|
and d.height_ratio(h) >= pcfg.danger_min_height_ratio
|
|
]
|
|
|
|
if pcfg.road_enabled:
|
|
road = self.road_detector.analyze(frame)
|
|
else:
|
|
road = RoadInfo(0.0, 0.0, 0.0, 0.0, None)
|
|
|
|
with self._lock:
|
|
self._seq += 1
|
|
seq = self._seq
|
|
result = PerceptionResult(
|
|
frame_w=w,
|
|
frame_h=h,
|
|
detections=dets,
|
|
persons=persons,
|
|
dangers=dangers,
|
|
road=road,
|
|
ts=time.monotonic(),
|
|
seq=seq,
|
|
frame=frame,
|
|
)
|
|
with self._lock:
|
|
self._latest = result
|
|
except Exception: # noqa: BLE001 - one bad frame must not kill us
|
|
_LOG.exception("PerceptionThread iteration failed; continuing")
|
|
|
|
# Throttle to roughly target_fps: sleep the remainder of the period.
|
|
elapsed = time.monotonic() - start
|
|
remaining = self._period - elapsed
|
|
if remaining > 0:
|
|
# Wait on the stop event so stop() wakes us promptly.
|
|
self._stop_event.wait(remaining)
|
|
|
|
def stop(self) -> None:
|
|
"""Signal the run loop to exit (thread-safe; idempotent)."""
|
|
self._stop_event.set()
|
|
|
|
def latest(self) -> Optional[PerceptionResult]:
|
|
"""Return the most recent :class:`PerceptionResult`, or ``None``.
|
|
|
|
Thread-safe; the returned snapshot is the object published by the last
|
|
completed iteration.
|
|
"""
|
|
with self._lock:
|
|
return self._latest
|