feat(lib): add private _coord.py with numpy transforms
Copy of coord.py (which is already pure numpy) into the private underscored module to be consumed by upcoming mouse.py rewrite. Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
This commit is contained in:
81
src/ai_mouse/_coord.py
Normal file
81
src/ai_mouse/_coord.py
Normal file
@@ -0,0 +1,81 @@
|
||||
"""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)
|
||||
30
tests/unit/test__coord.py
Normal file
30
tests/unit/test__coord.py
Normal file
@@ -0,0 +1,30 @@
|
||||
"""Test the private numpy coordinate transforms."""
|
||||
from __future__ import annotations
|
||||
|
||||
import numpy as np
|
||||
|
||||
from ai_mouse._coord import decode_trajectory, encode_trajectory
|
||||
|
||||
|
||||
def test_encode_decode_roundtrip() -> None:
|
||||
points = np.array([[100.0, 200.0], [300.0, 250.0], [500.0, 300.0]])
|
||||
start = (100, 200)
|
||||
end = (500, 300)
|
||||
encoded = encode_trajectory(points, start, end)
|
||||
decoded = decode_trajectory(encoded, start, end)
|
||||
assert np.allclose(decoded, points, atol=1e-6)
|
||||
|
||||
|
||||
def test_encode_endpoints() -> None:
|
||||
"""Start should encode to (0,0); end should encode to (1,0)."""
|
||||
points = np.array([[100.0, 200.0], [500.0, 300.0]])
|
||||
encoded = encode_trajectory(points, (100, 200), (500, 300))
|
||||
assert np.allclose(encoded[0], [0.0, 0.0], atol=1e-6)
|
||||
assert np.allclose(encoded[1], [1.0, 0.0], atol=1e-6)
|
||||
|
||||
|
||||
def test_zero_distance_returns_zeros() -> None:
|
||||
points = np.array([[100.0, 200.0]])
|
||||
encoded = encode_trajectory(points, (100, 200), (100, 200))
|
||||
assert encoded.shape == (1, 2)
|
||||
assert np.all(encoded == 0)
|
||||
Reference in New Issue
Block a user