feat(behavior): image-embedding lane for the behaviour classifier

src/glider/analysis/behavior/classify/batch.py

@@ -280,14 +280,16 @@ def classify_pose_data(
     trajectory at all, so for those models this is the only way to apply
     them. Motion needs ``video``, the recording ``pose`` was tracked from;
     ``pose_csv``, when given, lets it reuse the per-session cache training
-    wrote instead of decoding the video again.
+    wrote instead of decoding the video again. The image lane
+    (``model.image_lane``) needs ``video`` too, and caches the same way.
     """
     from glider.analysis.behavior.classify.features_stream import derive_stream_columns
     from glider.analysis.behavior.classify.smoothing import (
         MajorityVoteSmoother,
         centered_majority_vote,
     )
     from glider.analysis.behavior.features import compute_features
+    from glider.analysis.behavior.image_lane import required_columns
     from glider.analysis.behavior.trajectory import apply_trajectory_rolling
     from glider.analysis.behavior.windowing import apply_rolling, apply_spectral_rolling
 
@@ -298,6 +300,12 @@ def classify_pose_data(
     per_frame_names, spectral = derive_stream_columns(model)
     uses_motion = any(c.startswith("motion_") for c in per_frame_names)
     uses_traj = any(c.startswith("traj_") for c in model.feature_names)
+    lane = getattr(model, "image_lane", None)
+    if lane is not None and video is None:
+        raise NotImplementedError(
+            "this model uses the image lane, which is read from the source video, "
+            "and none was given"
+        )
     if uses_motion and video is None:
         # motion_* is frame differencing on the source video. Scoring without
         # it would blank every row, so decline and let the caller say why.
@@ -337,6 +345,12 @@ def classify_pose_data(
         traj = apply_trajectory_rolling(pose.xy, body_axis=ba, window=model.window).iloc[lag:]
         traj.index = rolled.index
         rolled = pd.concat([rolled, traj], axis=1)
+    if lane is not None:
+        # Row i of `rolled` describes frame i + lag; the lane is windowed at
+        # that same frame, from frames up to it (past-only lanes) -- frames
+        # the live path has already seen when it emits this row.
+        raw = _image_lane_raw(pose, model.spec, lane.config, video, pose_csv)
+        rolled = pd.concat([rolled, lane.frame_features(raw, query=rolled.index)], axis=1)
     windowed = rolled.reindex(columns=model.feature_names)
 
     # Rows are tagged with the middle frame they describe, so row i of `windowed` is
@@ -374,9 +388,13 @@ def classify_pose_data(
 
     values = windowed.to_numpy(dtype=np.float64)
     scored = values[emit_rows]
-    # Blank on any NaN, exactly as predict_one does — a partially-known row is
-    # not a weak prediction, it is no prediction.
-    usable = ~np.isnan(scored).any(axis=1)
+    # Blank on any NaN keypoint feature, exactly as predict_one does — a
+    # partially-known row is not a weak prediction, it is no prediction. The
+    # image lane is the exception: a missed mask leaves its columns NaN and
+    # the classifier routes them, as it did in training.
+    required = set(required_columns(model.feature_names))
+    req_idx = [i for i, c in enumerate(model.feature_names) if c in required]
+    usable = ~np.isnan(scored[:, req_idx]).any(axis=1)
 
     raw = np.full(len(scored), "", dtype=object)
     if usable.any():
@@ -435,6 +453,18 @@ def _motion_features(pose, spec, video, pose_csv, index) -> pd.DataFrame:
     return motion
 
 
+def _image_lane_raw(pose, spec, config, video, pose_csv) -> np.ndarray:
+    """Raw image-lane rows for ``pose``: the training cache when there is a pose CSV."""
+    from pathlib import Path
+
+    from glider.analysis.behavior.image_lane import compute_raw_for_video, load_or_compute_raw
+
+    ba = spec.with_resolved_body_axis(pose.n_keypoints).body_axis
+    if pose_csv is not None:
+        return load_or_compute_raw(Path(pose_csv), Path(video), pose.xy, ba, config)
+    return compute_raw_for_video(Path(video), pose.xy, ba, config)
+
+
 def classify_pose_tracks(tracks: PoseTracks, model, **kw) -> dict[int, EthogramRows]:
     """Score every animal in *tracks* with *model* — one slot, one call.
 

src/glider/analysis/behavior/classify/pipeline.py

@@ -31,6 +31,7 @@
     BehaviorClassifier,
     DisplayConsumer,
     FeatureEngine,
+    ImageLaneWorker,
     LatestLabel,
     PoseReplay,
     PoseTracker,
@@ -111,6 +112,8 @@ def _unstreamable_feature_families(feature_names) -> list[str]:
 TRACKED_QUEUE_MAX = 4
 DISPLAY_QUEUE_MAX = 4
 CLASSIFIER_QUEUE_MAX = 8
+# In front of the image-lane worker, which encodes what has queued in one batch.
+LANE_QUEUE_MAX = 16
 
 
 @dataclass
@@ -340,6 +343,34 @@ def __init__(self, config: LiveInferenceConfig, model=None):
         # extraction matches at training and inference time).
         self.spec = self.model.spec
 
+        # Image lane: an extra stage between the tracker and the feature
+        # engine. Refused up front for a centred lane, which would need frames
+        # the stream has not reached.
+        self.image_lane = getattr(self.model, "image_lane", None)
+        self.lane_worker = None
+        engine_queue = self.tracked_queue
+        if self.image_lane is not None:
+            if self.image_lane.config.center:
+                raise RuntimeError(
+                    "this model's image lane uses a centred window, which reads "
+                    "frames after the one it describes, so it cannot run live. "
+                    "Retrain with the default past-only lane."
+                )
+            # A deeper queue in front of the lane: the worker micro-batches
+            # whatever has queued, so a short stall is absorbed rather than
+            # dropped. Replaced before the tracker is built, so it feeds this.
+            self.tracked_queue = queue.Queue(maxsize=LANE_QUEUE_MAX)
+            self.lane_queue: queue.Queue = queue.Queue(maxsize=TRACKED_QUEUE_MAX)
+            self.lane_worker = ImageLaneWorker(
+                in_queue=self.tracked_queue,
+                out_queue=self.lane_queue,
+                stop_event=self.stop_event,
+                artifact=self.image_lane,
+                body_axis=self.spec.with_resolved_body_axis(len(config.keypoint_names)).body_axis,
+                device=config.device,
+            )
+            engine_queue = self.lane_queue
+
         # Pre-build the threads (started later via run()).
         self.producer = VideoProducer(
             source=config.source,
@@ -352,7 +383,7 @@ def __init__(self, config: LiveInferenceConfig, model=None):
             config, self.raw_queue, self.tracked_queue, self.display_queue, self.stop_event
         )
         self.feature_engine = FeatureEngine(
-            tracked_queue=self.tracked_queue,
+            tracked_queue=engine_queue,
             classifier_queue=self.classifier_queue,
             stop_event=self.stop_event,
             spec=self.spec,
@@ -366,6 +397,7 @@ def __init__(self, config: LiveInferenceConfig, model=None):
             dart_threshold=config.dart_threshold,
             freeze_min_frames=config.freeze_min_frames,
             dart_min_frames=config.dart_min_frames,
+            image_lane=self.image_lane,
         )
         self.classifier = BehaviorClassifier(
             classifier_queue=self.classifier_queue,
@@ -466,6 +498,7 @@ def run(self) -> None:
             threads = [
                 self.producer,
                 self.tracker,
+                *([self.lane_worker] if self.lane_worker is not None else []),
                 self.feature_engine,
                 self.classifier,
                 self.display,
@@ -490,7 +523,7 @@ def run(self) -> None:
                 t.join(timeout=2.0)
             # Surface any per-thread errors so the user knows what
             # broke instead of staring at a silent exit.
-            for t in (self.producer, self.tracker):
+            for t in (self.producer, self.tracker, getattr(self, "lane_worker", None)):
                 err = getattr(t, "error", None)
                 if err:
                     print(f"[{t.name}] error: {err}")

src/glider/analysis/behavior/classify/threads.py

@@ -430,6 +430,111 @@ def _propagate_eos(self) -> None:
                 pass
 
 
+class ImageLaneWorker(threading.Thread):
+    """Crops, encodes and reduces each tracked frame for the image lane.
+
+    Sits between the tracker and the :class:`FeatureEngine`, so the networks
+    run alongside the pose model rather than after it: per-frame throughput is
+    set by the slower of the two stages, not their sum. Payloads pass through
+    with the frame's reduced lane values appended (``None`` before the first
+    keypoint fix). The crop uses only this frame's and earlier keypoints
+    (:class:`~glider.analysis.behavior.image_lane.CropTracker`), the same
+    policy the offline extraction applies.
+
+    ``encode_ms`` keeps recent per-frame timings so a run can report whether
+    the lane kept up with the camera.
+    """
+
+    def __init__(
+        self,
+        in_queue: queue.Queue,
+        out_queue: queue.Queue,
+        stop_event: threading.Event,
+        artifact,
+        body_axis: tuple[int, int],
+        device: str | None = None,
+        max_batch: int = 8,
+    ):
+        super().__init__(name="ImageLaneWorker", daemon=True)
+        self.in_queue = in_queue
+        self.out_queue = out_queue
+        self.stop_event = stop_event
+        self.artifact = artifact
+        self.body_axis = body_axis
+        self.device = device
+        self.encode_ms: deque[float] = deque(maxlen=300)
+        self.max_batch = int(max_batch)
+        self.error: str | None = None
+
+    def run(self) -> None:
+        from glider.analysis.behavior.image_lane import CropTracker, crop_at, get_encoder
+
+        config = self.artifact.config
+        try:
+            encoder = get_encoder(config, device=self.device)
+        except Exception as e:  # noqa: BLE001 - surfaced through self.error
+            self.error = f"could not load the image-lane models: {e}"
+            self._end()
+            return
+        if not str(encoder.device).startswith("cuda"):
+            # Measured on a 32-core desktop: 76 ms/frame unbatched, 47 ms at
+            # batch 8 -- short of a 33 ms frame at any batch size.
+            logger.warning(
+                "the image lane is running on %s; it needs a CUDA GPU to keep up "
+                "with 30 fps, so frames will be dropped and windows broken",
+                encoder.device,
+            )
+        tracker = CropTracker(config, self.body_axis)
+        try:
+            while not self.stop_event.is_set():
+                try:
+                    batch = [self.in_queue.get(timeout=0.1)]
+                except queue.Empty:
+                    continue
+                # Micro-batch: whatever queued up while the last batch ran is
+                # encoded in one call. Per-frame cost falls steeply with batch
+                # size, so a stall (a GC pause, the pose model's turn on the
+                # GPU) is absorbed instead of overflowing the queue and
+                # dropping frames, which would break every window spanning them.
+                while len(batch) < self.max_batch and batch[-1] is not END_OF_STREAM:
+                    try:
+                        batch.append(self.in_queue.get_nowait())
+                    except queue.Empty:
+                        break
+                ended = batch[-1] is END_OF_STREAM
+                items = batch[:-1] if ended else batch
+                crops, slots = [], []
+                for i, (_idx, frame_bgr, keypoints, _conf) in enumerate(items):
+                    fix = tracker.push(keypoints)
+                    if fix is not None and frame_bgr is not None:
+                        crops.append(crop_at(frame_bgr, fix[0], fix[1], config.crop_px))
+                        slots.append(i)
+                values: list = [None] * len(items)
+                if crops:
+                    t0 = time.perf_counter()
+                    try:
+                        reduced = self.artifact.transform(encoder.encode(crops))
+                    except Exception as e:  # noqa: BLE001
+                        self.error = f"image lane failed on frame {items[slots[0]][0]}: {e}"
+                        return
+                    per_frame = (time.perf_counter() - t0) * 1000.0 / len(crops)
+                    self.encode_ms.extend([per_frame] * len(crops))
+                    for i, row in zip(slots, reduced, strict=True):
+                        values[i] = row
+                for item, v in zip(items, values, strict=True):
+                    _put_or_drop(self.out_queue, (*item, v), self.stop_event)
+                if ended:
+                    return
+        finally:
+            self._end()
+
+    def _end(self) -> None:
+        try:
+            self.out_queue.put(END_OF_STREAM, timeout=1.0)
+        except queue.Full:
+            pass
+
+
 class FeatureEngine(threading.Thread):
     """Per-frame features → sliding buffer → emit rolling row to classifier.
 
@@ -462,6 +567,7 @@ def __init__(
         dart_threshold: float | None = None,
         freeze_min_frames: int = 30,
         dart_min_frames: int = 3,
+        image_lane=None,
     ):
         super().__init__(name="FeatureEngine", daemon=True)
         self.tracked_queue = tracked_queue
@@ -490,6 +596,15 @@ def __init__(
             spectral_features=spectral_features,
         )
         self._tick: int = 0
+        # Image lane: per-frame values arrive on the payload from the
+        # ImageLaneWorker; each emitted row gets the lane window at the same
+        # (lagged) frame it describes.
+        self._lane = None
+        if image_lane is not None:
+            from glider.analysis.behavior.image_lane import StreamingImageLane
+
+            self._lane = StreamingImageLane(image_lane)
+            self._lane_columns = list(image_lane.columns)
         # Centered-feature lag: the emitted windowed row describes the MIDDLE
         # frame, which trails the current frame by this many frames.
         self._lag = self._extractor.lag
@@ -527,7 +642,9 @@ def run(self) -> None:
                 if item is END_OF_STREAM:
                     self._end()
                     return
-                frame_idx, _frame_bgr, keypoints, _confidences = item
+                frame_idx, _frame_bgr, keypoints, _confidences = item[:4]
+                if self._lane is not None:
+                    self._lane.push(frame_idx, item[4] if len(item) > 4 else None)
                 if self._speed_axis:
                     spd = self._causal_speed.push(keypoints)
                     self._speed_labels.append(self._freeze_dart.push(spd))
@@ -550,6 +667,9 @@ def run(self) -> None:
                 # current one by self._lag. Tag it with that index so the
                 # ethogram timeline stays aligned.
                 mid_idx = max(0, frame_idx - self._lag)
+                if self._lane is not None:
+                    column_names = list(column_names) + self._lane_columns
+                    row = np.concatenate([np.asarray(row, float), self._lane.row(mid_idx)])
                 if self._speed_axis:
                     # Speed label for that same middle frame = the oldest entry
                     # in the lag-sized buffer (once it has filled).
@@ -848,6 +968,10 @@ def run(self) -> None:
                 if item is END_OF_STREAM:
                     return
                 _frame_idx, frame_bgr, keypoints, confidences = item
+                # Draw on a copy: the tracker hands this same array to the
+                # feature path, and the image lane crops its pixels. Drawing
+                # in place put the skeleton and badge into the lane's input.
+                frame_bgr = frame_bgr.copy()
 
                 # Update render FPS estimate with EMA so the overlay
                 # doesn't flicker. Quick coefficient = 0.1.

src/glider/analysis/behavior/evaluation.py

@@ -133,6 +133,13 @@ def _windowed_for(model: BehaviorModel, pose_csv: Path, fps: float) -> pd.DataFr
         traj_frame = apply_trajectory_rolling(pose.xy, body_axis=axis, window=model.window)
         traj_frame.index = windowed.index
         windowed = pd.concat([windowed, traj_frame], axis=1)
+    lane = getattr(model, "image_lane", None)
+    if lane is not None:
+        from glider.analysis.behavior.pipeline import _image_lane_raw
+
+        img = lane.frame_features(_image_lane_raw(Path(pose_csv), pose, model.spec, lane.config))
+        img.index = windowed.index
+        windowed = pd.concat([windowed, img], axis=1)
     return windowed