intial commit

src/chicken_counter/motion.py

@@ -0,0 +1,99 @@
+"""Detect backward cart motion using sparse optical flow on the background."""
+
+from __future__ import annotations
+
+import cv2
+import numpy as np
+
+from chicken_counter.config import MotionConfig, RoiConfig
+from chicken_counter.types import MotionState, TrackObservation
+
+
+class BackwardMotionDetector:
+    def __init__(self, config: MotionConfig, roi: RoiConfig) -> None:
+        self.config = config
+        self.roi = roi
+        self.previous_gray: np.ndarray | None = None
+        self.state = MotionState()
+        self._roi_bounds = self._compute_roi_bounds()
+
+    def _compute_roi_bounds(self) -> tuple[int, int, int, int]:
+        x_values = [point[0] for point in self.roi.points]
+        y_values = [point[1] for point in self.roi.points]
+        return min(x_values), min(y_values), max(x_values), max(y_values)
+
+    def update(
+        self,
+        frame: np.ndarray,
+        tracks: list[TrackObservation],
+        frame_index: int,
+    ) -> MotionState:
+        if not self.config.enabled:
+            return self.state
+
+        stride = max(1, self.config.stride_frames)
+        if frame_index % stride != 0:
+            return self.state
+
+        gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
+        x_min, y_min, x_max, y_max = self._roi_bounds
+        gray = gray[y_min:y_max, x_min:x_max]
+
+        scale = self.config.flow_scale
+        if scale < 1.0:
+            target_width = max(1, int(gray.shape[1] * scale))
+            target_height = max(1, int(gray.shape[0] * scale))
+            gray = cv2.resize(gray, (target_width, target_height), interpolation=cv2.INTER_AREA)
+        else:
+            scale = 1.0
+
+        mask = np.full(gray.shape, 255, dtype=np.uint8)
+        for track in tracks:
+            x1, y1, x2, y2 = track.bbox_xyxy
+            r = self.config.block_radius
+            local_x1 = int((max(0, x1 - r) - x_min) * scale)
+            local_y1 = int((max(0, y1 - r) - y_min) * scale)
+            local_x2 = int((min(x_max, x2 + r) - x_min) * scale)
+            local_y2 = int((min(y_max, y2 + r) - y_min) * scale)
+            if local_x2 <= local_x1 or local_y2 <= local_y1:
+                continue
+            cv2.rectangle(mask, (local_x1, local_y1), (local_x2, local_y2), 0, -1)
+
+        points = cv2.goodFeaturesToTrack(
+            gray,
+            maxCorners=self.config.max_corners,
+            qualityLevel=self.config.quality_level,
+            minDistance=self.config.min_distance,
+            mask=mask,
+        )
+
+        if self.previous_gray is None or points is None or len(points) < self.config.min_features:
+            self.previous_gray = gray
+            return self.state
+
+        next_points, status, _ = cv2.calcOpticalFlowPyrLK(self.previous_gray, gray, points, None)
+        self.previous_gray = gray
+        if next_points is None or status is None:
+            return self.state
+
+        valid_prev = points[status.flatten() == 1]
+        valid_next = next_points[status.flatten() == 1]
+        if len(valid_prev) < self.config.min_features:
+            return self.state
+
+        flow = valid_next - valid_prev
+        axis_values = flow[:, 0, 1] if self.config.axis == "vertical" else flow[:, 0, 0]
+        median_axis_speed = float(np.median(axis_values)) * self.config.forward_sign
+        alpha = self.config.ema_alpha
+        self.state.smoothed_speed = alpha * median_axis_speed + (1.0 - alpha) * self.state.smoothed_speed
+
+        if self.state.smoothed_speed <= self.config.reverse_enter_threshold:
+            self.state.consecutive_reverse_frames += 1
+        elif self.state.smoothed_speed > self.config.reverse_exit_threshold:
+            self.state.consecutive_reverse_frames = 0
+            self.state.backward_active = False
+
+        if self.state.consecutive_reverse_frames >= self.config.debounce_frames:
+            self.state.backward_active = True
+
+        return self.state