Merge pull request #24 from clemense/clemense/add_viz_config_and_animate

Add configuration and animation to viz

src/yourdfpy/urdf.py

@@ -834,7 +834,7 @@ def clear_errors(self):
         """Clear the validation error log."""
         self._errors = []
 
-    def show(self, collision_geometry=False):
+    def show(self, collision_geometry=False, callback=None):
         """Open a simpler viewer displaying the URDF model.
 
         Args:
@@ -846,7 +846,7 @@ def show(self, collision_geometry=False):
                     "No collision scene available. Use build_collision_scene_graph=True and load_collision_meshes=True during loading."
                 )
             else:
-                self._scene_collision.show()
+                self._scene_collision.show(callback=callback)
         else:
             if self._scene is None:
                 raise ValueError(
@@ -857,7 +857,7 @@ def show(self, collision_geometry=False):
                     "Scene is empty, maybe meshes failed to load? Use build_scene_graph=True and load_meshes=True during loading."
                 )
             else:
-                self._scene.show()
+                self._scene.show(callback=callback)
 
     def validate(self, validation_fn=None) -> bool:
         """Validate URDF model.

src/yourdfpy/viz.py

@@ -2,9 +2,12 @@
 Script for visualizing a robot from a URDF.
 """
 
-import argparse
-import logging
 import sys
+import time
+import logging
+import argparse
+import numpy as np
+from functools import partial
 
 from yourdfpy import __version__
 from yourdfpy import URDF
@@ -48,6 +51,11 @@ def parse_args(args):
         action="store_true",
         help="Use collision geometry for the visualized URDF model.",
     )
+    parser.add_argument(
+        "--animate",
+        action="store_true",
+        help="Animate model by interpolating through all actuated joint limits.",
+    )
     parser.add_argument(
         "-v",
         "--verbose",
@@ -79,6 +87,71 @@ def setup_logging(loglevel):
     )
 
 
+def generate_joint_limit_trajectory(urdf_model, loop_time):
+    """Generate a trajectory for all actuated joints that interpolates between joint limits.
+    For continuous joint interpolate between [0, 2 * pi].
+
+    Args:
+        urdf_model (yourdfpy.URDF): _description_
+        loop_time (float): Time in seconds to loop through the trajectory.
+
+    Returns:
+        dict: A dictionary over all actuated joints with list of configuration values.
+    """
+    trajectory_via_points = {}
+    for joint_name in urdf_model.actuated_joint_names:
+        if urdf_model.joint_map[joint_name].type.lower() == "continuous":
+            via_point_0 = 0.0
+            via_point_2 = 2.0 * np.pi
+            via_point_1 = (via_point_2 - via_point_0) / 2.0
+        else:
+            limit_lower = (
+                urdf_model.joint_map[joint_name].limit.lower
+                if urdf_model.joint_map[joint_name].limit.lower is not None
+                else -np.pi
+            )
+            limit_upper = (
+                urdf_model.joint_map[joint_name].limit.upper
+                if urdf_model.joint_map[joint_name].limit.upper is not None
+                else +np.pi
+            )
+            via_point_0 = limit_lower
+            via_point_1 = limit_upper
+            via_point_2 = limit_lower
+
+        trajectory_via_points[joint_name] = np.array(
+            [
+                via_point_0,
+                via_point_1,
+                via_point_2,
+            ]
+        )
+    times = np.linspace(0.0, 1.0, int(loop_time * 100.0))
+    bins = np.arange(3) / 2.0
+
+    # Compute alphas for each time
+    inds = np.digitize(times, bins, right=True)
+    inds[inds == 0] = 1
+    alphas = (bins[inds] - times) / (bins[inds] - bins[inds - 1])
+
+    # Create the new interpolated trajectory
+    trajectory = {}
+    for k in trajectory_via_points:
+        trajectory[k] = (
+            alphas * trajectory_via_points[k][inds - 1]
+            + (1.0 - alphas) * trajectory_via_points[k][inds]
+        )
+
+    return trajectory
+
+
+def viewer_callback(scene, urdf_model, trajectory, loop_time):
+    frame = int(100.0 * (time.time() % loop_time))
+    cfg = {k: trajectory[k][frame] for k in trajectory}
+
+    urdf_model.update_cfg(configuration=cfg)
+
+
 def main(args):
     """Wrapper allowing string arguments in a CLI fashion.
 
@@ -95,7 +168,26 @@ def main(args):
         )
     else:
         urdf_model = URDF.load(args.input)
-    urdf_model.show(collision_geometry=args.collision)
+
+    if args.configuration:
+        urdf_model.update_cfg(args.configuration)
+
+    callback = None
+    if args.animate:
+        loop_time = 6.0
+        callback = partial(
+            viewer_callback,
+            urdf_model=urdf_model,
+            loop_time=loop_time,
+            trajectory=generate_joint_limit_trajectory(
+                urdf_model=urdf_model, loop_time=loop_time
+            ),
+        )
+
+    urdf_model.show(
+        collision_geometry=args.collision,
+        callback=callback,
+    )
 
 
 def run():