Merge pull request #599 from iory/primitives

Add convert-urdf-to-primitives command

Author: iory

skrobot/apps/cli.py

@@ -14,6 +14,7 @@ def get_available_apps():
     app_metadata = {
         'visualize_urdf': ('Visualize URDF model', None),
         'convert_urdf_mesh': ('Convert URDF mesh files', (3, 6)),
+        'convert_urdf_to_primitives': ('Convert URDF meshes to primitive shapes', None),
         'modularize_urdf': ('Modularize URDF files', None),
         'change_urdf_root': ('Change URDF root link', None),
         'transform_urdf': ('Add world link with transform to URDF', None),

skrobot/apps/convert_urdf_to_primitives.py

@@ -0,0 +1,177 @@
+#!/usr/bin/env python
+"""
+Convert URDF mesh geometries to primitive shapes.
+
+This tool analyzes mesh geometries in a URDF file and converts them to
+primitive shapes (box, cylinder, sphere) for more efficient physics
+simulation and simplified robot models.
+"""
+
+import argparse
+import logging
+from pathlib import Path
+import sys
+
+from skrobot.urdf import convert_meshes_to_primitives
+
+
+def main():
+    parser = argparse.ArgumentParser(
+        description='Convert URDF mesh geometries to primitive shapes',
+        formatter_class=argparse.RawDescriptionHelpFormatter,
+        epilog="""
+Examples:
+  # Convert both visual and collision meshes to primitives
+  convert-urdf-to-primitives robot.urdf
+
+  # Specify output file
+  convert-urdf-to-primitives robot.urdf -o robot_primitives.urdf
+
+  # Preserve visual meshes, only convert collision
+  convert-urdf-to-primitives robot.urdf --preserve-visual
+
+  # Force all meshes to be boxes
+  convert-urdf-to-primitives robot.urdf --primitive-type box
+
+  # Modify the input file in place
+  convert-urdf-to-primitives robot.urdf --inplace
+
+  # Process with verbose output
+  convert-urdf-to-primitives robot.urdf -v
+        """
+    )
+
+    parser.add_argument(
+        'urdf_file',
+        type=str,
+        help='Path to the input URDF file'
+    )
+
+    parser.add_argument(
+        '-o', '--output',
+        type=str,
+        default=None,
+        help='Path to the output URDF file (default: input_primitives.urdf)'
+    )
+
+    parser.add_argument(
+        '--inplace',
+        action='store_true',
+        help='Modify the input URDF file in place'
+    )
+
+    parser.add_argument(
+        '--preserve-visual',
+        action='store_true',
+        help='Preserve visual meshes, only convert collision meshes'
+    )
+
+    parser.add_argument(
+        '--preserve-collision',
+        action='store_true',
+        help='Preserve collision meshes, only convert visual meshes'
+    )
+
+    parser.add_argument(
+        '--primitive-type',
+        type=str,
+        choices=['box', 'cylinder', 'sphere', 'auto'],
+        default='auto',
+        help='Force a specific primitive type for all meshes (default: auto-detect)'
+    )
+
+    parser.add_argument(
+        '-v', '--verbose',
+        action='store_true',
+        help='Enable verbose output'
+    )
+
+    parser.add_argument(
+        '--force',
+        action='store_true',
+        help='Overwrite output file if it exists'
+    )
+
+    args = parser.parse_args()
+
+    if args.verbose:
+        logging.basicConfig(level=logging.INFO, format='%(message)s')
+    else:
+        logging.basicConfig(level=logging.WARNING, format='%(message)s')
+
+    input_path = Path(args.urdf_file)
+    if not input_path.exists():
+        print(f"Error: Input file '{args.urdf_file}' does not exist", file=sys.stderr)
+        return 1
+
+    if args.inplace and args.output:
+        print("Error: Cannot use both --output and --inplace options", file=sys.stderr)
+        return 1
+
+    if args.preserve_visual and args.preserve_collision:
+        print("Error: Cannot use both --preserve-visual and --preserve-collision", file=sys.stderr)
+        return 1
+
+    if args.inplace:
+        output_path = None
+    elif args.output:
+        output_path = Path(args.output)
+        if output_path.exists() and not args.force:
+            print(f"Error: Output file '{output_path}' already exists. Use --force to overwrite.", file=sys.stderr)
+            return 1
+    else:
+        output_path = input_path.parent / f"{input_path.stem}_primitives.urdf"
+        if output_path.exists() and not args.force:
+            print(f"Error: Output file '{output_path}' already exists. Use --force to overwrite.", file=sys.stderr)
+            return 1
+
+    convert_visual = not args.preserve_visual
+    convert_collision = not args.preserve_collision
+
+    primitive_type = None if args.primitive_type == 'auto' else args.primitive_type
+
+    try:
+        print(f"Loading URDF from: {input_path}")
+        if args.preserve_visual:
+            print("Preserving visual meshes, converting collision meshes to primitives...")
+        elif args.preserve_collision:
+            print("Preserving collision meshes, converting visual meshes to primitives...")
+        else:
+            print("Converting both visual and collision meshes to primitives...")
+
+        if primitive_type:
+            print(f"Forcing primitive type: {primitive_type}")
+        else:
+            print("Auto-detecting best primitive type for each mesh")
+
+        modified_count = convert_meshes_to_primitives(
+            str(input_path),
+            str(output_path) if output_path else None,
+            convert_visual=convert_visual,
+            convert_collision=convert_collision,
+            primitive_type=primitive_type
+        )
+
+        if modified_count > 0:
+            print(f"Converted {modified_count} geometry elements to primitives")
+        else:
+            print("No mesh geometries found to convert")
+
+        if args.inplace:
+            print(f"Modified URDF saved in place: {input_path}")
+        elif output_path:
+            print(f"Modified URDF saved to: {output_path}")
+
+        print("Conversion completed successfully!")
+        return 0
+
+    except Exception as e:
+        print(f"Error: {e}", file=sys.stderr)
+        if args.verbose:
+            import traceback
+            traceback.print_exc()
+        return 1
+
+
+if __name__ == '__main__':
+    sys.exit(main())

skrobot/coordinates/math.py

@@ -1541,6 +1541,77 @@ def rotation_matrix_from_axis(
     return np.vstack([e1, e2, e3])[np.argsort(indices)].T
 
 
+def rotation_matrix_z_to_axis(axis):
+    """Create rotation matrix that aligns Z-axis to target axis direction.
+
+    This is a convenience wrapper around rotation_matrix_from_axis that
+    automatically selects an appropriate second axis (X or Y) to fully determine
+    the rotation matrix. The resulting rotation matrix will transform the Z-axis
+    [0, 0, 1] to point in the target axis direction.
+
+    Difference from rotation_matrix_from_axis:
+    - rotation_matrix_from_axis: Requires explicit specification of both axes
+    - rotation_matrix_z_to_axis: Only requires target axis, automatically selects second axis
+
+    This function is particularly useful for URDF cylinder primitives, which are
+    defined with their central axis along the Z-axis.
+
+    Parameters
+    ----------
+    axis : numpy.ndarray or list or tuple
+        Target axis direction for Z-axis alignment (will be normalized).
+        The Z-axis [0, 0, 1] of the resulting rotation matrix will be aligned
+        with this direction.
+
+    Returns
+    -------
+    rotation : numpy.ndarray
+        3x3 rotation matrix R such that R @ [0, 0, 1] is parallel to axis.
+
+    Examples
+    --------
+    >>> from skrobot.coordinates.math import rotation_matrix_z_to_axis
+    >>> import numpy as np
+    >>> # Align Z-axis with X-axis [1, 0, 0]
+    >>> R = rotation_matrix_z_to_axis([1, 0, 0])
+    >>> np.allclose(R @ [0, 0, 1], [1, 0, 0])
+    True
+    >>> # Align Z-axis with diagonal direction [0, 1, 1]
+    >>> R = rotation_matrix_z_to_axis([0, 1, 1])
+    >>> np.allclose(R @ [0, 0, 1], [0, 1/np.sqrt(2), 1/np.sqrt(2)])
+    True
+
+    See Also
+    --------
+    rotation_matrix_from_axis : Create rotation matrix from two explicit axes
+
+    Notes
+    -----
+    The second axis is automatically chosen to avoid gimbal lock:
+    - If the target axis is close to the X-axis, Y-axis is used as a hint
+    - Otherwise, X-axis is used as a hint
+
+    This ensures a stable rotation even when the target axis is aligned with
+    or close to a cardinal axis.
+    """
+    axis = np.asarray(axis, dtype=float)
+    axis = axis / np.linalg.norm(axis)
+
+    # Select second axis hint to avoid gimbal lock
+    # If axis is close to X-axis, use Y-axis as hint; otherwise use X-axis
+    if abs(axis[0]) < 0.9:
+        second_axis = np.array([1, 0, 0])  # Use X-axis as hint
+    else:
+        second_axis = np.array([0, 1, 0])  # Use Y-axis as hint
+
+    # Create rotation matrix with Z-axis aligned to target and X-axis hint
+    return rotation_matrix_from_axis(
+        first_axis=axis,
+        second_axis=second_axis,
+        axes='zx'
+    )
+
+
 def rodrigues(axis, theta=None, skip_normalization=False):
     """Rodrigues formula.
 

skrobot/urdf/__init__.py

@@ -1,6 +1,7 @@
 from .aggregate import aggregate_urdf_mesh_files
 from .modularize_urdf import find_root_link
 from .modularize_urdf import transform_urdf_to_macro
+from .primitives_converter import convert_meshes_to_primitives
 from .transform_urdf import transform_urdf_with_world_link
 from .wheel_collision_converter import convert_wheel_collisions_to_cylinders
 from .wheel_collision_converter import get_mesh_dimensions
@@ -16,5 +17,6 @@
     'transform_urdf_with_world_link',
     'aggregate_urdf_mesh_files',
     'convert_wheel_collisions_to_cylinders',
+    'convert_meshes_to_primitives',
     'get_mesh_dimensions',
 ]