410 lines
16 KiB
Python
410 lines
16 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)
|
||
|
||
# ── Marker-Eckpunkte (Mehrpunkt-Residuen, doc/Homing_5_Pose_MultiPoint_Weighted.md Schritt 3) ──
|
||
|
||
@staticmethod
|
||
def _shortest_arc_R(normal: Sequence[float]) -> np.ndarray:
|
||
"""
|
||
Rotationsmatrix [0,0,1] → normal (kürzester Bogen).
|
||
|
||
Repliziert THREE.Quaternion.setFromUnitVectors(vFrom=[0,0,1], vTo=n)
|
||
exakt (inkl. antiparallelem Sonderfall), damit die hier erzeugten
|
||
Ecken zur visuell verifizierten Orientierungs-Konvention in
|
||
boardViewer.html passen (qNormalLoc dort).
|
||
"""
|
||
n = np.asarray(normal, dtype=float)
|
||
nn = float(np.linalg.norm(n))
|
||
n = n / nn if nn > 1e-12 else np.array([0.0, 0.0, 1.0])
|
||
# three.js: quat = (cross([0,0,1], n), dot([0,0,1], n) + 1), dann normalisieren.
|
||
r = float(n[2]) + 1.0
|
||
if r < 1e-12:
|
||
# antiparallel (n == [0,0,-1]): three.js else-Zweig für vFrom=[0,0,1] → (0,-1,0,0)
|
||
qx, qy, qz, qw = 0.0, -1.0, 0.0, 0.0
|
||
else:
|
||
qx, qy, qz, qw = -float(n[1]), float(n[0]), 0.0, r # cross([0,0,1], n) = (-ny, nx, 0)
|
||
ql = math.sqrt(qx * qx + qy * qy + qz * qz + qw * qw)
|
||
qx, qy, qz, qw = qx / ql, qy / ql, qz / ql, qw / ql
|
||
return np.array([
|
||
[1 - 2 * (qy * qy + qz * qz), 2 * (qx * qy - qz * qw), 2 * (qx * qz + qy * qw)],
|
||
[2 * (qx * qy + qz * qw), 1 - 2 * (qx * qx + qz * qz), 2 * (qy * qz - qx * qw)],
|
||
[2 * (qx * qz - qy * qw), 2 * (qy * qz + qx * qw), 1 - 2 * (qx * qx + qy * qy)],
|
||
])
|
||
|
||
@staticmethod
|
||
def _marker_plane_corners(half: float) -> np.ndarray:
|
||
"""
|
||
Die 4 Eckpunkte in der Marker-Ebene (+Z = Normale), in DERSELBEN
|
||
Reihenfolge wie die von 3b_corner_marker_poses.py triangulierten
|
||
`corners_m` (ArUco 0..3, im Uhrzeigersinn von der Vorderseite gesehen).
|
||
|
||
Ecke 0 zeigt in Richtung (+h, +h) — dieselbe Konvention wie der visuell
|
||
kalibrierte Orientierungszeiger (1,1,0) in boardViewer.html. Damit passt
|
||
sie zu den manuell/visuell gesetzten `spin`-Werten der ARM-Marker, die
|
||
im Homing die Gelenkwinkel bestimmen (gegen echte corners_m am
|
||
Seed-Pose verifiziert, test_robot_fk_corners.py, RMS ~1 mm).
|
||
|
||
Hinweis: Die Board-Referenzmarker (Set A0) sind ~90° anders kalibriert,
|
||
ihre Eckreihenfolge passt unter dieser Konvention NICHT — egal, weil
|
||
Board der Root-Link ist: ihr Eck-Residuum ist konstant bzgl. der
|
||
Gelenkvariablen und beeinflusst die Schätzung nicht (der robuste
|
||
Huber-Verlust dämpft es als Ausreißer). Siehe Doc-Notiz.
|
||
|
||
Die Drehrichtung 0→1→2→3 ist so, dass 3b's `corner_plane_normal()`
|
||
(outward = -cross(e01,e02)) wieder +Z liefert — identisch zur
|
||
Beobachtungs-Konvention.
|
||
"""
|
||
h = float(half)
|
||
return np.array([
|
||
[ h, h, 0.0], # 0
|
||
[ h, -h, 0.0], # 1
|
||
[-h, -h, 0.0], # 2
|
||
[-h, h, 0.0], # 3
|
||
])
|
||
|
||
@classmethod
|
||
def marker_corners_local(cls,
|
||
position: Sequence[float],
|
||
normal: Sequence[float],
|
||
size_mm: float,
|
||
spin_deg: float = 0.0) -> np.ndarray:
|
||
"""
|
||
Die 4 Marker-Ecken im LINK-lokalen Frame (mm), Reihenfolge wie `corners_m`.
|
||
|
||
Orientierung = Spin um die Normale ∘ Minimal-Rotation [0,0,1]→Normale,
|
||
exakt wie boardViewer.html (qSpinLoc.multiply(qNormalLoc)).
|
||
"""
|
||
n = np.asarray(normal, dtype=float)
|
||
R = _rot_axis_angle(n, float(spin_deg)) @ cls._shortest_arc_R(n)
|
||
plane = cls._marker_plane_corners(float(size_mm) / 2.0) # (4,3)
|
||
return np.asarray(position, dtype=float) + (R @ plane.T).T # (4,3)
|
||
|
||
@classmethod
|
||
def marker_corners_world(cls,
|
||
transforms: Dict[str, np.ndarray],
|
||
link_name: str,
|
||
position: Sequence[float],
|
||
normal: Sequence[float],
|
||
size_mm: float,
|
||
spin_deg: float = 0.0) -> np.ndarray:
|
||
"""Die 4 Marker-Ecken im Weltframe (mm), Reihenfolge wie `corners_m`."""
|
||
T = transforms.get(link_name, np.eye(4))
|
||
local = cls.marker_corners_local(position, normal, size_mm, spin_deg)
|
||
return np.array([transform_point(T, c) for c in local])
|
||
|
||
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, corners_world}
|
||
|
||
corners_world: (4,3) Welt-mm in `corners_m`-Reihenfolge (für den
|
||
marker_observation="corner_points"-Modus in 5_pose_estimation.py).
|
||
"""
|
||
default_size = float((self.robot.get("markerDefaults", {}) or {}).get("size", 25.0))
|
||
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)
|
||
size_mm = float(m.get("size", default_size))
|
||
spin_deg = float(m.get("spin", 0.0))
|
||
local_corners = self.marker_corners_local(loc, nor, size_mm, spin_deg)
|
||
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),
|
||
"corners_world": np.array([transform_point(T, c) for c in local_corners]),
|
||
}
|
||
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
|