901 lines
34 KiB
Python
901 lines
34 KiB
Python
|
#!/usr/bin/env python3
|
||
|
|
"""Fit primitive colliders to a robot's visual meshes and write URDF or USD.
|
||
|
|
|
||
|
|
The input and output formats are selected from their file extensions. URDF
|
||
|
|
output is a new complete document; USD output is an overlay of the input USD.
|
||
|
|
The source asset is never modified.
|
||
|
|
"""
|
||
|
|
|
||
|
|
from __future__ import annotations
|
||
|
|
|
||
|
|
import argparse
|
||
|
|
import os
|
||
|
|
from dataclasses import dataclass
|
||
|
|
from pathlib import Path
|
||
|
|
import sys
|
||
|
|
import tomllib
|
||
|
|
import traceback
|
||
|
|
import xml.etree.ElementTree as ET
|
||
|
|
from typing import Any
|
||
|
|
|
||
|
|
import numpy as np
|
||
|
|
import trimesh
|
||
|
|
|
||
|
|
|
||
|
|
GENERATOR_TAG = "trimesh-primitives-v1"
|
||
|
|
GENERATED_PREFIX = "AUTO_COLLISION_"
|
||
|
|
PRIMITIVE_TYPES = ("box", "sphere", "cylinder")
|
||
|
|
|
||
|
|
|
||
|
|
@dataclass(frozen=True)
|
||
|
|
class PrimitiveFit:
|
||
|
|
kind: str
|
||
|
|
transform: np.ndarray
|
||
|
|
dimensions: tuple[float, ...]
|
||
|
|
volume: float
|
||
|
|
|
||
|
|
|
||
|
|
def _points_array(points: np.ndarray) -> np.ndarray:
|
||
|
|
result = np.asarray(points, dtype=np.float64)
|
||
|
|
if result.ndim != 2 or result.shape[1] != 3 or len(result) < 4:
|
||
|
|
raise ValueError("at least four 3D points are required")
|
||
|
|
if not np.isfinite(result).all():
|
||
|
|
raise ValueError("points contain NaN or infinity")
|
||
|
|
return result
|
||
|
|
|
||
|
|
|
||
|
|
def fit_box(points: np.ndarray, padding: float = 0.0, scale: float = 1.0) -> PrimitiveFit:
|
||
|
|
points = _points_array(points)
|
||
|
|
to_box, extents = trimesh.bounds.oriented_bounds(points)
|
||
|
|
extents = np.asarray(extents, dtype=np.float64) * scale + 2.0 * padding
|
||
|
|
transform = np.linalg.inv(np.asarray(to_box, dtype=np.float64))
|
||
|
|
return PrimitiveFit("box", transform, tuple(extents), float(np.prod(extents)))
|
||
|
|
|
||
|
|
|
||
|
|
def fit_link_aligned_box(
|
||
|
|
points: np.ndarray, padding: float = 0.0, scale: float = 1.0
|
||
|
|
) -> PrimitiveFit:
|
||
|
|
points = _points_array(points)
|
||
|
|
lower = points.min(axis=0)
|
||
|
|
upper = points.max(axis=0)
|
||
|
|
extents = (upper - lower) * scale + 2.0 * padding
|
||
|
|
transform = np.eye(4)
|
||
|
|
transform[:3, 3] = (lower + upper) / 2.0
|
||
|
|
return PrimitiveFit("box", transform, tuple(extents), float(np.prod(extents)))
|
||
|
|
|
||
|
|
|
||
|
|
def fit_sphere(points: np.ndarray, padding: float = 0.0, scale: float = 1.0) -> PrimitiveFit:
|
||
|
|
points = _points_array(points)
|
||
|
|
center, radius = trimesh.nsphere.minimum_nsphere(points)
|
||
|
|
radius = float(radius) * scale + padding
|
||
|
|
transform = np.eye(4)
|
||
|
|
transform[:3, 3] = center
|
||
|
|
return PrimitiveFit("sphere", transform, (radius,), float(4.0 * np.pi * radius**3 / 3.0))
|
||
|
|
|
||
|
|
|
||
|
|
def fit_cylinder(
|
||
|
|
points: np.ndarray,
|
||
|
|
padding: float = 0.0,
|
||
|
|
scale: float = 1.0,
|
||
|
|
sample_count: int = 6,
|
||
|
|
angle_tol: float = 0.001,
|
||
|
|
) -> PrimitiveFit:
|
||
|
|
points = _points_array(points)
|
||
|
|
result = trimesh.bounds.minimum_cylinder(
|
||
|
|
points, sample_count=sample_count, angle_tol=angle_tol
|
||
|
|
)
|
||
|
|
radius = float(result["radius"]) * scale + padding
|
||
|
|
height = float(result["height"]) * scale + 2.0 * padding
|
||
|
|
transform = np.asarray(result["transform"], dtype=np.float64)
|
||
|
|
volume = float(np.pi * radius**2 * height)
|
||
|
|
return PrimitiveFit("cylinder", transform, (radius, height), volume)
|
||
|
|
|
||
|
|
|
||
|
|
def fit_axis_aligned_cylinder(
|
||
|
|
points: np.ndarray,
|
||
|
|
axis: str,
|
||
|
|
padding: float = 0.0,
|
||
|
|
scale: float = 1.0,
|
||
|
|
) -> PrimitiveFit:
|
||
|
|
points = _points_array(points)
|
||
|
|
axis = axis.lower()
|
||
|
|
if axis not in "xyz":
|
||
|
|
raise ValueError(f"cylinder axis must be x, y, or z: {axis}")
|
||
|
|
axis_index = "xyz".index(axis)
|
||
|
|
radial_indices = [index for index in range(3) if index != axis_index]
|
||
|
|
radial_center, radius = trimesh.nsphere.minimum_nsphere(
|
||
|
|
points[:, radial_indices]
|
||
|
|
)
|
||
|
|
axial_min = float(points[:, axis_index].min())
|
||
|
|
axial_max = float(points[:, axis_index].max())
|
||
|
|
|
||
|
|
center = np.zeros(3)
|
||
|
|
center[axis_index] = (axial_min + axial_max) / 2.0
|
||
|
|
center[radial_indices] = radial_center
|
||
|
|
radius = float(radius) * scale + padding
|
||
|
|
height = (axial_max - axial_min) * scale + 2.0 * padding
|
||
|
|
|
||
|
|
transform = np.eye(4)
|
||
|
|
if axis == "x":
|
||
|
|
transform[:3, :3] = np.array(
|
||
|
|
[[0.0, 0.0, 1.0], [0.0, 1.0, 0.0], [-1.0, 0.0, 0.0]]
|
||
|
|
)
|
||
|
|
elif axis == "y":
|
||
|
|
transform[:3, :3] = np.array(
|
||
|
|
[[1.0, 0.0, 0.0], [0.0, 0.0, 1.0], [0.0, -1.0, 0.0]]
|
||
|
|
)
|
||
|
|
transform[:3, 3] = center
|
||
|
|
volume = float(np.pi * radius**2 * height)
|
||
|
|
return PrimitiveFit("cylinder", transform, (radius, height), volume)
|
||
|
|
|
||
|
|
|
||
|
|
def fit_primitive(
|
||
|
|
points: np.ndarray,
|
||
|
|
kind: str,
|
||
|
|
*,
|
||
|
|
allowed: list[str],
|
||
|
|
padding: float,
|
||
|
|
scale: float,
|
||
|
|
cylinder_sample_count: int,
|
||
|
|
cylinder_angle_tol: float,
|
||
|
|
alignment: str = "oriented",
|
||
|
|
axis: str | None = None,
|
||
|
|
) -> PrimitiveFit:
|
||
|
|
def fit(candidate: str) -> PrimitiveFit:
|
||
|
|
if candidate == "box":
|
||
|
|
if alignment == "link":
|
||
|
|
return fit_link_aligned_box(points, padding, scale)
|
||
|
|
return fit_box(points, padding, scale)
|
||
|
|
if candidate == "sphere":
|
||
|
|
return fit_sphere(points, padding, scale)
|
||
|
|
if candidate == "cylinder":
|
||
|
|
if axis:
|
||
|
|
return fit_axis_aligned_cylinder(points, axis, padding, scale)
|
||
|
|
return fit_cylinder(
|
||
|
|
points,
|
||
|
|
padding,
|
||
|
|
scale,
|
||
|
|
cylinder_sample_count,
|
||
|
|
cylinder_angle_tol,
|
||
|
|
)
|
||
|
|
raise ValueError(f"unsupported primitive type: {candidate}")
|
||
|
|
|
||
|
|
if kind != "auto":
|
||
|
|
return fit(kind)
|
||
|
|
|
||
|
|
candidates: list[PrimitiveFit] = []
|
||
|
|
failures: list[str] = []
|
||
|
|
for candidate in allowed:
|
||
|
|
try:
|
||
|
|
candidates.append(fit(candidate))
|
||
|
|
except Exception as exc: # A degenerate mesh may fail one fitter only.
|
||
|
|
failures.append(f"{candidate}: {exc}")
|
||
|
|
if not candidates:
|
||
|
|
raise RuntimeError("all primitive fits failed: " + "; ".join(failures))
|
||
|
|
return min(candidates, key=lambda candidate: candidate.volume)
|
||
|
|
|
||
|
|
|
||
|
|
def load_config(path: Path | None) -> dict[str, Any]:
|
||
|
|
if path is None:
|
||
|
|
return {}
|
||
|
|
with path.open("rb") as stream:
|
||
|
|
config = tomllib.load(stream)
|
||
|
|
config["_config_dir"] = str(path.parent.resolve())
|
||
|
|
return config
|
||
|
|
|
||
|
|
|
||
|
|
def link_settings(config: dict[str, Any], link_name: str) -> dict[str, Any]:
|
||
|
|
settings = dict(config.get("defaults", {}))
|
||
|
|
settings.update(config.get("links", {}).get(link_name, {}))
|
||
|
|
return settings
|
||
|
|
|
||
|
|
|
||
|
|
def _find_robot_root(stage: Any, requested_path: str | None) -> Any:
|
||
|
|
if requested_path:
|
||
|
|
prim = stage.GetPrimAtPath(requested_path)
|
||
|
|
if not prim:
|
||
|
|
raise ValueError(f"robot root does not exist: {requested_path}")
|
||
|
|
return prim
|
||
|
|
|
||
|
|
default_prim = stage.GetDefaultPrim()
|
||
|
|
if default_prim and default_prim.GetRelationship("isaac:physics:robotLinks").IsValid():
|
||
|
|
return default_prim
|
||
|
|
|
||
|
|
for prim in stage.Traverse():
|
||
|
|
if prim.GetRelationship("isaac:physics:robotLinks").IsValid():
|
||
|
|
return prim
|
||
|
|
raise RuntimeError("could not find an Isaac robotLinks relationship; pass --robot-root")
|
||
|
|
|
||
|
|
|
||
|
|
def find_robot_links(stage: Any, robot_root_path: str | None) -> list[Any]:
|
||
|
|
root = _find_robot_root(stage, robot_root_path)
|
||
|
|
targets = root.GetRelationship("isaac:physics:robotLinks").GetTargets()
|
||
|
|
links = [stage.GetPrimAtPath(path) for path in targets]
|
||
|
|
links = [prim for prim in links if prim]
|
||
|
|
if not links:
|
||
|
|
raise RuntimeError(f"robot has no resolved links: {root.GetPath()}")
|
||
|
|
return links
|
||
|
|
|
||
|
|
|
||
|
|
def _computed_purpose(prim: Any, UsdGeom: Any) -> str:
|
||
|
|
imageable = UsdGeom.Imageable(prim)
|
||
|
|
if not imageable:
|
||
|
|
return ""
|
||
|
|
return str(imageable.ComputePurpose())
|
||
|
|
|
||
|
|
|
||
|
|
def collect_visual_points(link: Any, link_paths: set[str], Usd: Any, UsdGeom: Any, UsdPhysics: Any) -> tuple[np.ndarray, set[str]]:
|
||
|
|
"""Return visual vertices in link coordinates and direct mesh-collider roots."""
|
||
|
|
cache = UsdGeom.XformCache()
|
||
|
|
link_to_world = cache.GetLocalToWorldTransform(link)
|
||
|
|
world_to_link = np.asarray(link_to_world.GetInverse(), dtype=np.float64)
|
||
|
|
point_sets: list[np.ndarray] = []
|
||
|
|
mesh_collision_roots: set[str] = set()
|
||
|
|
|
||
|
|
for child in link.GetChildren():
|
||
|
|
if str(child.GetPath()) in link_paths or child.GetName().startswith(GENERATED_PREFIX):
|
||
|
|
continue
|
||
|
|
|
||
|
|
child_has_mesh_collision = False
|
||
|
|
for prim in Usd.PrimRange(child, Usd.TraverseInstanceProxies()):
|
||
|
|
if not prim.IsA(UsdGeom.Mesh):
|
||
|
|
continue
|
||
|
|
purpose = _computed_purpose(prim, UsdGeom)
|
||
|
|
is_collision = prim.HasAPI(UsdPhysics.CollisionAPI) or purpose == str(UsdGeom.Tokens.guide)
|
||
|
|
if is_collision:
|
||
|
|
child_has_mesh_collision = True
|
||
|
|
continue
|
||
|
|
if purpose not in ("", str(UsdGeom.Tokens.default_), str(UsdGeom.Tokens.render)):
|
||
|
|
continue
|
||
|
|
|
||
|
|
points = np.asarray(UsdGeom.Mesh(prim).GetPointsAttr().Get(), dtype=np.float64)
|
||
|
|
if not len(points):
|
||
|
|
continue
|
||
|
|
mesh_to_world = np.asarray(cache.GetLocalToWorldTransform(prim), dtype=np.float64)
|
||
|
|
mesh_to_link = mesh_to_world @ world_to_link
|
||
|
|
local_points = points @ mesh_to_link[:3, :3] + mesh_to_link[3, :3]
|
||
|
|
point_sets.append(local_points)
|
||
|
|
|
||
|
|
if child_has_mesh_collision:
|
||
|
|
mesh_collision_roots.add(str(child.GetPath()))
|
||
|
|
|
||
|
|
if not point_sets:
|
||
|
|
raise RuntimeError(f"no visual mesh vertices found below {link.GetPath()}")
|
||
|
|
return np.concatenate(point_sets), mesh_collision_roots
|
||
|
|
|
||
|
|
|
||
|
|
def collect_mesh_file_points(path: Path) -> np.ndarray:
|
||
|
|
"""Load a URDF visual mesh whose vertices are already link-local."""
|
||
|
|
if not path.is_file():
|
||
|
|
raise FileNotFoundError(path)
|
||
|
|
loaded = trimesh.load(path, force="mesh")
|
||
|
|
if isinstance(loaded, trimesh.Scene):
|
||
|
|
meshes = tuple(loaded.geometry.values())
|
||
|
|
if not meshes:
|
||
|
|
raise RuntimeError(f"mesh file has no geometry: {path}")
|
||
|
|
loaded = trimesh.util.concatenate(meshes)
|
||
|
|
return _points_array(np.asarray(loaded.vertices, dtype=np.float64))
|
||
|
|
|
||
|
|
|
||
|
|
def _set_transform(prim: Any, transform: np.ndarray, scale: tuple[float, float, float] | None, Gf: Any, UsdGeom: Any) -> None:
|
||
|
|
xformable = UsdGeom.Xformable(prim)
|
||
|
|
translation = transform[:3, 3]
|
||
|
|
quaternion = trimesh.transformations.quaternion_from_matrix(transform)
|
||
|
|
xformable.AddTranslateOp().Set(Gf.Vec3d(*translation.tolist()))
|
||
|
|
xformable.AddOrientOp(UsdGeom.XformOp.PrecisionDouble).Set(
|
||
|
|
Gf.Quatd(float(quaternion[0]), Gf.Vec3d(*quaternion[1:4].tolist()))
|
||
|
|
)
|
||
|
|
if scale is not None:
|
||
|
|
xformable.AddScaleOp().Set(Gf.Vec3d(*scale))
|
||
|
|
|
||
|
|
|
||
|
|
def prepare_authoring_links(stage: Any, link_paths: list[str]) -> None:
|
||
|
|
"""De-instance only branches that contain a link requiring a new child."""
|
||
|
|
instance_roots: set[str] = set()
|
||
|
|
for link_path in link_paths:
|
||
|
|
link = stage.GetPrimAtPath(link_path)
|
||
|
|
if not link:
|
||
|
|
continue
|
||
|
|
if link.IsInstance():
|
||
|
|
instance_roots.add(link_path)
|
||
|
|
continue
|
||
|
|
if not link.IsInstanceProxy():
|
||
|
|
continue
|
||
|
|
instance_root = link
|
||
|
|
while instance_root.IsInstanceProxy():
|
||
|
|
instance_root = instance_root.GetParent()
|
||
|
|
if not instance_root or not instance_root.IsInstance():
|
||
|
|
raise RuntimeError(f"could not find an instance root for collision link: {link_path}")
|
||
|
|
instance_roots.add(str(instance_root.GetPath()))
|
||
|
|
|
||
|
|
if not instance_roots:
|
||
|
|
return
|
||
|
|
for instance_root in instance_roots:
|
||
|
|
stage.OverridePrim(instance_root).SetInstanceable(False)
|
||
|
|
stage.GetRootLayer().Save()
|
||
|
|
stage.Reload()
|
||
|
|
|
||
|
|
still_proxies = [
|
||
|
|
link_path
|
||
|
|
for link_path in link_paths
|
||
|
|
if stage.GetPrimAtPath(link_path).IsInstanceProxy()
|
||
|
|
]
|
||
|
|
if still_proxies:
|
||
|
|
raise RuntimeError(f"links remained instance proxies: {still_proxies}")
|
||
|
|
|
||
|
|
|
||
|
|
def author_primitive(stage: Any, link_path: str, fit: PrimitiveFit, Gf: Any, Sdf: Any, UsdGeom: Any, UsdPhysics: Any) -> str:
|
||
|
|
prim_path = f"{link_path}/{GENERATED_PREFIX}{fit.kind.upper()}"
|
||
|
|
if fit.kind == "box":
|
||
|
|
shape = UsdGeom.Cube.Define(stage, prim_path)
|
||
|
|
shape.CreateSizeAttr(1.0)
|
||
|
|
scale = tuple(float(value) for value in fit.dimensions)
|
||
|
|
elif fit.kind == "sphere":
|
||
|
|
shape = UsdGeom.Sphere.Define(stage, prim_path)
|
||
|
|
shape.CreateRadiusAttr(float(fit.dimensions[0]))
|
||
|
|
scale = None
|
||
|
|
elif fit.kind == "cylinder":
|
||
|
|
shape = UsdGeom.Cylinder.Define(stage, prim_path)
|
||
|
|
shape.CreateAxisAttr(UsdGeom.Tokens.z)
|
||
|
|
shape.CreateRadiusAttr(float(fit.dimensions[0]))
|
||
|
|
shape.CreateHeightAttr(float(fit.dimensions[1]))
|
||
|
|
scale = None
|
||
|
|
else:
|
||
|
|
raise AssertionError(fit.kind)
|
||
|
|
|
||
|
|
prim = shape.GetPrim()
|
||
|
|
_set_transform(prim, fit.transform, scale, Gf, UsdGeom)
|
||
|
|
UsdPhysics.CollisionAPI.Apply(prim).CreateCollisionEnabledAttr(True)
|
||
|
|
UsdGeom.Imageable(prim).CreatePurposeAttr(UsdGeom.Tokens.guide)
|
||
|
|
prim.SetCustomDataByKey("collisionGenerator", GENERATOR_TAG)
|
||
|
|
prim.CreateAttribute("collision:primitiveType", Sdf.ValueTypeNames.Token, custom=True).Set(fit.kind)
|
||
|
|
return prim_path
|
||
|
|
|
||
|
|
|
||
|
|
def create_overlay_stage(
|
||
|
|
input_path: Path,
|
||
|
|
temporary_output: Path,
|
||
|
|
default_prim_path: str,
|
||
|
|
source_stage: Any,
|
||
|
|
Usd: Any,
|
||
|
|
) -> Any:
|
||
|
|
stage = Usd.Stage.CreateNew(str(temporary_output))
|
||
|
|
relative_input = os.path.relpath(input_path, temporary_output.parent)
|
||
|
|
stage.GetRootLayer().subLayerPaths = [relative_input]
|
||
|
|
for metadata_key in ("upAxis", "metersPerUnit", "kilogramsPerUnit"):
|
||
|
|
metadata_value = source_stage.GetMetadata(metadata_key)
|
||
|
|
if metadata_value is not None:
|
||
|
|
stage.SetMetadata(metadata_key, metadata_value)
|
||
|
|
source_default = stage.GetPrimAtPath(default_prim_path)
|
||
|
|
if not source_default:
|
||
|
|
raise RuntimeError(f"default prim did not compose into overlay: {default_prim_path}")
|
||
|
|
stage.SetDefaultPrim(source_default)
|
||
|
|
return stage
|
||
|
|
|
||
|
|
|
||
|
|
def validate_usd_output(path: Path, expected_count: int, Usd: Any, UsdGeom: Any, UsdPhysics: Any) -> None:
|
||
|
|
stage = Usd.Stage.Open(str(path))
|
||
|
|
if not stage.GetDefaultPrim():
|
||
|
|
raise RuntimeError("output USD has no default prim")
|
||
|
|
generated = [
|
||
|
|
prim
|
||
|
|
for prim in stage.Traverse()
|
||
|
|
if prim.GetCustomDataByKey("collisionGenerator") == GENERATOR_TAG
|
||
|
|
]
|
||
|
|
if len(generated) != expected_count:
|
||
|
|
raise RuntimeError(f"expected {expected_count} generated colliders, found {len(generated)}")
|
||
|
|
for prim in generated:
|
||
|
|
if prim.GetTypeName() not in ("Cube", "Sphere", "Cylinder"):
|
||
|
|
raise RuntimeError(f"generated collider is not a primitive: {prim.GetPath()}")
|
||
|
|
if not prim.HasAPI(UsdPhysics.CollisionAPI):
|
||
|
|
raise RuntimeError(f"CollisionAPI missing: {prim.GetPath()}")
|
||
|
|
if _computed_purpose(prim, UsdGeom) != str(UsdGeom.Tokens.guide):
|
||
|
|
raise RuntimeError(f"guide purpose missing: {prim.GetPath()}")
|
||
|
|
|
||
|
|
|
||
|
|
def run_usd(args: argparse.Namespace, Usd: Any, UsdGeom: Any, UsdPhysics: Any, Gf: Any, Sdf: Any) -> int:
|
||
|
|
input_path = args.input.resolve()
|
||
|
|
if not input_path.is_file():
|
||
|
|
raise FileNotFoundError(input_path)
|
||
|
|
if not args.dry_run and args.output is None:
|
||
|
|
raise ValueError("--output is required unless --dry-run is used")
|
||
|
|
|
||
|
|
config = load_config(args.config.resolve() if args.config else None)
|
||
|
|
defaults = config.get("defaults", {})
|
||
|
|
allowed = list(defaults.get("allowed_primitives", PRIMITIVE_TYPES))
|
||
|
|
invalid = set(allowed) - set(PRIMITIVE_TYPES)
|
||
|
|
if invalid:
|
||
|
|
raise ValueError(f"invalid allowed_primitives: {sorted(invalid)}")
|
||
|
|
|
||
|
|
source_stage = Usd.Stage.Open(str(input_path))
|
||
|
|
if not source_stage:
|
||
|
|
raise RuntimeError(f"could not open USD: {input_path}")
|
||
|
|
links = find_robot_links(source_stage, args.robot_root)
|
||
|
|
source_default = source_stage.GetDefaultPrim()
|
||
|
|
if not source_default:
|
||
|
|
raise RuntimeError("input USD has no default prim")
|
||
|
|
link_paths = {str(link.GetPath()) for link in links}
|
||
|
|
selected = set(args.only)
|
||
|
|
results: list[tuple[str, PrimitiveFit, set[str]]] = []
|
||
|
|
|
||
|
|
for link in links:
|
||
|
|
name = link.GetName()
|
||
|
|
if selected and name not in selected:
|
||
|
|
continue
|
||
|
|
settings = link_settings(config, name)
|
||
|
|
if not settings.get("enabled", True):
|
||
|
|
print(f"SKIP {name}: disabled by configuration")
|
||
|
|
continue
|
||
|
|
try:
|
||
|
|
points, collision_roots = collect_visual_points(
|
||
|
|
link, link_paths, Usd, UsdGeom, UsdPhysics
|
||
|
|
)
|
||
|
|
except RuntimeError as error:
|
||
|
|
mesh_file = settings.get("mesh_file")
|
||
|
|
if not mesh_file:
|
||
|
|
raise RuntimeError(
|
||
|
|
f"{error}; no mesh_file configured for link name {name!r}"
|
||
|
|
) from error
|
||
|
|
mesh_path = Path(mesh_file)
|
||
|
|
if not mesh_path.is_absolute():
|
||
|
|
mesh_path = Path(config["_config_dir"]) / mesh_path
|
||
|
|
points = collect_mesh_file_points(mesh_path.resolve())
|
||
|
|
collision_roots = set()
|
||
|
|
print(f"FALLBACK {name}: loaded {mesh_path}")
|
||
|
|
kind = str(settings.get("primitive", "auto"))
|
||
|
|
if kind not in (*PRIMITIVE_TYPES, "auto"):
|
||
|
|
raise ValueError(f"invalid primitive for {name}: {kind}")
|
||
|
|
fit = fit_primitive(
|
||
|
|
points,
|
||
|
|
kind,
|
||
|
|
allowed=list(settings.get("allowed_primitives", allowed)),
|
||
|
|
padding=float(settings.get("padding", 0.0)),
|
||
|
|
scale=float(settings.get("scale", 1.0)),
|
||
|
|
cylinder_sample_count=int(settings.get("cylinder_sample_count", defaults.get("cylinder_sample_count", 6))),
|
||
|
|
cylinder_angle_tol=float(settings.get("cylinder_angle_tol", defaults.get("cylinder_angle_tol", 0.001))),
|
||
|
|
alignment=str(settings.get("alignment", "oriented")),
|
||
|
|
axis=str(settings["axis"]) if "axis" in settings else None,
|
||
|
|
)
|
||
|
|
results.append((str(link.GetPath()), fit, collision_roots))
|
||
|
|
dimensions = ", ".join(f"{value:.6f}" for value in fit.dimensions)
|
||
|
|
print(f"FIT {name}: {fit.kind} ({dimensions}), vertices={len(points)}, volume={fit.volume:.8f}")
|
||
|
|
|
||
|
|
if selected:
|
||
|
|
found = {Path(path).name for path, _, _ in results}
|
||
|
|
missing = selected - found
|
||
|
|
if missing:
|
||
|
|
raise ValueError(f"selected links were not generated: {sorted(missing)}")
|
||
|
|
if args.dry_run:
|
||
|
|
print(f"Dry run complete: {len(results)} collider(s) fitted")
|
||
|
|
return 0
|
||
|
|
|
||
|
|
output_path = args.output.resolve()
|
||
|
|
if output_path == input_path:
|
||
|
|
raise ValueError("input and output must be different files")
|
||
|
|
if output_path.exists() and not args.replace:
|
||
|
|
raise FileExistsError(f"output exists; pass --replace: {output_path}")
|
||
|
|
output_path.parent.mkdir(parents=True, exist_ok=True)
|
||
|
|
temporary = output_path.with_name(f".{output_path.stem}.tmp{output_path.suffix}")
|
||
|
|
if temporary.exists():
|
||
|
|
temporary.unlink()
|
||
|
|
|
||
|
|
print(f"CREATE overlay: {temporary}", flush=True)
|
||
|
|
output_stage = create_overlay_stage(
|
||
|
|
input_path,
|
||
|
|
temporary,
|
||
|
|
str(source_default.GetPath()),
|
||
|
|
source_stage,
|
||
|
|
Usd,
|
||
|
|
)
|
||
|
|
print("CREATE overlay: composed", flush=True)
|
||
|
|
prepare_authoring_links(output_stage, [link_path for link_path, _, _ in results])
|
||
|
|
disable_meshes = bool(defaults.get("disable_existing_mesh_collisions", True))
|
||
|
|
for link_path, fit, collision_roots in results:
|
||
|
|
if disable_meshes:
|
||
|
|
for collision_root in collision_roots:
|
||
|
|
output_stage.OverridePrim(collision_root).SetActive(False)
|
||
|
|
authored = author_primitive(
|
||
|
|
output_stage, link_path, fit, Gf, Sdf, UsdGeom, UsdPhysics
|
||
|
|
)
|
||
|
|
print(f"WRITE {authored}")
|
||
|
|
output_stage.GetRootLayer().Save()
|
||
|
|
del output_stage
|
||
|
|
os.replace(temporary, output_path)
|
||
|
|
|
||
|
|
if args.validate:
|
||
|
|
validate_usd_output(output_path, len(results), Usd, UsdGeom, UsdPhysics)
|
||
|
|
print(f"Validated {len(results)} generated collider(s)")
|
||
|
|
print(f"Output: {output_path}")
|
||
|
|
return 0
|
||
|
|
|
||
|
|
|
||
|
|
def _parse_vector(
|
||
|
|
value: str | None, size: int, default: tuple[float, ...]
|
||
|
|
) -> np.ndarray:
|
||
|
|
if value is None:
|
||
|
|
return np.asarray(default, dtype=np.float64)
|
||
|
|
result = np.fromstring(value, sep=" ", dtype=np.float64)
|
||
|
|
if len(result) != size or not np.isfinite(result).all():
|
||
|
|
raise ValueError(f"expected {size} finite values, got {value!r}")
|
||
|
|
return result
|
||
|
|
|
||
|
|
|
||
|
|
def _urdf_origin_transform(origin: ET.Element | None) -> np.ndarray:
|
||
|
|
if origin is None:
|
||
|
|
return np.eye(4)
|
||
|
|
xyz = _parse_vector(origin.get("xyz"), 3, (0.0, 0.0, 0.0))
|
||
|
|
rpy = _parse_vector(origin.get("rpy"), 3, (0.0, 0.0, 0.0))
|
||
|
|
transform = trimesh.transformations.euler_matrix(*rpy, axes="sxyz")
|
||
|
|
transform[:3, 3] = xyz
|
||
|
|
return transform
|
||
|
|
|
||
|
|
|
||
|
|
def _resolve_urdf_mesh_path(filename: str, urdf_path: Path) -> Path:
|
||
|
|
if filename.startswith("file://"):
|
||
|
|
path = Path(filename.removeprefix("file://"))
|
||
|
|
elif filename.startswith("package://"):
|
||
|
|
package_path = Path(filename.removeprefix("package://"))
|
||
|
|
if len(package_path.parts) < 2:
|
||
|
|
raise ValueError(f"invalid package URI: {filename}")
|
||
|
|
package_name, relative_parts = package_path.parts[0], package_path.parts[1:]
|
||
|
|
candidates = [
|
||
|
|
parent / package_name / Path(*relative_parts)
|
||
|
|
for parent in (urdf_path.parent, *urdf_path.parents)
|
||
|
|
]
|
||
|
|
candidates.extend(
|
||
|
|
parent / Path(*relative_parts)
|
||
|
|
for parent in urdf_path.parents
|
||
|
|
if parent.name == package_name
|
||
|
|
)
|
||
|
|
for candidate in candidates:
|
||
|
|
if candidate.is_file():
|
||
|
|
return candidate.resolve()
|
||
|
|
raise FileNotFoundError(
|
||
|
|
f"could not resolve {filename!r} relative to {urdf_path}"
|
||
|
|
)
|
||
|
|
else:
|
||
|
|
path = Path(filename)
|
||
|
|
if not path.is_absolute():
|
||
|
|
path = urdf_path.parent / path
|
||
|
|
path = path.resolve()
|
||
|
|
if not path.is_file():
|
||
|
|
raise FileNotFoundError(path)
|
||
|
|
return path
|
||
|
|
|
||
|
|
|
||
|
|
def collect_urdf_visual_points(link: ET.Element, urdf_path: Path) -> np.ndarray:
|
||
|
|
"""Collect all visual mesh vertices in the link-local coordinate frame."""
|
||
|
|
point_sets: list[np.ndarray] = []
|
||
|
|
for visual in link.findall("visual"):
|
||
|
|
visual_transform = _urdf_origin_transform(visual.find("origin"))
|
||
|
|
geometry = visual.find("geometry")
|
||
|
|
mesh = geometry.find("mesh") if geometry is not None else None
|
||
|
|
if mesh is None:
|
||
|
|
continue
|
||
|
|
filename = mesh.get("filename")
|
||
|
|
if not filename:
|
||
|
|
raise ValueError(
|
||
|
|
f"visual mesh has no filename in link {link.get('name')!r}"
|
||
|
|
)
|
||
|
|
vertices = collect_mesh_file_points(
|
||
|
|
_resolve_urdf_mesh_path(filename, urdf_path)
|
||
|
|
)
|
||
|
|
mesh_scale = _parse_vector(mesh.get("scale"), 3, (1.0, 1.0, 1.0))
|
||
|
|
vertices = trimesh.transform_points(vertices * mesh_scale, visual_transform)
|
||
|
|
point_sets.append(vertices)
|
||
|
|
if not point_sets:
|
||
|
|
raise RuntimeError(f"no visual mesh found in link {link.get('name')!r}")
|
||
|
|
return np.concatenate(point_sets)
|
||
|
|
|
||
|
|
|
||
|
|
def _format_number(value: float) -> str:
|
||
|
|
if abs(value) < 5e-13:
|
||
|
|
value = 0.0
|
||
|
|
return f"{value:.12g}"
|
||
|
|
|
||
|
|
|
||
|
|
def _format_vector(values: np.ndarray | tuple[float, ...]) -> str:
|
||
|
|
return " ".join(_format_number(float(value)) for value in values)
|
||
|
|
|
||
|
|
|
||
|
|
def create_urdf_collision(fit: PrimitiveFit) -> ET.Element:
|
||
|
|
collision = ET.Element(
|
||
|
|
"collision", {"name": f"{GENERATED_PREFIX}{fit.kind.upper()}"}
|
||
|
|
)
|
||
|
|
translation = fit.transform[:3, 3]
|
||
|
|
rpy = trimesh.transformations.euler_from_matrix(fit.transform, axes="sxyz")
|
||
|
|
ET.SubElement(
|
||
|
|
collision,
|
||
|
|
"origin",
|
||
|
|
{"xyz": _format_vector(translation), "rpy": _format_vector(rpy)},
|
||
|
|
)
|
||
|
|
geometry = ET.SubElement(collision, "geometry")
|
||
|
|
if fit.kind == "box":
|
||
|
|
ET.SubElement(geometry, "box", {"size": _format_vector(fit.dimensions)})
|
||
|
|
elif fit.kind == "sphere":
|
||
|
|
ET.SubElement(
|
||
|
|
geometry, "sphere", {"radius": _format_number(fit.dimensions[0])}
|
||
|
|
)
|
||
|
|
elif fit.kind == "cylinder":
|
||
|
|
ET.SubElement(
|
||
|
|
geometry,
|
||
|
|
"cylinder",
|
||
|
|
{
|
||
|
|
"radius": _format_number(fit.dimensions[0]),
|
||
|
|
"length": _format_number(fit.dimensions[1]),
|
||
|
|
},
|
||
|
|
)
|
||
|
|
else:
|
||
|
|
raise AssertionError(fit.kind)
|
||
|
|
return collision
|
||
|
|
|
||
|
|
|
||
|
|
def _fit_urdf_link(
|
||
|
|
link: ET.Element,
|
||
|
|
urdf_path: Path,
|
||
|
|
config: dict[str, Any],
|
||
|
|
defaults: dict[str, Any],
|
||
|
|
allowed: list[str],
|
||
|
|
) -> tuple[PrimitiveFit, int]:
|
||
|
|
name = link.get("name", "")
|
||
|
|
settings = link_settings(config, name)
|
||
|
|
try:
|
||
|
|
points = collect_urdf_visual_points(link, urdf_path)
|
||
|
|
except RuntimeError as error:
|
||
|
|
mesh_file = settings.get("mesh_file")
|
||
|
|
if not mesh_file:
|
||
|
|
raise RuntimeError(str(error)) from error
|
||
|
|
mesh_path = Path(mesh_file)
|
||
|
|
if not mesh_path.is_absolute():
|
||
|
|
mesh_path = Path(config["_config_dir"]) / mesh_path
|
||
|
|
points = collect_mesh_file_points(mesh_path.resolve())
|
||
|
|
print(f"FALLBACK {name}: loaded {mesh_path}")
|
||
|
|
|
||
|
|
kind = str(settings.get("primitive", "auto"))
|
||
|
|
if kind not in (*PRIMITIVE_TYPES, "auto"):
|
||
|
|
raise ValueError(f"invalid primitive for {name}: {kind}")
|
||
|
|
link_allowed = list(settings.get("allowed_primitives", allowed))
|
||
|
|
invalid = set(link_allowed) - set(PRIMITIVE_TYPES)
|
||
|
|
if invalid:
|
||
|
|
raise ValueError(f"invalid allowed_primitives for {name}: {sorted(invalid)}")
|
||
|
|
fit = fit_primitive(
|
||
|
|
points,
|
||
|
|
kind,
|
||
|
|
allowed=link_allowed,
|
||
|
|
padding=float(settings.get("padding", 0.0)),
|
||
|
|
scale=float(settings.get("scale", 1.0)),
|
||
|
|
cylinder_sample_count=int(
|
||
|
|
settings.get(
|
||
|
|
"cylinder_sample_count", defaults.get("cylinder_sample_count", 6)
|
||
|
|
)
|
||
|
|
),
|
||
|
|
cylinder_angle_tol=float(
|
||
|
|
settings.get(
|
||
|
|
"cylinder_angle_tol", defaults.get("cylinder_angle_tol", 0.001)
|
||
|
|
)
|
||
|
|
),
|
||
|
|
alignment=str(settings.get("alignment", "oriented")),
|
||
|
|
axis=str(settings["axis"]) if "axis" in settings else None,
|
||
|
|
)
|
||
|
|
return fit, len(points)
|
||
|
|
|
||
|
|
|
||
|
|
def _parse_urdf(path: Path) -> ET.ElementTree:
|
||
|
|
parser = ET.XMLParser(target=ET.TreeBuilder(insert_comments=True))
|
||
|
|
tree = ET.parse(path, parser=parser)
|
||
|
|
root = tree.getroot()
|
||
|
|
if root.tag != "robot":
|
||
|
|
raise ValueError(f"URDF root must be <robot>, found <{root.tag}>")
|
||
|
|
return tree
|
||
|
|
|
||
|
|
|
||
|
|
def validate_urdf_output(path: Path, expected: dict[str, PrimitiveFit]) -> None:
|
||
|
|
root = _parse_urdf(path).getroot()
|
||
|
|
links = {link.get("name", ""): link for link in root.findall("link")}
|
||
|
|
missing = set(expected) - set(links)
|
||
|
|
if missing:
|
||
|
|
raise RuntimeError(f"output URDF is missing links: {sorted(missing)}")
|
||
|
|
for name, fit in expected.items():
|
||
|
|
generated = [
|
||
|
|
collision
|
||
|
|
for collision in links[name].findall("collision")
|
||
|
|
if collision.get("name", "").startswith(GENERATED_PREFIX)
|
||
|
|
]
|
||
|
|
if len(generated) != 1:
|
||
|
|
raise RuntimeError(
|
||
|
|
f"expected one generated collider on {name}, found {len(generated)}"
|
||
|
|
)
|
||
|
|
geometry = generated[0].find("geometry")
|
||
|
|
primitive_count = (
|
||
|
|
sum(geometry.find(kind) is not None for kind in PRIMITIVE_TYPES)
|
||
|
|
if geometry is not None
|
||
|
|
else 0
|
||
|
|
)
|
||
|
|
if primitive_count != 1:
|
||
|
|
raise RuntimeError(f"invalid generated collision geometry on {name}")
|
||
|
|
written_transform = _urdf_origin_transform(generated[0].find("origin"))
|
||
|
|
if not np.allclose(written_transform, fit.transform, atol=1e-9, rtol=1e-9):
|
||
|
|
error = float(np.max(np.abs(written_transform - fit.transform)))
|
||
|
|
raise RuntimeError(
|
||
|
|
f"generated collision transform changed on {name}: max error {error}"
|
||
|
|
)
|
||
|
|
primitive = geometry.find(fit.kind)
|
||
|
|
if primitive is None:
|
||
|
|
raise RuntimeError(f"expected {fit.kind} collision geometry on {name}")
|
||
|
|
if fit.kind == "box":
|
||
|
|
dimensions = _parse_vector(primitive.get("size"), 3, ())
|
||
|
|
elif fit.kind == "sphere":
|
||
|
|
dimensions = np.asarray([float(primitive.get("radius", "nan"))])
|
||
|
|
else:
|
||
|
|
dimensions = np.asarray(
|
||
|
|
[
|
||
|
|
float(primitive.get("radius", "nan")),
|
||
|
|
float(primitive.get("length", "nan")),
|
||
|
|
]
|
||
|
|
)
|
||
|
|
expected_dimensions = np.asarray(fit.dimensions)
|
||
|
|
if (
|
||
|
|
not np.isfinite(dimensions).all()
|
||
|
|
or (dimensions <= 0.0).any()
|
||
|
|
or not np.allclose(
|
||
|
|
dimensions, expected_dimensions, atol=1e-9, rtol=1e-9
|
||
|
|
)
|
||
|
|
):
|
||
|
|
raise RuntimeError(
|
||
|
|
f"generated collision dimensions changed on {name}: "
|
||
|
|
f"expected {expected_dimensions}, found {dimensions}"
|
||
|
|
)
|
||
|
|
|
||
|
|
|
||
|
|
def run_urdf(args: argparse.Namespace) -> int:
|
||
|
|
input_path = args.input.resolve()
|
||
|
|
if not input_path.is_file():
|
||
|
|
raise FileNotFoundError(input_path)
|
||
|
|
if not args.dry_run and args.output is None:
|
||
|
|
raise ValueError("--output is required unless --dry-run is used")
|
||
|
|
|
||
|
|
config = load_config(args.config.resolve() if args.config else None)
|
||
|
|
defaults = config.get("defaults", {})
|
||
|
|
allowed = list(defaults.get("allowed_primitives", PRIMITIVE_TYPES))
|
||
|
|
invalid = set(allowed) - set(PRIMITIVE_TYPES)
|
||
|
|
if invalid:
|
||
|
|
raise ValueError(f"invalid allowed_primitives: {sorted(invalid)}")
|
||
|
|
|
||
|
|
tree = _parse_urdf(input_path)
|
||
|
|
root = tree.getroot()
|
||
|
|
selected = set(args.only)
|
||
|
|
known_names = {link.get("name", "") for link in root.findall("link")}
|
||
|
|
unknown = selected - known_names
|
||
|
|
if unknown:
|
||
|
|
raise ValueError(f"selected links do not exist: {sorted(unknown)}")
|
||
|
|
|
||
|
|
results: list[tuple[ET.Element, PrimitiveFit]] = []
|
||
|
|
for link in root.findall("link"):
|
||
|
|
name = link.get("name", "")
|
||
|
|
if selected and name not in selected:
|
||
|
|
continue
|
||
|
|
settings = link_settings(config, name)
|
||
|
|
if not settings.get("enabled", True):
|
||
|
|
print(f"SKIP {name}: disabled by configuration")
|
||
|
|
continue
|
||
|
|
fit, vertex_count = _fit_urdf_link(
|
||
|
|
link, input_path, config, defaults, allowed
|
||
|
|
)
|
||
|
|
results.append((link, fit))
|
||
|
|
dimensions = ", ".join(f"{value:.6f}" for value in fit.dimensions)
|
||
|
|
print(
|
||
|
|
f"FIT {name}: {fit.kind} ({dimensions}), "
|
||
|
|
f"vertices={vertex_count}, volume={fit.volume:.8f}"
|
||
|
|
)
|
||
|
|
|
||
|
|
if args.dry_run:
|
||
|
|
print(f"Dry run complete: {len(results)} collider(s) fitted")
|
||
|
|
return 0
|
||
|
|
|
||
|
|
output_path = args.output.resolve()
|
||
|
|
if output_path == input_path:
|
||
|
|
raise ValueError("input and output must be different files")
|
||
|
|
if output_path.exists() and not args.replace:
|
||
|
|
raise FileExistsError(f"output exists; pass --replace: {output_path}")
|
||
|
|
|
||
|
|
expected: dict[str, PrimitiveFit] = {}
|
||
|
|
for link, fit in results:
|
||
|
|
if not args.keep_existing:
|
||
|
|
for collision in list(link.findall("collision")):
|
||
|
|
link.remove(collision)
|
||
|
|
link.append(create_urdf_collision(fit))
|
||
|
|
expected[link.get("name", "")] = fit
|
||
|
|
|
||
|
|
output_path.parent.mkdir(parents=True, exist_ok=True)
|
||
|
|
temporary = output_path.with_name(f".{output_path.stem}.tmp{output_path.suffix}")
|
||
|
|
if temporary.exists():
|
||
|
|
temporary.unlink()
|
||
|
|
ET.indent(tree, space=" ")
|
||
|
|
tree.write(temporary, encoding="utf-8", xml_declaration=True)
|
||
|
|
os.replace(temporary, output_path)
|
||
|
|
|
||
|
|
if args.validate:
|
||
|
|
validate_urdf_output(output_path, expected)
|
||
|
|
print(f"Validated {len(expected)} generated collider(s)")
|
||
|
|
print(f"Output: {output_path}")
|
||
|
|
return 0
|
||
|
|
|
||
|
|
|
||
|
|
def parse_args(argv: list[str]) -> argparse.Namespace:
|
||
|
|
parser = argparse.ArgumentParser(description=__doc__)
|
||
|
|
parser.add_argument("--input", required=True, type=Path, help="source URDF or USD")
|
||
|
|
parser.add_argument("--output", type=Path, help="generated URDF or overlay USD")
|
||
|
|
parser.add_argument("--config", type=Path, help="TOML fitting configuration")
|
||
|
|
parser.add_argument(
|
||
|
|
"--robot-root", help="USD robot root prim path, if auto-detection fails"
|
||
|
|
)
|
||
|
|
parser.add_argument(
|
||
|
|
"--only",
|
||
|
|
action="append",
|
||
|
|
default=[],
|
||
|
|
metavar="LINK",
|
||
|
|
help="generate only selected link (repeatable)",
|
||
|
|
)
|
||
|
|
parser.add_argument(
|
||
|
|
"--keep-existing",
|
||
|
|
action="store_true",
|
||
|
|
help="URDF only: append instead of replacing collisions on generated links",
|
||
|
|
)
|
||
|
|
parser.add_argument(
|
||
|
|
"--dry-run", action="store_true", help="fit and print without writing output"
|
||
|
|
)
|
||
|
|
parser.add_argument(
|
||
|
|
"--replace", action="store_true", help="atomically replace an existing output"
|
||
|
|
)
|
||
|
|
parser.add_argument(
|
||
|
|
"--validate", action="store_true", help="reopen and validate generated output"
|
||
|
|
)
|
||
|
|
return parser.parse_args(argv)
|
||
|
|
|
||
|
|
|
||
|
|
def _asset_format(path: Path) -> str:
|
||
|
|
suffix = path.suffix.lower()
|
||
|
|
if suffix == ".urdf":
|
||
|
|
return "urdf"
|
||
|
|
if suffix in (".usd", ".usda", ".usdc"):
|
||
|
|
return "usd"
|
||
|
|
raise ValueError(
|
||
|
|
f"unsupported file extension {path.suffix!r}; expected .urdf, .usd, .usda, or .usdc"
|
||
|
|
)
|
||
|
|
|
||
|
|
|
||
|
|
def _validate_format_options(args: argparse.Namespace, asset_format: str) -> None:
|
||
|
|
if args.output is not None and _asset_format(args.output) != asset_format:
|
||
|
|
raise ValueError("input and output formats must match")
|
||
|
|
if asset_format == "urdf" and args.robot_root:
|
||
|
|
raise ValueError("--robot-root is only valid for USD input")
|
||
|
|
if asset_format == "usd" and args.keep_existing:
|
||
|
|
raise ValueError("--keep-existing is only valid for URDF input")
|
||
|
|
|
||
|
|
|
||
|
|
def main(argv: list[str] | None = None) -> int:
|
||
|
|
try:
|
||
|
|
args = parse_args(sys.argv[1:] if argv is None else argv)
|
||
|
|
asset_format = _asset_format(args.input)
|
||
|
|
_validate_format_options(args, asset_format)
|
||
|
|
if asset_format == "urdf":
|
||
|
|
return run_urdf(args)
|
||
|
|
|
||
|
|
from isaacsim import SimulationApp
|
||
|
|
|
||
|
|
simulation_app = SimulationApp({"headless": True})
|
||
|
|
try:
|
||
|
|
from pxr import Gf, Sdf, Usd, UsdGeom, UsdPhysics
|
||
|
|
|
||
|
|
return run_usd(args, Usd, UsdGeom, UsdPhysics, Gf, Sdf)
|
||
|
|
finally:
|
||
|
|
simulation_app.close()
|
||
|
|
except Exception:
|
||
|
|
traceback.print_exc()
|
||
|
|
sys.stderr.flush()
|
||
|
|
return 1
|
||
|
|
|
||
|
|
|
||
|
|
if __name__ == "__main__":
|
||
|
|
raise SystemExit(main())
|