core/pose.py

"""Deep pose-estimation backend (wired, not bundled).

The intended pipeline: run a trained SLEAP / DeepLabCut model to get multi-fly
body-part keypoints per frame (head, thorax, abdomen, wings, legs), then derive
richer behaviors (wing extension, orientation, courtship) than centroid tracking.
Return the same result dict shape as `fast_track.analyze` so the viz is shared.

No redistributable fly pose model is bundled. Provide one via FLY_POSE_MODEL and
implement `_run_model` to enable this engine. Until then it raises clearly.
"""
from __future__ import annotations

import os

import numpy as np


def available() -> bool:
    path = os.environ.get("FLY_POSE_MODEL")
    return bool(path) and os.path.exists(path)


def analyze(movie: np.ndarray, fps: float = 15.0, **_) -> dict:
    if not available():
        raise RuntimeError(
            "pose engine is wired but no model is bundled. Set FLY_POSE_MODEL to a "
            "SLEAP/DeepLabCut model and implement core/pose._run_model. "
            "Use engine='fast' for the always-available centroid tracker."
        )
    return _run_model(movie, fps)  # pragma: no cover


def _run_model(movie: np.ndarray, fps: float) -> dict:  # pragma: no cover
    raise NotImplementedError("Run SLEAP/DLC inference and assemble the result dict here.")