appRobotRender/pipeline/marker_sets.py

#!/usr/bin/env python3
"""
marker_sets.py
==============
Klassifiziert die Marker aus robot.json in drei Rollen für die Szenen-Kalibrierung
(siehe doc/callibrate_scene_roadmap.md, Phase P0):

  arm   : Marker an BEWEGLICHEN Roboter-Links (Arm1, Ellbow, Arm2, Finger, ...).
          Ihre Lage je Link ist bekannt -> Welt-Referenz, wird NICHT kalibriert.
  set   : Marker an einem STATISCHEN Link (Welt-Wurzel) MIT "set"-Feld. Gleiches
          "set" = ein starres Objekt mit fixen relativen Bezügen -> es wird nur die
          6-DoF-Platzierung des ganzen Sets kalibriert.
  loose : Marker an einem statischen Link OHNE "set"-Feld. Lose, einzeln zu vermessen.

Die Set-Zugehörigkeit steht ausschließlich in robot.json (optionales Feld "set" am
Marker). Hier wird nichts hartkodiert — die Sets ergeben sich per Auto-Discovery aus
den vorhandenen "set"-Werten.

"Beweglich" = es liegt irgendwo auf der Kette bis zur Wurzel ein revolute/linear-Joint.
Damit zählt die Wurzel (Board) und jeder rein über "fixed"-Joints angebundene Link als
statisch (Welt), alles ab dem ersten Gelenk als Arm.

Public API
----------
data  = load_robot("data/robot/robot.json")
cls   = classify_markers(data)         # id -> MarkerInfo
sets  = get_sets(data)                 # set_name -> [MarkerInfo, ...]
loose = get_loose_markers(data)        # [MarkerInfo, ...]
arm   = get_arm_markers(data)          # id -> MarkerInfo
rep   = set_summary(data)              # serialisierbarer Report (für QA / --json)

CLI
---
python pipeline/marker_sets.py -robot data/robot/robot.json
python pipeline/marker_sets.py -robot data/robot/robot.json --json
"""
from __future__ import annotations

import argparse
import json
import sys
from dataclasses import dataclass, asdict
from typing import Any, Dict, List, Optional


# ──────────────────────────────────────────────────────────────
# Datentyp
# ──────────────────────────────────────────────────────────────

@dataclass
class MarkerInfo:
    id: int
    link: str
    role: str                       # "arm" | "set" | "loose"
    set_name: Optional[str]         # nur bei role == "set"
    position: List[float]           # im Link-Frame (mm), wie in robot.json
    normal: Optional[List[float]]
    spin: Optional[float]


# ──────────────────────────────────────────────────────────────
# Laden
# ──────────────────────────────────────────────────────────────

def load_robot(path: str) -> Dict[str, Any]:
    with open(path, "r", encoding="utf-8") as f:
        return json.load(f)


# ──────────────────────────────────────────────────────────────
# Link-Topologie: statisch (Welt) vs. beweglich (Arm)
# ──────────────────────────────────────────────────────────────

_MOVABLE_JOINTS = ("revolute", "linear")


def link_is_movable(links: Dict[str, Any], name: str) -> bool:
    """
    True, wenn auf der Kette von `name` bis zur Wurzel mindestens ein
    revolute/linear-Joint liegt (Marker dort gehören zum beweglichen Roboter).
    """
    seen = set()
    cur: Optional[str] = name
    while cur and cur in links and cur not in seen:
        seen.add(cur)
        joint = links[cur].get("jointToParent") or {}
        if str(joint.get("type", "")).lower() in _MOVABLE_JOINTS:
            return True
        cur = links[cur].get("parent")
    return False


def root_links(links: Dict[str, Any]) -> List[str]:
    return [n for n, l in links.items()
            if not l.get("parent") or l.get("parent") not in links]


# ──────────────────────────────────────────────────────────────
# Klassifizierung
# ──────────────────────────────────────────────────────────────

def _set_name_of(marker: Dict[str, Any]) -> Optional[str]:
    val = marker.get("set")
    if val is None:
        return None
    s = str(val).strip()
    return s or None


def classify_markers(robot_data: Dict[str, Any]) -> Dict[int, MarkerInfo]:
    """
    Liefert id -> MarkerInfo über alle Links. Bei doppelten IDs gewinnt das erste
    Vorkommen (zusätzlich als Warnung in set_summary sichtbar).
    """
    links = robot_data.get("links", {}) or {}
    out: Dict[int, MarkerInfo] = {}

    for link_name, link in links.items():
        movable = link_is_movable(links, link_name)
        for mk in link.get("markers", []) or []:
            if "id" not in mk or "position" not in mk:
                continue
            mid = int(mk["id"])
            if mid in out:
                continue  # erstes Vorkommen behalten
            set_name = _set_name_of(mk)
            if movable:
                role = "arm"
                set_name = None
            elif set_name is not None:
                role = "set"
            else:
                role = "loose"

            normal = mk.get("normal")
            spin = mk.get("spin")
            out[mid] = MarkerInfo(
                id=mid,
                link=link_name,
                role=role,
                set_name=set_name,
                position=[float(v) for v in mk["position"]],
                normal=[float(v) for v in normal] if normal is not None else None,
                spin=float(spin) if spin is not None else None,
            )
    return out


def get_arm_markers(robot_data: Dict[str, Any]) -> Dict[int, MarkerInfo]:
    return {m.id: m for m in classify_markers(robot_data).values() if m.role == "arm"}


def get_sets(robot_data: Dict[str, Any]) -> Dict[str, List[MarkerInfo]]:
    """set_name -> Liste der Marker (role == 'set'), gruppiert nach 'set'-Wert."""
    sets: Dict[str, List[MarkerInfo]] = {}
    for m in classify_markers(robot_data).values():
        if m.role == "set" and m.set_name is not None:
            sets.setdefault(m.set_name, []).append(m)
    return sets


def get_loose_markers(robot_data: Dict[str, Any]) -> List[MarkerInfo]:
    return [m for m in classify_markers(robot_data).values() if m.role == "loose"]


# ──────────────────────────────────────────────────────────────
# QA / Report
# ──────────────────────────────────────────────────────────────

def _duplicate_ids(robot_data: Dict[str, Any]) -> List[int]:
    links = robot_data.get("links", {}) or {}
    seen, dups = set(), set()
    for link in links.values():
        for mk in link.get("markers", []) or []:
            if "id" not in mk:
                continue
            mid = int(mk["id"])
            (dups if mid in seen else seen).add(mid)
    return sorted(dups)


def set_summary(robot_data: Dict[str, Any]) -> Dict[str, Any]:
    cls = classify_markers(robot_data)
    sets = get_sets(robot_data)
    loose = get_loose_markers(robot_data)
    arm = [m for m in cls.values() if m.role == "arm"]

    warnings: List[str] = []
    for mid in _duplicate_ids(robot_data):
        warnings.append(f"Marker-ID {mid} kommt mehrfach vor (erstes Vorkommen gewertet)")
    for name, members in sets.items():
        links_used = sorted({m.link for m in members})
        if len(members) < 3:
            warnings.append(
                f"Set '{name}' hat nur {len(members)} Marker — 6-DoF-Platzierung "
                f"nicht voll bestimmbar (>=3 nicht-kollineare nötig)")
        if len(links_used) > 1:
            warnings.append(f"Set '{name}' verteilt sich über mehrere Links {links_used} "
                            f"— ein Set sollte ein physisches Objekt sein")

    return {
        "counts": {
            "total": len(cls),
            "arm": len(arm),
            "set": sum(len(v) for v in sets.values()),
            "loose": len(loose),
            "num_sets": len(sets),
        },
        "root_links": root_links(robot_data.get("links", {}) or {}),
        "sets": {name: sorted(m.id for m in members) for name, members in sorted(sets.items())},
        "loose_ids": sorted(m.id for m in loose),
        "arm_ids": sorted(m.id for m in arm),
        "warnings": warnings,
    }


# ──────────────────────────────────────────────────────────────
# CLI
# ──────────────────────────────────────────────────────────────

def main() -> None:
    ap = argparse.ArgumentParser(description="Marker aus robot.json in arm/set/loose klassifizieren")
    ap.add_argument("-robot", "--robot", required=True, help="Pfad zu robot.json")
    ap.add_argument("--json", action="store_true", help="Report als JSON ausgeben")
    args = ap.parse_args()

    data = load_robot(args.robot)
    rep = set_summary(data)

    if args.json:
        print(json.dumps(rep, indent=2, ensure_ascii=False))
        return

    c = rep["counts"]
    print(f"robot.json: {args.robot}")
    print(f"Wurzel-Link(s): {rep['root_links']}")
    print(f"\nMarker gesamt: {c['total']}  |  arm: {c['arm']}  set: {c['set']}  "
          f"loose: {c['loose']}  (Sets: {c['num_sets']})")

    print("\nSets (starr, fixe interne Lage -> 6-DoF kalibrieren):")
    if rep["sets"]:
        for name, ids in rep["sets"].items():
            print(f"  {name:10s} ({len(ids):3d}): {ids}")
    else:
        print("  (keine)")

    print(f"\nLose Marker (einzeln zu vermessen): {rep['loose_ids'] or '(keine)'}")
    print(f"Arm-Marker (Welt-Referenz, nicht kalibriert): {rep['arm_ids']}")

    if rep["warnings"]:
        print("\n[WARN]")
        for w in rep["warnings"]:
            print(f"  - {w}")


if __name__ == "__main__":
    try:
        main()
    except Exception as exc:  # pragma: no cover
        print(f"[ERROR] {exc}", file=sys.stderr)
        sys.exit(1)