From 0e856496b139a54dfa963bbeb00b48b823027922 Mon Sep 17 00:00:00 2001 From: gbradham Date: Thu, 5 Nov 2026 11:19:06 -0600 Subject: [PATCH] feat(behavior): image-embedding lane for the behaviour classifier Adds an optional image lane beside the keypoint features: a causal crop around the animal, a YOLO segmentation mask, mask-shape descriptors, masked DINOv2 ViT-S/14 and the mmv2 pose-model neck, reduced by a PCA fitted on the training sessions only and windowed past-only (value, mean, std, diff). - image_lane.py: config, crop tracker, encoder, per-session raw cache, PCA artifact stored in the model bundle, one window function shared by training, offline scoring and the live stream. - Wired into training, cross-validation, batch scoring and the live pipeline (image_lane=True); lane columns are NaN-tolerant, so a missed mask leaves the keypoint features usable. - Live display now draws its overlay on a copy of the frame, so stages that read pixels downstream no longer see it. Reproduces the research pipeline: PyMouse leave-one-mouse-out 0.776 centred / 0.768 live (research 0.778 / 0.768); Blast leave-one-session-out 0.834 -> 0.842 with the lane. --- .../analysis/behavior/classify/batch.py | 38 +- .../analysis/behavior/classify/pipeline.py | 37 +- .../analysis/behavior/classify/threads.py | 126 ++- src/glider/analysis/behavior/evaluation.py | 7 + src/glider/analysis/behavior/image_lane.py | 715 ++++++++++++++++++ src/glider/analysis/behavior/model.py | 29 +- src/glider/analysis/behavior/pipeline.py | 164 +++- src/glider/gui/behavior/window.py | 59 +- .../unit/analysis/behavior/test_image_lane.py | 596 +++++++++++++++ tests/unit/gui/behavior/test_window.py | 27 + 10 files changed, 1751 insertions(+), 47 deletions(-) create mode 100644 src/glider/analysis/behavior/image_lane.py create mode 100644 tests/unit/analysis/behavior/test_image_lane.py diff --git a/src/glider/analysis/behavior/classify/batch.py b/src/glider/analysis/behavior/classify/batch.py index ce786f81..3c33dd65 100644 --- a/src/glider/analysis/behavior/classify/batch.py +++ b/src/glider/analysis/behavior/classify/batch.py @@ -280,7 +280,8 @@ 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 ( @@ -288,6 +289,7 @@ def classify_pose_data( 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. diff --git a/src/glider/analysis/behavior/classify/pipeline.py b/src/glider/analysis/behavior/classify/pipeline.py index 0a6b8922..8b626412 100644 --- a/src/glider/analysis/behavior/classify/pipeline.py +++ b/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}") diff --git a/src/glider/analysis/behavior/classify/threads.py b/src/glider/analysis/behavior/classify/threads.py index 0a37bb5c..9a662c98 100644 --- a/src/glider/analysis/behavior/classify/threads.py +++ b/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. diff --git a/src/glider/analysis/behavior/evaluation.py b/src/glider/analysis/behavior/evaluation.py index 0840ef4c..d1c484f5 100644 --- a/src/glider/analysis/behavior/evaluation.py +++ b/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 diff --git a/src/glider/analysis/behavior/image_lane.py b/src/glider/analysis/behavior/image_lane.py new file mode 100644 index 00000000..f6563312 --- /dev/null +++ b/src/glider/analysis/behavior/image_lane.py @@ -0,0 +1,715 @@ +"""The image lane: an embedding of the animal crop, beside the keypoint features. + +Keypoints say where the body parts are; they cannot say what the body looks +like. A rear and a hunched sit can put seven keypoints in nearly the same +place, and grooming forepaws are not tracked at all. The image lane adds what +the camera sees around the animal, in four steps per frame: + +1. **Crop.** A ``crop_px`` square around the body centre, spanning + ``crop_body_lengths`` body lengths (2.2 ≈ the 20 cm crop of the research + pipeline for an adult mouse). Both the centre and the scale are *causal*: + the centre carries forward over undetected frames and the scale is a + running median of the body length over the frames seen so far, so the live + path and an offline run crop every frame identically. +2. **Mask.** A YOLO segmentation model (``mouse_seg.pt``) masks the animal: + the mask under the crop centre, else the largest one. +3. **Encode.** 19 mask-shape descriptors; DINOv2 ViT-S/14 on the masked crop + (CLS + mean patch token, 768 values); and the neck of a YOLO pose model + (``mmv2.pt`` layer 19, 128 channels) pooled inside the mask. +4. **Reduce and window.** The two learned embeddings are standardised and + projected to ``dims`` principal components fitted on the *training* + sessions only (:class:`ImageLaneArtifact`, stored in the model bundle). + Each frame then gets its value, the mean and std over a ``window``-frame + window, and its first difference. + +Steps 1-3 are the expensive part and are cached per session +(:func:`load_or_compute_raw`), exactly as the motion lane caches its +features. Steps 4 are cheap, and the windows are computed by one function +(:func:`window_lane`) for training, offline scoring and the live stream, so +the three cannot drift apart. + +Window direction. ``center=False`` (the default) is past-only: frame ``f`` +uses frames ``f - window + 1 .. f``. That is the variant the manuscript calls +LIVE, and the only one the live pipeline accepts. ``center=True`` reproduces +the research's offline-best numbers but reads frames after ``f``. + +Lane columns are named ``img_*`` and are *NaN-tolerant*: a frame whose mask +was missed still has its keypoint features, and LightGBM routes the missing +image values down its default branch, as the research pipeline did. Rows are +only dropped (or scored blank) for NaN in the keypoint features — see +:func:`required_columns`. + +Reference: ``Data/behavior_embedding_data/scripts/pm_raster.py`` +(``img_lane``/``win``) and ``pm_extract.py`` in the GLIDER manuscript folder. +""" + +from __future__ import annotations + +import hashlib +import logging +import os +from collections import deque +from collections.abc import Sequence +from dataclasses import asdict, dataclass, field +from pathlib import Path + +import numpy as np +import pandas as pd + +logger = logging.getLogger(__name__) + +#: Every image-lane column starts with this. Columns that do NOT are the +#: keypoint features, whose NaN still blanks a row. +IMAGE_LANE_PREFIX = "img_" + +#: Number of mask-shape descriptors (:func:`shape_features`). +N_SHAPE = 19 + +#: Nominal width of the crop in cm, used only to put the shape descriptors in +#: the units the research pipeline used. The crop is scaled by body length, +#: so these are "cm for an adult mouse", not a calibration. +NOMINAL_CROP_CM = 20.0 + +#: Running-median history for the crop scale, in frames (30 s at 30 fps). +SCALE_HISTORY = 900 + +#: Stats appended to each per-frame lane value, in column order. +LANE_STATS: tuple[str, ...] = ("", "_m", "_s", "_d") + +CACHE_DIR = Path("image_lane_cache") + +#: Where :func:`default_config` looks for ``mouse_seg.pt`` and ``mmv2.pt`` +#: unless ``GLIDER_IMAGE_LANE_MODELS`` names another folder. +DEFAULT_MODEL_DIR = Path.home() / ".glider" / "models" / "image_lane" +SEG_MODEL_NAME = "mouse_seg.pt" +NECK_MODEL_NAME = "mmv2.pt" + +_BODY_KERNEL = (9, 9) # opening that strips the tail before measuring body area +_TAIL_KERNEL = (17, 17) # opening that strips the tail before the centre width +_GREY = (128, 128, 128) +_DINO_MEAN = (0.485, 0.456, 0.406) +_DINO_STD = (0.229, 0.224, 0.225) +#: CLS + mean patch token width of each DINOv2 backbone. +_DINO_WIDTH = { + "dinov2_vits14": 768, + "dinov2_vitb14": 1536, + "dinov2_vitl14": 2048, + "dinov2_vitg14": 3072, +} + + +@dataclass(frozen=True) +class ImageLaneConfig: + """Everything that decides what the image lane computes. + + ``seg_model`` and ``neck_model`` are paths to Ultralytics ``.pt`` files + (the lab's ``mouse_seg.pt`` and ``mmv2.pt``). DINOv2 comes from + ``torch.hub`` (downloaded once, then cached). + """ + + seg_model: str + neck_model: str + neck_layer: int = 19 + dino_repo: str = "facebookresearch/dinov2" + dino_model: str = "dinov2_vits14" + crop_px: int = 320 + crop_body_lengths: float = 2.2 + #: Keypoint the crop centres on; ``None`` = centroid of the visible ones. + center_keypoint: int | None = None + #: Principal components kept per learned embedding. + dims: int = 32 + #: Frames in the lane's rolling window. + window: int = 5 + #: Centred window (offline only) instead of past-only. + center: bool = False + + def to_dict(self) -> dict: + return asdict(self) + + @classmethod + def from_dict(cls, d: dict) -> ImageLaneConfig: + return cls(**{k: v for k, v in d.items() if k in cls.__dataclass_fields__}) + + def cache_tag(self) -> str: + """Short hash of everything that changes the raw (pre-PCA) embedding.""" + parts = [ + _file_ident(self.seg_model), + _file_ident(self.neck_model), + str(self.neck_layer), + self.dino_repo, + self.dino_model, + str(self.crop_px), + f"{self.crop_body_lengths:g}", + str(self.center_keypoint), + ] + return hashlib.sha1("|".join(parts).encode("utf-8")).hexdigest()[:10] + + +def default_config(**overrides) -> ImageLaneConfig: + """The manuscript's image lane, with the two YOLO models from the model folder. + + The folder is ``$GLIDER_IMAGE_LANE_MODELS`` or :data:`DEFAULT_MODEL_DIR`. + Raises ``FileNotFoundError`` naming the folder when a model is missing, so + a training run fails before it decodes anything. + """ + folder = Path(os.environ.get("GLIDER_IMAGE_LANE_MODELS") or DEFAULT_MODEL_DIR) + seg, neck = folder / SEG_MODEL_NAME, folder / NECK_MODEL_NAME + missing = [p.name for p in (seg, neck) if not p.is_file()] + if missing: + raise FileNotFoundError( + f"the image lane needs {' and '.join(missing)} in {folder} " + "(or set GLIDER_IMAGE_LANE_MODELS to the folder that holds them)" + ) + return ImageLaneConfig(seg_model=str(seg), neck_model=str(neck), **overrides) + + +def resolve_config(value) -> ImageLaneConfig | None: + """``None``/``False`` -> no lane; ``True`` -> :func:`default_config`; a config or dict as is.""" + if value is None or value is False: + return None + if value is True: + return default_config() + if isinstance(value, ImageLaneConfig): + return value + if isinstance(value, dict): + return ImageLaneConfig.from_dict(value) + raise TypeError(f"image_lane must be a bool, dict or ImageLaneConfig, not {type(value)!r}") + + +def _file_ident(path: str) -> str: + """Name + size of a model file, so a retrained model invalidates the cache.""" + p = Path(path) + try: + return f"{p.name}:{p.stat().st_size}" + except OSError: + return p.name + + +def required_columns(columns: Sequence[str]) -> list[str]: + """The columns whose NaN makes a row unusable: everything but the image lane.""" + return [c for c in columns if not str(c).startswith(IMAGE_LANE_PREFIX)] + + +def has_lane_columns(columns: Sequence[str]) -> bool: + return any(str(c).startswith(IMAGE_LANE_PREFIX) for c in columns) + + +# --------------------------------------------------------------------------- +# Crop +# --------------------------------------------------------------------------- + + +class CropTracker: + """Causal crop centre and scale, fed one frame's keypoints at a time. + + The same object drives the offline extraction and the live worker, which + is what makes their crops identical. The centre is the chosen keypoint (or + the centroid of the visible keypoints), carried forward when the frame has + none. The scale is the running median body length over the last + :data:`SCALE_HISTORY` frames that had one — a rear shortens the body, and a + per-frame scale would zoom in on exactly the frames the lane exists for. + """ + + def __init__(self, config: ImageLaneConfig, body_axis: tuple[int, int]): + self.config = config + self.head, self.tail = int(body_axis[0]), int(body_axis[1]) + self._lengths: deque[float] = deque(maxlen=SCALE_HISTORY) + self._center: np.ndarray | None = None + + def push(self, keypoints: np.ndarray) -> tuple[np.ndarray, float] | None: + """``(centre_xy, px_per_crop_px)`` for this frame, or None before the first fix.""" + kp = np.asarray(keypoints, dtype=np.float64) + ci = self.config.center_keypoint + if ci is None: + ok = np.isfinite(kp).all(axis=1) + c = kp[ok].mean(axis=0) if ok.any() else None + else: + c = kp[ci] if np.isfinite(kp[ci]).all() else None + if c is not None: + self._center = c + h, t = kp[self.head], kp[self.tail] + if np.isfinite(h).all() and np.isfinite(t).all(): + length = float(np.hypot(*(t - h))) + if length > 1e-6: + self._lengths.append(length) + if self._center is None or not self._lengths: + return None + body = float(np.median(self._lengths)) + return self._center, self.config.crop_px / (self.config.crop_body_lengths * body) + + +def crop_at(frame: np.ndarray, center: np.ndarray, scale: float, px: int) -> np.ndarray: + """``px`` x ``px`` crop of ``frame`` centred on ``center``, magnified by ``scale``.""" + import cv2 + + m = np.float32([[scale, 0, px / 2 - scale * center[0]], [0, scale, px / 2 - scale * center[1]]]) + return cv2.warpAffine( + frame, m, (px, px), flags=cv2.INTER_LINEAR, borderMode=cv2.BORDER_REPLICATE + ) + + +# --------------------------------------------------------------------------- +# Mask shape +# --------------------------------------------------------------------------- + + +def _center_width(mask: np.ndarray, uv: tuple[float, float]) -> float | None: + """Distance from the crop centre to the outline, tail removed (crop px).""" + import cv2 + + body = cv2.morphologyEx( + mask.astype(np.uint8), + cv2.MORPH_OPEN, + cv2.getStructuringElement(cv2.MORPH_ELLIPSE, _TAIL_KERNEL), + ) + if body.sum() < 500: + return None + y = int(np.clip(int(uv[1]), 0, mask.shape[0] - 1)) + x = int(np.clip(int(uv[0]), 0, mask.shape[1] - 1)) + return float(cv2.distanceTransform(body, cv2.DIST_L2, 5)[y, x]) + + +def shape_features(mask: np.ndarray) -> np.ndarray: + """19 descriptors of the animal mask: size, elongation, solidity, Hu moments. + + All NaN when there is no usable mask (fewer than 50 px). + """ + import cv2 + + px = mask.shape[0] + pc = px / NOMINAL_CROP_CM + f = np.full(N_SHAPE, np.nan, np.float32) + m = mask.astype(np.uint8) + contours, _ = cv2.findContours(m, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_NONE) + if not contours: + return f + c = max(contours, key=cv2.contourArea) + area = cv2.contourArea(c) + if area < 50 or len(c) < 5: + return f + body = cv2.morphologyEx( + m, cv2.MORPH_OPEN, cv2.getStructuringElement(cv2.MORPH_ELLIPSE, _BODY_KERNEL) + ) + (_, _), (a, b), _ = cv2.fitEllipse(c) + hull = cv2.convexHull(c) + x, y, w, h = cv2.boundingRect(c) + hu = cv2.HuMoments(cv2.moments(c)).ravel() + per = cv2.arcLength(c, True) + width = _center_width(mask, (px / 2, px / 2)) + f[:12] = [ + area / pc**2, + body.sum() / pc**2, + max(a, b) / pc, + min(a, b) / pc, + min(a, b) / max(a, b, 1e-6), + area / max(cv2.contourArea(hull), 1), + area / max(w * h, 1), + per**2 / max(area, 1), + (width if width else np.nan) / pc, + per / pc, + w / pc, + h / pc, + ] + f[12:] = -np.sign(hu) * np.log10(np.abs(hu) + 1e-30) + return f + + +# --------------------------------------------------------------------------- +# Encoder (the networks) +# --------------------------------------------------------------------------- + + +class ImageLaneEncoder: + """Mask + shape + DINOv2 + pose-model neck, for a batch of crops. + + Returns one raw row per crop: ``[shape (19) | dino (768) | neck (C)]``. + Loading takes a few seconds (three networks), so build one and reuse it. + """ + + def __init__(self, config: ImageLaneConfig, device: str | None = None): + import torch + from ultralytics import YOLO + + from glider.vision.pose.device import resolve_device + + self.config = config + self.device = resolve_device(device) + self._torch = torch + self._half = self.device.startswith("cuda") + self._seg = YOLO(config.seg_model) + self._net = YOLO(config.neck_model).model.to(self.device).eval() + self._feats: dict[str, object] = {} + self._net.model[config.neck_layer].register_forward_hook( + lambda _m, _i, out: self._feats.__setitem__("neck", out) + ) + self._dino = ( + torch.hub.load(config.dino_repo, config.dino_model, verbose=False) + .to(self.device) + .eval() + ) + self._mean = torch.tensor(_DINO_MEAN, device=self.device).view(1, 3, 1, 1) + self._std = torch.tensor(_DINO_STD, device=self.device).view(1, 3, 1, 1) + self.n_neck: int | None = None + + def masks(self, crops: list[np.ndarray]) -> list[np.ndarray]: + px = self.config.crop_px + out = [] + results = self._seg( + crops, + imgsz=px, + half=self._half, + device=self.device, + verbose=False, + retina_masks=True, + ) + for res in results: + if res.masks is None or len(res.masks) == 0: + out.append(np.zeros((px, px), bool)) + continue + ms = res.masks.data.cpu().numpy().astype(bool) + hit = [j for j in range(len(ms)) if ms[j][px // 2, px // 2]] + out.append(ms[hit[0]] if hit else ms[int(np.argmax(ms.sum((1, 2))))]) + return out + + def encode(self, crops: list[np.ndarray]) -> np.ndarray: + """``(len(crops), D)`` float32 raw lane rows for BGR crops.""" + import cv2 + + torch = self._torch + if not crops: + return np.empty((0, 0), np.float32) + masks = self.masks(crops) + with torch.no_grad(): + x = torch.from_numpy(np.stack(crops)).to(self.device) + x = x.permute(0, 3, 1, 2).float()[:, [2, 1, 0]] / 255 + self._net(x) + neck = self._feats["neck"].float() + mk = torch.from_numpy(np.stack(masks)).to(self.device)[:, None].float() + mk = torch.nn.functional.interpolate(mk, size=neck.shape[-2:], mode="area") + pooled = (neck * mk).sum((2, 3)) / mk.sum((2, 3)).clamp(min=1e-3) + + masked = [] + for c, m in zip(crops, masks, strict=True): + c2 = c.copy() + c2[~m] = _GREY + masked.append(cv2.resize(c2, (224, 224), interpolation=cv2.INTER_AREA)) + xd = torch.from_numpy(np.stack(masked)).to(self.device) + xd = xd.permute(0, 3, 1, 2).float()[:, [2, 1, 0]] / 255 + o = self._dino.forward_features((xd - self._mean) / self._std) + dino = torch.cat([o["x_norm_clstoken"], o["x_norm_patchtokens"].mean(1)], 1) + self.n_neck = int(pooled.shape[1]) + shape = np.stack([shape_features(m) for m in masks]) + # Through float16, as the research cache stored them, so a fresh run + # and a cached one feed the PCA the same numbers. + learned = torch.cat([dino, pooled], 1).half().float().cpu().numpy() + return np.concatenate([shape, learned], axis=1).astype(np.float32) + + +_ENCODERS: dict[tuple, ImageLaneEncoder] = {} + + +def get_encoder(config: ImageLaneConfig, device: str | None = None) -> ImageLaneEncoder: + """A shared encoder per (config, device): loading three networks per session is slow.""" + key = (config.seg_model, config.neck_model, config.neck_layer, config.dino_model, device) + if key not in _ENCODERS: + _ENCODERS[key] = ImageLaneEncoder(config, device=device) + return _ENCODERS[key] + + +# --------------------------------------------------------------------------- +# Per-session raw extraction + cache +# --------------------------------------------------------------------------- + + +def compute_raw_for_video( + video_path: Path, + xy: np.ndarray, + body_axis: tuple[int, int], + config: ImageLaneConfig, + *, + encoder: ImageLaneEncoder | None = None, + batch: int = 64, + progress=None, +) -> np.ndarray: + """Raw lane rows for every pose frame of ``video_path`` (NaN where none). + + Video frame ``i`` is pose frame ``i``, as for the motion lane. Frames + before the first keypoint fix, and past the end of the video, are NaN. + """ + import cv2 + + xy = np.asarray(xy, dtype=np.float64) + n = len(xy) + encoder = encoder or get_encoder(config) + tracker = CropTracker(config, body_axis) + cap = cv2.VideoCapture(str(video_path)) + if not cap.isOpened(): + raise FileNotFoundError(f"cannot open video {video_path}") + out: np.ndarray | None = None + idx: list[int] = [] + buf: list[np.ndarray] = [] + + def flush(): + nonlocal out + if not buf: + return + rows = encoder.encode(buf) + if out is None: + out = np.full((n, rows.shape[1]), np.nan, np.float32) + out[idx] = rows + idx.clear() + buf.clear() + + try: + for i in range(n): + ok, frame = cap.read() + if not ok: + break # video shorter than pose; the rest stays NaN + fix = tracker.push(xy[i]) + if fix is None: + continue + idx.append(i) + buf.append(crop_at(frame, fix[0], fix[1], config.crop_px)) + if len(buf) >= batch: + flush() + if progress and i and i % 5000 == 0: + progress(i, n) + flush() + finally: + cap.release() + if out is None: + raise ValueError(f"no frame of {Path(video_path).name} had a keypoint fix to crop") + return out + + +def load_or_compute_raw( + pose_csv: Path, + video_path: Path, + xy: np.ndarray, + body_axis: tuple[int, int], + config: ImageLaneConfig, + *, + cache_dir: Path = CACHE_DIR, + encoder: ImageLaneEncoder | None = None, + progress=None, +) -> np.ndarray: + """Cached :func:`compute_raw_for_video`; the key covers pose, video and config. + + The cache is a float16 ``.npy`` returned memory-mapped, so a cohort of long + sessions can be fitted and windowed without holding every embedding in RAM. + """ + pose_csv = Path(pose_csv) + cache_dir = Path(cache_dir) + ident = f"{pose_csv.resolve()}|{Path(video_path).resolve()}" + digest = hashlib.sha1(ident.encode("utf-8")).hexdigest()[:12] + head, tail = int(body_axis[0]), int(body_axis[1]) + cache_file = cache_dir / f"{pose_csv.stem}_{digest}_a{head}-{tail}_{config.cache_tag()}.npy" + n = len(np.asarray(xy)) + if cache_file.exists(): + raw = np.load(cache_file, mmap_mode="r") + if len(raw) == n: + return raw + del raw # stale (pose length changed): recompute + if not Path(video_path).exists(): + raise FileNotFoundError( + f"the image lane needs the source video, but {video_path} is missing " + f"(pose {pose_csv.name})." + ) + raw = compute_raw_for_video( + video_path, xy, body_axis, config, encoder=encoder, progress=progress + ) + cache_dir.mkdir(parents=True, exist_ok=True) + np.save(cache_file, raw.astype(np.float16)) + return np.load(cache_file, mmap_mode="r") + + +# --------------------------------------------------------------------------- +# Reducer (fitted, stored in the model bundle) +# --------------------------------------------------------------------------- + + +@dataclass +class _Block: + name: str + start: int + stop: int + mean: np.ndarray | None = None + scale: np.ndarray | None = None + components: np.ndarray | None = None # (D_block, dims); None = pass through + + @property + def width(self) -> int: + return self.stop - self.start if self.components is None else self.components.shape[1] + + +@dataclass +class ImageLaneArtifact: + """Fitted image-lane reducer: what a model bundle needs to rebuild the lane. + + Fitted on training sessions only — PCA fitted on a held-out session would + leak it into the features that score it. + """ + + config: ImageLaneConfig + blocks: list[_Block] = field(default_factory=list) + + @classmethod + def fit( + cls, raws: Sequence[np.ndarray], config: ImageLaneConfig, *, stride: int = 5 + ) -> ImageLaneArtifact: + """Standardise + PCA each learned block over every ``stride``-th finite frame.""" + if not raws: + raise ValueError("the image lane needs at least one training session to fit on") + width = raws[0].shape[1] + if width <= N_SHAPE: + raise ValueError(f"raw image-lane rows have {width} values; expected > {N_SHAPE}") + n_dino = _DINO_WIDTH.get(config.dino_model) + if n_dino is None or N_SHAPE + n_dino >= width: + raise ValueError(f"cannot split {width} raw values into shape | dino | neck") + spans = [ + ("shape", 0, N_SHAPE), + ("dino", N_SHAPE, N_SHAPE + n_dino), + ("neck", N_SHAPE + n_dino, width), + ] + blocks = [_Block("shape", 0, N_SHAPE)] + for name, a, b in spans[1:]: + sample = np.concatenate([r[::stride, a:b] for r in raws]).astype(np.float64) + sample = sample[np.isfinite(sample).all(axis=1)] + if len(sample) < 2: + raise ValueError("too few frames with an image embedding to fit the lane") + mu = sample.mean(axis=0) + sd = sample.std(axis=0) + 1e-6 + _u, _s, vt = np.linalg.svd((sample - mu) / sd, full_matrices=False) + k = min(config.dims, vt.shape[0]) + blocks.append(_Block(name, a, b, mu, sd, vt[:k].T.astype(np.float64))) + return cls(config=config, blocks=blocks) + + @property + def frame_columns(self) -> list[str]: + names = [] + for b in self.blocks: + names += [f"{IMAGE_LANE_PREFIX}{b.name}{i:02d}" for i in range(b.width)] + return names + + @property + def columns(self) -> list[str]: + base = self.frame_columns + return [f"{c}{s}" for s in LANE_STATS for c in base] + + def transform(self, raw: np.ndarray) -> np.ndarray: + """Raw rows -> per-frame lane values ``(n, sum of block widths)``.""" + raw = np.asarray(raw, dtype=np.float64) + if raw.ndim == 1: + raw = raw[None] + parts = [] + for b in self.blocks: + x = raw[:, b.start : b.stop] + if b.components is not None: + x = ((x - b.mean) / b.scale) @ b.components + parts.append(x) + return np.concatenate(parts, axis=1) + + def window(self, values: np.ndarray, frames: np.ndarray, query: np.ndarray) -> np.ndarray: + return window_lane( + values, frames, query, window=self.config.window, center=self.config.center + ) + + def frame_features(self, raw: np.ndarray, query: np.ndarray | None = None) -> pd.DataFrame: + """Windowed lane columns for a whole session's raw rows, at ``query`` frames.""" + values = self.transform(raw) + frames = np.arange(len(values)) + query = frames if query is None else np.asarray(query) + return pd.DataFrame(self.window(values, frames, query), index=query, columns=self.columns) + + +def window_lane( + values: np.ndarray, + frames: np.ndarray, + query: np.ndarray, + *, + window: int, + center: bool = False, + chunk: int = 20000, +) -> np.ndarray: + """Value, window mean, window std and first difference at each query frame. + + ``values[i]`` belongs to source frame ``frames[i]`` (sorted, unique); + frames that are absent or NaN are skipped, like pandas + ``rolling(min_periods=1)``. Std is the sample std (NaN with one value); + the difference is ``v[f] - v[f-1]``. Output columns follow + :data:`LANE_STATS`: ``[value | mean | std | diff]``. + """ + values = np.asarray(values, dtype=np.float64) + frames = np.asarray(frames, dtype=np.int64) + query = np.asarray(query, dtype=np.int64) + d = values.shape[1] + offsets = np.arange(window) - (window // 2 if center else window - 1) + out = np.full((len(query), 4 * d), np.nan) + + def lookup(targets: np.ndarray) -> np.ndarray: + pos = np.searchsorted(frames, targets).clip(0, max(len(frames) - 1, 0)) + got = np.full(targets.shape + (d,), np.nan) + if len(frames): + hit = frames[pos] == targets + got[hit] = values[pos[hit]] + return got + + for s in range(0, len(query), chunk): + q = query[s : s + chunk] + win = lookup(q[:, None] + offsets[None, :]) # (q, w, d) + ok = np.isfinite(win) + count = ok.sum(axis=1) + total = np.where(ok, win, 0.0).sum(axis=1) + with np.errstate(invalid="ignore", divide="ignore"): + mean = total / count + dev = np.where(ok, win - mean[:, None, :], 0.0) + std = np.sqrt((dev**2).sum(axis=1) / (count - 1)) + mean[count == 0] = np.nan + std[count < 2] = np.nan + cur = lookup(q) + prev = lookup(q - 1) + out[s : s + len(q)] = np.concatenate([cur, mean, std, cur - prev], axis=1) + return out + + +# --------------------------------------------------------------------------- +# Streaming (live) +# --------------------------------------------------------------------------- + + +class StreamingImageLane: + """Live counterpart of :meth:`ImageLaneArtifact.frame_features`. + + Push each frame's reduced lane values as they arrive; :meth:`row` gives + the windowed columns for any recent frame using only frames pushed so far. + It calls the same :func:`window_lane` the offline path does. + """ + + def __init__(self, artifact: ImageLaneArtifact, history: int = 32): + if artifact.config.center: + raise ValueError( + "this model's image lane uses a centred window, which reads frames " + "after the one it describes; retrain with a past-only lane for live use" + ) + self.artifact = artifact + self._frames: deque[int] = deque(maxlen=history) + self._values: deque[np.ndarray] = deque(maxlen=history) + self._width = len(artifact.frame_columns) + + def push(self, frame_idx: int, values: np.ndarray | None) -> None: + if self._frames and frame_idx <= self._frames[-1]: + self.reset() # the stream restarted; a window must not span it + v = np.full(self._width, np.nan) if values is None else np.asarray(values, float) + self._frames.append(int(frame_idx)) + self._values.append(v.reshape(-1)) + + def reset(self) -> None: + self._frames.clear() + self._values.clear() + + def row(self, frame_idx: int) -> np.ndarray: + if not self._frames: + return np.full(4 * self._width, np.nan) + return self.artifact.window( + np.stack(self._values), np.asarray(self._frames), np.array([frame_idx]) + )[0] diff --git a/src/glider/analysis/behavior/model.py b/src/glider/analysis/behavior/model.py index 5a26eff6..5eccc49c 100644 --- a/src/glider/analysis/behavior/model.py +++ b/src/glider/analysis/behavior/model.py @@ -85,6 +85,7 @@ def __init__( library_versions: dict[str, str] | None = None, glider_version: str | None = None, embedding: Any | None = None, + image_lane: Any | None = None, ): self.classifier = classifier self.feature_names = list(feature_names) @@ -100,6 +101,9 @@ def __init__( # for the live "galaxy" view. None for models trained without it # (and for all format_version-1 bundles). self.embedding = embedding + # Optional fitted ImageLaneArtifact: the PCA + windows that rebuild the + # img_* columns from the video. None for keypoint-only models. + self.image_lane = image_lane # ------------------------------------------------------------------ # Inference @@ -139,7 +143,7 @@ def predict( # same column names it was fit on (avoids a UserWarning about # missing feature names). df = windowed[self.feature_names] - valid_mask = ~df.isna().any(axis=1).to_numpy() + valid_mask = ~df[self.required_feature_names].isna().any(axis=1).to_numpy() labels = np.full(len(df), "", dtype=object) if valid_mask.any(): valid_df = df.loc[valid_mask] @@ -174,7 +178,7 @@ def posteriors(self, windowed: pd.DataFrame) -> tuple[np.ndarray, np.ndarray]: f"first few: {missing[:5]}" ) df = windowed[self.feature_names] - valid = ~df.isna().any(axis=1).to_numpy() + valid = ~df[self.required_feature_names].isna().any(axis=1).to_numpy() probs = ( self.classifier.predict_proba(df.loc[valid]) if valid.any() @@ -205,7 +209,7 @@ def predict_one( raise ValueError( f"feature_row shape {row.shape} != " f"({len(self.feature_names)},)" ) - if np.isnan(row).any(): + if np.isnan(row[self._required_index]).any(): return "" # Wrap in a 1-row DataFrame so sklearn sees the feature names it # was fit on (avoids a UserWarning that fires on bare arrays). @@ -243,6 +247,7 @@ def save(self, path: str | Path) -> Path: "library_versions": dict(self.library_versions), "glider_version": self.glider_version, "embedding": self.embedding, + "image_lane": self.image_lane, "format_version": 2, } joblib.dump(payload, path) @@ -276,11 +281,29 @@ def load(cls, path: str | Path) -> BehaviorModel: # A half-reconstructed embedding is worse than none: drop it so # nothing downstream tries to project through a stubbed reducer. embedding=payload.get("embedding") if embedding_usable else None, + image_lane=payload.get("image_lane"), ) # ------------------------------------------------------------------ # Helpers # ------------------------------------------------------------------ + @property + def required_feature_names(self) -> list[str]: + """Columns whose NaN blanks a prediction: all but the NaN-tolerant image lane.""" + from glider.analysis.behavior.image_lane import required_columns + + return required_columns(self.feature_names) + + @property + def _required_index(self) -> np.ndarray: + cached = getattr(self, "_required_index_cache", None) + if cached is None or len(cached[0]) != len(self.feature_names): + req = set(self.required_feature_names) + idx = np.array([i for i, c in enumerate(self.feature_names) if c in req], dtype=int) + cached = (list(self.feature_names), idx) + self._required_index_cache = cached + return cached[1] + def _is_fitted(self) -> bool: return self.classifier is not None and hasattr(self.classifier, "classes_") diff --git a/src/glider/analysis/behavior/pipeline.py b/src/glider/analysis/behavior/pipeline.py index a5f70ca0..154b995a 100644 --- a/src/glider/analysis/behavior/pipeline.py +++ b/src/glider/analysis/behavior/pipeline.py @@ -44,11 +44,13 @@ if TYPE_CHECKING: from glider.analysis.behavior.hybrid import HybridModel + from glider.analysis.behavior.image_lane import ImageLaneArtifact, ImageLaneConfig from glider.vision.pose.core import PoseData from glider.analysis.behavior.annotations import AnnotationStore from glider.analysis.behavior.benchmarks.metrics import macro_frame_f1 from glider.analysis.behavior.features import FeatureSpec, compute_features +from glider.analysis.behavior.image_lane import required_columns, resolve_config from glider.analysis.behavior.labels import ( AMBIGUOUS, build_label_and_group_series, @@ -230,6 +232,7 @@ def train_model( motion_features: bool = False, individuals: list[int | None] | None = None, holdout_individuals: list[int | None] | None = None, + image_lane=None, ) -> TrainResult: """Fit a behavior classifier from one or more (pose, annotations) pairs. @@ -335,6 +338,13 @@ def train_model( with ``holdout_sessions`` and independent of ``individuals`` and its length. ``None`` (default) is today's behaviour — every holdout session sees every zone. + image_lane + Add the image embedding of the animal crop beside the keypoint + features (:mod:`~glider.analysis.behavior.image_lane`). ``True`` uses + :func:`~glider.analysis.behavior.image_lane.default_config`; an + :class:`~glider.analysis.behavior.image_lane.ImageLaneConfig` sets it + explicitly. Needs each session's source video. The lane's PCA is + fitted on ``sessions`` only, never on ``holdout_sessions``. Returns ------- @@ -348,6 +358,15 @@ def train_model( spec = spec or FeatureSpec() fps = resolve_sessions_fps(sessions, fps, holdout_sessions=holdout_sessions) + lane_config = resolve_config(image_lane) + if lane_config is not None and mirror_augment: + raise ValueError(_LANE_MIRROR_REFUSAL) + # Fitted on the training sessions alone, before any holdout is assembled. + lane = ( + _fit_image_lane(sessions, spec=spec, fps=fps, config=lane_config) + if lane_config is not None + else None + ) # ---- 1-2. Assemble features/labels, drop unusable rows, subsample bg ---- assembled = _assemble_and_filter( @@ -367,6 +386,7 @@ def train_model( background_subsample_ratio=background_subsample_ratio, random_state=random_state, individuals=individuals, + image_lane=lane, ) x_kept, y_kept, g_kept = assembled.x_kept, assembled.y_kept, assembled.g_kept n_total = assembled.n_total @@ -413,6 +433,7 @@ def train_model( individuals=( [holdout_individuals[hi]] if holdout_individuals is not None else None ), + image_lane=lane, ) n_h = len(y_h) clean_h = _clean_window_rows( @@ -426,10 +447,7 @@ def train_model( # cross-validation and evaluate_model score under, so this holdout # accuracy is comparable to theirs. test_keep = ( - (y_test_all != "") - & (y_test_all != AMBIGUOUS) - & ~x_test_all.isna().any(axis=1) - & test_clean + (y_test_all != "") & (y_test_all != AMBIGUOUS) & ~_nan_rows(x_test_all) & test_clean ) # Note: if include_background is on, we don't add background to # the test set — the test set's labels should only be the @@ -552,6 +570,7 @@ def train_model( "lgbm_reg": ( dataclasses.asdict(lgbm_reg or LgbmReg()) if classifier_label == "lightgbm" else None ), + "image_lane": lane_config.to_dict() if lane_config is not None else None, } # Optional 3D embedding of the training feature space for the live @@ -586,6 +605,7 @@ def train_model( training_summary=summary, library_versions=capture_library_versions(), embedding=embedding_artifact, + image_lane=lane, ) return TrainResult(model=model, summary=summary) @@ -767,6 +787,7 @@ def _assemble_sessions( traj_features: bool = False, motion_features: bool = False, individuals: list[int | None] | None = None, + image_lane: ImageLaneArtifact | None = None, ) -> tuple[pd.DataFrame, pd.Series, pd.Series, list[dict[str, int]]]: """Load each (pose, annotations) pair, extract features + label series. @@ -782,7 +803,12 @@ def _assemble_sessions( zone in the file. A two-animal video appears as TWO sessions sharing one annotations CSV and differing only in this list — which is why the loop body needs no notion of multiple animals at all. + + ``image_lane`` is an already-fitted lane; its columns are appended after + the keypoint windows, since the lane has windows of its own. """ + if image_lane is not None and mirror_augment: + raise ValueError(_LANE_MIRROR_REFUSAL) if motion_features and mirror_augment: raise ValueError( "motion features are not supported with mirror augmentation yet " @@ -874,6 +900,11 @@ def _assemble_sessions( traj = apply_trajectory_rolling(variant.xy, body_axis=ba, window=window) traj.index = windowed.index windowed = pd.concat([windowed, traj], axis=1) + if image_lane is not None: + raw = _image_lane_raw(pose_csv, variant, spec, image_lane.config) + img = image_lane.frame_features(raw) + img.index = windowed.index + windowed = pd.concat([windowed, img], axis=1) if len(windowed) != len(labels): raise RuntimeError( f"windowed length {len(windowed)} != labels length {len(labels)} for {pose_csv}" @@ -925,6 +956,7 @@ def _assemble_and_filter( background_subsample_ratio: float, random_state: int, individuals: list[int | None] | None = None, + image_lane: ImageLaneArtifact | None = None, ) -> _AssembledData: """Assemble sessions, drop unusable rows, and subsample background. @@ -952,6 +984,7 @@ def _assemble_and_filter( traj_features=traj_features, motion_features=motion_features, individuals=individuals, + image_lane=image_lane, ) # ---- 2. Optionally promote unannotated frames to a background class ---- @@ -963,7 +996,7 @@ def _assemble_and_filter( raise ValueError("background_class_name must be non-empty") y_all = y_all.where(y_all != "", background_class_name) n_total = len(x_all) - keep_mask = (y_all != "") & (y_all != AMBIGUOUS) & ~x_all.isna().any(axis=1) + keep_mask = (y_all != "") & (y_all != AMBIGUOUS) & ~_nan_rows(x_all) n_kept = int(keep_mask.sum()) if n_kept < 2: raise ValueError( @@ -1066,6 +1099,76 @@ def _append_motion(feats: pd.DataFrame, pose_csv: Path, pose, spec: FeatureSpec) return pd.concat([feats, motion], axis=1) +_LANE_MIRROR_REFUSAL = ( + "the image lane is not supported with mirror augmentation (the pose is " + "mirrored but the video the lane is read from is not)" +) + + +def _nan_rows(x: pd.DataFrame) -> pd.Series: + """Rows with a NaN keypoint feature. Image-lane NaN (a missed mask) does not count.""" + return x[required_columns(x.columns)].isna().any(axis=1) + + +def _image_lane_raw(pose_csv: Path, pose, spec: FeatureSpec, config: ImageLaneConfig): + """The session's raw lane rows (memory-mapped cache), read from its source video.""" + from glider.analysis.behavior.image_lane import load_or_compute_raw + from glider.analysis.behavior.motion import video_path_for + + ba = spec.with_resolved_body_axis(pose.n_keypoints).body_axis + return load_or_compute_raw(Path(pose_csv), video_path_for(Path(pose_csv)), pose.xy, ba, config) + + +#: Every this-many frames of each session goes into the lane's PCA fit. +LANE_FIT_STRIDE = 5 + + +def _fit_image_lane( + sessions: list[SessionPair], *, spec: FeatureSpec, fps: float, config: ImageLaneConfig +) -> ImageLaneArtifact: + """Fit the lane's standardisation + PCA on these (training) sessions.""" + from glider.analysis.behavior.image_lane import ImageLaneArtifact + + samples = [] + for pose_csv, _ann in sessions: + pose = from_dlc_csv(Path(pose_csv), fps=fps) + raw = _image_lane_raw(pose_csv, pose, spec, config) + samples.append(np.asarray(raw[::LANE_FIT_STRIDE], dtype=np.float32)) + return ImageLaneArtifact.fit(samples, config, stride=1) + + +@dataclass +class _CvLane: + """Raw lane rows per CV session, so each fold can fit its own PCA. + + ``raws[i]`` is session ``i``'s memory-mapped cache. A fold fits on its + training sessions and projects every row, which keeps the held-out + session out of the PCA that scores it, as in the research protocol. + """ + + config: ImageLaneConfig + raws: list[np.ndarray] + + def fit(self, train_sessions) -> ImageLaneArtifact: + from glider.analysis.behavior.image_lane import ImageLaneArtifact + + samples = [ + np.asarray(self.raws[int(i)][::LANE_FIT_STRIDE], dtype=np.float32) + for i in train_sessions + ] + return ImageLaneArtifact.fit(samples, self.config, stride=1) + + def join(self, x: pd.DataFrame, artifact, sess: np.ndarray, frame: np.ndarray) -> pd.DataFrame: + """``x`` with the lane columns for each row's (session, frame) appended.""" + lane = np.full((len(x), len(artifact.columns)), np.nan) + for i in np.unique(sess): + rows = np.flatnonzero(sess == i) + values = artifact.transform(self.raws[int(i)]) + lane[rows] = artifact.window(values, np.arange(len(values)), frame[rows]) + img = pd.DataFrame(lane, columns=artifact.columns, index=x.index) + return pd.concat([x, img], axis=1) + + def _mirror_pose(pose): """Return a copy of ``pose`` with x-coords mirrored + L/R keypoints swapped. @@ -1134,6 +1237,7 @@ def cross_validate_sessions( freq_features: bool = False, traj_features: bool = False, motion_features: bool = False, + image_lane=None, ) -> dict[str, Any]: """Session-grouped K-fold cross-validation — a robust cross-session generalization estimate. @@ -1165,6 +1269,9 @@ def cross_validate_sessions( ``classifier_type`` defaults to ``"lightgbm"`` (matching :func:`train_model`) so CV measures the same backend that gets deployed; pass ``"rf"`` to evaluate the RandomForest backend instead. + + ``image_lane`` is as for :func:`train_model`; each fold fits the lane's + PCA on its own training sessions. """ spec = spec or FeatureSpec() fps = resolve_sessions_fps(sessions, fps) @@ -1184,6 +1291,7 @@ def cross_validate_sessions( background_class_name=background_class_name, background_ratio=background_ratio, random_state=random_state, + image_lane=resolve_config(image_lane), ) return _run_cv_folds( data, @@ -1223,6 +1331,7 @@ def cross_validate_and_train( traj_features: bool = False, motion_features: bool = False, embedding: str = "none", + image_lane=None, ) -> tuple[dict[str, Any], TrainResult]: """Cross-validate, then refit on everything — one feature pass for both. @@ -1261,6 +1370,7 @@ def cross_validate_and_train( background_class_name=background_class_name, background_ratio=background_ratio, random_state=random_state, + image_lane=resolve_config(image_lane), ) cv_result = _run_cv_folds( data, @@ -1285,7 +1395,10 @@ def cross_validate_and_train( class_weight=class_weight, lgbm_reg=lgbm_reg, ) - clf.fit(data.x, data.y) + # The deployed lane is fitted on every session, like the classifier. + lane = data.lane.fit(np.unique(data.sess)) if data.lane is not None else None + x_all = data.lane.join(data.x, lane, data.sess, data.frame) if lane is not None else data.x + clf.fit(x_all, data.y) summary = { "split_strategy": "cross_validated", @@ -1298,7 +1411,7 @@ def cross_validate_and_train( # only say how many rows were assembled, leaving the reader to infer # the real denominator from the per-class supports. "n_rows_scored": cv_result.get("n_rows_scored"), - "train_accuracy": float(clf.score(data.x, data.y)), + "train_accuracy": float(clf.score(x_all, data.y)), # The honest headline: measured across folds, not on the fit itself. "test_accuracy": cv_result.get("mean_accuracy"), # The names every reader looks for -- run_report.TrainingRun, @@ -1317,7 +1430,7 @@ def cross_validate_and_train( "bout_metrics": cv_result.get("bout_metrics") or {}, "false_alarm_rate": cv_result.get("false_alarm_rate"), "classes": [str(c) for c in clf.classes_], - "n_features": int(data.x.shape[1]), + "n_features": int(x_all.shape[1]), "kept_label_counts": split_label_counts(data.y), "window": int(window), "stats": list(stats), @@ -1330,12 +1443,13 @@ def cross_validate_and_train( if _classifier_label_for(clf) == "lightgbm" else None ), + "image_lane": lane.config.to_dict() if lane is not None else None, } if hasattr(clf, "feature_importances_"): importances = np.asarray(clf.feature_importances_, dtype=float) order = np.argsort(-importances) summary["top_features"] = [ - {"feature": data.x.columns[int(i)], "importance": float(importances[int(i)])} + {"feature": x_all.columns[int(i)], "importance": float(importances[int(i)])} for i in order[:20] ] @@ -1349,7 +1463,7 @@ def cross_validate_and_train( try: embedding_artifact = _embedding_mod.fit_embedding( - data.x, data.y, method=embedding, random_state=random_state + x_all, data.y, method=embedding, random_state=random_state ) except Exception as e: # noqa: BLE001 - never lose a fit over a chart warnings.warn( @@ -1359,7 +1473,7 @@ def cross_validate_and_train( model = BehaviorModel( classifier=clf, - feature_names=list(data.x.columns), + feature_names=list(x_all.columns), spec=spec, window=window, stats=tuple(stats), @@ -1368,6 +1482,7 @@ def cross_validate_and_train( training_summary=summary, library_versions=capture_library_versions(), embedding=embedding_artifact, + image_lane=lane, ) return cv_result, TrainResult(model=model, summary=summary) @@ -1396,6 +1511,9 @@ class _CvData: #: same run counting evaluation uses). ``None`` recomputes it over the #: kept rows. clean: np.ndarray | None = None + #: Raw image-lane rows per session; ``None`` when the lane is off. The lane + #: columns are not in ``x``: each fold appends its own (see _CvLane). + lane: _CvLane | None = None def _assemble_for_cv( @@ -1415,6 +1533,7 @@ def _assemble_for_cv( background_class_name: str, background_ratio: float, random_state: int, + image_lane: ImageLaneConfig | None = None, ) -> _CvData: """Feature matrix + per-row session / mirror / frame metadata. @@ -1426,6 +1545,8 @@ def _assemble_for_cv( """ if not sessions: raise ValueError("cross-validation requires at least one session") + if image_lane is not None and mirror_augment: + raise ValueError(_LANE_MIRROR_REFUSAL) if motion_features and mirror_augment: raise ValueError( "motion features are not supported with mirror augmentation yet " @@ -1464,9 +1585,12 @@ def _assemble_for_cv( sess_ids: list[np.ndarray] = [] is_mirror: list[np.ndarray] = [] frame_ids: list[np.ndarray] = [] + lane_raws: list[np.ndarray] = [] for sidx, (pose_csv, ann_csv) in enumerate(sessions): pose = from_dlc_csv(Path(pose_csv), fps=fps) store = AnnotationStore.load_csv(Path(ann_csv)) + if image_lane is not None: + lane_raws.append(_image_lane_raw(pose_csv, pose, spec, image_lane)) variants = [(pose, False)] if mirror_augment: variants.append((_mirror_pose(pose), True)) @@ -1517,7 +1641,7 @@ def _assemble_for_cv( y_all.to_numpy(), sess_all, frame_all_rows, window, mirror=mirror_all ) - keep = ((y_all != AMBIGUOUS) & ~x_all.isna().any(axis=1)).to_numpy() + keep = ((y_all != AMBIGUOUS) & ~_nan_rows(x_all)).to_numpy() if not include_background: keep = keep & (y_all != "").to_numpy() x = x_all.loc[keep].reset_index(drop=True) @@ -1548,7 +1672,10 @@ def _assemble_for_cv( frame = frame[sel] clean = clean[sel] - return _CvData(x=x, y=y, sess=sess, mirror=mirror, frame=frame, window=window, clean=clean) + lane = _CvLane(image_lane, lane_raws) if image_lane is not None else None + return _CvData( + x=x, y=y, sess=sess, mirror=mirror, frame=frame, window=window, clean=clean, lane=lane + ) def _run_cv_folds( @@ -1614,8 +1741,13 @@ def _run_cv_folds( class_weight=class_weight, lgbm_reg=lgbm_reg, ) - clf.fit(x.iloc[train_idx], y.iloc[train_idx]) - real_pred = clf.predict(x.iloc[real_idx]) + xf = x + if data.lane is not None: + # This fold's PCA, fitted without the held-out session(s). + fold_lane = data.lane.fit(np.unique(sess[train_idx])) + xf = data.lane.join(x, fold_lane, sess, frame) + clf.fit(xf.iloc[train_idx], y.iloc[train_idx]) + real_pred = clf.predict(xf.iloc[real_idx]) y_true = y.iloc[eval_idx].to_numpy() y_pred = real_pred[clean_sel] fold_acc.append(float((y_pred == y_true).mean())) @@ -1626,7 +1758,7 @@ def _run_cv_folds( bout_true.append(y.iloc[real_idx].to_numpy()) bout_pred.append(real_pred) if threshold_curve: - all_proba.append((clf.predict_proba(x.iloc[eval_idx]), [str(c) for c in clf.classes_])) + all_proba.append((clf.predict_proba(xf.iloc[eval_idx]), [str(c) for c in clf.classes_])) # Aggregate per-class metrics + confusion over ALL folds. GroupKFold # test sets are disjoint and cover every kept row exactly once, so diff --git a/src/glider/gui/behavior/window.py b/src/glider/gui/behavior/window.py index b35a257c..2865775f 100644 --- a/src/glider/gui/behavior/window.py +++ b/src/glider/gui/behavior/window.py @@ -1339,6 +1339,16 @@ def __init__(self, parent=None): "mirrors the pose but not the video, so the two would disagree.\n" "train_model refuses the combination outright." ) + self._image_lane_check = QCheckBox("Image lane") + self._image_lane_check.setToolTip( + "Adds an embedding of the animal crop from the source video: mask\n" + "shape, DINOv2 and the pose model's neck features (the classifier\n" + "the GLIDER paper reports). Needs mouse_seg.pt and mmv2.pt in\n" + "~/.glider/models/image_lane (or $GLIDER_IMAGE_LANE_MODELS), and a\n" + "GPU to keep up in real time.\n\n" + "Cannot be combined with mirror augmentation, for the same reason\n" + "as motion features: the video is not mirrored." + ) self._freq_check = QCheckBox("Frequency features") self._freq_check.setToolTip( "Dominant frequency and spectral flatness per kinematic column —\n" @@ -1362,18 +1372,20 @@ def __init__(self, parent=None): options.add(self._traj_check) options.add(self._motion_check) + options.add(self._image_lane_check) options.add(self._freq_check) # train_model raises on motion + mirror together, and the raise lands # minutes into a fit with the model lost. Enforce it here, where the # click happens, rather than letting the run start and fail. self._mirror_check.toggled.connect(self._on_motion_mirror_toggled) self._motion_check.toggled.connect(self._on_motion_mirror_toggled) + self._image_lane_check.toggled.connect(self._on_motion_mirror_toggled) self._motion_mirror_hint = hint("") options.add(self._motion_mirror_hint) self._on_motion_mirror_toggled() options.add( hint( - "These three are computed but off by default, so no existing " + "These four are computed but off by default, so no existing " "model used them. Turning one on changes the feature set, and " "the resulting model is not comparable to one trained without it." ) @@ -1411,31 +1423,35 @@ def _on_remove_holdout(self) -> None: self._refresh_counts() def _on_motion_mirror_toggled(self, *_args) -> None: - """Keep motion features and mirror augmentation mutually exclusive. + """Keep the video-read features and mirror augmentation mutually exclusive. The pipeline mirrors the pose but not the source video, so motion - features read from the video would disagree with the mirrored pose; - ``train_model`` refuses the pair. Whichever the operator picked most - recently wins, and the other is cleared and greyed out with the reason - on screen — a disabled box with no explanation reads as a bug. + features and the image lane, both read from the video, would disagree + with the mirrored pose; ``train_model`` refuses either pair. Whichever + the operator picked most recently wins, and the other is cleared and + greyed out with the reason on screen — a disabled box with no + explanation reads as a bug. """ - mirror, motion = self._mirror_check, self._motion_check - # Signals are blocked while clearing so the two handlers do not - # bounce the state between them. - if mirror.isChecked() and motion.isChecked(): - loser = motion if self.sender() is mirror else mirror - loser.blockSignals(True) - loser.setChecked(False) - loser.blockSignals(False) - - motion.setEnabled(not mirror.isChecked()) - mirror.setEnabled(not motion.isChecked()) + mirror = self._mirror_check + video = [self._motion_check, self._image_lane_check] + # Signals are blocked while clearing so the handlers do not bounce the + # state between them. + if mirror.isChecked() and any(box.isChecked() for box in video): + losers = [b for b in video if b.isChecked()] if self.sender() is mirror else [mirror] + for loser in losers: + loser.blockSignals(True) + loser.setChecked(False) + loser.blockSignals(False) + + for box in video: + box.setEnabled(not mirror.isChecked()) + mirror.setEnabled(not any(box.isChecked() for box in video)) if mirror.isChecked(): - note = "Motion features are unavailable while mirror augment is on." - elif motion.isChecked(): - note = "Mirror augment is unavailable while motion features are on." + note = "Motion features and the image lane are unavailable while mirror augment is on." + elif any(box.isChecked() for box in video): + note = "Mirror augment is unavailable while a video feature is on." else: - note = "Motion features and mirror augment cannot be used together." + note = "Motion features and the image lane cannot be used with mirror augment." self._motion_mirror_hint.setText(note) def _refresh_window_hint(self) -> None: @@ -1561,6 +1577,7 @@ def _shared_options(self) -> dict: "random_state": int(self._seed_spin.value()), "traj_features": self._traj_check.isChecked(), "motion_features": self._motion_check.isChecked(), + "image_lane": self._image_lane_check.isChecked(), "freq_features": self._freq_check.isChecked(), "embedding": self._embedding_combo.currentData(), # 0 is the spin box's "auto": leave fps unset and let the pipeline diff --git a/tests/unit/analysis/behavior/test_image_lane.py b/tests/unit/analysis/behavior/test_image_lane.py new file mode 100644 index 00000000..34d613c6 --- /dev/null +++ b/tests/unit/analysis/behavior/test_image_lane.py @@ -0,0 +1,596 @@ +"""The image lane: crop embedding beside the keypoint features. + +The networks are not run here (they need model files and minutes of video); +every test that needs raw embeddings gets synthetic ones through the same +seam the pipeline uses, ``_image_lane_raw``. What is pinned is everything +GLIDER itself computes: the causal crop, the mask shape descriptors, the PCA +fitted on training sessions only, the windows (identical to the research +pandas recipe), the live stream reproducing the offline columns frame for +frame, and the NaN tolerance of the lane columns. +""" + +from __future__ import annotations + +import queue +import threading +from unittest.mock import MagicMock + +import numpy as np +import pandas as pd +import pytest + +from glider.analysis.behavior import FeatureSpec +from glider.analysis.behavior import image_lane as il +from glider.analysis.behavior.image_lane import ( + IMAGE_LANE_PREFIX, + N_SHAPE, + CropTracker, + ImageLaneArtifact, + ImageLaneConfig, + ImageLaneEncoder, + StreamingImageLane, + required_columns, + window_lane, +) + +pytest.importorskip("lightgbm") + +N_NECK = 16 +RAW_WIDTH = N_SHAPE + 768 + N_NECK +CONFIG = ImageLaneConfig(seg_model="seg.pt", neck_model="neck.pt", dims=4) +SPEC = FeatureSpec(body_axis=(0, 4)) +COMMON = {"window": 8, "n_estimators": 20, "random_state": 42, "fps": 30.0} + + +# --------------------------------------------------------------------------- +# Synthetic sessions + raw lane rows +# --------------------------------------------------------------------------- + + +def _pose(seed: int, n_frames: int = 300): + from glider.vision.pose.core import PoseData + + rng = np.random.default_rng(seed) + names = ["snout", "left_ear", "right_ear", "neck", "tail_base"] + offsets = np.array([[0, -20], [-10, -10], [10, -10], [0, 0], [0, 30]], dtype=float) + xy = np.empty((n_frames, len(names), 2)) + per = n_frames // 3 + for regime, slc in enumerate([slice(0, per), slice(per, 2 * per), slice(2 * per, n_frames)]): + n = slc.stop - slc.start + t = np.arange(n) + if regime == 0: + cx, cy, jitter = 50 + 0.5 * t, np.full(n, 200.0), 0.4 + elif regime == 1: + cx, cy, jitter = 350 + 3 * np.sin(0.5 * t), 200 + 3 * np.cos(0.5 * t), 1.0 + else: + cx, cy, jitter = np.full(n, 360.0), np.full(n, 200.0), 0.15 + for k in range(len(names)): + xy[slc, k, 0] = cx + offsets[k, 0] + rng.normal(0, jitter, n) + xy[slc, k, 1] = cy + offsets[k, 1] + rng.normal(0, jitter, n) + return PoseData( + xy=xy, confidence=np.full((n_frames, len(names)), 0.95), keypoint_names=names, fps=30.0 + ) + + +def _raw_for(n_frames: int, seed: int) -> np.ndarray: + """Raw lane rows whose level tracks the regime, so the lane carries signal.""" + rng = np.random.default_rng(seed) + regime = np.minimum(np.arange(n_frames) // (n_frames // 3), 2) + raw = rng.normal(size=(n_frames, RAW_WIDTH)) + regime[:, None] + raw[::17, :N_SHAPE] = np.nan # the seg model misses the odd frame + return raw.astype(np.float32) + + +@pytest.fixture +def sessions(tmp_path): + from glider.analysis.behavior.annotations import AnnotationStore, BehaviorZone + from glider.vision.pose.dlc import to_dlc_csv + + pairs = [] + for i in range(4): + pose = _pose(seed=i) + pose_csv = tmp_path / f"s{i}.csv" + to_dlc_csv(pose, pose_csv) + store = AnnotationStore() + third = pose.n_frames // 3 + store.add(BehaviorZone("walk", 0, third)) + store.add(BehaviorZone("mill", third, 2 * third)) + store.add(BehaviorZone("still", 2 * third, pose.n_frames)) + ann = tmp_path / f"s{i}_annotations.csv" + store.save_csv(ann) + pairs.append((pose_csv, ann)) + return pairs + + +@pytest.fixture +def fake_raw(monkeypatch): + """Route the pipeline's raw-lane seam to synthetic rows; record who asked.""" + from glider.analysis.behavior import pipeline + + asked: list[str] = [] + + def raw(pose_csv, pose, spec, config): + asked.append(str(pose_csv)) + return _raw_for(pose.n_frames, seed=int(str(pose_csv)[-5])) + + monkeypatch.setattr(pipeline, "_image_lane_raw", raw) + return asked + + +# --------------------------------------------------------------------------- +# Windows +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize("center", [False, True]) +def test_windows_match_the_research_pandas_recipe(center): + """value, rolling mean/std (min_periods=1) and diff, as pm_raster.win() builds them.""" + rng = np.random.default_rng(0) + vals = rng.normal(size=(120, 5)) + vals[[3, 40, 41, 90]] = np.nan + df = pd.DataFrame(vals) + r = df.rolling(5, min_periods=1, center=center) + ref = pd.concat([df, r.mean(), r.std(), df.diff()], axis=1).to_numpy() + + got = window_lane(vals, np.arange(120), np.arange(120), window=5, center=center) + + np.testing.assert_allclose(got, ref, atol=1e-12) + assert (np.isnan(got) == np.isnan(ref)).all() + + +def test_absent_frames_count_as_missing_not_as_neighbours(): + """A sparse frame set must give what the dense, NaN-filled one does.""" + rng = np.random.default_rng(1) + dense = rng.normal(size=(50, 3)) + keep = np.setdiff1d(np.arange(50), [10, 11, 30]) + dense_nan = dense.copy() + dense_nan[[10, 11, 30]] = np.nan + + sparse = window_lane(dense[keep], keep, np.arange(50), window=5) + full = window_lane(dense_nan, np.arange(50), np.arange(50), window=5) + + np.testing.assert_allclose(sparse, full, equal_nan=True) + + +def test_past_window_never_reads_a_later_frame(): + rng = np.random.default_rng(2) + a = rng.normal(size=(40, 3)) + b = a.copy() + b[21:] += 100.0 + ra = window_lane(a, np.arange(40), np.arange(21), window=5) + rb = window_lane(b, np.arange(40), np.arange(21), window=5) + np.testing.assert_array_equal(ra, rb) + + +# --------------------------------------------------------------------------- +# Reducer +# --------------------------------------------------------------------------- + + +def test_fit_projects_each_learned_block_to_dims_and_passes_shape_through(): + raws = [_raw_for(300, 0), _raw_for(300, 1)] + art = ImageLaneArtifact.fit(raws, CONFIG) + + per_frame = art.transform(raws[0]) + assert per_frame.shape == (300, N_SHAPE + 2 * CONFIG.dims) + np.testing.assert_array_equal(per_frame[:, :N_SHAPE], raws[0][:, :N_SHAPE]) + assert len(art.columns) == 4 * per_frame.shape[1] + assert all(c.startswith(IMAGE_LANE_PREFIX) for c in art.columns) + # No lane column may look like a keypoint rolling stat, or the live + # engine would try to stream it as a per-frame keypoint feature. + assert not any("__" in c for c in art.columns) + + +def test_fit_uses_only_the_sessions_it_is_given(): + """The held-out session must not shape the PCA that scores it.""" + a, b = _raw_for(300, 0), _raw_for(300, 1) + held_out = _raw_for(300, 2) * 50.0 + art1 = ImageLaneArtifact.fit([a, b], CONFIG) + art2 = ImageLaneArtifact.fit([a, b], CONFIG) + art3 = ImageLaneArtifact.fit([a, b, held_out], CONFIG) + np.testing.assert_array_equal(art1.transform(held_out), art2.transform(held_out)) + assert not np.allclose(art1.transform(a), art3.transform(a)) + + +def test_fit_refuses_rows_it_cannot_split(): + with pytest.raises(ValueError, match="shape"): + ImageLaneArtifact.fit([np.zeros((10, N_SHAPE + 5))], CONFIG) + + +# --------------------------------------------------------------------------- +# Crop + mask shape +# --------------------------------------------------------------------------- + + +def test_crop_tracker_carries_the_centre_and_takes_a_running_median_scale(): + tracker = CropTracker(CONFIG, body_axis=(0, 1)) + assert tracker.push(np.full((2, 2), np.nan)) is None # nothing seen yet + + center, scale = tracker.push(np.array([[0.0, 0.0], [100.0, 0.0]])) + np.testing.assert_allclose(center, [50.0, 0.0]) + assert scale == pytest.approx(CONFIG.crop_px / (CONFIG.crop_body_lengths * 100.0)) + + # A missed frame keeps the last centre; a short (rearing) body barely moves the scale. + center2, scale2 = tracker.push(np.full((2, 2), np.nan)) + np.testing.assert_allclose(center2, center) + tracker.push(np.array([[0.0, 0.0], [100.0, 0.0]])) + _c, scale3 = tracker.push(np.array([[0.0, 0.0], [40.0, 0.0]])) + assert scale3 == pytest.approx(scale) + + +def test_crop_tracker_can_centre_on_a_named_keypoint(): + cfg = ImageLaneConfig(seg_model="s", neck_model="n", center_keypoint=1) + center, _ = CropTracker(cfg, (0, 1)).push(np.array([[0.0, 0.0], [10.0, 4.0]])) + np.testing.assert_allclose(center, [10.0, 4.0]) + + +def test_shape_features_describe_an_ellipse_and_blank_an_empty_mask(): + cv2 = pytest.importorskip("cv2") + mask = np.zeros((320, 320), np.uint8) + cv2.ellipse(mask, (160, 160), (80, 30), 0, 0, 360, 1, -1) + f = il.shape_features(mask.astype(bool)) + assert np.isfinite(f).all() + pc = 320 / il.NOMINAL_CROP_CM + assert f[0] == pytest.approx(np.pi * 80 * 30 / pc**2, rel=0.02) # area, nominal cm^2 + assert f[4] == pytest.approx(30 / 80, rel=0.05) # minor/major axis + assert np.isnan(il.shape_features(np.zeros((320, 320), bool))).all() + + +# --------------------------------------------------------------------------- +# Config +# --------------------------------------------------------------------------- + + +def test_default_config_names_the_folder_when_a_model_is_missing(tmp_path, monkeypatch): + monkeypatch.setenv("GLIDER_IMAGE_LANE_MODELS", str(tmp_path)) + with pytest.raises(FileNotFoundError, match="mouse_seg.pt and mmv2.pt"): + il.default_config() + (tmp_path / "mouse_seg.pt").write_bytes(b"x") + (tmp_path / "mmv2.pt").write_bytes(b"x") + cfg = il.resolve_config(True) + assert cfg.seg_model == str(tmp_path / "mouse_seg.pt") + assert cfg.center is False # the shipped lane is the live (past-only) one + + +def test_resolve_config_accepts_off_dicts_and_configs(): + assert il.resolve_config(None) is None + assert il.resolve_config(False) is None + assert il.resolve_config(CONFIG) is CONFIG + assert il.resolve_config(CONFIG.to_dict()) == CONFIG + with pytest.raises(TypeError): + il.resolve_config("yes") + + +def test_cache_tag_changes_with_anything_that_changes_the_raw_rows(): + assert CONFIG.cache_tag() != ImageLaneConfig("seg.pt", "neck.pt", crop_px=256).cache_tag() + # dims / window / center only change the reducer, so the cache is shared. + assert CONFIG.cache_tag() == ImageLaneConfig("seg.pt", "neck.pt", dims=8, window=9).cache_tag() + + +# --------------------------------------------------------------------------- +# Training + cross-validation +# --------------------------------------------------------------------------- + + +def test_train_model_appends_the_lane_and_bundles_its_reducer(sessions, fake_raw, tmp_path): + from glider.analysis.behavior.model import BehaviorModel + from glider.analysis.behavior.pipeline import train_model + + result = train_model(sessions[:3], spec=SPEC, image_lane=CONFIG, **COMMON) + model = result.model + + assert isinstance(model.image_lane, ImageLaneArtifact) + lane_cols = [c for c in model.feature_names if c.startswith(IMAGE_LANE_PREFIX)] + assert lane_cols == model.image_lane.columns + assert result.summary["image_lane"]["dims"] == CONFIG.dims + + path = model.save(tmp_path / "m.pkl") + loaded = BehaviorModel.load(path) + np.testing.assert_array_equal( + loaded.image_lane.transform(_raw_for(50, 9)), model.image_lane.transform(_raw_for(50, 9)) + ) + + +def test_rows_with_a_missed_mask_still_train(sessions, fake_raw): + """Lane NaN is routed by the classifier, not a reason to drop the frame.""" + from glider.analysis.behavior.pipeline import train_model + + with_lane = train_model(sessions[:3], spec=SPEC, image_lane=CONFIG, **COMMON) + without = train_model(sessions[:3], spec=SPEC, **COMMON) + assert with_lane.summary["n_rows_kept"] == without.summary["n_rows_kept"] + + +def test_the_holdout_session_is_kept_out_of_the_lane_fit(sessions, fake_raw, monkeypatch): + from glider.analysis.behavior import pipeline + + fitted_on: list[int] = [] + real_fit = ImageLaneArtifact.fit.__func__ + + def spy(cls, raws, config, **kw): + fitted_on.append(len(raws)) + return real_fit(cls, raws, config, **kw) + + monkeypatch.setattr(ImageLaneArtifact, "fit", classmethod(spy)) + result = pipeline.train_model( + sessions[:3], spec=SPEC, image_lane=CONFIG, holdout_sessions=[sessions[3]], **COMMON + ) + assert fitted_on == [3] + assert result.summary["test_size"] > 0 + + +def test_image_lane_with_mirror_is_refused(sessions, fake_raw): + from glider.analysis.behavior.pipeline import cross_validate_sessions, train_model + + with pytest.raises(ValueError, match="mirror"): + train_model(sessions[:2], spec=SPEC, image_lane=CONFIG, mirror_augment=True, **COMMON) + with pytest.raises(ValueError, match="mirror"): + cross_validate_sessions( + sessions, spec=SPEC, image_lane=CONFIG, mirror_augment=True, n_folds=2, **COMMON + ) + + +def test_each_cv_fold_fits_the_lane_without_its_held_out_session(sessions, fake_raw, monkeypatch): + from glider.analysis.behavior import pipeline + + folds: list[set[int]] = [] + real = pipeline._CvLane.fit + + def spy(self, train_sessions): + folds.append({int(i) for i in train_sessions}) + return real(self, train_sessions) + + monkeypatch.setattr(pipeline._CvLane, "fit", spy) + cv, trained = pipeline.cross_validate_and_train( + sessions, spec=SPEC, image_lane=CONFIG, n_folds=4, **COMMON + ) + + assert len(folds) == 5 + assert all(len(f) == 3 for f in folds[:4]) # one session held out per fold + assert len({frozenset(f) for f in folds[:4]}) == 4 # a different one each time + assert folds[4] == {0, 1, 2, 3} # the deployed lane sees every session + assert cv["n_folds"] == 4 + assert trained.model.image_lane is not None + assert any(c.startswith(IMAGE_LANE_PREFIX) for c in trained.model.feature_names) + + +# --------------------------------------------------------------------------- +# Inference: NaN tolerance, offline scoring, live stream +# --------------------------------------------------------------------------- + + +@pytest.fixture +def lane_model(sessions, fake_raw): + from glider.analysis.behavior.pipeline import train_model + + return train_model(sessions[:3], spec=SPEC, image_lane=CONFIG, **COMMON).model + + +def test_required_columns_leave_out_only_the_lane(): + cols = ["dist_0_1__mean", "img_shape00", "img_dino01_m", "speed__std"] + assert required_columns(cols) == ["dist_0_1__mean", "speed__std"] + + +def test_predict_one_ignores_lane_nan_but_blanks_on_keypoint_nan(lane_model): + row = np.zeros(len(lane_model.feature_names)) + lane = [i for i, c in enumerate(lane_model.feature_names) if c.startswith(IMAGE_LANE_PREFIX)] + keypoint = [i for i in range(len(row)) if i not in set(lane)] + + row[lane] = np.nan + assert lane_model.predict_one(row) != "" + row[keypoint[0]] = np.nan + assert lane_model.predict_one(row) == "" + + +def test_offline_scoring_needs_the_video(lane_model): + from glider.analysis.behavior.classify.batch import classify_pose_data + + with pytest.raises(NotImplementedError, match="image lane"): + classify_pose_data(_pose(seed=7), lane_model) + + +def test_offline_scoring_appends_the_lane_at_each_row_frame(lane_model, monkeypatch): + from glider.analysis.behavior.classify import batch + + pose = _pose(seed=7) + raw = _raw_for(pose.n_frames, 7) + monkeypatch.setattr(batch, "_image_lane_raw", lambda *a, **k: raw) + captured = {} + real_predict = batch._predict + + def capture(model, frame, *a): + captured["frame"] = frame + return real_predict(model, frame, *a) + + monkeypatch.setattr(batch, "_predict", capture) + rows = batch.classify_pose_data(pose, lane_model, video="v.mp4", predict_every=1) + + # Only the keypoint warm-up (a one-row std is NaN) blanks a row; the + # frames whose mask was missed (every 17th) are still scored. + blank = [f for f, label in zip(rows.frames, rows.labels, strict=True) if not label] + assert len(rows) > 250 and all(f <= 3 for f in blank) + lane_part = captured["frame"][lane_model.image_lane.columns].to_numpy() + scored = [f for f, label in zip(rows.frames, rows.labels, strict=True) if label] + expected = lane_model.image_lane.frame_features(raw, query=np.asarray(scored)).to_numpy() + np.testing.assert_allclose(lane_part, expected, equal_nan=True) + + +def test_the_live_stream_reproduces_the_offline_lane_columns(lane_model): + """Frame by frame, the streaming lane equals the offline one — causally.""" + art = lane_model.image_lane + raw = _raw_for(120, 3) + values = art.transform(raw) + offline = art.frame_features(raw).to_numpy() + + stream = StreamingImageLane(art) + lag = 2 # the keypoint lane emits the frame two behind the newest + for f in range(120): + stream.push(f, values[f]) + if f >= lag: + np.testing.assert_allclose(stream.row(f - lag), offline[f - lag], equal_nan=True) + + +def test_the_stream_refuses_a_centred_lane(): + art = ImageLaneArtifact.fit([_raw_for(100, 0)], ImageLaneConfig("s", "n", center=True)) + with pytest.raises(ValueError, match="centred"): + StreamingImageLane(art) + + +def test_feature_engine_appends_the_lane_to_every_emitted_row(lane_model): + from glider.analysis.behavior.classify.features_stream import derive_stream_columns + from glider.analysis.behavior.classify.threads import END_OF_STREAM, FeatureEngine + + pose = _pose(seed=5, n_frames=60) + raw = _raw_for(60, 5) + values = lane_model.image_lane.transform(raw) + per_frame, spectral = derive_stream_columns(lane_model) + tracked, out = queue.Queue(), queue.Queue() + engine = FeatureEngine( + tracked_queue=tracked, + classifier_queue=out, + stop_event=threading.Event(), + spec=lane_model.spec, + keypoint_names=pose.keypoint_names, + window=lane_model.window, + stats=lane_model.stats, + per_frame_feature_names=per_frame, + predict_every=1, + spectral_features=spectral, + image_lane=lane_model.image_lane, + ) + for f in range(60): + tracked.put((f, None, pose.xy[f], pose.confidence[f], values[f])) + tracked.put(END_OF_STREAM) + engine.run() + + offline = lane_model.image_lane.frame_features(raw) + rows = [] + while (item := out.get_nowait()) is not END_OF_STREAM: + rows.append(item) + assert rows + for frame_idx, names, row in rows: + assert names[-len(lane_model.image_lane.columns) :] == lane_model.image_lane.columns + np.testing.assert_allclose( + row[-len(lane_model.image_lane.columns) :], + offline.loc[frame_idx].to_numpy(), + equal_nan=True, + ) + + +def test_lane_worker_encodes_each_fixed_frame_and_passes_none_before_one(lane_model, monkeypatch): + from glider.analysis.behavior.classify.threads import END_OF_STREAM, ImageLaneWorker + + encoder = MagicMock(spec=ImageLaneEncoder) + encoder.device = "cuda:0" + encoder.encode.side_effect = lambda crops: _raw_for(len(crops) * 3, 1)[: len(crops)] + monkeypatch.setattr(il, "get_encoder", lambda config, device=None: encoder) + + pose = _pose(seed=1, n_frames=6) + inq, outq = queue.Queue(), queue.Queue() + worker = ImageLaneWorker( + inq, outq, threading.Event(), lane_model.image_lane, body_axis=(0, 4), device="cpu" + ) + frame = np.zeros((480, 640, 3), np.uint8) + inq.put((0, frame, np.full((5, 2), np.nan), np.zeros(5))) + for f in range(1, 6): + inq.put((f, frame, pose.xy[f], pose.confidence[f])) + inq.put(END_OF_STREAM) + worker.run() + + items = [] + while (item := outq.get_nowait()) is not END_OF_STREAM: + items.append(item) + assert worker.error is None + assert [it[0] for it in items] == list(range(6)) # every frame, in order + assert items[0][4] is None # no keypoint fix yet: nothing to crop + assert all(it[4].shape == (len(lane_model.image_lane.frame_columns),) for it in items[1:]) + # Everything was already queued, so it went through as one micro-batch. + batches = [len(c.args[0]) for c in encoder.encode.call_args_list] + assert batches == [5] + crop = encoder.encode.call_args.args[0][0] + assert crop.shape == (CONFIG.crop_px, CONFIG.crop_px, 3) + assert len(worker.encode_ms) == 5 + + +def test_lane_worker_batches_no_more_than_max_batch(lane_model, monkeypatch): + from glider.analysis.behavior.classify.threads import END_OF_STREAM, ImageLaneWorker + + encoder = MagicMock(spec=ImageLaneEncoder) + encoder.device = "cuda:0" + encoder.encode.side_effect = lambda crops: _raw_for(len(crops) * 3, 1)[: len(crops)] + monkeypatch.setattr(il, "get_encoder", lambda config, device=None: encoder) + pose = _pose(seed=1, n_frames=12) + inq, outq = queue.Queue(), queue.Queue() + worker = ImageLaneWorker( + inq, outq, threading.Event(), lane_model.image_lane, (0, 4), max_batch=4 + ) + frame = np.zeros((480, 640, 3), np.uint8) + for f in range(12): + inq.put((f, frame, pose.xy[f], pose.confidence[f])) + inq.put(END_OF_STREAM) + worker.run() + assert [len(c.args[0]) for c in encoder.encode.call_args_list] == [4, 4, 4] + + +def test_the_display_draws_on_a_copy_not_the_frame_the_lane_crops(): + """The tracker shares one array between display and features; the lane reads its pixels.""" + from glider.analysis.behavior.classify.threads import ( + END_OF_STREAM, + DisplayConsumer, + LatestLabel, + ) + + frame = np.zeros((120, 160, 3), np.uint8) + q: queue.Queue = queue.Queue() + pose = _pose(seed=0, n_frames=1) + q.put((0, frame, pose.xy[0] / 4, pose.confidence[0])) + q.put(END_OF_STREAM) + display = DisplayConsumer( + display_queue=q, + latest_label=LatestLabel(), + stop_event=threading.Event(), + keypoint_names=pose.keypoint_names, + display_window=False, + ) + display.run() + assert not frame.any() + + +def test_live_pipeline_refuses_a_centred_lane(sessions, fake_raw): + from glider.analysis.behavior.classify.pipeline import ( + LiveInferenceConfig, + LiveInferencePipeline, + ) + from glider.analysis.behavior.pipeline import train_model + + centred = ImageLaneConfig(seg_model="seg.pt", neck_model="neck.pt", dims=4, center=True) + model = train_model(sessions[:2], spec=SPEC, image_lane=centred, **COMMON).model + config = LiveInferenceConfig( + source="x.mp4", + keypoint_names=["snout", "left_ear", "right_ear", "neck", "tail_base"], + yolo_model_path=None, + behavior_model_path=None, + display=False, + ) + with pytest.raises(RuntimeError, match="centred"): + LiveInferencePipeline(config, model=model) + + +def test_live_pipeline_puts_the_lane_worker_between_tracker_and_engine(lane_model): + from glider.analysis.behavior.classify.pipeline import ( + LiveInferenceConfig, + LiveInferencePipeline, + ) + from glider.analysis.behavior.classify.threads import ImageLaneWorker + + config = LiveInferenceConfig( + source="x.mp4", + keypoint_names=["snout", "left_ear", "right_ear", "neck", "tail_base"], + yolo_model_path="pose.pt", + behavior_model_path=None, + display=False, + ) + pipe = LiveInferencePipeline(config, model=lane_model) + assert isinstance(pipe.lane_worker, ImageLaneWorker) + assert pipe.lane_worker.in_queue is pipe.tracked_queue + assert pipe.feature_engine.tracked_queue is pipe.lane_worker.out_queue + assert pipe.lane_worker.body_axis == (0, 4) diff --git a/tests/unit/gui/behavior/test_window.py b/tests/unit/gui/behavior/test_window.py index 761a1c84..c4925a83 100644 --- a/tests/unit/gui/behavior/test_window.py +++ b/tests/unit/gui/behavior/test_window.py @@ -1740,6 +1740,33 @@ def click_both(tab): assert not (options.get("motion_features") and options.get("mirror_augment")) +def test_image_lane_defaults_off_and_is_sent_when_ticked(qtbot, tmp_path, monkeypatch): + _tab, options = _fit_options(qtbot, tmp_path, monkeypatch) + assert options["image_lane"] is False + _tab, options = _fit_options( + qtbot, tmp_path, monkeypatch, configure=lambda t: t._image_lane_check.setChecked(True) + ) + assert options["image_lane"] is True + + +def test_image_lane_and_mirror_are_mutually_exclusive(qtbot, tmp_path, monkeypatch): + """The lane reads the unmirrored video, so train_model refuses it with mirror.""" + tab, _ = _fit_options(qtbot, tmp_path, monkeypatch) + tab._image_lane_check.setChecked(True) + assert tab._mirror_check.isEnabled() is False + tab._image_lane_check.setChecked(False) + tab._mirror_check.setChecked(True) + assert tab._image_lane_check.isEnabled() is False + assert tab._image_lane_check.isChecked() is False + + def click_both(t): + t._mirror_check.setChecked(True) + t._image_lane_check.setChecked(True) + + _tab, options = _fit_options(qtbot, tmp_path, monkeypatch, configure=click_both) + assert not (options.get("image_lane") and options.get("mirror_augment")) + + # -------------------------------------------------------------------------- # Train tab: cross-validation #