From 8b79cc0fa5e30d9a2cff7fe81c360ab4864cd9f2 Mon Sep 17 00:00:00 2001
From: Paul Gesel <paul.gesel@picknik.ai>
Date: Mon, 13 Jan 2025 10:24:06 -0700
Subject: [PATCH] Use core segmentation behavior in lab_sim

Signed-off-by: Paul Gesel <paul.gesel@picknik.ai>
---
 src/example_behaviors/CMakeLists.txt          |  2 -
 src/lab_sim/CMakeLists.txt                    |  1 +
 .../models/decoder.onnx                       |  0
 .../models/sam2_hiera_large_encoder.onnx      |  0
 src/lab_sim/objectives/run_sam2_onnx.xml      | 41 ---------
 ...segment_point_cloud_from_clicked_point.xml | 88 +++++++++++++++++++
 6 files changed, 89 insertions(+), 43 deletions(-)
 rename src/{example_behaviors => lab_sim}/models/decoder.onnx (100%)
 rename src/{example_behaviors => lab_sim}/models/sam2_hiera_large_encoder.onnx (100%)
 delete mode 100644 src/lab_sim/objectives/run_sam2_onnx.xml
 create mode 100644 src/lab_sim/objectives/segment_point_cloud_from_clicked_point.xml

diff --git a/src/example_behaviors/CMakeLists.txt b/src/example_behaviors/CMakeLists.txt
index 2c7cd6a0..6f54a740 100644
--- a/src/example_behaviors/CMakeLists.txt
+++ b/src/example_behaviors/CMakeLists.txt
@@ -15,7 +15,6 @@ example_interfaces)
 foreach(package IN ITEMS ${THIS_PACKAGE_INCLUDE_DEPENDS})
   find_package(${package} REQUIRED)
 endforeach()
-find_package(moveit_pro_ml REQUIRED)
 
 add_library(
   example_behaviors
@@ -41,7 +40,6 @@ target_include_directories(
   PRIVATE ${PCL_INCLUDE_DIRS})
 ament_target_dependencies(example_behaviors
                           ${THIS_PACKAGE_INCLUDE_DEPENDS})
-target_link_libraries(example_behaviors onnx_sam2)
 
 # Install Libraries
 install(
diff --git a/src/lab_sim/CMakeLists.txt b/src/lab_sim/CMakeLists.txt
index a2db2edd..2b665a3f 100644
--- a/src/lab_sim/CMakeLists.txt
+++ b/src/lab_sim/CMakeLists.txt
@@ -33,6 +33,7 @@ install(
     config
     description
     launch
+    models
     objectives
     waypoints
   DESTINATION
diff --git a/src/example_behaviors/models/decoder.onnx b/src/lab_sim/models/decoder.onnx
similarity index 100%
rename from src/example_behaviors/models/decoder.onnx
rename to src/lab_sim/models/decoder.onnx
diff --git a/src/example_behaviors/models/sam2_hiera_large_encoder.onnx b/src/lab_sim/models/sam2_hiera_large_encoder.onnx
similarity index 100%
rename from src/example_behaviors/models/sam2_hiera_large_encoder.onnx
rename to src/lab_sim/models/sam2_hiera_large_encoder.onnx
diff --git a/src/lab_sim/objectives/run_sam2_onnx.xml b/src/lab_sim/objectives/run_sam2_onnx.xml
deleted file mode 100644
index f05b6821..00000000
--- a/src/lab_sim/objectives/run_sam2_onnx.xml
+++ /dev/null
@@ -1,41 +0,0 @@
-<?xml version="1.0" encoding="utf-8" ?>
-<root
-  BTCPP_format="4"
-  main_tree_to_execute="Segment Point Cloud from Clicked Point"
->
-  <!--//////////-->
-  <BehaviorTree
-    ID="Segment Point Cloud from Clicked Point"
-    _description="Captures a point cloud and requests the user to click an object in the image to be segmented. The point cloud is then filtered to only include the selected object."
-    _favorite="true"
-  >
-    <Control ID="Sequence">
-      <Action ID="ClearSnapshot" />
-      <Action ID="GetImage" topic_name="/wrist_camera/color" />
-      <Action
-        ID="GetPointsFromUser"
-        point_prompts="Select the object to be segmented;"
-        point_names="Point1;"
-        view_name="/wrist_camera/color"
-      />
-      <Action ID="ExampleSAM2Segmentation" />
-      <Action ID="GetPointCloud" topic_name="/wrist_camera/points" />
-      <Action
-        ID="GetCameraInfo"
-        topic_name="/wrist_camera/camera_info"
-        message_out="{camera_info}"
-        timeout_sec="5.000000"
-      />
-      <Action ID="GetMasks3DFromMasks2D" />
-      <Decorator ID="ForEach" vector_in="{masks3d}" out="{mask3d}">
-        <Action ID="GetPointCloudFromMask3D" point_cloud="{point_cloud}" />
-      </Decorator>
-      <Action ID="SendPointCloudToUI" point_cloud="{point_cloud_fragment}" />
-      <Action ID="PublishPointCloud" point_cloud="{point_cloud_fragment}" />
-      <Action ID="SwitchUIPrimaryView" />
-    </Control>
-  </BehaviorTree>
-  <TreeNodesModel>
-    <SubTree ID="Segment Point Cloud from Clicked Point" />
-  </TreeNodesModel>
-</root>
diff --git a/src/lab_sim/objectives/segment_point_cloud_from_clicked_point.xml b/src/lab_sim/objectives/segment_point_cloud_from_clicked_point.xml
new file mode 100644
index 00000000..87a6908d
--- /dev/null
+++ b/src/lab_sim/objectives/segment_point_cloud_from_clicked_point.xml
@@ -0,0 +1,88 @@
+<?xml version="1.0" encoding="utf-8" ?>
+<root
+  BTCPP_format="4"
+  main_tree_to_execute="Segment Point Cloud from Clicked Point"
+>
+  <!--//////////-->
+  <BehaviorTree
+    ID="Segment Point Cloud from Clicked Point"
+    _description="Captures a point cloud and requests the user to click an object in the image to be segmented. The point cloud is then filtered to only include the selected object."
+    _favorite="true"
+  >
+    <Control ID="Sequence">
+      <Action ID="ClearSnapshot" />
+      <Action
+        ID="GetImage"
+        topic_name="/wrist_camera/color"
+        timeout_sec="5.000000"
+        message_out="{image}"
+      />
+      <Action
+        ID="GetPointsFromUser"
+        point_prompts="Select the object to be segmented;"
+        point_names="Point1;"
+        view_name="/wrist_camera/color"
+        pixel_coords="{pixel_coords}"
+      />
+      <Action
+        ID="GetMasks2DFromPointQuery"
+        image="{image}"
+        pixel_coords="{pixel_coords}"
+        masks2d="{masks2d}"
+        decoder_model_path="models/decoder.onnx"
+        encoder_model_path="models/sam2_hiera_large_encoder.onnx"
+        model_package="lab_sim"
+      />
+      <Action
+        ID="GetPointCloud"
+        topic_name="/wrist_camera/points"
+        timeout_sec="5.000000"
+        message_out="{point_cloud}"
+      />
+      <Action
+        ID="GetCameraInfo"
+        topic_name="/wrist_camera/camera_info"
+        message_out="{camera_info}"
+        timeout_sec="5.000000"
+      />
+      <Action
+        ID="GetMasks3DFromMasks2D"
+        camera_info="{camera_info}"
+        masks2d="{masks2d}"
+        point_cloud="{point_cloud}"
+        masks3d="{masks3d}"
+      />
+      <Decorator
+        ID="ForEach"
+        vector_in="{masks3d}"
+        out="{mask3d}"
+        index="{index}"
+      >
+        <Action
+          ID="GetPointCloudFromMask3D"
+          point_cloud="{point_cloud}"
+          mask3d="{mask3d}"
+          point_cloud_fragment="{point_cloud_fragment}"
+        />
+      </Decorator>
+      <Action
+        ID="SendPointCloudToUI"
+        point_cloud="{point_cloud_fragment}"
+        pcd_topic="/pcd_pointcloud_captures"
+      />
+      <Action
+        ID="PublishPointCloud"
+        point_cloud="{point_cloud_fragment}"
+        point_cloud_topic="/my_point_cloud"
+      />
+      <Action ID="SwitchUIPrimaryView" primary_view_name="Visualization" />
+    </Control>
+  </BehaviorTree>
+  <TreeNodesModel>
+    <SubTree ID="Segment Point Cloud from Clicked Point">
+      <MetadataFields>
+        <Metadata runnable="true" />
+      </MetadataFields>
+    </SubTree>
+  </TreeNodesModel>
+</root>