From 154e98f21c5150ab74503b501c55f837d36752d0 Mon Sep 17 00:00:00 2001
From: Amal Nanavati <amaln@uw.edu>
Date: Mon, 6 May 2024 16:50:11 -0700
Subject: [PATCH] Allow users to preload a collision mesh

---
 examples/ex_collision_mesh.py | 11 +++++++++++
 pymoveit2/moveit2.py          | 18 ++++++++++++++----
 2 files changed, 25 insertions(+), 4 deletions(-)

diff --git a/examples/ex_collision_mesh.py b/examples/ex_collision_mesh.py
index f0bea29..fa6ec0d 100755
--- a/examples/ex_collision_mesh.py
+++ b/examples/ex_collision_mesh.py
@@ -12,6 +12,7 @@
 from threading import Thread
 
 import rclpy
+import trimesh
 from rclpy.callback_groups import ReentrantCallbackGroup
 from rclpy.node import Node
 
@@ -41,6 +42,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])
+    node.declare_parameter("preload_mesh", False)
 
     # Create callback group that allows execution of callbacks in parallel without restrictions
     callback_group = ReentrantCallbackGroup()
@@ -68,6 +70,7 @@ def main():
     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
+    preload_mesh = node.get_parameter("preload_mesh").get_parameter_value().bool_value
 
     # Use the default example mesh if invalid
     if not filepath:
@@ -91,12 +94,20 @@ def main():
             f"Adding collision mesh '{filepath}' "
             f"{{position: {list(position)}, quat_xyzw: {list(quat_xyzw)}}}"
         )
+
+        # Load the mesh if specified
+        mesh = None
+        if preload_mesh:
+            mesh = trimesh.load(filepath)
+            filepath = None
+
         moveit2.add_collision_mesh(
             filepath=filepath,
             id=object_id,
             position=position,
             quat_xyzw=quat_xyzw,
             scale=scale,
+            mesh=mesh,
         )
     elif action == "remove":
         # Remove collision mesh
diff --git a/pymoveit2/moveit2.py b/pymoveit2/moveit2.py
index 5cbc753..d5ceedf 100644
--- a/pymoveit2/moveit2.py
+++ b/pymoveit2/moveit2.py
@@ -1,7 +1,7 @@
 import copy
 import threading
 from enum import Enum
-from typing import List, Optional, Tuple, Union
+from typing import Any, List, Optional, Tuple, Union
 
 import numpy as np
 from action_msgs.msg import GoalStatus
@@ -1611,7 +1611,7 @@ def add_collision_cone(
 
     def add_collision_mesh(
         self,
-        filepath: str,
+        filepath: Optional[str],
         id: str,
         pose: Optional[Union[PoseStamped, Pose]] = None,
         position: Optional[Union[Point, Tuple[float, float, float]]] = None,
@@ -1621,9 +1621,11 @@ def add_collision_mesh(
         frame_id: Optional[str] = None,
         operation: int = CollisionObject.ADD,
         scale: Union[float, Tuple[float, float, float]] = 1.0,
+        mesh: Optional[Any] = None,
     ):
         """
-        Add collision object with a mesh geometry specified by `filepath`.
+        Add collision object with a mesh geometry. Either `filepath` must be
+        specified or `mesh` must be provided.
         Note: This function required 'trimesh' Python module to be installed.
         """
 
@@ -1636,10 +1638,17 @@ def add_collision_mesh(
                 "to add collision objects into the MoveIt 2 planning scene."
             ) from err
 
+        # Check the parameters
         if (pose is None) and (position is None or quat_xyzw is None):
             raise ValueError(
                 "Either `pose` or `position` and `quat_xyzw` must be specified!"
             )
+        if (filepath is None and mesh is None) or (
+            filepath is not None and mesh is not None
+        ):
+            raise ValueError("Exactly one of `filepath` or `mesh` must be specified!")
+        if mesh is not None and not isinstance(mesh, trimesh.Trimesh):
+            raise ValueError("`mesh` must be an instance of `trimesh.Trimesh`!")
 
         if isinstance(pose, PoseStamped):
             pose_stamped = pose
@@ -1682,7 +1691,8 @@ def add_collision_mesh(
             pose=pose_stamped.pose,
         )
 
-        mesh = trimesh.load(filepath)
+        if filepath is not None:
+            mesh = trimesh.load(filepath)
 
         # Scale the mesh
         if isinstance(scale, float):