from __future__ import annotations
import warnings
from functools import partial
from pathlib import Path
from typing import TYPE_CHECKING, List, Tuple
import numpy as np
import trimesh
from trimesh.scene import Scene
from typing_extensions import assert_never
import viser
from .. import transforms as tf
if TYPE_CHECKING:
import yourdfpy
[docs]
class ViserUrdf:
"""Helper for rendering URDFs in Viser. This is a self-contained example
that uses only basic Viser features. It can be copied and modified if you
need more fine-grained control.
To move or control visibility of the entire robot, you can create a
parent frame that the URDF will be attached to. This is because
ViserUrdf creates the robot's geometry as children of the specified
`root_node_name`, but doesn't create the root node itself.
.. code-block:: python
import time
import numpy as np
import viser
from viser.extras import ViserUrdf
from robot_descriptions.loaders.yourdfpy import load_robot_description
server = viser.ViserServer()
# Create a parent frame for the robot.
# ViserUrdf will attach the robot's geometry as children of this frame.
robot_base = server.scene.add_frame("/robot", show_axes=False)
# Load a URDF from robot_descriptions package.
urdf = ViserUrdf(
server,
load_robot_description("panda_description"),
root_node_name="/robot"
)
# Move the entire robot by updating the base frame.
robot_base.position = (1.0, 0.0, 0.5) # Move to (x=1, y=0, z=0.5).
# Update joint configuration.
urdf.update_cfg(np.array([0.0, 0.5, -0.5, 0.0, 0.0, 0.0, 0.0, 0.0]))
# Make the robot blink.
while True:
robot_base.visible = False
time.sleep(0.2)
robot_base.visible = True
time.sleep(3.0)
Args:
target: ViserServer or ClientHandle object to add URDF to.
urdf_or_path: Either a path to a URDF file or a yourdfpy URDF object.
scale: Scale factor to apply to resize the URDF.
root_node_name: Viser scene tree name for the root of the URDF geometry.
mesh_color_override: Optional color to override the URDF's visual mesh colors. RGB or RGBA tuple.
collision_mesh_color_override: Optional color to override the URDF's collision mesh colors. RGB or RGBA tuple.
load_meshes: If true, shows the URDF's visual meshes. If a yourdfpy
URDF object is used as input, this expects the URDF to have a visual
scene graph configured.
load_collision_meshes: If true, shows the URDF's collision meshes. If a
yourdfpy URDF object is used as input, this expects the URDF to have a
collision scene graph configured.
"""
def __init__(
self,
target: viser.ViserServer | viser.ClientHandle,
urdf_or_path: yourdfpy.URDF | Path,
scale: float = 1.0,
root_node_name: str = "/",
mesh_color_override: tuple[float, float, float]
| tuple[float, float, float, float]
| None = None,
collision_mesh_color_override: tuple[float, float, float]
| tuple[float, float, float, float]
| None = None,
load_meshes: bool = True,
load_collision_meshes: bool = False,
) -> None:
assert root_node_name.startswith("/")
assert len(root_node_name) == 1 or not root_node_name.endswith("/")
try:
import yourdfpy
except ImportError as e:
raise ImportError(
"yourdfpy is required for ViserUrdf but is not installed. "
"Install it with `pip install yourdfpy` or `pip install viser[urdf]`."
) from e
if isinstance(urdf_or_path, Path):
urdf = yourdfpy.URDF.load(
urdf_or_path,
build_scene_graph=load_meshes,
build_collision_scene_graph=load_collision_meshes,
load_meshes=load_meshes,
load_collision_meshes=load_collision_meshes,
filename_handler=partial(
yourdfpy.filename_handler_magic,
dir=urdf_or_path.parent,
),
)
else:
urdf = urdf_or_path
assert isinstance(urdf, yourdfpy.URDF)
self._target = target
self._urdf = urdf
self._scale = scale
self._root_node_name = root_node_name
self._load_meshes = load_meshes
self._collision_root_frame: viser.FrameHandle | None = None
self._visual_root_frame: viser.FrameHandle | None = None
self._joint_frames: List[viser.SceneNodeHandle] = []
self._meshes: List[viser.SceneNodeHandle] = []
num_joints_to_repeat = 0
if load_meshes:
if urdf.scene is not None:
num_joints_to_repeat += 1
self._visual_root_frame = self._add_joint_frames_and_meshes(
urdf.scene,
root_node_name,
collision_geometry=False,
mesh_color_override=mesh_color_override,
)
else:
warnings.warn(
"load_meshes is enabled but the URDF model does not have a visual scene configured. Not displaying."
)
if load_collision_meshes:
if urdf.collision_scene is not None:
num_joints_to_repeat += 1
self._collision_root_frame = self._add_joint_frames_and_meshes(
urdf.collision_scene,
root_node_name,
collision_geometry=True,
mesh_color_override=collision_mesh_color_override,
)
else:
warnings.warn(
"load_collision_meshes is enabled but the URDF model does not have a collision scene configured. Not displaying."
)
self._joint_map_values = [*self._urdf.joint_map.values()] * num_joints_to_repeat
@property
def show_visual(self) -> bool:
"""Returns whether the visual meshes are currently visible."""
return self._visual_root_frame is not None and self._visual_root_frame.visible
@show_visual.setter
def show_visual(self, visible: bool) -> None:
"""Set whether the visual meshes are currently visible."""
if self._visual_root_frame is not None:
self._visual_root_frame.visible = visible
else:
warnings.warn(
"Cannot set `.show_visual`, since no visual meshes were loaded."
)
@property
def show_collision(self) -> bool:
"""Returns whether the collision meshes are currently visible."""
return (
self._collision_root_frame is not None
and self._collision_root_frame.visible
)
@show_collision.setter
def show_collision(self, visible: bool) -> None:
"""Set whether the collision meshes are currently visible."""
if self._collision_root_frame is not None:
self._collision_root_frame.visible = visible
else:
warnings.warn(
"Cannot set `.show_collision`, since no collision meshes were loaded."
)
[docs]
def remove(self) -> None:
"""Remove URDF from scene."""
# Some of this will be redundant, since children are removed when
# parents are removed.
for frame in self._joint_frames:
frame.remove()
for mesh in self._meshes:
mesh.remove()
[docs]
def update_cfg(self, configuration: np.ndarray) -> None:
"""Update the joint angles of the visualized URDF."""
import yourdfpy
self._urdf.update_cfg(configuration)
for joint, frame_handle in zip(self._joint_map_values, self._joint_frames):
assert isinstance(joint, yourdfpy.Joint)
T_parent_child = self._urdf.get_transform(
joint.child, joint.parent, collision_geometry=not self._load_meshes
)
frame_handle.wxyz = tf.SO3.from_matrix(T_parent_child[:3, :3]).wxyz
frame_handle.position = T_parent_child[:3, 3] * self._scale
[docs]
def get_actuated_joint_limits(
self,
) -> dict[str, tuple[float | None, float | None]]:
"""Returns an ordered mapping from actuated joint names to position limits."""
import yourdfpy
out: dict[str, tuple[float | None, float | None]] = {}
for joint_name, joint in zip(
self._urdf.actuated_joint_names, self._urdf.actuated_joints
):
assert isinstance(joint_name, str)
assert isinstance(joint, yourdfpy.Joint)
if joint.limit is None:
out[joint_name] = (-np.pi, np.pi)
else:
out[joint_name] = (joint.limit.lower, joint.limit.upper)
return out
[docs]
def get_actuated_joint_names(self) -> Tuple[str, ...]:
"""Returns a tuple of actuated joint names, in order."""
return tuple(self._urdf.actuated_joint_names)
def _add_joint_frames_and_meshes(
self,
scene: Scene,
root_node_name: str,
collision_geometry: bool,
mesh_color_override: tuple[float, float, float]
| tuple[float, float, float, float]
| None,
) -> viser.FrameHandle:
"""
Helper function to add joint frames and meshes to the ViserUrdf object.
"""
import yourdfpy
prefix = "collision" if collision_geometry else "visual"
prefixed_root_node_name = (f"{root_node_name}/{prefix}").replace("//", "/")
root_frame = self._target.scene.add_frame(
prefixed_root_node_name, show_axes=False
)
# Add coordinate frame for each joint.
for joint in self._urdf.joint_map.values():
assert isinstance(joint, yourdfpy.Joint)
self._joint_frames.append(
self._target.scene.add_frame(
_viser_name_from_frame(
scene,
joint.child,
prefixed_root_node_name,
),
show_axes=False,
)
)
# Add the URDF's meshes/geometry to viser.
for link_name, mesh in scene.geometry.items():
assert isinstance(mesh, trimesh.Trimesh)
T_parent_child = self._urdf.get_transform(
link_name,
scene.graph.transforms.parents[link_name],
collision_geometry=collision_geometry,
)
name = _viser_name_from_frame(scene, link_name, prefixed_root_node_name)
# Scale + transform the mesh. (these will mutate it!)
#
# It's important that we use apply_transform() instead of unpacking
# the rotation/translation terms, since the scene graph transform
# can also contain scale and reflection terms.
mesh = mesh.copy()
mesh.apply_scale(self._scale)
mesh.apply_transform(T_parent_child)
if mesh_color_override is None:
self._meshes.append(self._target.scene.add_mesh_trimesh(name, mesh))
elif len(mesh_color_override) == 3:
self._meshes.append(
self._target.scene.add_mesh_simple(
name,
mesh.vertices,
mesh.faces,
color=mesh_color_override,
)
)
elif len(mesh_color_override) == 4:
self._meshes.append(
self._target.scene.add_mesh_simple(
name,
mesh.vertices,
mesh.faces,
color=mesh_color_override[:3],
opacity=mesh_color_override[3],
)
)
else:
assert_never(mesh_color_override)
return root_frame
def _viser_name_from_frame(
scene: Scene,
frame_name: str,
root_node_name: str = "/",
) -> str:
"""Given the (unique) name of a frame in our URDF's kinematic tree, return a
scene node name for viser.
For a robot manipulator with four frames, that looks like:
((shoulder)) == ((elbow))
/ / |X|
/ / ((wrist))
____/ /____ |X|
[ ] [=======]
[ base_link ] [] []
[___________]
this would map a name like "elbow" to "base_link/shoulder/elbow".
"""
assert root_node_name.startswith("/")
assert len(root_node_name) == 1 or not root_node_name.endswith("/")
frames = []
while frame_name != scene.graph.base_frame:
frames.append(frame_name)
frame_name = scene.graph.transforms.parents[frame_name]
if root_node_name != "/":
frames.append(root_node_name)
return "/".join(frames[::-1])