Huang Qi
4d414fd180
- Add .gitignore for Python/data/models - Add matplotlib>=3.8.0 for eval plots - Add PretrainConfig, FinetuneConfig, BalabitAdapterConfig, EvalConfig dataclasses
82 lines
2.1 KiB
Python
82 lines
2.1 KiB
Python
"""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)
|