From 31fd884dfdb4e12c720f54ebf5fd68787da48f64 Mon Sep 17 00:00:00 2001 From: Huang Qi Date: Tue, 12 May 2026 01:23:52 +0800 Subject: [PATCH] refactor(lib): remove legacy generator.py / coord.py / scroll module MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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). --- scripts/build_golden_mouse.py | 49 ---- scripts/build_golden_scroll.py | 48 ---- src/ai_mouse/_scroll_legacy.py | 148 ------------ src/ai_mouse/coord.py | 81 ------- src/ai_mouse/generator.py | 353 ---------------------------- tests/unit/test_coord.py | 113 --------- tests/unit/test_generator.py | 158 ------------- tests/unit/test_scroll_generator.py | 50 ---- tools/eval/__main__.py | 13 +- tools/server/routes_scroll.py | 12 +- tools/server/routes_verify.py | 11 +- tools/trainer.py | 2 +- 12 files changed, 20 insertions(+), 1018 deletions(-) delete mode 100644 scripts/build_golden_mouse.py delete mode 100644 scripts/build_golden_scroll.py delete mode 100644 src/ai_mouse/_scroll_legacy.py delete mode 100644 src/ai_mouse/coord.py delete mode 100644 src/ai_mouse/generator.py delete mode 100644 tests/unit/test_coord.py delete mode 100644 tests/unit/test_generator.py delete mode 100644 tests/unit/test_scroll_generator.py diff --git a/scripts/build_golden_mouse.py b/scripts/build_golden_mouse.py deleted file mode 100644 index a30dfa3..0000000 --- a/scripts/build_golden_mouse.py +++ /dev/null @@ -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() diff --git a/scripts/build_golden_scroll.py b/scripts/build_golden_scroll.py deleted file mode 100644 index 4937775..0000000 --- a/scripts/build_golden_scroll.py +++ /dev/null @@ -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() diff --git a/src/ai_mouse/_scroll_legacy.py b/src/ai_mouse/_scroll_legacy.py deleted file mode 100644 index f066d1e..0000000 --- a/src/ai_mouse/_scroll_legacy.py +++ /dev/null @@ -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 diff --git a/src/ai_mouse/coord.py b/src/ai_mouse/coord.py deleted file mode 100644 index ff3d528..0000000 --- a/src/ai_mouse/coord.py +++ /dev/null @@ -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) diff --git a/src/ai_mouse/generator.py b/src/ai_mouse/generator.py deleted file mode 100644 index 2d6fefd..0000000 --- a/src/ai_mouse/generator.py +++ /dev/null @@ -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), - ] diff --git a/tests/unit/test_coord.py b/tests/unit/test_coord.py deleted file mode 100644 index 0fd376f..0000000 --- a/tests/unit/test_coord.py +++ /dev/null @@ -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) diff --git a/tests/unit/test_generator.py b/tests/unit/test_generator.py deleted file mode 100644 index 64ef61a..0000000 --- a/tests/unit/test_generator.py +++ /dev/null @@ -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) diff --git a/tests/unit/test_scroll_generator.py b/tests/unit/test_scroll_generator.py deleted file mode 100644 index b143433..0000000 --- a/tests/unit/test_scroll_generator.py +++ /dev/null @@ -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 diff --git a/tools/eval/__main__.py b/tools/eval/__main__.py index c09f99d..429b7e7 100644 --- a/tools/eval/__main__.py +++ b/tools/eval/__main__.py @@ -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 diff --git a/tools/server/routes_scroll.py b/tools/server/routes_scroll.py index 0c9af75..9e4760f 100644 --- a/tools/server/routes_scroll.py +++ b/tools/server/routes_scroll.py @@ -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} diff --git a/tools/server/routes_verify.py b/tools/server/routes_verify.py index 0430b0e..b17a1d4 100644 --- a/tools/server/routes_verify.py +++ b/tools/server/routes_verify.py @@ -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 diff --git a/tools/trainer.py b/tools/trainer.py index 65ecf5c..47cbdd3 100644 --- a/tools/trainer.py +++ b/tools/trainer.py @@ -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