mirror of
https://github.com/allaunthefox/Research-Stack.git
synced 2026-07-31 03:05:21 +00:00
278 lines
10 KiB
Python
278 lines
10 KiB
Python
from __future__ import annotations
|
|
|
|
from dataclasses import dataclass
|
|
from math import degrees
|
|
from pathlib import Path
|
|
import xml.etree.ElementTree as ET
|
|
|
|
URDF_SUFFIX = ".urdf"
|
|
SUPPORTED_JOINT_TYPES = {"fixed", "continuous", "revolute", "prismatic"}
|
|
SUPPORTED_MESH_SUFFIXES = {".stl"}
|
|
|
|
|
|
class UrdfSourceError(ValueError):
|
|
pass
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class UrdfJoint:
|
|
name: str
|
|
joint_type: str
|
|
parent_link: str
|
|
child_link: str
|
|
min_value_deg: float | None
|
|
max_value_deg: float | None
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class UrdfSource:
|
|
file_ref: str
|
|
source_path: Path
|
|
robot_name: str
|
|
root_link: str
|
|
links: tuple[str, ...]
|
|
joints: tuple[UrdfJoint, ...]
|
|
mesh_paths: tuple[Path, ...]
|
|
|
|
|
|
def file_ref_from_urdf_path(urdf_path: Path) -> str:
|
|
resolved = urdf_path.resolve()
|
|
if resolved.suffix.lower() != URDF_SUFFIX:
|
|
raise UrdfSourceError(f"{resolved} is not a URDF source file")
|
|
return _relative_to_repo(resolved)
|
|
|
|
|
|
def read_urdf_source(urdf_path: Path) -> UrdfSource:
|
|
resolved_path = urdf_path.resolve()
|
|
if resolved_path.suffix.lower() != URDF_SUFFIX:
|
|
raise UrdfSourceError(f"{resolved_path} is not a URDF source file")
|
|
|
|
try:
|
|
root = ET.fromstring(resolved_path.read_text(encoding="utf-8"))
|
|
except (OSError, ET.ParseError) as exc:
|
|
raise UrdfSourceError(f"{_relative_to_repo(resolved_path)} could not be parsed as URDF XML") from exc
|
|
|
|
if root.tag != "robot":
|
|
raise UrdfSourceError(f"{_relative_to_repo(resolved_path)} root element must be <robot>")
|
|
robot_name = str(root.attrib.get("name") or "").strip()
|
|
if not robot_name:
|
|
raise UrdfSourceError(f"{_relative_to_repo(resolved_path)} robot name is required")
|
|
|
|
link_names = []
|
|
for link_element in root.findall("link"):
|
|
name = str(link_element.attrib.get("name") or "").strip()
|
|
if not name:
|
|
raise UrdfSourceError(f"{_relative_to_repo(resolved_path)} link name is required")
|
|
link_names.append(name)
|
|
if not link_names:
|
|
raise UrdfSourceError(f"{_relative_to_repo(resolved_path)} must define at least one link")
|
|
_raise_on_duplicates(link_names, source_path=resolved_path, label="link")
|
|
link_name_set = set(link_names)
|
|
|
|
mesh_paths: list[Path] = []
|
|
for link_element in root.findall("link"):
|
|
for visual_element in link_element.findall("visual"):
|
|
geometry_element = visual_element.find("geometry")
|
|
if geometry_element is None:
|
|
continue
|
|
mesh_element = geometry_element.find("mesh")
|
|
if mesh_element is None:
|
|
raise UrdfSourceError(
|
|
f"{_relative_to_repo(resolved_path)} only mesh visual geometry is supported"
|
|
)
|
|
filename = str(mesh_element.attrib.get("filename") or "").strip()
|
|
if not filename:
|
|
raise UrdfSourceError(f"{_relative_to_repo(resolved_path)} mesh filename is required")
|
|
mesh_path = _resolve_mesh_path(filename, source_path=resolved_path)
|
|
if mesh_path.suffix.lower() not in SUPPORTED_MESH_SUFFIXES:
|
|
raise UrdfSourceError(
|
|
f"{_relative_to_repo(resolved_path)} only STL mesh geometry is supported: {filename!r}"
|
|
)
|
|
if not mesh_path.is_file():
|
|
raise UrdfSourceError(
|
|
f"{_relative_to_repo(resolved_path)} references missing mesh file: {filename!r}"
|
|
)
|
|
mesh_paths.append(mesh_path)
|
|
|
|
joints = []
|
|
joint_names = []
|
|
parent_by_child: dict[str, str] = {}
|
|
children = set()
|
|
joints_by_parent: dict[str, list[str]] = {}
|
|
for joint_element in root.findall("joint"):
|
|
name = str(joint_element.attrib.get("name") or "").strip()
|
|
if not name:
|
|
raise UrdfSourceError(f"{_relative_to_repo(resolved_path)} joint name is required")
|
|
joint_names.append(name)
|
|
joint_type = str(joint_element.attrib.get("type") or "").strip().lower()
|
|
if joint_type not in SUPPORTED_JOINT_TYPES:
|
|
raise UrdfSourceError(
|
|
f"{_relative_to_repo(resolved_path)} joint {name!r} uses unsupported type {joint_type!r}"
|
|
)
|
|
parent_element = joint_element.find("parent")
|
|
child_element = joint_element.find("child")
|
|
parent_link = str(parent_element.attrib.get("link") if parent_element is not None else "").strip()
|
|
child_link = str(child_element.attrib.get("link") if child_element is not None else "").strip()
|
|
if not parent_link or not child_link:
|
|
raise UrdfSourceError(
|
|
f"{_relative_to_repo(resolved_path)} joint {name!r} must define parent and child links"
|
|
)
|
|
if parent_link not in link_name_set:
|
|
raise UrdfSourceError(
|
|
f"{_relative_to_repo(resolved_path)} joint {name!r} references missing parent link {parent_link!r}"
|
|
)
|
|
if child_link not in link_name_set:
|
|
raise UrdfSourceError(
|
|
f"{_relative_to_repo(resolved_path)} joint {name!r} references missing child link {child_link!r}"
|
|
)
|
|
if child_link in parent_by_child:
|
|
raise UrdfSourceError(
|
|
f"{_relative_to_repo(resolved_path)} link {child_link!r} has multiple parents"
|
|
)
|
|
parent_by_child[child_link] = parent_link
|
|
children.add(child_link)
|
|
joints_by_parent.setdefault(parent_link, []).append(child_link)
|
|
min_value_deg, max_value_deg = _joint_limits_deg(joint_element, joint_type=joint_type, source_path=resolved_path)
|
|
joints.append(
|
|
UrdfJoint(
|
|
name=name,
|
|
joint_type=joint_type,
|
|
parent_link=parent_link,
|
|
child_link=child_link,
|
|
min_value_deg=min_value_deg,
|
|
max_value_deg=max_value_deg,
|
|
)
|
|
)
|
|
|
|
_raise_on_duplicates(joint_names, source_path=resolved_path, label="joint")
|
|
root_candidates = [link_name for link_name in link_names if link_name not in children]
|
|
if len(root_candidates) != 1:
|
|
raise UrdfSourceError(
|
|
f"{_relative_to_repo(resolved_path)} must form a single rooted tree; found roots {root_candidates!r}"
|
|
)
|
|
root_link = root_candidates[0]
|
|
|
|
visited: set[str] = set()
|
|
visiting: set[str] = set()
|
|
|
|
def visit(link_name: str) -> None:
|
|
if link_name in visited:
|
|
return
|
|
if link_name in visiting:
|
|
raise UrdfSourceError(f"{_relative_to_repo(resolved_path)} joint graph contains a cycle")
|
|
visiting.add(link_name)
|
|
for child_link in joints_by_parent.get(link_name, ()):
|
|
visit(child_link)
|
|
visiting.remove(link_name)
|
|
visited.add(link_name)
|
|
|
|
visit(root_link)
|
|
if visited != link_name_set:
|
|
missing_links = sorted(link_name_set - visited)
|
|
raise UrdfSourceError(
|
|
f"{_relative_to_repo(resolved_path)} leaves links disconnected from the root: {missing_links!r}"
|
|
)
|
|
if len(joints) != len(link_names) - 1:
|
|
raise UrdfSourceError(
|
|
f"{_relative_to_repo(resolved_path)} must form a tree with exactly links-1 joints"
|
|
)
|
|
|
|
_validate_with_yourdfpy(resolved_path)
|
|
|
|
return UrdfSource(
|
|
file_ref=file_ref_from_urdf_path(resolved_path),
|
|
source_path=resolved_path,
|
|
robot_name=robot_name,
|
|
root_link=root_link,
|
|
links=tuple(link_names),
|
|
joints=tuple(joints),
|
|
mesh_paths=tuple(mesh_paths),
|
|
)
|
|
|
|
|
|
def _validate_with_yourdfpy(source_path: Path) -> None:
|
|
try:
|
|
from yourdfpy import URDF
|
|
except ModuleNotFoundError as exc:
|
|
raise UrdfSourceError(
|
|
f"{_relative_to_repo(source_path)} requires yourdfpy for URDF validation; "
|
|
"install it in the CAD Python environment"
|
|
) from exc
|
|
|
|
try:
|
|
urdf = URDF.load(
|
|
str(source_path),
|
|
build_scene_graph=False,
|
|
build_collision_scene_graph=False,
|
|
load_meshes=False,
|
|
load_collision_meshes=False,
|
|
)
|
|
is_valid = urdf.validate()
|
|
except Exception as exc:
|
|
raise UrdfSourceError(f"{_relative_to_repo(source_path)} failed yourdfpy validation: {exc}") from exc
|
|
|
|
if not is_valid:
|
|
errors = "; ".join(str(error) for error in urdf.errors) or "unknown error"
|
|
raise UrdfSourceError(
|
|
f"{_relative_to_repo(source_path)} failed yourdfpy validation: {errors}"
|
|
)
|
|
|
|
|
|
def _joint_limits_deg(
|
|
joint_element: ET.Element,
|
|
*,
|
|
joint_type: str,
|
|
source_path: Path,
|
|
) -> tuple[float | None, float | None]:
|
|
if joint_type == "fixed":
|
|
return 0.0, 0.0
|
|
if joint_type == "continuous":
|
|
return -180.0, 180.0
|
|
limit_element = joint_element.find("limit")
|
|
if limit_element is None:
|
|
raise UrdfSourceError(
|
|
f"{_relative_to_repo(source_path)} {joint_type} joint {joint_element.attrib.get('name', '')!r} requires <limit>"
|
|
)
|
|
try:
|
|
lower = float(limit_element.attrib["lower"])
|
|
upper = float(limit_element.attrib["upper"])
|
|
except KeyError as exc:
|
|
raise UrdfSourceError(
|
|
f"{_relative_to_repo(source_path)} {joint_type} joint {joint_element.attrib.get('name', '')!r} requires lower and upper limits"
|
|
) from exc
|
|
except ValueError as exc:
|
|
raise UrdfSourceError(
|
|
f"{_relative_to_repo(source_path)} {joint_type} joint {joint_element.attrib.get('name', '')!r} has invalid limits"
|
|
) from exc
|
|
if joint_type == "prismatic":
|
|
return lower, upper
|
|
return degrees(lower), degrees(upper)
|
|
|
|
|
|
def _resolve_mesh_path(filename: str, *, source_path: Path) -> Path:
|
|
if filename.startswith("package://"):
|
|
relative = filename.removeprefix("package://").lstrip("/")
|
|
return (Path.cwd() / relative).resolve()
|
|
return (source_path.parent / filename).resolve()
|
|
|
|
|
|
def _raise_on_duplicates(values: list[str], *, source_path: Path, label: str) -> None:
|
|
seen: set[str] = set()
|
|
duplicates: set[str] = set()
|
|
for value in values:
|
|
if value in seen:
|
|
duplicates.add(value)
|
|
continue
|
|
seen.add(value)
|
|
if duplicates:
|
|
duplicate_text = ", ".join(repr(item) for item in sorted(duplicates))
|
|
raise UrdfSourceError(
|
|
f"{_relative_to_repo(source_path)} {label} names contain duplicates {duplicate_text}"
|
|
)
|
|
|
|
|
|
def _relative_to_repo(path: Path) -> str:
|
|
try:
|
|
return path.resolve().relative_to(Path.cwd().resolve()).as_posix()
|
|
except ValueError:
|
|
return path.resolve().as_posix()
|