311 lines
12 KiB
Python
311 lines
12 KiB
Python
#!/usr/bin/env python3
|
||
"""
|
||
robot_fk.py
|
||
-----------
|
||
Minimal forward kinematics engine for the robot.json format.
|
||
|
||
Matches the Blender hierarchy used by render_robot.py exactly:
|
||
world_T_link = world_T_parent
|
||
@ Translate(mountPosition) @ Rotate_xyz(mountRotation)
|
||
@ Translate(jointOrigin) @ Rotate_xyz(joint.rotation)
|
||
@ T_motion
|
||
|
||
T_motion = Rotate(axis, value_deg) for revolute joints
|
||
Translate(axis * value_mm) for linear joints
|
||
|
||
Units throughout: millimetres, degrees.
|
||
|
||
Public API
|
||
----------
|
||
fk = RobotFK.from_file("robot.json")
|
||
|
||
transforms = fk.compute({"x": 180, "y": 86, "z": -120,
|
||
"a": -60, "b": 22, "c": 91, "e": 10})
|
||
# → dict link_name -> 4×4 np.ndarray (world frame, mm)
|
||
|
||
p_world = fk.marker_world(transforms, "Arm1", [0, -160, 35])
|
||
# → np.ndarray shape (3,), in mm
|
||
|
||
all_m = fk.all_markers_world(transforms)
|
||
# → dict marker_id -> {"world_mm", "link", "local_mm"}
|
||
|
||
# Cumulative x-offset for a link at all-zero state
|
||
# (useful for 4a: x_slider = world_x - local_x - link_x_at_zero)
|
||
x0 = fk.link_x_at_zero_state("Arm1") # → float mm
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import json
|
||
import math
|
||
from pathlib import Path
|
||
from typing import Any, Dict, List, Optional, Sequence, Tuple
|
||
|
||
import numpy as np
|
||
|
||
STATE_KEYS = ("x", "y", "z", "a", "b", "c", "e")
|
||
|
||
|
||
# ──────────────────────────────────────────────────────────────
|
||
# Low-level matrix helpers
|
||
# ──────────────────────────────────────────────────────────────
|
||
|
||
def _rot_axis_angle(axis: Sequence[float], angle_deg: float) -> np.ndarray:
|
||
"""3×3 rotation matrix via Rodrigues (axis need not be normalised)."""
|
||
ax = np.asarray(axis, dtype=float)
|
||
n = float(np.linalg.norm(ax))
|
||
if n < 1e-12:
|
||
return np.eye(3)
|
||
ax = ax / n
|
||
c = math.cos(math.radians(angle_deg))
|
||
s = math.sin(math.radians(angle_deg))
|
||
t = 1.0 - c
|
||
x, y, z = ax
|
||
return np.array([
|
||
[t*x*x + c, t*x*y - s*z, t*x*z + s*y],
|
||
[t*x*y + s*z, t*y*y + c, t*y*z - s*x],
|
||
[t*x*z - s*y, t*y*z + s*x, t*z*z + c ],
|
||
])
|
||
|
||
|
||
def _rot_xyz_euler(rx: float, ry: float, rz: float) -> np.ndarray:
|
||
"""XYZ Euler angles (degrees) → 3×3 — matches Blender XYZ Euler mode."""
|
||
return (_rot_axis_angle([0, 0, 1], rz)
|
||
@ _rot_axis_angle([0, 1, 0], ry)
|
||
@ _rot_axis_angle([1, 0, 0], rx))
|
||
|
||
|
||
def make_T(R: np.ndarray, t: Sequence[float]) -> np.ndarray:
|
||
"""4×4 homogeneous transform."""
|
||
T = np.eye(4)
|
||
T[:3, :3] = R
|
||
T[:3, 3] = np.asarray(t, dtype=float)
|
||
return T
|
||
|
||
|
||
def transform_point(T: np.ndarray, p: Sequence[float]) -> np.ndarray:
|
||
"""Apply 4×4 transform to a 3-D point."""
|
||
h = np.array([p[0], p[1], p[2], 1.0])
|
||
return (T @ h)[:3]
|
||
|
||
|
||
# ──────────────────────────────────────────────────────────────
|
||
# FK engine
|
||
# ──────────────────────────────────────────────────────────────
|
||
|
||
class RobotFK:
|
||
"""Forward kinematics for the robot.json format."""
|
||
|
||
def __init__(self, robot_data: Dict[str, Any]):
|
||
self.robot = robot_data
|
||
self.links: Dict[str, Any] = robot_data.get("links", {})
|
||
self._topo: List[str] = self._topological_sort()
|
||
|
||
# ── construction ─────────────────────────────────────────
|
||
|
||
@classmethod
|
||
def from_file(cls, path) -> "RobotFK":
|
||
with open(path, "r", encoding="utf-8") as f:
|
||
return cls(json.load(f))
|
||
|
||
def _topological_sort(self) -> List[str]:
|
||
parent_map = {n: d.get("parent") for n, d in self.links.items()}
|
||
visited, order = set(), []
|
||
|
||
def visit(name: str) -> None:
|
||
if name in visited:
|
||
return
|
||
visited.add(name)
|
||
p = parent_map.get(name)
|
||
if p and p in self.links:
|
||
visit(p)
|
||
order.append(name)
|
||
|
||
for name in self.links:
|
||
visit(name)
|
||
return order
|
||
|
||
# ── core computation ──────────────────────────────────────
|
||
|
||
def compute(self, joint_values: Dict[str, float]) -> Dict[str, np.ndarray]:
|
||
"""
|
||
Compute link world transforms for the given joint state.
|
||
|
||
Parameters
|
||
----------
|
||
joint_values : dict variable_name -> value
|
||
Linear joints (x, e): mm
|
||
Revolute joints (y, z, a, b, c): degrees
|
||
|
||
Returns
|
||
-------
|
||
dict link_name -> 4×4 np.ndarray (world frame, mm)
|
||
"""
|
||
state = {k: 0.0 for k in STATE_KEYS}
|
||
for k, v in joint_values.items():
|
||
if k in state:
|
||
state[k] = float(v)
|
||
|
||
transforms: Dict[str, np.ndarray] = {}
|
||
|
||
for link_name in self._topo:
|
||
d = self.links[link_name]
|
||
parent = d.get("parent")
|
||
T_parent = transforms.get(parent, np.eye(4)) if parent else np.eye(4)
|
||
|
||
# 1 · Mount (static in parent frame)
|
||
mp = d.get("mountPosition", [0, 0, 0])
|
||
mr = d.get("mountRotation", [0, 0, 0])
|
||
T_m = make_T(_rot_xyz_euler(*mr), mp)
|
||
|
||
# 2 · Joint origin (pivot point in mount frame)
|
||
ji = d.get("jointToParent", {}) or {}
|
||
jp = ji.get("origin", [0, 0, 0])
|
||
jr = ji.get("rotation", [0, 0, 0])
|
||
T_j = make_T(_rot_xyz_euler(*jr), jp)
|
||
|
||
# 3 · Motion
|
||
jtype = str(ji.get("type", "fixed")).lower()
|
||
var = str(ji.get("variable", "")).lower()
|
||
axis = np.asarray(ji.get("axis", [1, 0, 0]), dtype=float)
|
||
val = state.get(var, 0.0)
|
||
|
||
if jtype == "revolute":
|
||
T_mot = make_T(_rot_axis_angle(axis, val), [0, 0, 0])
|
||
elif jtype == "linear":
|
||
n = float(np.linalg.norm(axis))
|
||
u = axis / n if n > 1e-12 else np.array([1.0, 0, 0])
|
||
T_mot = make_T(np.eye(3), u * val)
|
||
else:
|
||
T_mot = np.eye(4)
|
||
|
||
transforms[link_name] = T_parent @ T_m @ T_j @ T_mot
|
||
|
||
return transforms
|
||
|
||
# ── marker helpers ────────────────────────────────────────
|
||
|
||
@staticmethod
|
||
def marker_world(transforms: Dict[str, np.ndarray],
|
||
link_name: str,
|
||
local_pos: Sequence[float]) -> np.ndarray:
|
||
"""Transform a local marker position → world (mm)."""
|
||
return transform_point(transforms.get(link_name, np.eye(4)), local_pos)
|
||
|
||
def all_markers_world(self,
|
||
transforms: Dict[str, np.ndarray]
|
||
) -> Dict[int, Dict[str, Any]]:
|
||
"""
|
||
Returns
|
||
-------
|
||
dict marker_id -> {world_mm, local_mm, link, normal_world}
|
||
"""
|
||
result: Dict[int, Dict[str, Any]] = {}
|
||
for lname, ldata in self.links.items():
|
||
T = transforms.get(lname, np.eye(4))
|
||
R = T[:3, :3]
|
||
for m in ldata.get("markers", []):
|
||
mid = int(m.get("id", -1))
|
||
if mid < 0 or "position" not in m:
|
||
continue
|
||
loc = np.array(m["position"], dtype=float)
|
||
nor = np.array(m.get("normal", [0, 0, 1]), dtype=float)
|
||
result[mid] = {
|
||
"world_mm": transform_point(T, loc),
|
||
"local_mm": loc,
|
||
"link": lname,
|
||
"normal_world": (R @ nor) / max(np.linalg.norm(R @ nor), 1e-12),
|
||
}
|
||
return result
|
||
|
||
# ── x-axis invariant helpers (used by 4a) ────────────────
|
||
|
||
def link_x_at_zero_state(self, link_name: str) -> float:
|
||
"""
|
||
Return the world x-coordinate of the link-frame origin
|
||
when ALL joint variables are zero.
|
||
|
||
For any link reachable via only x-axis rotations (Arm1, Ellbow, Arm2),
|
||
this value is constant regardless of the actual revolute angles.
|
||
Adding the slider value x_mm gives the true link origin x:
|
||
link_origin_world_x = x_mm + link_x_at_zero_state(link_name)
|
||
|
||
And for any marker in that link:
|
||
marker_world_x = x_mm + link_x_at_zero_state(link_name) + marker_local_x
|
||
"""
|
||
T = self.compute({k: 0.0 for k in STATE_KEYS})
|
||
return float(T[link_name][0, 3])
|
||
|
||
def joint_origin_world(self,
|
||
link_name: str,
|
||
joint_state: Dict[str, float]) -> np.ndarray:
|
||
"""
|
||
World position of the pivot that link_name rotates / slides around.
|
||
"""
|
||
d = self.links[link_name]
|
||
parent = d.get("parent")
|
||
T_all = self.compute(joint_state)
|
||
T_parent = T_all.get(parent, np.eye(4)) if parent else np.eye(4)
|
||
|
||
mp = d.get("mountPosition", [0, 0, 0])
|
||
mr = d.get("mountRotation", [0, 0, 0])
|
||
T_m = make_T(_rot_xyz_euler(*mr), mp)
|
||
|
||
ji = d.get("jointToParent", {}) or {}
|
||
jp = ji.get("origin", [0, 0, 0])
|
||
jr = ji.get("rotation", [0, 0, 0])
|
||
T_j = make_T(_rot_xyz_euler(*jr), jp)
|
||
|
||
return transform_point(T_parent @ T_m @ T_j, [0, 0, 0])
|
||
|
||
def joint_axis_world(self,
|
||
link_name: str,
|
||
joint_state: Dict[str, float]) -> np.ndarray:
|
||
"""
|
||
Joint axis of link_name expressed in world frame.
|
||
"""
|
||
d = self.links[link_name]
|
||
parent = d.get("parent")
|
||
T_all = self.compute(joint_state)
|
||
T_parent = T_all.get(parent, np.eye(4)) if parent else np.eye(4)
|
||
|
||
mp = d.get("mountPosition", [0, 0, 0])
|
||
mr = d.get("mountRotation", [0, 0, 0])
|
||
T_m = make_T(_rot_xyz_euler(*mr), mp)
|
||
|
||
ji = d.get("jointToParent", {}) or {}
|
||
jp = ji.get("origin", [0, 0, 0])
|
||
jr = ji.get("rotation", [0, 0, 0])
|
||
T_j = make_T(_rot_xyz_euler(*jr), jp)
|
||
|
||
R_to_pivot = (T_parent @ T_m @ T_j)[:3, :3]
|
||
axis_local = np.asarray(ji.get("axis", [1, 0, 0]), dtype=float)
|
||
world = R_to_pivot @ axis_local
|
||
n = float(np.linalg.norm(world))
|
||
return world / n if n > 1e-12 else world
|
||
|
||
# ── utility ───────────────────────────────────────────────
|
||
|
||
def chain(self, link_name: str) -> List[str]:
|
||
"""Return chain from root to link_name (inclusive)."""
|
||
out, cur = [], link_name
|
||
while cur:
|
||
out.append(cur)
|
||
cur = self.links.get(cur, {}).get("parent")
|
||
return list(reversed(out))
|
||
|
||
def board_marker_world_positions(self, length_scale: float = 1.0) -> Dict[int, np.ndarray]:
|
||
"""
|
||
Return known world positions for all Board markers (in mm).
|
||
Board is the root, so its marker positions ARE world positions.
|
||
length_scale: 1/1000 if robot.json uses mm.
|
||
"""
|
||
board = self.links.get("Board", {})
|
||
result: Dict[int, np.ndarray] = {}
|
||
for m in board.get("markers", []):
|
||
mid = int(m.get("id", -1))
|
||
if mid >= 0 and "position" in m:
|
||
p = np.array(m["position"], dtype=float) * length_scale
|
||
result[mid] = p
|
||
return result
|