From 1d705a5bb8b6b5f8ac88a55caf0a2b87f18a0f06 Mon Sep 17 00:00:00 2001
From: Amal Nanavati <amaln@uw.edu>
Date: Fri, 26 Apr 2024 04:26:13 -0700
Subject: [PATCH] Allow Scaling Collision Meshes (#62)

* Allow scaling of collision meshes

* Added trimesh to rosdeps for easier installation

* Ran tests with example

* Remove trimesh as a required dependency
---
 examples/ex_collision_mesh.py |  9 +++++++--
 pymoveit2/moveit2.py          | 11 +++++++++++
 2 files changed, 18 insertions(+), 2 deletions(-)

diff --git a/examples/ex_collision_mesh.py b/examples/ex_collision_mesh.py
index 0fa16fe..f0bea29 100755
--- a/examples/ex_collision_mesh.py
+++ b/examples/ex_collision_mesh.py
@@ -2,7 +2,7 @@
 """
 Example of adding and removing a collision object with a mesh geometry.
 Note: Python module `trimesh` is required for this example (`pip install trimesh`).
-- ros2 run pymoveit2 ex_collision_mesh.py --ros-args -p position:="[0.5, 0.0, 0.5]" -p quat_xyzw:="[0.0, 0.0, -0.7071, 0.7071]"
+- ros2 run pymoveit2 ex_collision_mesh.py --ros-args -p position:="[0.5, 0.0, 0.5]" -p quat_xyzw:="[0.0, 0.0, -0.7071, 0.7071]" -p scale:="[1.0, 1.0, 1.0]"
 - ros2 run pymoveit2 ex_collision_mesh.py --ros-args -p action:="move" -p position:="[0.2, 0.0, 0.2]"
 - ros2 run pymoveit2 ex_collision_mesh.py --ros-args -p filepath:="./my_favourity_mesh.stl"
 - ros2 run pymoveit2 ex_collision_mesh.py --ros-args -p action:="remove"
@@ -40,6 +40,7 @@ def main():
     )
     node.declare_parameter("position", [0.5, 0.0, 0.5])
     node.declare_parameter("quat_xyzw", [0.0, 0.0, -0.7071, 0.7071])
+    node.declare_parameter("scale", [1.0, 1.0, 1.0])
 
     # Create callback group that allows execution of callbacks in parallel without restrictions
     callback_group = ReentrantCallbackGroup()
@@ -66,10 +67,13 @@ def main():
     action = node.get_parameter("action").get_parameter_value().string_value
     position = node.get_parameter("position").get_parameter_value().double_array_value
     quat_xyzw = node.get_parameter("quat_xyzw").get_parameter_value().double_array_value
+    scale = node.get_parameter("scale").get_parameter_value().double_array_value
 
     # Use the default example mesh if invalid
     if not filepath:
-        node.get_logger().info(f"Using the default example mesh file")
+        node.get_logger().info(
+            f"Using the default example mesh file {DEFAULT_EXAMPLE_MESH}"
+        )
         filepath = DEFAULT_EXAMPLE_MESH
 
     # Make sure the mesh file exists
@@ -92,6 +96,7 @@ def main():
             id=object_id,
             position=position,
             quat_xyzw=quat_xyzw,
+            scale=scale,
         )
     elif action == "remove":
         # Remove collision mesh
diff --git a/pymoveit2/moveit2.py b/pymoveit2/moveit2.py
index 90ed568..ceb3670 100644
--- a/pymoveit2/moveit2.py
+++ b/pymoveit2/moveit2.py
@@ -3,6 +3,7 @@
 from enum import Enum
 from typing import List, Optional, Tuple, Union
 
+import numpy as np
 from action_msgs.msg import GoalStatus
 from geometry_msgs.msg import Point, Pose, PoseStamped, Quaternion
 from moveit_msgs.action import ExecuteTrajectory, MoveGroup
@@ -1618,12 +1619,14 @@ def add_collision_mesh(
         ] = None,
         frame_id: Optional[str] = None,
         operation: int = CollisionObject.ADD,
+        scale: Union[float, Tuple[float, float, float]] = 1.0,
     ):
         """
         Add collision object with a mesh geometry specified by `filepath`.
         Note: This function required 'trimesh' Python module to be installed.
         """
 
+        # Load the mesh
         try:
             import trimesh
         except ImportError as err:
@@ -1679,6 +1682,14 @@ def add_collision_mesh(
         )
 
         mesh = trimesh.load(filepath)
+
+        # Scale the mesh
+        if isinstance(scale, float):
+            scale = (scale, scale, scale)
+        transform = np.eye(4)
+        np.fill_diagonal(transform, scale)
+        mesh.apply_transform(transform)
+
         msg.meshes.append(
             Mesh(
                 triangles=[MeshTriangle(vertex_indices=face) for face in mesh.faces],