Refactor GetGrippingPointTool and introduce unit tests and manual tests

Author: Juliaj

pyproject.toml

@@ -71,7 +71,8 @@ build-backend = "poetry.core.masonry.api"
 markers = [
     "billable: marks test as billable  (deselect with '-m \"not billable\"')",
     "ci_only: marks test as cli only (deselect with '-m \"not ci_only\"')",
+    "manual: marks tests as manual (may require demo app to be running)",
 ]
-addopts = "-m 'not billable and not ci_only' --ignore=src"
+addopts = "-m 'not billable and not ci_only and not manual' --ignore=src"
 log_cli = true
 log_cli_level = "INFO"

src/rai_core/rai/__init__.py

@@ -20,7 +20,7 @@
     get_llm_model_direct,
     get_tracing_callbacks,
 )
-from .utils import timeout
+from .tools import timeout
 
 __all__ = [
     "AgentRunner",

src/rai_core/rai/tools/ros2/detection/pcl.py

@@ -11,6 +11,7 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+import time
 from typing import List, Literal, Optional, cast
 
 import numpy as np
@@ -46,6 +47,77 @@ def depth_to_point_cloud(
     return points.astype(np.float32, copy=False)
 
 
+def _publish_gripping_point_debug_data(
+    connector: ROS2Connector,
+    obj_points_xyz: NDArray[np.float32],
+    gripping_points_xyz: list[NDArray[np.float32]],
+    base_frame_id: str = "egoarm_base_link",
+    publish_duration: float = 10.0,
+) -> None:
+    """Publish the gripping point debug data for visualization in RVIZ via point cloud and marker array.
+
+    Args:
+        connector: The ROS2 connector.
+        obj_points_xyz: The list of objects found in the image.
+        gripping_points_xyz: The list of gripping points in the base frame.
+        base_frame_id: The base frame id.
+        publish_duration: Duration in seconds to publish the data (default: 10.0).
+    """
+
+    from geometry_msgs.msg import Point, Point32, Pose, Vector3
+    from sensor_msgs.msg import PointCloud
+    from std_msgs.msg import Header
+    from visualization_msgs.msg import Marker, MarkerArray
+
+    points = (
+        np.concatenate(obj_points_xyz, axis=0)
+        if obj_points_xyz
+        else np.zeros((0, 3), dtype=np.float32)
+    )
+
+    msg = PointCloud()  # type: ignore[reportUnknownArgumentType]
+    msg.header.frame_id = base_frame_id  # type: ignore[reportUnknownMemberType]
+    msg.points = [Point32(x=float(p[0]), y=float(p[1]), z=float(p[2])) for p in points]  # type: ignore[reportUnknownArgumentType]
+    pub = connector.node.create_publisher(  # type: ignore[reportUnknownMemberType]
+        PointCloud, "/debug_gripping_points_pointcloud", 10
+    )
+
+    marker_pub = connector.node.create_publisher(  # type: ignore[reportUnknownMemberType]
+        MarkerArray, "/debug_gripping_points_markerarray", 10
+    )
+    marker_array = MarkerArray()
+    header = Header()
+    header.frame_id = base_frame_id
+    header.stamp = connector.node.get_clock().now().to_msg()
+    markers = []
+    for i, p in enumerate(gripping_points_xyz):
+        m = Marker()
+        m.header = header
+        m.type = Marker.SPHERE
+        m.action = Marker.ADD
+        m.pose = Pose(position=Point(x=float(p[0]), y=float(p[1]), z=float(p[2])))
+        m.scale = Vector3(x=0.04, y=0.04, z=0.04)
+        m.id = i
+        m.color.r = 1.0  # type: ignore[reportUnknownMemberType]
+        m.color.g = 0.0  # type: ignore[reportUnknownMemberType]
+        m.color.b = 0.0  # type: ignore[reportUnknownMemberType]
+        m.color.a = 1.0  # type: ignore[reportUnknownMemberType]
+
+        # m.ns = str(i)
+
+        markers.append(m)  # type: ignore[reportUnknownArgumentType]
+    marker_array.markers = markers
+
+    start_time = time.time()
+    publish_rate = 10.0  # Hz
+    sleep_duration = 1.0 / publish_rate
+
+    while time.time() - start_time < publish_duration:
+        marker_pub.publish(marker_array)
+        pub.publish(msg)
+        time.sleep(sleep_duration)
+
+
 class PointCloudFromSegmentation:
     """Generate a masked point cloud for an object and transform it to a target frame.
 
@@ -511,96 +583,3 @@ def run(
                 f = pts
             filtered.append(f.astype(np.float32, copy=False))
         return filtered
-
-
-import time
-
-from rai.communication.ros2 import ROS2Context
-
-ROS2Context()
-
-
-def main():
-    from rai.communication.ros2.connectors import ROS2Connector
-
-    connector = ROS2Connector()
-    connector.node.declare_parameter("conversion_ratio", 1.0)
-    time.sleep(5)
-    est = GrippingPointEstimator(
-        strategy="biggest_plane", ransac_iterations=400, distance_threshold_m=0.008
-    )
-
-    pc_gen = PointCloudFromSegmentation(
-        connector=connector,
-        camera_topic="/rgbd_camera/camera_image_color",
-        depth_topic="/rgbd_camera/camera_image_depth",
-        camera_info_topic="/rgbd_camera/camera_info",
-        source_frame="egofront_rgbd_camera_depth_optical_frame",
-        target_frame="egoarm_base_link",
-    )
-    points_xyz = pc_gen.run(
-        object_name="box"
-    )  # ndarray of shape (N, 3) in target frame
-    print(points_xyz)
-    filt = PointCloudFilter(strategy="dbscan", dbscan_eps=0.02, dbscan_min_samples=10)
-    points_xyz = filt.run(points_xyz)  # same list shape, each np.float32 (N,3)
-    grip_points = est.run(points_xyz)
-
-    print(grip_points)
-    from geometry_msgs.msg import Point32
-    from sensor_msgs.msg import PointCloud
-
-    points = (
-        np.concatenate(points_xyz, axis=0)
-        if points_xyz
-        else np.zeros((0, 3), dtype=np.float32)
-    )
-
-    msg = PointCloud()  # type: ignore[reportUnknownArgumentType]
-    msg.header.frame_id = "egoarm_base_link"  # type: ignore[reportUnknownMemberType]
-    msg.points = [Point32(x=float(p[0]), y=float(p[1]), z=float(p[2])) for p in points]  # type: ignore[reportUnknownArgumentType]
-    pub = connector.node.create_publisher(  # type: ignore[reportUnknownMemberType]
-        PointCloud, "/debug/get_grabbing_point_pointcloud", 10
-    )
-    from geometry_msgs.msg import Point, Pose, Vector3
-    from std_msgs.msg import Header
-    from visualization_msgs.msg import Marker, MarkerArray
-
-    marker_pub = connector.node.create_publisher(  # type: ignore[reportUnknownMemberType]
-        MarkerArray, "/debug/get_grabbing_point_marker_array", 10
-    )
-    marker_array = MarkerArray()
-    header = Header()
-    header.frame_id = "egoarm_base_link"
-    # header.stamp = connector.node.get_clock().now().to_msg()
-    markers = []
-    for i, p in enumerate(grip_points):
-        m = Marker()
-        m.header = header
-        m.type = Marker.SPHERE
-        m.action = Marker.ADD
-        m.pose = Pose(position=Point(x=float(p[0]), y=float(p[1]), z=float(p[2])))
-        m.scale = Vector3(x=0.04, y=0.04, z=0.04)
-        m.id = i
-        m.color.r = 1.0  # type: ignore[reportUnknownMemberType]
-        m.color.g = 0.0  # type: ignore[reportUnknownMemberType]
-        m.color.b = 0.0  # type: ignore[reportUnknownMemberType]
-        m.color.a = 1.0  # type: ignore[reportUnknownMemberType]
-
-        # m.ns = str(i)
-
-        markers.append(m)  # type: ignore[reportUnknownArgumentType]
-    marker_array.markers = markers
-
-    while True:
-        connector.node.get_logger().info(  # type: ignore[reportUnknownMemberType]
-            f"publishing pointcloud to /debug/get_grabbing_point_pointcloud: {len(msg.points)} points, mean: {np.array(points.mean(axis=0))}."
-        )
-
-        marker_pub.publish(marker_array)
-        pub.publish(msg)
-        time.sleep(0.1)
-
-
-if __name__ == "__main__":
-    main()

src/rai_core/rai/tools/ros2/detection/tools.py

@@ -12,11 +12,10 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import Type
+from typing import Any, Optional, Type
 
 from pydantic import BaseModel, Field
 
-from rai.tools import timeout
 from rai.tools.ros2.base import BaseROS2Tool
 from rai.tools.ros2.detection.pcl import (
     GrippingPointEstimator,
@@ -33,27 +32,51 @@ class GetGrippingPointToolInput(BaseModel):
 
 
 # TODO(maciejmajek): Configuration system configurable with namespacing
+# TODO(juliajia): Question for Maciej: for comments above on configuration system with namespacing, can you provide an use case for this?
 class GetGrippingPointTool(BaseROS2Tool):
     name: str = "get_gripping_point"
     description: str = "Get gripping points for specified object/objects. Returns 3D coordinates where a robot gripper can grasp the object."
 
-    point_cloud_from_segmentation: PointCloudFromSegmentation
+    target_frame: str
+    source_frame: str
+    camera_topic: str  # rgb camera topic
+    depth_topic: str
+    camera_info_topic: str  # rgb camera info topic
+
     gripping_point_estimator: GrippingPointEstimator
     point_cloud_filter: PointCloudFilter
 
+    # Auto-initialized in model_post_init
+    point_cloud_from_segmentation: Optional[PointCloudFromSegmentation] = None
+
     timeout_sec: float = Field(
         default=10.0, description="Timeout in seconds to get the gripping point"
     )
 
     args_schema: Type[GetGrippingPointToolInput] = GetGrippingPointToolInput
 
-    def _run(self, object_name: str) -> str:
-        @timeout(
-            self.timeout_sec,
-            f"Gripping point detection for object '{object_name}' exceeded {self.timeout_sec} seconds",
+    def model_post_init(self, __context: Any) -> None:
+        """Initialize PointCloudFromSegmentation with the provided camera parameters."""
+        self.point_cloud_from_segmentation = PointCloudFromSegmentation(
+            connector=self.connector,
+            camera_topic=self.camera_topic,
+            depth_topic=self.depth_topic,
+            camera_info_topic=self.camera_info_topic,
+            source_frame=self.source_frame,
+            target_frame=self.target_frame,
         )
+
+    def _run(self, object_name: str) -> str:
+        # this will be not work in agent scenario because signal need to be run in main thread, comment out for now
+        # @timeout(
+        #     self.timeout_sec,
+        #     f"Gripping point detection for object '{object_name}' exceeded {self.timeout_sec} seconds",
+        # )
         def _run_with_timeout():
             pcl = self.point_cloud_from_segmentation.run(object_name)
+            if len(pcl) == 0:
+                return f"No {object_name}s detected."
+
             pcl = self.point_cloud_filter.run(pcl)
             gps = self.gripping_point_estimator.run(pcl)
 

src/rai_extensions/rai_open_set_vision/rai_open_set_vision/tools/segmentation_tools.py

@@ -293,23 +293,24 @@ def _process_mask(
 
         points = pcd
         # publish resulting pointcloud
-        import time
-
-        from geometry_msgs.msg import Point32
-        from sensor_msgs.msg import PointCloud
-
-        msg = PointCloud()
-        msg.header.frame_id = "egofront_rgbd_camera_depth_optical_frame"
-        msg.points = [Point32(x=p[0], y=p[1], z=p[2]) for p in points]
-        pub = self.connector.node.create_publisher(
-            PointCloud, "/debug/get_grabbing_point_pointcloud", 10
-        )
-        while True:
-            self.connector.node.get_logger().info(
-                f"publishing pointcloud to /debug/get_grabbing_point_pointcloud: {len(msg.points)} points, mean: {np.array(points.mean(axis=0))}."
-            )
-            pub.publish(msg)
-            time.sleep(0.1)
+        # TODO(juliajia): remove this after debugging
+        # import time
+
+        # from geometry_msgs.msg import Point32
+        # from sensor_msgs.msg import PointCloud
+
+        # msg = PointCloud()
+        # msg.header.frame_id = "egofront_rgbd_camera_depth_optical_frame"
+        # msg.points = [Point32(x=p[0], y=p[1], z=p[2]) for p in points]
+        # pub = self.connector.node.create_publisher(
+        #     PointCloud, "/debug/get_grabbing_point_pointcloud", 10
+        # )
+        # while True:
+        #     self.connector.node.get_logger().info(
+        #         f"publishing pointcloud to /debug/get_grabbing_point_pointcloud: {len(msg.points)} points, mean: {np.array(points.mean(axis=0))}."
+        #     )
+        #     pub.publish(msg)
+        #     time.sleep(0.1)
 
         # https://github.com/ycheng517/tabletop-handybot/blob/6d401e577e41ea86529d091b406fbfc936f37a8d/tabletop_handybot/tabletop_handybot/tabletop_handybot_node.py#L413-L424
         grasp_z = points[:, 2].max()