refactor(lib): remove legacy generator.py / coord.py / scroll module
Drop the pre-migration PyTorch inference pipeline now that the ONNX-backed
MouseModel/ScrollModel in mouse.py and scroll.py are wired up through the
public ai_mouse API.
Deleted:
* src/ai_mouse/generator.py (legacy torch flow ODE + post-processing)
* src/ai_mouse/coord.py (legacy public coord transforms,
superseded by ai_mouse._coord)
* src/ai_mouse/_scroll_legacy.py (legacy torch scroll VAE inference)
* scripts/build_golden_*.py (one-shot capture scripts, no longer
needed once goldens are committed)
* tests/unit/test_generator.py (legacy module gone)
* tests/unit/test_scroll_generator.py (legacy module gone)
* tests/unit/test_coord.py (legacy module gone; ai_mouse._coord is
tested by test__coord.py)
* scripts/ (empty, removed)
Tools migrations:
* tools/trainer.py: import encode_trajectory from ai_mouse._coord
instead of the deleted ai_mouse.coord
* tools/server/routes_verify.py, tools/server/routes_scroll.py: route to
the public ai_mouse.generate / generate_scroll. They no longer accept
a model_dir override — the bundled ONNX is the source of truth, and a
fresh export goes through `python -m tools.export_onnx`.
* tools/eval/__main__.py: same migration; model_dir CLI arg retained as
a deprecation shim but ignored.
Final src/ai_mouse/ layout (matches plan):
__init__.py, _assets.py, _coord.py, _model_cache.py, _postprocess.py,
errors.py, mouse.py, py.typed, scroll.py, assets/
Test suite: 188 passed (was 188 before deletion; obsolete suites cleaned
out alongside the modules they covered).
This commit is contained in:
@@ -1,49 +0,0 @@
|
||||
"""One-shot script to capture golden mouse trajectories from the current torch
|
||||
implementation. Run BEFORE the migration so we can verify the numpy/ORT rewrite
|
||||
in Phase 4 produces equivalent output.
|
||||
|
||||
Output: tests/unit/data/golden_mouse.npz
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import random
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
# Allow running as `uv run python scripts/build_golden_mouse.py` from project root.
|
||||
sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
|
||||
from ai_mouse import generate
|
||||
|
||||
CASES: list[tuple[tuple[int, int], tuple[int, int]]] = [
|
||||
((100, 200), (900, 400)), # horizontal 800px
|
||||
((500, 500), (500, 100)), # vertical 400px upward
|
||||
((200, 600), (800, 200)), # 720px diagonal
|
||||
((100, 100), (130, 110)), # very short 31px
|
||||
((50, 50), (1500, 900)), # very long 1700px
|
||||
((400, 300), (500, 300)), # short horizontal 100px
|
||||
((300, 300), (700, 700)), # 45° diagonal
|
||||
((600, 400), (200, 100)), # reverse diagonal
|
||||
]
|
||||
SEEDS = (0, 1, 2, 3)
|
||||
|
||||
|
||||
def main() -> None:
|
||||
out: dict[str, np.ndarray] = {}
|
||||
for case_idx, (start, end) in enumerate(CASES):
|
||||
for seed in SEEDS:
|
||||
random.seed(seed)
|
||||
np.random.seed(seed)
|
||||
torch.manual_seed(seed)
|
||||
pts = generate(start=start, end=end)
|
||||
out[f"case{case_idx}_seed{seed}"] = np.array(pts, dtype=np.int64)
|
||||
out_path = Path("tests/unit/data/golden_mouse.npz")
|
||||
np.savez_compressed(out_path, **out)
|
||||
print(f"Wrote {len(out)} golden traces to {out_path}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -1,48 +0,0 @@
|
||||
"""Capture golden scroll event sequences from current torch implementation."""
|
||||
from __future__ import annotations
|
||||
|
||||
import random
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
# Allow running as `uv run python scripts/build_golden_scroll.py` from project root.
|
||||
sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
|
||||
from ai_mouse import generate_scroll
|
||||
|
||||
CASES: list[tuple[int, int, str]] = [
|
||||
(0, 1500, "target"),
|
||||
(0, 500, "precise"),
|
||||
(0, 5000, "fast"),
|
||||
(2000, 0, "target"), # upward
|
||||
(0, 800, "precise"),
|
||||
(0, 3500, "fast"),
|
||||
(1000, 1200, "precise"), # tiny scroll
|
||||
(0, 10000, "fast"), # very long
|
||||
]
|
||||
SEEDS = (0, 1, 2, 3)
|
||||
|
||||
|
||||
def main() -> None:
|
||||
out: dict[str, np.ndarray] = {}
|
||||
for case_idx, (start_y, end_y, mode) in enumerate(CASES):
|
||||
for seed in SEEDS:
|
||||
random.seed(seed)
|
||||
np.random.seed(seed)
|
||||
torch.manual_seed(seed)
|
||||
events = generate_scroll(start_y, end_y, mode=mode)
|
||||
arr = np.array(
|
||||
[[e["deltaY"], e["deltaMode"], e["t"]] for e in events],
|
||||
dtype=np.int64,
|
||||
)
|
||||
out[f"case{case_idx}_seed{seed}"] = arr
|
||||
out_path = Path("tests/unit/data/golden_scroll.npz")
|
||||
np.savez_compressed(out_path, **out)
|
||||
print(f"Wrote {len(out)} scroll golden traces to {out_path}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -1,148 +0,0 @@
|
||||
"""Scroll wheel event sequence generator."""
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import logging
|
||||
import math
|
||||
from pathlib import Path
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
|
||||
from tools.scroll.models import ScrollCVAE
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
_BUNDLED_SCROLL_MODELS = Path(__file__).resolve().parent.parent.parent / "data" / "scroll_models"
|
||||
|
||||
|
||||
def _build_condition(
|
||||
distance: float,
|
||||
direction: int,
|
||||
mode: str,
|
||||
viewport_height: float = 900.0,
|
||||
) -> np.ndarray:
|
||||
"""Build 7-dim condition vector matching the trainer layout.
|
||||
|
||||
Dims: [dist/5000, log(dist/500), direction, viewport_norm, mode_onehot*3]
|
||||
"""
|
||||
mode_onehot = [0.0, 0.0, 0.0]
|
||||
if mode == "target":
|
||||
mode_onehot[0] = 1.0
|
||||
elif mode == "fast":
|
||||
mode_onehot[1] = 1.0
|
||||
elif mode == "precise":
|
||||
mode_onehot[2] = 1.0
|
||||
|
||||
viewport_norm = viewport_height / 1000.0
|
||||
|
||||
return np.array([
|
||||
distance / 5000.0,
|
||||
math.log(max(distance, 1.0) / 500.0),
|
||||
float(direction),
|
||||
viewport_norm,
|
||||
*mode_onehot,
|
||||
], dtype=np.float32)
|
||||
|
||||
|
||||
def generate_scroll(
|
||||
start_scrollY: int,
|
||||
target_scrollY: int,
|
||||
mode: str = "target",
|
||||
model_dir: str | None = None,
|
||||
) -> list[dict]:
|
||||
"""Generate a realistic scroll event sequence.
|
||||
|
||||
Args:
|
||||
start_scrollY: Current scroll position (px from top).
|
||||
target_scrollY: Target scroll position.
|
||||
mode: "target" | "fast" | "precise"
|
||||
model_dir: Path to scroll model files. None = bundled.
|
||||
|
||||
Returns:
|
||||
List of {"deltaY": int, "deltaMode": 0, "t": int}.
|
||||
Positive deltaY = scroll down, negative = scroll up.
|
||||
"""
|
||||
model_dir_path = Path(model_dir) if model_dir else _BUNDLED_SCROLL_MODELS
|
||||
model_pt = model_dir_path / "scroll_model.pt"
|
||||
config_json = model_dir_path / "scroll_config.json"
|
||||
|
||||
if not model_pt.exists():
|
||||
raise FileNotFoundError(f"Scroll model not found at {model_pt}")
|
||||
|
||||
seq_len = 32
|
||||
if config_json.exists():
|
||||
cfg = json.loads(config_json.read_text())
|
||||
seq_len = cfg.get("seq_len", 32)
|
||||
|
||||
model = ScrollCVAE(seq_len=seq_len)
|
||||
model.load_state_dict(torch.load(model_pt, map_location="cpu", weights_only=True))
|
||||
model.eval()
|
||||
|
||||
distance = abs(target_scrollY - start_scrollY)
|
||||
direction = 1 if target_scrollY > start_scrollY else -1
|
||||
distance = max(distance, 10)
|
||||
|
||||
cond = _build_condition(float(distance), direction, mode)
|
||||
cond_t = torch.from_numpy(cond).unsqueeze(0)
|
||||
|
||||
with torch.no_grad():
|
||||
z = torch.randn(1, model.latent_dim)
|
||||
decoded = model.decode(z, cond_t).squeeze(0).numpy()
|
||||
|
||||
delta_norm = decoded[:, 0]
|
||||
log_dt = decoded[:, 1]
|
||||
|
||||
# De-normalise delta: use softmax-like normalisation so they sum to ~1
|
||||
delta_weights = np.exp(delta_norm)
|
||||
delta_weights = delta_weights / delta_weights.sum()
|
||||
delta_px = delta_weights * distance * direction
|
||||
|
||||
# Quantise to realistic wheel increments
|
||||
quantum = 40 if mode == "precise" else 120
|
||||
|
||||
delta_quantised = np.round(delta_px / quantum) * quantum
|
||||
for i in range(len(delta_quantised)):
|
||||
if delta_quantised[i] == 0:
|
||||
delta_quantised[i] = quantum * direction
|
||||
|
||||
# Adjust last event so total matches target distance
|
||||
current_total = delta_quantised.sum()
|
||||
diff = (distance * direction) - current_total
|
||||
delta_quantised[-1] += diff
|
||||
|
||||
# Timestamps from log_dt
|
||||
if len(log_dt) > 3:
|
||||
median_log = float(np.median(log_dt))
|
||||
log_dt[:2] = np.clip(log_dt[:2], None, median_log + 0.5)
|
||||
log_dt[-2:] = np.clip(log_dt[-2:], None, median_log + 0.5)
|
||||
|
||||
dt_ms = np.exp(log_dt).clip(5, 80)
|
||||
|
||||
# Scale to realistic total duration
|
||||
if mode == "fast":
|
||||
expected_duration = distance * 0.2 + 100
|
||||
elif mode == "precise":
|
||||
expected_duration = distance * 1.5 + 300
|
||||
else:
|
||||
expected_duration = distance * 0.4 + 200
|
||||
|
||||
dt_ms = dt_ms * (expected_duration / max(dt_ms.sum(), 1.0))
|
||||
dt_ms = dt_ms.clip(5, 80)
|
||||
|
||||
t_abs = np.cumsum(dt_ms).astype(int)
|
||||
t_abs = np.concatenate([[0], t_abs[:-1]])
|
||||
|
||||
# Ensure monotonic
|
||||
for i in range(1, len(t_abs)):
|
||||
if t_abs[i] <= t_abs[i - 1]:
|
||||
t_abs[i] = t_abs[i - 1] + 5
|
||||
|
||||
# Build events, removing zero-delta (keep at least 5)
|
||||
events = []
|
||||
for i in range(seq_len):
|
||||
dy = int(delta_quantised[i])
|
||||
if dy != 0 or len(events) < 5:
|
||||
events.append({"deltaY": dy, "deltaMode": 0, "t": int(t_abs[i])})
|
||||
|
||||
return events
|
||||
@@ -1,81 +0,0 @@
|
||||
"""Rotated coordinate system for angle-invariant trajectory encoding.
|
||||
|
||||
All trajectories are normalised into a frame where:
|
||||
- start → (0, 0)
|
||||
- end → (1, 0)
|
||||
- lateral displacement is perpendicular to start→end axis
|
||||
|
||||
This makes the model angle-invariant: a 45° diagonal move and a horizontal
|
||||
move look identical in the rotated frame (just "forward from 0 to 1").
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import math
|
||||
|
||||
import numpy as np
|
||||
|
||||
|
||||
def encode_trajectory(
|
||||
points: np.ndarray,
|
||||
start: tuple[int, int],
|
||||
end: tuple[int, int],
|
||||
) -> np.ndarray:
|
||||
"""Transform pixel coordinates to rotated normalised frame.
|
||||
|
||||
Args:
|
||||
points: (N, 2) array of (x, y) pixel coordinates.
|
||||
start: (x, y) start position.
|
||||
end: (x, y) end position.
|
||||
|
||||
Returns:
|
||||
(N, 2) array of (forward, lateral) in normalised rotated frame.
|
||||
"""
|
||||
sx, sy = float(start[0]), float(start[1])
|
||||
ex, ey = float(end[0]), float(end[1])
|
||||
dist = math.hypot(ex - sx, ey - sy)
|
||||
if dist < 1e-8:
|
||||
return np.zeros_like(points)
|
||||
|
||||
ux, uy = (ex - sx) / dist, (ey - sy) / dist
|
||||
vx, vy = -uy, ux
|
||||
|
||||
dx = points[:, 0] - sx
|
||||
dy = points[:, 1] - sy
|
||||
|
||||
forward = (dx * ux + dy * uy) / dist
|
||||
lateral = (dx * vx + dy * vy) / dist
|
||||
|
||||
return np.stack([forward, lateral], axis=1)
|
||||
|
||||
|
||||
def decode_trajectory(
|
||||
normalised: np.ndarray,
|
||||
start: tuple[int, int],
|
||||
end: tuple[int, int],
|
||||
) -> np.ndarray:
|
||||
"""Transform rotated normalised frame back to pixel coordinates.
|
||||
|
||||
Args:
|
||||
normalised: (N, 2) array of (forward, lateral).
|
||||
start: (x, y) start position.
|
||||
end: (x, y) end position.
|
||||
|
||||
Returns:
|
||||
(N, 2) array of (x, y) pixel coordinates.
|
||||
"""
|
||||
sx, sy = float(start[0]), float(start[1])
|
||||
ex, ey = float(end[0]), float(end[1])
|
||||
dist = math.hypot(ex - sx, ey - sy)
|
||||
if dist < 1e-8:
|
||||
return np.full_like(normalised, [sx, sy])
|
||||
|
||||
ux, uy = (ex - sx) / dist, (ey - sy) / dist
|
||||
vx, vy = -uy, ux
|
||||
|
||||
forward = normalised[:, 0]
|
||||
lateral = normalised[:, 1]
|
||||
|
||||
px = sx + forward * dist * ux + lateral * dist * vx
|
||||
py = sy + forward * dist * uy + lateral * dist * vy
|
||||
|
||||
return np.stack([px, py], axis=1)
|
||||
@@ -1,353 +0,0 @@
|
||||
"""Inference layer: Flow Matching trajectory generation.
|
||||
|
||||
Pipeline:
|
||||
1. Load model from model_dir (flow_model.pt, click_dist.json,
|
||||
duration_dist.json, train_config.json).
|
||||
2. Compute condition vector: [dist/2000, log(dist/100), log(total_dur/500)].
|
||||
3. Sample total_duration from duration_dist.json by distance bin (log-normal).
|
||||
4. 10-step Euler ODE: start from noise, integrate velocity field to get trajectory.
|
||||
5. Spatial post-processing:
|
||||
a. Endpoint snapping: force first=(0,0), last=(1,0), lerp last 6 points.
|
||||
b. Smooth start: dampen lateral near start (first 4 points).
|
||||
c. Enforce forward monotonicity (prevent x-axis jitter).
|
||||
d. 5-point gaussian smooth on lateral (preserve endpoints).
|
||||
6. Temporal post-processing:
|
||||
a. Clip log_dt to [0, 5] to prevent exponential explosion.
|
||||
(speed profile and median±1.1 hard clip removed in 2026-05 refactor —
|
||||
let the model's learned timing distribution come through naturally.)
|
||||
7. Decode to pixels via decode_trajectory.
|
||||
8. Resample to n_points if n_points != model seq_len.
|
||||
9. Convert log_dt → ms timestamps, scale to total_duration, clip [2, 150].
|
||||
10. Ensure timestamps monotonically increasing.
|
||||
11. Append click events sampled from truncated normal.
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import logging
|
||||
import math
|
||||
from pathlib import Path
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
from scipy.stats import truncnorm
|
||||
|
||||
from ai_mouse.coord import decode_trajectory
|
||||
from tools.config import GenerateConfig
|
||||
from tools.models import TrajectoryFlowModel
|
||||
from tools.utils import resample_arc
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
_BUNDLED_MODELS_DIR = Path(__file__).parent.parent / "data" / "models_v2"
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Duration sampling helper
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _sample_duration(duration_dist: dict, dist: float) -> float:
|
||||
"""Sample a total movement duration (ms) for the given pixel distance.
|
||||
|
||||
Uses per-distance-bin log-normal parameters from duration_dist.
|
||||
|
||||
Args:
|
||||
duration_dist: dict with "bins" and "params" keys.
|
||||
dist: pixel distance between start and end.
|
||||
|
||||
Returns:
|
||||
Sampled duration in milliseconds.
|
||||
"""
|
||||
bins = duration_dist["bins"]
|
||||
params = duration_dist["params"]
|
||||
# Find bin for this distance
|
||||
bin_idx = len(bins) - 1
|
||||
for i in range(len(bins) - 1):
|
||||
if dist < bins[i + 1]:
|
||||
bin_idx = i
|
||||
break
|
||||
# Clamp to valid params index
|
||||
bin_idx = min(bin_idx, len(params) - 1)
|
||||
mu_log = params[bin_idx]["mu_log"]
|
||||
sigma_log = params[bin_idx]["sigma_log"]
|
||||
return float(np.exp(np.random.normal(mu_log, sigma_log)))
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Smoothing helper
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _gaussian_smooth(x: np.ndarray, sigma: float = 1.0) -> np.ndarray:
|
||||
"""5-point gaussian smoothing along a 1-D array, preserving endpoints.
|
||||
|
||||
Args:
|
||||
x: 1-D input array.
|
||||
sigma: Gaussian std (px); larger = more smoothing. Default 1.0 gives
|
||||
weights ≈ [0.054, 0.244, 0.403, 0.244, 0.054].
|
||||
|
||||
Returns:
|
||||
Smoothed array of the same shape. x[0] and x[-1] are unchanged.
|
||||
If len(x) < 5, returns x unchanged (kernel won't fit).
|
||||
"""
|
||||
if len(x) < 5:
|
||||
return x.copy()
|
||||
kernel = np.exp(-0.5 * (np.arange(-2, 3) / sigma) ** 2)
|
||||
kernel /= kernel.sum()
|
||||
# Pad with edge values to avoid boundary artifacts, then slice back
|
||||
padded = np.pad(x, pad_width=2, mode="edge")
|
||||
smoothed = np.convolve(padded, kernel, mode="valid")
|
||||
smoothed[0] = x[0]
|
||||
smoothed[-1] = x[-1]
|
||||
return smoothed
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Main generate function
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def generate(
|
||||
start: tuple[int, int],
|
||||
end: tuple[int, int],
|
||||
n_points: int = 64,
|
||||
speed: float | None = None,
|
||||
model_dir: str | None = None,
|
||||
config: GenerateConfig | None = None,
|
||||
) -> list[tuple[int, int, int]]:
|
||||
"""Generate a human-like mouse trajectory from start to end.
|
||||
|
||||
Uses a Flow Matching model with 4-step Euler ODE integration.
|
||||
|
||||
Args:
|
||||
start: (x, y) starting pixel coordinate.
|
||||
end: (x, y) target pixel coordinate.
|
||||
n_points: number of movement points in the path (default 64).
|
||||
speed: optional speed multiplier; speed=2 halves the duration.
|
||||
model_dir: directory containing flow_model.pt, click_dist.json,
|
||||
duration_dist.json, train_config.json.
|
||||
None → use bundled pre-trained weights.
|
||||
config: GenerateConfig instance; None → use defaults.
|
||||
|
||||
Returns:
|
||||
List of (x, y, t_ms) tuples. All values are ints.
|
||||
Last two entries are the mouse-down and mouse-up click events.
|
||||
"""
|
||||
if config is None:
|
||||
config = GenerateConfig()
|
||||
|
||||
model_dir_path = Path(model_dir) if model_dir else _BUNDLED_MODELS_DIR
|
||||
|
||||
flow_pt = model_dir_path / "flow_model.pt"
|
||||
click_json = model_dir_path / "click_dist.json"
|
||||
duration_json = model_dir_path / "duration_dist.json"
|
||||
config_json = model_dir_path / "train_config.json"
|
||||
|
||||
if not flow_pt.exists():
|
||||
if model_dir is not None:
|
||||
raise FileNotFoundError(
|
||||
f"Model weights not found in {model_dir_path}. "
|
||||
"Run training first or omit model_dir to use bundled weights."
|
||||
)
|
||||
raise FileNotFoundError(
|
||||
f"Bundled model weights missing at {_BUNDLED_MODELS_DIR}. "
|
||||
"Run training first."
|
||||
)
|
||||
|
||||
# Load train config for model architecture params
|
||||
seq_len = config.seq_len
|
||||
d_model = 128
|
||||
nhead = 4
|
||||
num_layers = 4
|
||||
dim_feedforward = 256
|
||||
cond_dim = 3
|
||||
if config_json.exists():
|
||||
cfg = json.loads(config_json.read_text())
|
||||
seq_len = int(cfg.get("seq_len", seq_len))
|
||||
d_model = int(cfg.get("d_model", d_model))
|
||||
nhead = int(cfg.get("nhead", nhead))
|
||||
num_layers = int(cfg.get("num_layers", num_layers))
|
||||
dim_feedforward = int(cfg.get("dim_feedforward", dim_feedforward))
|
||||
cond_dim = int(cfg.get("cond_dim", cond_dim))
|
||||
|
||||
# Load model
|
||||
model = TrajectoryFlowModel(
|
||||
seq_len=seq_len,
|
||||
d_model=d_model,
|
||||
nhead=nhead,
|
||||
num_layers=num_layers,
|
||||
dim_feedforward=dim_feedforward,
|
||||
cond_dim=cond_dim,
|
||||
)
|
||||
model.load_state_dict(
|
||||
torch.load(flow_pt, map_location="cpu", weights_only=True)
|
||||
)
|
||||
model.eval()
|
||||
|
||||
# Load auxiliary distributions
|
||||
click_params: dict = {"mu": 80.0, "sigma": 30.0, "low": 20.0, "high": 300.0}
|
||||
if click_json.exists():
|
||||
click_params = json.loads(click_json.read_text())
|
||||
|
||||
duration_dist: dict | None = None
|
||||
if duration_json.exists():
|
||||
duration_dist = json.loads(duration_json.read_text())
|
||||
|
||||
# Compute pixel distance
|
||||
sx, sy = float(start[0]), float(start[1])
|
||||
ex, ey = float(end[0]), float(end[1])
|
||||
dist = math.hypot(ex - sx, ey - sy)
|
||||
dist = max(dist, 1.0)
|
||||
|
||||
# Sample total duration
|
||||
if duration_dist is not None:
|
||||
total_duration = _sample_duration(duration_dist, dist)
|
||||
else:
|
||||
# Fallback: simple heuristic ~2px/ms
|
||||
total_duration = dist / 2.0
|
||||
if speed is not None and speed > 0:
|
||||
total_duration /= speed
|
||||
total_duration = max(total_duration, 10.0)
|
||||
|
||||
# Build condition vector: [dist_norm, log_dist, log_total_dur]
|
||||
cond_arr = np.array(
|
||||
[
|
||||
dist / 2000.0,
|
||||
math.log(dist / 100.0),
|
||||
math.log(total_duration / 500.0),
|
||||
],
|
||||
dtype=np.float32,
|
||||
)
|
||||
cond_t = torch.from_numpy(cond_arr).unsqueeze(0) # (1, 3)
|
||||
|
||||
# -----------------------------------------------------------------------
|
||||
# 4-step Euler ODE integration
|
||||
# -----------------------------------------------------------------------
|
||||
n_steps = config.n_steps
|
||||
dt = 1.0 / n_steps
|
||||
|
||||
with torch.no_grad():
|
||||
x = torch.randn(1, seq_len, 3) # start from noise
|
||||
for step in range(n_steps):
|
||||
t_val = step * dt
|
||||
t_tensor = torch.tensor([t_val])
|
||||
v = model(x, t_tensor, cond_t)
|
||||
x = x + v * dt
|
||||
# x is now the generated trajectory in (forward, lateral, log_dt) space
|
||||
|
||||
decoded = x.squeeze(0).numpy() # (seq_len, 3)
|
||||
|
||||
forward = decoded[:, 0].copy() # (seq_len,)
|
||||
lateral = decoded[:, 1].copy() # (seq_len,)
|
||||
log_dt = decoded[:, 2].copy() # (seq_len,)
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Spatial post-processing
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
# Endpoint snapping: lerp last 6 points towards (1.0, 0.0)
|
||||
n_snap = min(6, seq_len // 4)
|
||||
for i in range(n_snap):
|
||||
alpha = ((i + 1) / n_snap) ** 2 # quadratic ease-in
|
||||
k = seq_len - n_snap + i
|
||||
forward[k] = forward[k] * (1.0 - alpha) + 1.0 * alpha
|
||||
lateral[k] = lateral[k] * (1.0 - alpha) + 0.0 * alpha
|
||||
|
||||
# Force first and last points to canonical values
|
||||
forward[0], lateral[0] = 0.0, 0.0
|
||||
forward[-1], lateral[-1] = 1.0, 0.0
|
||||
|
||||
# Smooth start: dampen lateral near start (first 4 points)
|
||||
n_start_fix = min(4, seq_len // 4)
|
||||
for i in range(1, n_start_fix + 1):
|
||||
blend = i / (n_start_fix + 1) # 0.2, 0.4, 0.6, 0.8
|
||||
forward[i] = max(forward[i], forward[i - 1]) # ensure monotonic start
|
||||
lateral[i] = lateral[i] * blend # dampen lateral near start
|
||||
|
||||
# Enforce forward monotonicity with soft correction (prevent x-jitter)
|
||||
for i in range(1, seq_len - 1): # skip last point (already snapped to 1.0)
|
||||
if forward[i] < forward[i - 1]:
|
||||
forward[i] = forward[i - 1] + 0.001
|
||||
|
||||
# Clamp forward to [0, 1] and re-force endpoints after monotonicity fix
|
||||
forward = np.clip(forward, 0.0, 1.0)
|
||||
forward[0] = 0.0
|
||||
forward[-1] = 1.0
|
||||
|
||||
# Lateral 5-point gaussian smoothing (endpoints preserved)
|
||||
lateral = _gaussian_smooth(lateral, sigma=1.0)
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Temporal post-processing (log_dt)
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
# Clip log_dt to prevent extreme values after exp()
|
||||
# Training data log_dt = log(Δt_ms + 1), typical range [0, 4.5]
|
||||
# (e.g., Δt=1ms → 0.69, Δt=10ms → 2.40, Δt=80ms → 4.39)
|
||||
log_dt = np.clip(log_dt, 0.0, 5.0)
|
||||
|
||||
# First point has no interval (padding from training)
|
||||
log_dt[0] = 0.0
|
||||
|
||||
# Decode spatial coordinates to pixels
|
||||
normalised = np.stack([forward, lateral], axis=1) # (seq_len, 2)
|
||||
pixels = decode_trajectory(normalised, start, end) # (seq_len, 2)
|
||||
|
||||
# Resample to n_points if needed
|
||||
if n_points != seq_len:
|
||||
pixels = resample_arc(pixels, n_points)
|
||||
# Also resample log_dt via linear interpolation in uniform arc
|
||||
log_dt = np.interp(
|
||||
np.linspace(0, 1, n_points),
|
||||
np.linspace(0, 1, seq_len),
|
||||
log_dt,
|
||||
)
|
||||
|
||||
xs = pixels[:, 0]
|
||||
ys = pixels[:, 1]
|
||||
|
||||
# Convert log_dt → dt (ms), scale to total_duration, clip
|
||||
dt_raw = np.exp(log_dt)
|
||||
dt_raw = np.clip(dt_raw, 0.0, None)
|
||||
dt_sum = dt_raw.sum()
|
||||
if dt_sum > 1e-6:
|
||||
scale = total_duration / dt_sum
|
||||
else:
|
||||
scale = total_duration / max(n_points, 1)
|
||||
dt_ms = np.clip(
|
||||
dt_raw * scale,
|
||||
config.dt_clip_min_ms,
|
||||
config.dt_clip_max_ms,
|
||||
)
|
||||
|
||||
# Cumulative timestamps (start at 0)
|
||||
t_abs = np.cumsum(dt_ms)
|
||||
t_abs = np.concatenate([[0.0], t_abs[:-1]]) # shift so first point = 0
|
||||
|
||||
# Ensure monotonically increasing
|
||||
for i in range(1, len(t_abs)):
|
||||
if t_abs[i] <= t_abs[i - 1]:
|
||||
t_abs[i] = t_abs[i - 1] + 1.0
|
||||
|
||||
move_points: list[tuple[int, int, int]] = [
|
||||
(int(round(xs[i])), int(round(ys[i])), int(round(t_abs[i])))
|
||||
for i in range(n_points)
|
||||
]
|
||||
|
||||
# Sample click duration from truncated normal
|
||||
mu = float(click_params["mu"])
|
||||
sigma = float(click_params["sigma"])
|
||||
low = float(click_params["low"])
|
||||
high = float(click_params["high"])
|
||||
a, b = (low - mu) / sigma, (high - mu) / sigma
|
||||
click_duration = int(truncnorm.rvs(a, b, loc=mu, scale=sigma))
|
||||
click_duration = max(click_duration, int(low))
|
||||
|
||||
last_t = move_points[-1][2]
|
||||
click_x = int(round(xs[-1]))
|
||||
click_y = int(round(ys[-1]))
|
||||
return move_points + [
|
||||
(click_x, click_y, last_t),
|
||||
(click_x, click_y, last_t + click_duration),
|
||||
]
|
||||
@@ -1,113 +0,0 @@
|
||||
"""Tests for rotated coordinate system transforms."""
|
||||
from __future__ import annotations
|
||||
|
||||
import math
|
||||
|
||||
import numpy as np
|
||||
import pytest
|
||||
|
||||
from ai_mouse.coord import encode_trajectory, decode_trajectory
|
||||
|
||||
|
||||
class TestEncodeTrajectory:
|
||||
"""Test pixel → rotated normalised frame."""
|
||||
|
||||
def test_start_maps_to_origin(self):
|
||||
start = (100, 200)
|
||||
end = (400, 500)
|
||||
points = np.array([[100, 200]], dtype=float)
|
||||
result = encode_trajectory(points, start, end)
|
||||
np.testing.assert_allclose(result[0], [0.0, 0.0], atol=1e-10)
|
||||
|
||||
def test_end_maps_to_one_zero(self):
|
||||
start = (100, 200)
|
||||
end = (400, 500)
|
||||
points = np.array([[400, 500]], dtype=float)
|
||||
result = encode_trajectory(points, start, end)
|
||||
np.testing.assert_allclose(result[0], [1.0, 0.0], atol=1e-10)
|
||||
|
||||
def test_midpoint_maps_to_half_zero(self):
|
||||
start = (0, 0)
|
||||
end = (200, 0)
|
||||
points = np.array([[100, 0]], dtype=float)
|
||||
result = encode_trajectory(points, start, end)
|
||||
np.testing.assert_allclose(result[0], [0.5, 0.0], atol=1e-10)
|
||||
|
||||
def test_lateral_offset_positive(self):
|
||||
"""Point at (100, 50) with horizontal start→end has lateral = 50/200 = 0.25."""
|
||||
start = (0, 0)
|
||||
end = (200, 0)
|
||||
# For horizontal u=(1,0), v=(-0, 1)=(0,1).
|
||||
# Point (100, 50): forward = 100/200=0.5, lateral = 50/200=0.25
|
||||
points = np.array([[100, 50]], dtype=float)
|
||||
result = encode_trajectory(points, start, end)
|
||||
np.testing.assert_allclose(result[0], [0.5, 0.25], atol=1e-10)
|
||||
|
||||
def test_various_angles(self):
|
||||
"""Encode/decode round-trip works for various angles."""
|
||||
angles = [0, 45, 90, 135, 180, -45, -90, -135]
|
||||
for deg in angles:
|
||||
rad = math.radians(deg)
|
||||
start = (400, 300)
|
||||
dist = 200
|
||||
end = (int(400 + dist * math.cos(rad)), int(300 + dist * math.sin(rad)))
|
||||
# Create a curved path
|
||||
t = np.linspace(0, 1, 20)
|
||||
px = start[0] + t * (end[0] - start[0]) + 20 * np.sin(t * math.pi)
|
||||
py = start[1] + t * (end[1] - start[1]) + 20 * np.cos(t * math.pi)
|
||||
points = np.stack([px, py], axis=1)
|
||||
|
||||
encoded = encode_trajectory(points, start, end)
|
||||
assert encoded[0, 0] == pytest.approx(0.0, abs=0.2)
|
||||
assert encoded[-1, 0] == pytest.approx(1.0, abs=0.2)
|
||||
|
||||
|
||||
class TestDecodeTrajectory:
|
||||
"""Test rotated normalised frame → pixel."""
|
||||
|
||||
def test_origin_maps_to_start(self):
|
||||
start = (100, 200)
|
||||
end = (400, 500)
|
||||
normalised = np.array([[0.0, 0.0]], dtype=float)
|
||||
result = decode_trajectory(normalised, start, end)
|
||||
np.testing.assert_allclose(result[0], [100, 200], atol=1e-10)
|
||||
|
||||
def test_one_zero_maps_to_end(self):
|
||||
start = (100, 200)
|
||||
end = (400, 500)
|
||||
normalised = np.array([[1.0, 0.0]], dtype=float)
|
||||
result = decode_trajectory(normalised, start, end)
|
||||
np.testing.assert_allclose(result[0], [400, 500], atol=1e-10)
|
||||
|
||||
|
||||
class TestRoundTrip:
|
||||
"""Encode then decode should return original points."""
|
||||
|
||||
def test_round_trip_horizontal(self):
|
||||
start = (50, 100)
|
||||
end = (350, 100)
|
||||
points = np.array([[50, 100], [150, 130], [250, 90], [350, 100]], dtype=float)
|
||||
encoded = encode_trajectory(points, start, end)
|
||||
decoded = decode_trajectory(encoded, start, end)
|
||||
np.testing.assert_allclose(decoded, points, atol=1e-8)
|
||||
|
||||
def test_round_trip_diagonal(self):
|
||||
start = (100, 100)
|
||||
end = (500, 400)
|
||||
rng = np.random.default_rng(42)
|
||||
points = np.column_stack([
|
||||
np.linspace(100, 500, 30) + rng.normal(0, 10, 30),
|
||||
np.linspace(100, 400, 30) + rng.normal(0, 10, 30),
|
||||
])
|
||||
encoded = encode_trajectory(points, start, end)
|
||||
decoded = decode_trajectory(encoded, start, end)
|
||||
np.testing.assert_allclose(decoded, points, atol=1e-8)
|
||||
|
||||
def test_round_trip_vertical(self):
|
||||
"""Vertical movement (angle=90°) doesn't collapse."""
|
||||
start = (300, 50)
|
||||
end = (300, 450)
|
||||
points = np.array([[300, 50], [310, 200], [295, 350], [300, 450]], dtype=float)
|
||||
encoded = encode_trajectory(points, start, end)
|
||||
decoded = decode_trajectory(encoded, start, end)
|
||||
np.testing.assert_allclose(decoded, points, atol=1e-8)
|
||||
@@ -1,158 +0,0 @@
|
||||
"""Tests for Flow Matching trajectory generator."""
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
from pathlib import Path
|
||||
|
||||
import numpy as np
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from ai_mouse.generator import generate
|
||||
from tools.models import TrajectoryFlowModel
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def model_dir(tmp_path):
|
||||
"""Create temp dir with Flow model artifacts."""
|
||||
model = TrajectoryFlowModel(seq_len=64)
|
||||
torch.save(model.state_dict(), tmp_path / "flow_model.pt")
|
||||
|
||||
click_dist = {"mu": 80.0, "sigma": 30.0, "low": 20.0, "high": 300.0}
|
||||
(tmp_path / "click_dist.json").write_text(json.dumps(click_dist))
|
||||
|
||||
duration_dist = {
|
||||
"bins": [0, 50, 100, 200, 400, 600, 800, 1200, float("inf")],
|
||||
"params": [
|
||||
{"mu_log": 5.5, "sigma_log": 0.3},
|
||||
{"mu_log": 5.8, "sigma_log": 0.3},
|
||||
{"mu_log": 6.0, "sigma_log": 0.3},
|
||||
{"mu_log": 6.2, "sigma_log": 0.3},
|
||||
{"mu_log": 6.5, "sigma_log": 0.3},
|
||||
{"mu_log": 6.7, "sigma_log": 0.3},
|
||||
{"mu_log": 6.9, "sigma_log": 0.3},
|
||||
{"mu_log": 7.0, "sigma_log": 0.3},
|
||||
],
|
||||
}
|
||||
(tmp_path / "duration_dist.json").write_text(json.dumps(duration_dist))
|
||||
|
||||
train_config = {
|
||||
"seq_len": 64,
|
||||
"d_model": 128,
|
||||
"nhead": 4,
|
||||
"num_layers": 4,
|
||||
"dim_feedforward": 256,
|
||||
"cond_dim": 3,
|
||||
}
|
||||
(tmp_path / "train_config.json").write_text(json.dumps(train_config))
|
||||
|
||||
return tmp_path
|
||||
|
||||
|
||||
class TestGenerate:
|
||||
def test_returns_list_of_tuples(self, model_dir):
|
||||
result = generate(start=(100, 200), end=(500, 400), model_dir=str(model_dir))
|
||||
assert isinstance(result, list)
|
||||
assert all(isinstance(p, tuple) and len(p) == 3 for p in result)
|
||||
# All elements are ints
|
||||
for p in result:
|
||||
assert all(isinstance(v, int) for v in p)
|
||||
|
||||
def test_timestamps_monotonically_increasing(self, model_dir):
|
||||
result = generate(start=(100, 200), end=(500, 400), model_dir=str(model_dir))
|
||||
times = [p[2] for p in result]
|
||||
for i in range(1, len(times)):
|
||||
assert times[i] >= times[i - 1]
|
||||
|
||||
def test_starts_near_start(self, model_dir):
|
||||
start = (100, 200)
|
||||
result = generate(start=start, end=(500, 400), model_dir=str(model_dir))
|
||||
first = result[0]
|
||||
assert abs(first[0] - start[0]) < 30
|
||||
assert abs(first[1] - start[1]) < 30
|
||||
|
||||
def test_ends_near_end(self, model_dir):
|
||||
end = (500, 400)
|
||||
result = generate(start=(100, 200), end=end, model_dir=str(model_dir))
|
||||
# Last two are click events; the one before is last movement point
|
||||
last_move = result[-3]
|
||||
assert abs(last_move[0] - end[0]) < 30
|
||||
assert abs(last_move[1] - end[1]) < 30
|
||||
|
||||
def test_last_two_are_click_events(self, model_dir):
|
||||
result = generate(start=(100, 200), end=(500, 400), model_dir=str(model_dir))
|
||||
down = result[-2]
|
||||
up = result[-1]
|
||||
# Same x, y for click down and up
|
||||
assert down[0] == up[0]
|
||||
assert down[1] == up[1]
|
||||
# Up timestamp > down timestamp
|
||||
assert up[2] > down[2]
|
||||
# Click duration within bounds
|
||||
assert 20 <= up[2] - down[2] <= 300
|
||||
|
||||
def test_different_z_gives_different_paths(self, model_dir):
|
||||
r1 = generate(start=(100, 200), end=(500, 400), model_dir=str(model_dir))
|
||||
r2 = generate(start=(100, 200), end=(500, 400), model_dir=str(model_dir))
|
||||
points1 = [(p[0], p[1]) for p in r1[:-2]]
|
||||
points2 = [(p[0], p[1]) for p in r2[:-2]]
|
||||
assert points1 != points2
|
||||
|
||||
def test_n_points_parameter(self, model_dir):
|
||||
result = generate(
|
||||
start=(100, 200), end=(500, 400), n_points=32, model_dir=str(model_dir)
|
||||
)
|
||||
# 32 move points + 2 click events = 34
|
||||
assert len(result) == 34
|
||||
|
||||
|
||||
class TestPostProcessing:
|
||||
def test_dt_diversity_preserved(self, model_dir):
|
||||
"""After removing speed_profile + median clip, multiple generations
|
||||
should differ in their Δt sequences (not all identical)."""
|
||||
results = [generate(start=(100, 200), end=(500, 400), model_dir=str(model_dir))
|
||||
for _ in range(5)]
|
||||
# Extract Δt sequences (only move events, not click events)
|
||||
dts = []
|
||||
for r in results:
|
||||
moves = r[:-2]
|
||||
dt_seq = [moves[i+1][2] - moves[i][2] for i in range(len(moves)-1)]
|
||||
dts.append(dt_seq)
|
||||
# At least 2 of the 5 sequences should differ at any given index
|
||||
for i in range(min(len(d) for d in dts)):
|
||||
values = {tuple([d[i]]) for d in dts}
|
||||
if len(values) > 1:
|
||||
return # at least one position has variation — pass
|
||||
pytest.fail("All 5 Δt sequences are identical at every position — diversity collapsed")
|
||||
|
||||
|
||||
class TestGaussianSmooth:
|
||||
def test_endpoints_preserved(self):
|
||||
from ai_mouse.generator import _gaussian_smooth
|
||||
x = np.array([1.0, 5.0, 3.0, 7.0, 2.0], dtype=np.float64)
|
||||
smoothed = _gaussian_smooth(x, sigma=1.0)
|
||||
assert smoothed[0] == 1.0
|
||||
assert smoothed[-1] == 2.0
|
||||
|
||||
def test_smooths_high_frequency(self):
|
||||
"""A high-frequency square wave should have reduced amplitude after smoothing."""
|
||||
from ai_mouse.generator import _gaussian_smooth
|
||||
x = np.array([0, 1, 0, 1, 0, 1, 0, 1, 0, 1], dtype=np.float64)
|
||||
smoothed = _gaussian_smooth(x, sigma=1.0)
|
||||
# Interior amplitude should be reduced
|
||||
interior_orig = x[2:-2]
|
||||
interior_smooth = smoothed[2:-2]
|
||||
assert interior_smooth.std() < interior_orig.std()
|
||||
|
||||
def test_constant_signal_unchanged(self):
|
||||
from ai_mouse.generator import _gaussian_smooth
|
||||
x = np.full(20, 0.5, dtype=np.float64)
|
||||
smoothed = _gaussian_smooth(x, sigma=1.0)
|
||||
np.testing.assert_allclose(smoothed, x, rtol=1e-6)
|
||||
|
||||
def test_short_array_returns_unchanged(self):
|
||||
"""Arrays shorter than the kernel are returned unchanged."""
|
||||
from ai_mouse.generator import _gaussian_smooth
|
||||
x = np.array([1.0, 2.0, 3.0], dtype=np.float64)
|
||||
smoothed = _gaussian_smooth(x, sigma=1.0)
|
||||
np.testing.assert_allclose(smoothed, x, rtol=1e-6)
|
||||
@@ -1,50 +0,0 @@
|
||||
"""Tests for scroll generator."""
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
from pathlib import Path
|
||||
|
||||
import torch
|
||||
import pytest
|
||||
|
||||
from ai_mouse._scroll_legacy import generate_scroll
|
||||
from tools.scroll.models import ScrollCVAE
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def scroll_model_dir(tmp_path):
|
||||
model = ScrollCVAE(seq_len=32)
|
||||
torch.save(model.state_dict(), tmp_path / "scroll_model.pt")
|
||||
config = {"seq_len": 32, "epochs": 100}
|
||||
(tmp_path / "scroll_config.json").write_text(json.dumps(config))
|
||||
return tmp_path
|
||||
|
||||
|
||||
class TestGenerateScroll:
|
||||
def test_returns_list_of_dicts(self, scroll_model_dir):
|
||||
result = generate_scroll(1000, 3000, mode="target", model_dir=str(scroll_model_dir))
|
||||
assert isinstance(result, list)
|
||||
assert len(result) > 0
|
||||
assert all("deltaY" in e and "t" in e and "deltaMode" in e for e in result)
|
||||
|
||||
def test_timestamps_monotonic(self, scroll_model_dir):
|
||||
result = generate_scroll(1000, 3000, mode="target", model_dir=str(scroll_model_dir))
|
||||
times = [e["t"] for e in result]
|
||||
for i in range(1, len(times)):
|
||||
assert times[i] >= times[i - 1]
|
||||
|
||||
def test_total_scroll_approximately_matches_distance(self, scroll_model_dir):
|
||||
result = generate_scroll(1000, 3000, mode="target", model_dir=str(scroll_model_dir))
|
||||
total = sum(e["deltaY"] for e in result)
|
||||
# Should be within 30% of target distance (2000px)
|
||||
assert abs(total - 2000) < 2000 * 0.4
|
||||
|
||||
def test_deltaY_are_integers(self, scroll_model_dir):
|
||||
result = generate_scroll(1000, 3000, mode="target", model_dir=str(scroll_model_dir))
|
||||
assert all(isinstance(e["deltaY"], int) for e in result)
|
||||
|
||||
def test_direction_up(self, scroll_model_dir):
|
||||
result = generate_scroll(3000, 1000, mode="target", model_dir=str(scroll_model_dir))
|
||||
total = sum(e["deltaY"] for e in result)
|
||||
# Negative total for scrolling up
|
||||
assert total < 0
|
||||
@@ -48,8 +48,15 @@ def _load_reference_jsonl(path: Path, n_samples: int) -> list[dict]:
|
||||
def _generate_n_samples(
|
||||
model_dir: str, n_samples: int, seed: int = 0
|
||||
) -> list[dict]:
|
||||
"""Call the project's generator N times with random start/end pairs."""
|
||||
from ai_mouse.generator import generate
|
||||
"""Call the project's generator N times with random start/end pairs.
|
||||
|
||||
``model_dir`` is accepted for CLI backward compatibility but is no longer
|
||||
used — generation goes through the public ai_mouse API which loads the
|
||||
bundled ONNX model. Export a fresh .onnx via ``python -m tools.export_onnx``
|
||||
to refresh.
|
||||
"""
|
||||
del model_dir # legacy arg, unused
|
||||
from ai_mouse import generate
|
||||
|
||||
rng = random.Random(seed)
|
||||
out: list[dict] = []
|
||||
@@ -63,7 +70,7 @@ def _generate_n_samples(
|
||||
ex = max(0, min(800, ex))
|
||||
ey = max(0, min(600, ey))
|
||||
try:
|
||||
pts = generate(start=(sx, sy), end=(ex, ey), model_dir=model_dir)
|
||||
pts = generate(start=(sx, sy), end=(ex, ey))
|
||||
except Exception as exc: # noqa: BLE001
|
||||
logger.warning("generate() failed at i=%d: %s", i, exc)
|
||||
continue
|
||||
|
||||
@@ -182,21 +182,17 @@ async def scroll_train(req: ScrollTrainRequest) -> StreamingResponse:
|
||||
|
||||
@router.post("/verify")
|
||||
def scroll_verify(req: ScrollVerifyRequest) -> dict:
|
||||
from ai_mouse.scroll.generator import generate_scroll
|
||||
# Uses the bundled ONNX scroll model exposed via the public ai_mouse API.
|
||||
# The legacy scroll_model.pt path is no longer wired in; export a fresh
|
||||
# scroll_decoder.onnx via `python -m tools.export_onnx` to update.
|
||||
from ai_mouse import generate_scroll
|
||||
|
||||
_, models_dir = _paths()
|
||||
if not (models_dir / "scroll_model.pt").exists():
|
||||
raise HTTPException(
|
||||
status_code=400,
|
||||
detail="滚轮模型尚未训练，请先在「训练模型 → 滚轮模型」中完成训练。",
|
||||
)
|
||||
paths = []
|
||||
for _ in range(min(req.n_paths, 12)):
|
||||
events = generate_scroll(
|
||||
req.start_scrollY,
|
||||
req.target_scrollY,
|
||||
mode=req.mode,
|
||||
model_dir=str(models_dir),
|
||||
)
|
||||
paths.append(events)
|
||||
return {"paths": paths}
|
||||
|
||||
@@ -7,8 +7,6 @@ import logging
|
||||
from fastapi import APIRouter, HTTPException
|
||||
from pydantic import BaseModel
|
||||
|
||||
from .deps import get_data_dir
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
router = APIRouter()
|
||||
@@ -32,18 +30,19 @@ class VerifyRequest(BaseModel):
|
||||
|
||||
@router.post("/verify")
|
||||
def verify(req: VerifyRequest) -> dict:
|
||||
from ai_mouse.generator import generate
|
||||
# Uses the bundled ONNX model exposed via the public ai_mouse API.
|
||||
# The legacy req.model_dir / data/models_v2 .pt path is no longer wired
|
||||
# in; export a fresh .onnx via `python -m tools.export_onnx` to update.
|
||||
from ai_mouse import generate
|
||||
|
||||
n = max(1, min(req.n_paths, 12))
|
||||
models_dir = get_data_dir() / "models_v2"
|
||||
model_dir_arg = req.model_dir if req.model_dir else str(models_dir)
|
||||
start = tuple(req.start) # type: ignore[arg-type]
|
||||
end = tuple(req.end) # type: ignore[arg-type]
|
||||
|
||||
paths = []
|
||||
try:
|
||||
for _ in range(n):
|
||||
pts = generate(start=start, end=end, model_dir=model_dir_arg)
|
||||
pts = generate(start=start, end=end)
|
||||
paths.append([[x, y, t] for x, y, t in pts])
|
||||
except FileNotFoundError as exc:
|
||||
raise HTTPException(status_code=404, detail=str(exc)) from exc
|
||||
|
||||
@@ -23,7 +23,7 @@ import numpy as np
|
||||
import torch
|
||||
from torch.utils.data import DataLoader
|
||||
|
||||
from ai_mouse.coord import encode_trajectory
|
||||
from ai_mouse._coord import encode_trajectory
|
||||
from tools.config import TrainConfig
|
||||
from tools.models import TrajectoryFlowModel
|
||||
from tools.utils import resample_arc
|
||||
|
||||
Reference in New Issue
Block a user