Add example SAM 2 behavior and objective

.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,76 @@
+#pragma once
+
+
+#include <moveit_studio_behavior_interface/async_behavior_base.hpp>
+#include <moveit_pro_ml/onnx_sam2.hpp>
+#include <sensor_msgs/msg/image.hpp>
+#include <moveit_studio_vision_msgs/msg/mask2_d.hpp>
+#include <fmt/format.h>
+
+namespace example_behaviors
+{
+/**
+ * @brief Segment an image using the SAM 2 model
+ */
+class SAM2Segmentation : public moveit_studio::behaviors::AsyncBehaviorBase
+{
+public:
+/**
+   * @brief Constructor for the sam2_segmentation 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 sam2_segmentation 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 sam2_segmentation does not use expose any ports, so this function returns an empty list.
+  */
+ 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:
+  /**
+  * @brief Convert a ROS image message to the ONNX image format used by the SAM 2 model.
+   * @param image_msg The ROS message to be converted
+   * @param onnx_image The ONNX image
+   */
+ void set_onnx_image_from_ros_image(const sensor_msgs::msg::Image& image_msg, moveit_pro_ml::ONNXImage& onnx_image);
+
+  /**
+  * @brief Convert an ONNX image to a ROS image message.
+   * @param onnx_image The ONNX image to be converted
+   * @param image_msg The ROS message
+   */
+ void set_ros_image_from_onnx_image(const moveit_pro_ml::ONNXImage& onnx_image,  sensor_msgs::msg::Image& image_msg);
+
+ std::shared_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_;
+
+  /** @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,137 @@
+#include <spdlog/spdlog.h>
+
+#include <example_behaviors/sam2_segmentation.hpp>
+
+#include <geometry_msgs/msg/point_stamped.hpp>
+#include <moveit_studio_behavior_interface/get_required_ports.hpp>
+#include <ament_index_cpp/get_package_share_directory.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
+{
+  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)
+  {
+
+    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_shared<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.")
+
+
+    };
+  }
+
+  void SAM2Segmentation::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);
+     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;
+    }
+  }
+
+
+  void SAM2Segmentation::set_ros_image_from_onnx_image(const moveit_pro_ml::ONNXImage& onnx_image, sensor_msgs::msg::Image& image_msg)
+  {
+    image_msg.height = static_cast<uint32_t>(onnx_image.shape[0]);
+    image_msg.width = static_cast<uint32_t>(onnx_image.shape[1]);
+    image_msg.encoding = "mono8";
+    image_msg.data.resize(image_msg.height * image_msg.width);
+    image_msg.step = image_msg.width;
+    for (size_t i = 0; i < onnx_image.data.size(); ++i)
+    {
+      image_msg.data[i] = onnx_image.data[i] > 0.5 ? 255: 0;
+    }
+   }
+
+  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 point 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
+    {
+      auto masks = sam2_->predict(onnx_image_, point_prompts);
+
+      mask_image_msg_.header = image_msg.header;
+      set_ros_image_from_onnx_image(masks, mask_image_msg_);
+      mask_.pixels = mask_image_msg_;
+      mask_.x = 0;
+      mask_.y = 0;
+      setOutput<std::vector<moveit_studio_vision_msgs::msg::Mask2D>>(kPortMasks, {mask_});
+    }
+    catch (std::invalid_argument& e)
+    {
+      auto error_message = fmt::format("Invalid argument: {}", e.what());
+      return tl::make_unexpected(error_message);
+    }
+
+    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>