Add example SAM 2 behavior and objective (#23)

* Add SAM2 behavior and ONNX models
---------

Signed-off-by: Paul Gesel <[email protected]>
Co-authored-by: Griswald Brooks <[email protected]>

.gitattributes

@@ -0,0 +1 @@
+*.onnx filter=lfs diff=lfs merge=lfs -text

src/example_behaviors/CMakeLists.txt

@@ -15,22 +15,24 @@ 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
   SHARED 
   src/add_two_ints_service_client.cpp
   src/convert_mtc_solution_to_joint_trajectory.cpp
   src/delayed_message.cpp
-  src/get_string_from_topic.cpp
   src/fibonacci_action_client.cpp
+  src/get_string_from_topic.cpp
   src/hello_world.cpp
+  src/ndt_registration.cpp
   src/publish_color_rgba.cpp
+  src/ransac_registration.cpp
+  src/sam2_segmentation.cpp
   src/setup_mtc_pick_from_pose.cpp
   src/setup_mtc_place_from_pose.cpp
   src/setup_mtc_wave_hand.cpp
-  src/ndt_registration.cpp
-  src/ransac_registration.cpp
   src/register_behaviors.cpp)
 target_include_directories(
   example_behaviors
@@ -39,6 +41,7 @@ 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(
@@ -50,7 +53,7 @@ install(
   INCLUDES
   DESTINATION include)
 
-install(DIRECTORY config DESTINATION share/${PROJECT_NAME})
+install(DIRECTORY config models DESTINATION share/${PROJECT_NAME})
 
 if(BUILD_TESTING)
   moveit_pro_behavior_test(example_behaviors)

src/example_behaviors/include/example_behaviors/sam2_segmentation.hpp

@@ -0,0 +1,67 @@
+#pragma once
+
+#include <future>
+#include <memory>
+#include <string>
+
+#include <moveit_pro_ml/onnx_sam2.hpp>
+#include <moveit_pro_ml/onnx_sam2_types.hpp>
+#include <moveit_studio_behavior_interface/async_behavior_base.hpp>
+#include <moveit_studio_vision_msgs/msg/mask2_d.hpp>
+#include <sensor_msgs/msg/image.hpp>
+#include <tl_expected/expected.hpp>
+
+
+namespace example_behaviors
+{
+/**
+ * @brief Segment an image using the SAM 2 model
+ */
+class SAM2Segmentation : public moveit_studio::behaviors::AsyncBehaviorBase
+{
+public:
+/**
+   * @brief Constructor for the SAM2Segmentation behavior.
+   * @param name The name of a particular instance of this Behavior. This will be set by the behavior tree factory when this Behavior is created within a new behavior tree.
+   * @param config This contains runtime configuration info for this Behavior, such as the mapping between the Behavior's data ports on the behavior tree's blackboard. This will be set by the behavior tree factory when this Behavior is created within a new behavior tree.
+   * @details An important limitation is that the members of the base Behavior class are not instantiated until after the initialize() function is called, so these classes should not be used within the constructor.
+   */
+ SAM2Segmentation(const std::string& name, const BT::NodeConfiguration& config,
+                                     const std::shared_ptr<moveit_studio::behaviors::BehaviorContext>& shared_resources);
+
+ /**
+  * @brief Implementation of the required providedPorts() function for the Behavior.
+  * @details The BehaviorTree.CPP library requires that Behaviors must implement a static function named providedPorts() which defines their input and output ports. If the Behavior does not use any ports, this function must return an empty BT::PortsList.
+  * This function returns a list of ports with their names and port info, which is used internally by the behavior tree.
+  * @return List of ports for the behavior.
+  */
+ static BT::PortsList providedPorts();
+
+ /**
+  * @brief Implementation of the metadata() function for displaying metadata, such as Behavior description and
+  * subcategory, in the MoveIt Studio Developer Tool.
+  * @return A BT::KeyValueVector containing the Behavior metadata.
+  */
+ static BT::KeyValueVector metadata();
+
+protected:
+  tl::expected<bool, std::string> doWork() override;
+
+
+private:
+ std::unique_ptr<moveit_pro_ml::SAM2> sam2_;
+ moveit_pro_ml::ONNXImage onnx_image_;
+ sensor_msgs::msg::Image mask_image_msg_;
+ moveit_studio_vision_msgs::msg::Mask2D mask_msg_;
+
+  /** @brief Classes derived from AsyncBehaviorBase must implement getFuture() so that it returns a shared_future class member */
+  std::shared_future<tl::expected<bool, std::string>>& getFuture() override
+  {
+    return future_;
+  }
+
+  /** @brief Classes derived from AsyncBehaviorBase must have this shared_future as a class member */
+  std::shared_future<tl::expected<bool, std::string>> future_;
+
+};
+}  // namespace sam2_segmentation

src/example_behaviors/src/register_behaviors.cpp

@@ -2,18 +2,19 @@
 #include <moveit_studio_behavior_interface/behavior_context.hpp>
 #include <moveit_studio_behavior_interface/shared_resources_node_loader.hpp>
 
-#include <example_behaviors/hello_world.hpp>
+#include <example_behaviors/add_two_ints_service_client.hpp>
 #include <example_behaviors/convert_mtc_solution_to_joint_trajectory.hpp>
 #include <example_behaviors/delayed_message.hpp>
-#include <example_behaviors/setup_mtc_wave_hand.hpp>
-#include <example_behaviors/add_two_ints_service_client.hpp>
 #include <example_behaviors/fibonacci_action_client.hpp>
 #include <example_behaviors/get_string_from_topic.hpp>
+#include <example_behaviors/hello_world.hpp>
+#include <example_behaviors/ndt_registration.hpp>
 #include <example_behaviors/publish_color_rgba.hpp>
+#include <example_behaviors/ransac_registration.hpp>
+#include <example_behaviors/sam2_segmentation.hpp>
 #include <example_behaviors/setup_mtc_pick_from_pose.hpp>
 #include <example_behaviors/setup_mtc_place_from_pose.hpp>
-#include <example_behaviors/ndt_registration.hpp>
-#include <example_behaviors/ransac_registration.hpp>
+#include <example_behaviors/setup_mtc_wave_hand.hpp>
 
 #include <pluginlib/class_list_macros.hpp>
 
@@ -35,6 +36,7 @@ class ExampleBehaviorsLoader : public moveit_studio::behaviors::SharedResourcesN
     moveit_studio::behaviors::registerBehavior<FibonacciActionClient>(factory, "FibonacciActionClient",
                                                                       shared_resources);
     moveit_studio::behaviors::registerBehavior<PublishColorRGBA>(factory, "PublishColorRGBA", shared_resources);
+    moveit_studio::behaviors::registerBehavior<SAM2Segmentation>(factory, "SAM2Segmentation", shared_resources);
     moveit_studio::behaviors::registerBehavior<SetupMtcPickFromPose>(factory, "SetupMtcPickFromPose", shared_resources);
     moveit_studio::behaviors::registerBehavior<SetupMtcPlaceFromPose>(factory, "SetupMtcPlaceFromPose",
                                                                       shared_resources);

src/example_behaviors/src/sam2_segmentation.cpp

@@ -0,0 +1,142 @@
+#include <future>
+#include <memory>
+#include <string>
+
+#include <ament_index_cpp/get_package_share_directory.hpp>
+#include <example_behaviors/sam2_segmentation.hpp>
+#include <geometry_msgs/msg/point_stamped.hpp>
+#include <moveit_pro_ml/onnx_sam2.hpp>
+#include <moveit_studio_behavior_interface/async_behavior_base.hpp>
+#include <moveit_studio_behavior_interface/get_required_ports.hpp>
+#include <moveit_studio_vision_msgs/msg/mask2_d.hpp>
+#include <sensor_msgs/msg/image.hpp>
+#include <tl_expected/expected.hpp>
+
+namespace
+{
+  constexpr auto kPortImage = "image";
+  constexpr auto kPortImageDefault = "{image}";
+  constexpr auto kPortPoint = "pixel_coords";
+  constexpr auto kPortPointDefault = "{pixel_coords}";
+  constexpr auto kPortMasks = "masks2d";
+  constexpr auto kPortMasksDefault = "{masks2d}";
+
+  constexpr auto kImageInferenceWidth = 1024;
+  constexpr auto kImageInferenceHeight = 1024;
+} // namespace
+
+namespace example_behaviors
+{
+  // Convert a ROS image message to the ONNX image format used by the SAM 2 model
+  void set_onnx_image_from_ros_image(const sensor_msgs::msg::Image& image_msg,
+                                                       moveit_pro_ml::ONNXImage& onnx_image)
+  {
+    onnx_image.shape = {1, image_msg.height, image_msg.width, 3};
+    onnx_image.data.resize(image_msg.height * image_msg.width * 3);
+    const int stride = image_msg.encoding != "rgb8" ? 3: 4;
+    for (size_t i = 0; i < onnx_image.data.size(); i+=stride)
+    {
+      onnx_image.data[i] = static_cast<float>(image_msg.data[i]) / 255.0f;
+      onnx_image.data[i+1] = static_cast<float>(image_msg.data[i+1]) / 255.0f;
+      onnx_image.data[i+2] = static_cast<float>(image_msg.data[i+2]) / 255.0f;
+    }
+  }
+
+  // Converts a single channel ONNX image mask to a ROS mask message.
+  void set_ros_mask_from_onnx_mask(const moveit_pro_ml::ONNXImage& onnx_image, sensor_msgs::msg::Image& mask_image_msg, moveit_studio_vision_msgs::msg::Mask2D& mask_msg)
+  {
+    mask_image_msg.height = static_cast<uint32_t>(onnx_image.shape[0]);
+    mask_image_msg.width = static_cast<uint32_t>(onnx_image.shape[1]);
+    mask_image_msg.encoding = "mono8";
+    mask_image_msg.data.resize(mask_image_msg.height * mask_image_msg.width);
+    mask_image_msg.step = mask_image_msg.width;
+    for (size_t i = 0; i < onnx_image.data.size(); ++i)
+    {
+      mask_image_msg.data[i] = onnx_image.data[i] > 0.5 ? 255: 0;
+    }
+    mask_msg.pixels = mask_image_msg;
+    mask_msg.x = 0;
+    mask_msg.y = 0;
+  }
+
+  SAM2Segmentation::SAM2Segmentation(const std::string& name, const BT::NodeConfiguration& config,
+                                     const std::shared_ptr<moveit_studio::behaviors::BehaviorContext>& shared_resources)
+    : moveit_studio::behaviors::AsyncBehaviorBase(name, config, shared_resources)
+  {
+
+    const std::filesystem::path package_path = ament_index_cpp::get_package_share_directory("example_behaviors");
+    const std::filesystem::path encoder_onnx_file = package_path / "models" / "sam2_hiera_large_encoder.onnx";
+    const std::filesystem::path decoder_onnx_file = package_path / "models" / "decoder.onnx";
+    sam2_ = std::make_unique<moveit_pro_ml::SAM2>(encoder_onnx_file, decoder_onnx_file);
+  }
+
+  BT::PortsList SAM2Segmentation::providedPorts()
+  {
+    return {
+      BT::InputPort<sensor_msgs::msg::Image>(kPortImage, kPortImageDefault,
+                                             "The Image to run segmentation on."),
+      BT::InputPort<std::vector<geometry_msgs::msg::PointStamped>>(kPortPoint, kPortPointDefault,
+                                                                   "The input points, as a vector of <code>geometry_msgs/PointStamped</code> messages to be used for segmentation."),
+
+  BT::OutputPort<std::vector<moveit_studio_vision_msgs::msg::Mask2D>>(kPortMasks, kPortMasksDefault,
+                                                                   "The masks contained in a vector of <code>moveit_studio_vision_msgs::msg::Mask2D</code> messages.")
+    };
+  }
+
+  tl::expected<bool, std::string> SAM2Segmentation::doWork()
+  {
+    const auto ports = moveit_studio::behaviors::getRequiredInputs(getInput<sensor_msgs::msg::Image>(kPortImage),
+                                                                   getInput<std::vector<
+                                                                     geometry_msgs::msg::PointStamped>>(kPortPoint));
+
+    // Check that all required input data ports were set.
+    if (!ports.has_value())
+    {
+      auto error_message = fmt::format("Failed to get required values from input data ports:\n{}", ports.error());
+      return tl::make_unexpected(error_message);
+    }
+    const auto& [image_msg, points_2d] = ports.value();
+
+    if (image_msg.encoding != "rgb8" && image_msg.encoding != "rgba8")
+    {
+      auto error_message = fmt::format("Invalid image message format. Expected `(rgb8, rgba8)` got :\n{}", image_msg.encoding);
+      return tl::make_unexpected(error_message);
+    }
+
+    // Create ONNX formatted image tensor from ROS image
+    set_onnx_image_from_ros_image(image_msg, onnx_image_);
+
+    std::vector<moveit_pro_ml::PointPrompt> point_prompts;
+    for (auto const& point : points_2d)
+    {
+      // Assume all points are the same label
+      point_prompts.push_back({{kImageInferenceWidth*static_cast<float>(point.point.x), kImageInferenceHeight*static_cast<float>(point.point.y)}, {1.0f}});
+    }
+
+    try
+    {
+      const auto masks = sam2_->predict(onnx_image_, point_prompts);
+
+      mask_image_msg_.header = image_msg.header;
+      set_ros_mask_from_onnx_mask(masks, mask_image_msg_, mask_msg_);
+
+      setOutput<std::vector<moveit_studio_vision_msgs::msg::Mask2D>>(kPortMasks, {mask_msg_});
+    }
+    catch (const std::invalid_argument& e)
+    {
+      return tl::make_unexpected(fmt::format("Invalid argument: {}", e.what()));
+    }
+
+    return true;
+  }
+
+  BT::KeyValueVector SAM2Segmentation::metadata()
+  {
+    return {
+      {
+        "description",
+        "Segments a ROS image message using the provided points represented as a vector of <code>geometry_msgs/PointStamped</code> messages."
+      }
+    };
+  }
+} // namespace sam2_segmentation

src/lab_sim/objectives/run_sam2_onnx.xml

@@ -0,0 +1,41 @@
+<?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="SAM2Segmentation" />
+      <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="ForEachMask3D" 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>