feat(lib): add MouseModel (numpy + ONNX Runtime)

src/ai_mouse/mouse.py

@@ -0,0 +1,166 @@
+"""MouseModel — ONNX Runtime-backed mouse trajectory generation."""
+from __future__ import annotations
+
+import json
+import math
+from collections.abc import Sequence
+from pathlib import Path
+from typing import Optional
+
+import numpy as np
+import onnxruntime as ort
+
+from ai_mouse._assets import resolve
+from ai_mouse._coord import decode_trajectory
+from ai_mouse._postprocess import (
+    build_timestamps,
+    enforce_forward_monotonic,
+    gaussian_smooth,
+    resample_arc,
+    sample_duration,
+    smooth_start,
+    snap_endpoints,
+    truncnorm_sample,
+)
+from ai_mouse.errors import GenerationError, ModelLoadError
+
+_N_EULER_STEPS = 10
+
+
+class MouseModel:
+    """Persistent ONNX Runtime session for mouse trajectory generation."""
+
+    def __init__(
+        self,
+        model_path: str | Path | None = None,
+        providers: Sequence[str] | None = None,
+        seed: int | None = None,
+    ) -> None:
+        path_obj: Optional[Path] = Path(model_path) if model_path is not None else None
+
+        onnx_path = resolve(path_obj, "flow_model.onnx")
+        cfg_path = resolve(path_obj, "train_config.json")
+        click_path = resolve(path_obj, "click_dist.json")
+        dur_path = resolve(path_obj, "duration_dist.json")
+
+        cfg = json.loads(cfg_path.read_text())
+        self._seq_len = int(cfg["seq_len"])
+        self._cond_dim = int(cfg.get("cond_dim", 3))
+        self._click_params = json.loads(click_path.read_text())
+        self._duration_dist = json.loads(dur_path.read_text())
+
+        try:
+            self._session = ort.InferenceSession(
+                str(onnx_path),
+                providers=list(providers) if providers else ["CPUExecutionProvider"],
+            )
+        except Exception as exc:
+            raise ModelLoadError(f"Failed to load ONNX session: {exc}") from exc
+
+        self._default_seed = seed
+        self._rng = np.random.default_rng(seed)
+
+    def generate(
+        self,
+        start: tuple[int, int],
+        end: tuple[int, int],
+        n_points: int = 64,
+        speed: float | None = None,
+        click: bool = True,
+        seed: int | None = None,
+    ) -> list[tuple[int, int, int]]:
+        rng = np.random.default_rng(seed) if seed is not None else self._rng
+
+        sx, sy = float(start[0]), float(start[1])
+        ex, ey = float(end[0]), float(end[1])
+        dist = max(math.hypot(ex - sx, ey - sy), 1.0)
+
+        total_duration = sample_duration(self._duration_dist, dist, rng)
+        if speed is not None and speed > 0:
+            total_duration /= speed
+        total_duration = max(total_duration, 10.0)
+
+        cond = np.array(
+            [
+                dist / 2000.0,
+                math.log(dist / 100.0),
+                math.log(total_duration / 500.0),
+            ],
+            dtype=np.float32,
+        )[None]
+
+        x = rng.standard_normal((1, self._seq_len, 3)).astype(np.float32)
+        dt = 1.0 / _N_EULER_STEPS
+        for step in range(_N_EULER_STEPS):
+            t = np.full((1,), step * dt, dtype=np.float32)
+            v = self._session.run(["v"], {"x_t": x, "t": t, "cond": cond})[0]
+            x = x + v * dt
+
+        if not np.all(np.isfinite(x)):
+            raise GenerationError("Trajectory contains NaN/Inf after Euler integration")
+
+        forward = x[0, :, 0].copy()
+        lateral = x[0, :, 1].copy()
+        log_dt = x[0, :, 2].copy()
+
+        forward, lateral = snap_endpoints(forward, lateral, self._seq_len)
+        forward, lateral = smooth_start(forward, lateral)
+        forward = enforce_forward_monotonic(forward)
+        lateral = gaussian_smooth(lateral, sigma=1.0)
+
+        log_dt = np.clip(log_dt, 0.0, 5.0)
+        log_dt[0] = 0.0
+
+        normalised = np.stack([forward, lateral], axis=1)
+        pixels = decode_trajectory(normalised, start, end)
+
+        if n_points != self._seq_len:
+            pixels = resample_arc(pixels, n_points)
+            log_dt = np.interp(
+                np.linspace(0, 1, n_points),
+                np.linspace(0, 1, self._seq_len),
+                log_dt,
+            )
+
+        ts = build_timestamps(log_dt, total_duration)
+
+        moves: list[tuple[int, int, int]] = [
+            (int(round(pixels[i, 0])), int(round(pixels[i, 1])), int(round(ts[i])))
+            for i in range(n_points)
+        ]
+        if not click:
+            return moves
+
+        click_dur = int(
+            truncnorm_sample(
+                float(self._click_params["mu"]),
+                float(self._click_params["sigma"]),
+                float(self._click_params["low"]),
+                float(self._click_params["high"]),
+                rng,
+            )
+        )
+        click_dur = max(click_dur, int(float(self._click_params["low"])))
+        last_t = moves[-1][2]
+        cx, cy = moves[-1][0], moves[-1][1]
+        return moves + [(cx, cy, last_t), (cx, cy, last_t + click_dur)]
+
+    def sample_click_duration_ms(self, seed: int | None = None) -> int:
+        rng = np.random.default_rng(seed) if seed is not None else self._rng
+        v = truncnorm_sample(
+            float(self._click_params["mu"]),
+            float(self._click_params["sigma"]),
+            float(self._click_params["low"]),
+            float(self._click_params["high"]),
+            rng,
+        )
+        return max(int(v), int(float(self._click_params["low"])))
+
+    def close(self) -> None:
+        self._session = None  # type: ignore[assignment]
+
+    def __enter__(self) -> "MouseModel":
+        return self
+
+    def __exit__(self, *exc) -> None:
+        self.close()

tests/unit/test_mouse.py

@@ -0,0 +1,45 @@
+"""Tests for MouseModel and ai_mouse.generate()."""
+from __future__ import annotations
+
+import pytest
+
+
+def test_mouse_model_init_default() -> None:
+    from ai_mouse.mouse import MouseModel
+    m = MouseModel()
+    assert m._seq_len > 0
+    assert m._session is not None
+    m.close()
+
+
+def test_mouse_model_generate_returns_correct_shape() -> None:
+    from ai_mouse.mouse import MouseModel
+    m = MouseModel()
+    pts = m.generate((100, 200), (900, 400))
+    assert len(pts) == 66  # 64 moves + 2 clicks
+    for x, y, t in pts:
+        assert isinstance(x, int)
+        assert isinstance(y, int)
+        assert isinstance(t, int)
+
+
+def test_mouse_model_click_false_omits_clicks() -> None:
+    from ai_mouse.mouse import MouseModel
+    m = MouseModel()
+    pts = m.generate((100, 200), (900, 400), click=False)
+    assert len(pts) == 64
+
+
+def test_mouse_model_seed_reproducibility() -> None:
+    from ai_mouse.mouse import MouseModel
+    m = MouseModel()
+    a = m.generate((100, 200), (900, 400), seed=42)
+    b = m.generate((100, 200), (900, 400), seed=42)
+    assert a == b
+
+
+def test_mouse_model_invalid_path_raises_model_load_error() -> None:
+    from ai_mouse.mouse import MouseModel
+    from ai_mouse.errors import ModelLoadError
+    with pytest.raises(ModelLoadError):
+        MouseModel(model_path="/nonexistent/path/here")