diff --git a/src/ai_mouse/mouse.py b/src/ai_mouse/mouse.py new file mode 100644 index 0000000..80629f8 --- /dev/null +++ b/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() diff --git a/tests/unit/test_mouse.py b/tests/unit/test_mouse.py new file mode 100644 index 0000000..0488af7 --- /dev/null +++ b/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")