core/processing.py

"""The ONE entrypoint a tool author edits.

`process(video_path, engine, fps, px_per_mm)` tracks flies and computes behavior,
returning a trajectory+ethogram summary image and a report. `simulate_full` also
returns the raw result + movie so the UI can build the annotated video.
"""
from __future__ import annotations

import numpy as np

from . import fast_track, viz
from .io import read_video

ENGINES = ["fast", "pose"]


def simulate_full(video_path: str, engine: str = "fast",
                  fps: float = 15.0, px_per_mm: float = 4.0) -> dict:
    if engine not in ENGINES:
        raise ValueError(f"Unknown engine '{engine}'. Choose one of {ENGINES}.")
    movie = read_video(video_path)

    if engine == "pose":
        from . import pose
        res = pose.analyze(movie, fps=float(fps))
    else:
        res = fast_track.analyze(movie, fps=float(fps), px_per_mm=float(px_per_mm))

    return {
        "summary": viz.summary_image(movie, res),
        "report": res["report"],
        "res": res,
        "movie": movie,
    }


def process(video_path: str, engine: str = "fast",
            fps: float = 15.0, px_per_mm: float = 4.0) -> tuple[np.ndarray, dict]:
    """Returns (trajectory+ethogram summary RGB, report dict)."""
    r = simulate_full(video_path, engine, fps, px_per_mm)
    return r["summary"], r["report"]