feat: ROI-crop face detection to cut worker CPU

Author: alozowski

src/reachy_mini/vision/face_tracking.py

@@ -20,6 +20,9 @@
 
 _logger = logging.getLogger("reachymini-face-tracker")
 
+# Full-frame detect cadence while locked, so a moved or new face is re-found within ~1 s.
+_FULL_SCAN_INTERVAL = 12
+
 
 @dataclass
 class FaceObservation:
@@ -48,6 +51,18 @@ def _dist2(a: tuple[float, float], b: tuple[float, float]) -> float:
     return (a[0] - b[0]) ** 2 + (a[1] - b[1]) ** 2
 
 
+def _offset_faces(faces: "list[Face]", dx: int, dy: int) -> "list[Face]":
+    """Shift faces detected in a crop back to full-frame pixel coordinates."""
+    return [
+        type(f)(
+            bbox=(f.bbox[0] + dx, f.bbox[1] + dy, f.bbox[2], f.bbox[3]),
+            right_eye=(f.right_eye[0] + dx, f.right_eye[1] + dy),
+            left_eye=(f.left_eye[0] + dx, f.left_eye[1] + dy),
+        )
+        for f in faces
+    ]
+
+
 class Tracker:
     """Greedy single-face track that rejects spurious detections.
 
@@ -62,11 +77,13 @@ def __init__(
         min_area_frac: float = 0.003,
         max_jump: float = 0.5,
         max_misses: int = 20,
+        crop: int = 256,
     ) -> None:
-        """Create a tracker with the given acquisition and association gates."""
+        """Create a tracker with the given gates and ROI crop size (pixels)."""
         self._min_area_frac = min_area_frac
         self._max_jump = max_jump
         self._max_misses = max_misses
+        self._crop = crop
         self._center: tuple[float, float] | None = None
         self._misses = 0
 
@@ -92,6 +109,20 @@ def select(self, faces: "list[Face]", width: int, height: int) -> "Face | None":
         self._misses = 0
         return face
 
+    def roi(self, width: int, height: int) -> tuple[int, int, int, int] | None:
+        """Crop box around the current track for ROI detection, or None when unlocked."""
+        if self._center is None:
+            return None
+        cx = round((self._center[0] + 1.0) * 0.5 * max(width - 1, 1))
+        cy = round((self._center[1] + 1.0) * 0.5 * max(height - 1, 1))
+        half = self._crop // 2
+        return (
+            max(0, cx - half),
+            max(0, cy - half),
+            min(width, cx + half),
+            min(height, cy + half),
+        )
+
     def _miss(self) -> None:
         self._misses += 1
         if self._misses > self._max_misses:
@@ -169,6 +200,7 @@ def run(conn: Connection, stop: EventType, camera_specs: "CameraSpecs") -> None:
 
     crop_scale = camera_specs.default_resolution.value[3]
     camera_matrix: NDArray[np.float64] | None = None
+    since_full_scan = 0
     try:
         while not stop.is_set():
             sample = appsink.emit("try-pull-sample", 200_000_000)
@@ -185,7 +217,16 @@ def run(conn: Connection, stop: EventType, camera_specs: "CameraSpecs") -> None:
                 camera_matrix = intrinsics_for_size(
                     camera_specs.K, crop_scale, (width, height)
                 )
-            face = tracker.select(detector.detect(frame), width, height)
+            roi = tracker.roi(width, height)
+            if roi is None or since_full_scan >= _FULL_SCAN_INTERVAL:
+                faces = detector.detect(frame)
+                since_full_scan = 0
+            else:
+                x0, y0, x1, y1 = roi
+                window = np.ascontiguousarray(frame[y0:y1, x0:x1])
+                faces = _offset_faces(detector.detect(window), x0, y0)
+                since_full_scan += 1
+            face = tracker.select(faces, width, height)
             obs = to_observation(
                 face,
                 width,

tests/unit_tests/test_face_tracking.py

@@ -4,7 +4,7 @@
 
 import numpy as np
 
-from reachy_mini.vision.face_tracking import Tracker, to_observation
+from reachy_mini.vision.face_tracking import Tracker, _offset_faces, to_observation
 
 
 def _face(
@@ -71,3 +71,28 @@ def test_tracker_drops_track_after_misses_then_reacquires() -> None:
     assert tracker.select([far], 200, 200) is None  # miss 1, track held
     assert tracker.select([far], 200, 200) is None  # miss 2, track dropped
     assert tracker.select([far], 200, 200) is far  # re-acquired
+
+
+def test_tracker_roi_is_none_until_locked_then_crops_around_track() -> None:
+    """ROI is full-frame (None) until a lock, then a crop box centered on the track."""
+    tracker = Tracker(min_area_frac=0.0, crop=100)
+    assert tracker.roi(200, 200) is None
+
+    tracker.select(
+        [_face((90.0, 88.0, 20.0, 20.0), (95.0, 95.0), (105.0, 95.0))], 200, 200
+    )
+    roi = tracker.roi(200, 200)
+    assert roi is not None
+    x0, y0, x1, y1 = roi
+    assert (x1 - x0, y1 - y0) == (100, 100)
+    assert x0 <= 100 <= x1 and y0 <= 95 <= y1  # contains the eye center
+
+
+def test_offset_faces_maps_crop_coords_to_full_frame() -> None:
+    """A face found in a crop is shifted by the crop origin; its size is unchanged."""
+    shifted = _offset_faces(
+        [_face((10.0, 20.0, 30.0, 40.0), (15.0, 25.0), (35.0, 25.0))], 100, 200
+    )
+    assert shifted[0].bbox == (110.0, 220.0, 30.0, 40.0)
+    assert shifted[0].right_eye == (115.0, 225.0)
+    assert shifted[0].left_eye == (135.0, 225.0)