#!/usr/bin/env python3 """ 4b_revolute_angle.py -------------------- Generic revolute-joint angle estimator. For each movable link (Arm1, Ellbow, Arm2 …) whose joint type is 'revolute', this script estimates the rotation angle using the pairwise-vector method: For every PAIR (m1, m2) of markers belonging to the target link: v_model = local_pos_m2 - local_pos_m1 (in link's own frame) v_obs = world_pos_m2 - world_pos_m1 (in world frame) Both vectors are projected perpendicular to the joint axis (in world frame), and the signed angle from v_model_perp to v_obs_perp is measured. The joint axis in world frame is computed via FK using the already-known joint values (from 4a, 4b-prev …), so it is ALWAYS correct — even for deeply-nested joints whose world-frame axis differs from their local axis. Pair weights = baseline_model × baseline_obs (longer baselines → more reliable). How to use sequentially ----------------------- Run 4b once per revolute joint, from root to tip: python 4b_revolute_angle.py --robot r.json --aruco obs.json --link Arm1 --x-mm 180 python 4b_revolute_angle.py --robot r.json --aruco obs.json --link Ellbow --from-state state.json python 4b_revolute_angle.py --robot r.json --aruco obs.json --link Arm2 --from-state state.json The --from-state flag reads the accumulated joint state JSON so you don't have to pass every preceding value on the command line. Output JSON ----------- { "link": "Arm1", "joint": "y", "mean_angle_deg": 86.3, "circular_std_deg": 0.7, "num_pairs": 6, "joint_origin_world_mm": [290, 108, 61], "joint_axis_world": [-1, 0, 0], ... } The file also contains the full accumulated state so the next 4b invocation can read it via --from-state. """ from __future__ import annotations import argparse import json import math from itertools import combinations from pathlib import Path from typing import Dict, List, Optional, Sequence, Tuple import numpy as np from robot_fk import RobotFK STATE_KEYS = ("x", "y", "z", "a", "b", "c", "e") # ────────────────────────────────────────────────────────────── # I/O # ────────────────────────────────────────────────────────────── def _load_json(path: Path) -> dict: with path.open("r", encoding="utf-8") as f: return json.load(f) def _load_observed_mm(aruco_json: dict) -> Dict[int, np.ndarray]: result: Dict[int, np.ndarray] = {} for m in aruco_json.get("markers", []): mid = int(m.get("marker_id", m.get("id", -1))) if mid < 0: continue if "position_mm" in m: result[mid] = np.array(m["position_mm"], dtype=float) elif "position_m" in m: result[mid] = np.array(m["position_m"], dtype=float) * 1000.0 return result def _load_state(path: Path) -> Dict[str, float]: """Load accumulated joint state from a previous 4b output JSON.""" raw = _load_json(path) state = raw.get("accumulated_state", raw.get("state", {})) return {k: float(v) for k, v in state.items() if k in STATE_KEYS} # ────────────────────────────────────────────────────────────── # Angle maths # ────────────────────────────────────────────────────────────── def _project_perp(v: np.ndarray, axis: np.ndarray) -> np.ndarray: """Remove the component of v along axis.""" a = axis / (np.linalg.norm(axis) + 1e-15) return v - np.dot(v, a) * a def _signed_angle_rad(v_from: np.ndarray, v_to: np.ndarray, axis: np.ndarray) -> float: """Signed angle rotating v_from onto v_to around axis (radians).""" a = axis / (np.linalg.norm(axis) + 1e-15) return math.atan2(float(np.dot(a, np.cross(v_from, v_to))), float(np.dot(v_from, v_to))) def _wrap(angle: float) -> float: """Wrap to (−π, π].""" return (angle + math.pi) % (2.0 * math.pi) - math.pi def _circular_mean_deg(angles_rad: np.ndarray, weights: np.ndarray) -> Tuple[float, float, float]: """Returns mean_deg, circular_variance, circular_std_deg.""" w = np.clip(weights, 0, None) if w.sum() < 1e-15: w = np.ones_like(w) s, c = np.sum(w * np.sin(angles_rad)), np.sum(w * np.cos(angles_rad)) mean = math.atan2(s, c) R = math.sqrt(s*s + c*c) / w.sum() R = float(np.clip(R, 0, 1)) c_var = 1.0 - R c_std = math.sqrt(max(0.0, -2.0 * math.log(max(R, 1e-15)))) return math.degrees(mean), c_var, math.degrees(c_std) # ────────────────────────────────────────────────────────────── # Core estimator (generic — works for any revolute joint) # ────────────────────────────────────────────────────────────── def estimate_revolute_angle( fk: RobotFK, observed_mm: Dict[int, np.ndarray], link_name: str, known_state: Dict[str, float], min_baseline_mm: float = 15.0, verbose: bool = True, ) -> dict: """ Estimate the revolute joint angle for `link_name`. Parameters ---------- fk : RobotFK observed_mm : world marker positions from step 3 link_name : e.g. "Arm1", "Ellbow", "Arm2" known_state : already-estimated joint values (e.g. {"x": 180.0, "y": 86.0}) The target joint variable should NOT be in this dict. min_baseline_mm : skip pairs with model or observed baseline shorter than this verbose : print report Returns ------- dict with estimation results + updated accumulated_state """ # ── sanity checks ───────────────────────────────────────── links = fk.links if link_name not in links: return {"status": "failed", "reason": f"Link '{link_name}' not in robot.json"} ji = links[link_name].get("jointToParent", {}) or {} jtype = str(ji.get("type", "")).lower() if jtype != "revolute": return {"status": "failed", "reason": f"Joint of '{link_name}' is not revolute (type='{jtype}')"} var = str(ji.get("variable", "")).lower() # ── FK: joint origin and axis in world ─────────────────── # Use known_state with the TARGET joint at 0 zero_state = dict(known_state) zero_state[var] = 0.0 origin_world = fk.joint_origin_world(link_name, zero_state) axis_world = fk.joint_axis_world(link_name, zero_state) # ── collect matched markers ─────────────────────────────── model_local: Dict[int, np.ndarray] = {} for m in links[link_name].get("markers", []): mid = int(m.get("id", -1)) if mid >= 0 and "position" in m: model_local[mid] = np.array(m["position"], dtype=float) matched = {mid: (model_local[mid], observed_mm[mid]) for mid in model_local if mid in observed_mm} if len(matched) < 2: return { "status": "failed", "reason": (f"Need ≥2 matched markers, found {len(matched)}: " f"{list(matched.keys())}. " f"Model marker IDs: {list(model_local.keys())}"), } # ── pairwise estimation ─────────────────────────────────── ids = sorted(matched.keys()) angle_rad_list: List[float] = [] weight_list: List[float] = [] per_pair: List[dict] = [] for id1, id2 in combinations(ids, 2): l1, o1 = matched[id1] l2, o2 = matched[id2] v_model = l2 - l1 # local frame, both in same link v_obs = o2 - o1 # world frame v_model_perp = _project_perp(v_model, axis_world) v_obs_perp = _project_perp(v_obs, axis_world) bl_model = float(np.linalg.norm(v_model_perp)) bl_obs = float(np.linalg.norm(v_obs_perp)) if bl_model < min_baseline_mm or bl_obs < min_baseline_mm: per_pair.append({ "marker_ids": [id1, id2], "skipped": True, "reason": f"bl_model={bl_model:.1f} bl_obs={bl_obs:.1f} < {min_baseline_mm}", }) continue angle = _wrap(_signed_angle_rad(v_model_perp, v_obs_perp, axis_world)) w = bl_model * bl_obs angle_rad_list.append(angle) weight_list.append(w) per_pair.append({ "marker_ids": [id1, id2], "skipped": False, "angle_deg": math.degrees(angle), "baseline_model_mm": bl_model, "baseline_obs_mm": bl_obs, "weight": w, }) if not angle_rad_list: return { "status": "failed", "reason": "All pairs below min_baseline_mm. " "Try --min-baseline 5 or check step-3 output.", } mean_deg, c_var, c_std_deg = _circular_mean_deg( np.array(angle_rad_list), np.array(weight_list) ) # ── verbose report ──────────────────────────────────────── if verbose: print(f"\n── 4b: '{link_name}' angle ({var}) ──────────────────────") print(f" Joint origin (world): [{', '.join(f'{v:.1f}' for v in origin_world)}] mm") print(f" Joint axis (world): [{', '.join(f'{v:.3f}' for v in axis_world)}]") print(f" Matched markers: {list(matched.keys())}") print(f" Pairs used: {len(angle_rad_list)} / {len(list(combinations(ids, 2)))}") print(f" Angle: {mean_deg:+.2f} ° circular_σ {c_std_deg:.2f} °") if c_std_deg > 5.0: print(f" [WARN] high spread – step-3 errors or marker overlap") print(f"\n Pair detail:") for pp in per_pair: if pp["skipped"]: print(f" M{pp['marker_ids'][0]}↔M{pp['marker_ids'][1]}: SKIPPED – {pp['reason']}") else: print(f" M{pp['marker_ids'][0]}↔M{pp['marker_ids'][1]}: " f"{pp['angle_deg']:+7.2f}° " f"bl_model={pp['baseline_model_mm']:.1f} " f"bl_obs={pp['baseline_obs_mm']:.1f}") # ── build accumulated state ─────────────────────────────── accumulated = dict(known_state) accumulated[var] = mean_deg return { "status": "ok", "link": link_name, "joint": var, "joint_origin_world_mm": origin_world.tolist(), "joint_axis_world": axis_world.tolist(), "mean_angle_deg": mean_deg, "circular_variance": c_var, "circular_std_deg": c_std_deg, "num_pairs_used": len(angle_rad_list), "num_markers_matched": len(matched), "per_pair": per_pair, "accumulated_state": accumulated, } # ────────────────────────────────────────────────────────────── # CLI # ────────────────────────────────────────────────────────────── def main() -> int: p = argparse.ArgumentParser( description="4b: estimate revolute joint angle for any link") p.add_argument("--robot", type=Path, default=Path("robot.json")) p.add_argument("--aruco", type=Path, default=Path("aruco_positions_optimized.json")) p.add_argument("--link", type=str, required=True, help="Link name, e.g. Arm1, Ellbow, Arm2") p.add_argument("--from-state", type=Path, default=None, help="JSON from a previous 4b run (carries accumulated state)") # Manual joint-value overrides (for use without --from-state) for k in STATE_KEYS: p.add_argument(f"--{k}-mm" if k in ("x", "e") else f"--{k}-deg", type=float, default=None, help=f"Known value for joint '{k}'" + (" (mm)" if k in ("x", "e") else " (deg)")) p.add_argument("--min-baseline", type=float, default=15.0, help="Skip pairs with perpendicular baseline < this (mm)") p.add_argument("--output", type=Path, default=None, help="Save result JSON (readable by next 4b as --from-state)") args = p.parse_args() # Assemble known state if args.from_state: known_state = _load_state(args.from_state) print(f" Loaded state from {args.from_state}: {known_state}") else: known_state = {} # CLI overrides for k in STATE_KEYS: attr = f"{k}_mm" if k in ("x", "e") else f"{k}_deg" v = getattr(args, attr, None) if v is not None: known_state[k] = float(v) fk = RobotFK.from_file(args.robot) aruco_json = _load_json(args.aruco) observed_mm = _load_observed_mm(aruco_json) result = estimate_revolute_angle( fk = fk, observed_mm = observed_mm, link_name = args.link, known_state = known_state, min_baseline_mm = args.min_baseline, verbose = True, ) if args.output and result["status"] == "ok": args.output.parent.mkdir(parents=True, exist_ok=True) with args.output.open("w", encoding="utf-8") as f: json.dump(result, f, indent=2) print(f"\n Saved → {args.output}") return 0 if result["status"] == "ok" else 1 if __name__ == "__main__": raise SystemExit(main())