Hackathon Paul, David, Shaur: add L2G network in ONNX

src/example_behaviors/CMakeLists.txt

@@ -33,6 +33,7 @@ add_library(
   src/example_ndt_registration.cpp
   src/example_ransac_registration.cpp
   src/example_sam2_segmentation.cpp
+  src/l2g.cpp
   src/register_behaviors.cpp)
 target_include_directories(
   example_behaviors
@@ -41,7 +42,7 @@ target_include_directories(
   PRIVATE ${PCL_INCLUDE_DIRS})
 ament_target_dependencies(example_behaviors
                           ${THIS_PACKAGE_INCLUDE_DEPENDS})
-target_link_libraries(example_behaviors onnx_sam2)
+target_link_libraries(example_behaviors onnx_sam2 moveit_pro_ml)
 
 # Install Libraries
 install(

src/example_behaviors/include/example_behaviors/l2g.hpp

@@ -0,0 +1,69 @@
+#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 moveit_pro_ml
+{
+  class L2GModel;
+}
+
+namespace example_behaviors
+{
+/**
+ * @brief Segment an image using the SAM 2 model
+ */
+class L2GBehavior : public moveit_studio::behaviors::AsyncBehaviorBase
+{
+public:
+/**
+   * @brief Constructor for the L2GBehavior 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.
+   */
+ L2GBehavior(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::shared_ptr<moveit_pro_ml::L2GModel> l2g_;
+
+  /** @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/l2g.cpp

@@ -0,0 +1,184 @@
+#include <future>
+#include <memory>
+#include <string>
+
+#include <ament_index_cpp/get_package_share_directory.hpp>
+#include <geometry_msgs/msg/point_stamped.hpp>
+#include <geometry_msgs/msg/pose_stamped.hpp>
+#include <moveit_pro_ml/onnx_model.hpp>
+#include <moveit_studio_behavior_interface/async_behavior_base.hpp>
+#include <moveit_studio_behavior_interface/get_required_ports.hpp>
+#include <moveit_studio_vision/pointcloud/point_cloud_tools.hpp>
+#include <sensor_msgs/msg/point_cloud2.hpp>
+#include <tl_expected/expected.hpp>
+
+#include <pcl/filters/filter.h>
+#include <pcl/point_cloud.h>
+#include <pcl_conversions/pcl_conversions.h>
+
+#include <example_behaviors/l2g.hpp>
+
+namespace moveit_pro_ml
+{
+  class L2GModel
+  {
+  public:
+    L2GModel(const std::filesystem::path& onnx_file)
+      : model{std::make_shared<ONNXTensorModel>(onnx_file)}
+    {
+    }
+
+    [[nodiscard]] std::vector<float> predict(const std::vector<std::array<float, 3>>& points) const
+    {
+      // move image into tensor
+      const int num_points = points.size();
+      const auto point_data = model->create_dynamic_tensor<float>(model->dynamic_inputs.at("point_cloud"),
+                                                                  {1, num_points, 3});
+      std::copy_n(points.data()->data(), num_points * 3, point_data);
+
+      auto pred = model->predict_base(model->inputs, model->dynamic_inputs);
+
+      // copy out predicted_grasps data
+      const auto shape = pred.at("predicted_grasps").onnx_shape;
+      const auto predicted_grasps_data = model->get_tensor_data<float>(pred.at("predicted_grasps"));
+
+      // output format [center_x, center_y, center_z, qw, qx, qy, qz, grasp_width] x N
+      return {predicted_grasps_data, predicted_grasps_data + get_size(shape)};
+    }
+
+    std::shared_ptr<ONNXTensorModel> model;
+  };
+} // namespace moveit_pro_ml
+
+
+namespace
+{
+  constexpr auto kPortPointCloud = "point_cloud";
+  constexpr auto kPortPointCloudDefault = "{point_cloud}";
+  constexpr auto kPortGrasps = "grasps";
+  constexpr auto kPortGraspsDefault = "{grasps}";
+  constexpr auto kNumberOfPoints = 5000;
+  constexpr auto kMaxGrasps = 10;
+} // namespace
+
+namespace example_behaviors
+{
+  L2GBehavior::L2GBehavior(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 onnx_file = package_path / "models" / "l2g.onnx";
+    l2g_ = std::make_unique<moveit_pro_ml::L2GModel>(onnx_file);
+  }
+
+  BT::PortsList L2GBehavior::providedPorts()
+  {
+    return {
+      BT::InputPort<sensor_msgs::msg::PointCloud2>(kPortPointCloud, kPortPointCloudDefault,
+                                                   "The point cloud to grasp. The point cloud should be approximately fit in a 0.22 meter cube."),
+      BT::OutputPort<std::vector<moveit_studio_vision_msgs::msg::Mask2D>>(kPortGrasps, kPortGraspsDefault,
+                                                                          "The grasp poses in a vector of <code>geometry_msgs::msg::PoseStamped</code> messages.")
+    };
+  }
+
+  tl::expected<bool, std::string> L2GBehavior::doWork()
+  {
+    const auto ports = moveit_studio::behaviors::getRequiredInputs(
+      getInput<sensor_msgs::msg::PointCloud2>(kPortPointCloud));
+
+    // 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& [point_cloud_msg] = ports.value();
+
+    const auto cloud = std::make_shared<pcl::PointCloud<pcl::PointXYZ>>();
+    pcl::fromROSMsg(point_cloud_msg, *cloud);
+    auto filtered_cloud = std::make_shared<pcl::PointCloud<pcl::PointXYZ>>();
+    pcl::Indices index;
+    pcl::removeNaNFromPointCloud(*cloud, *filtered_cloud, index);
+    if (filtered_cloud->size() < kNumberOfPoints)
+    {
+      auto error_message = fmt::format("Point cloud must have at least {} points in it.", 5000);
+      return tl::make_unexpected(error_message);
+    }
+    double fraction = std::min(static_cast<double>(kNumberOfPoints) / filtered_cloud->points.size(), 1.0);
+    auto downsampled_cloud = moveit_studio::point_cloud_tools::downsampleRandom(filtered_cloud, fraction);
+
+    // Subtract centroid from the point cloud
+    pcl::PointXYZ centroid;
+    computeCentroid(*downsampled_cloud, centroid);
+
+    std::vector<std::array<float, 3>> points_centered;
+    points_centered.reserve(downsampled_cloud->points.size());
+    for (auto& point : downsampled_cloud->points)
+    {
+      points_centered.push_back({point.x - centroid.x, point.y - centroid.y, point.z - centroid.z});
+    }
+
+    // Filter out points outside the bounds defines by L2G.
+    std::vector<std::array<float, 3>> points;
+    points.reserve(downsampled_cloud->points.size());
+    for (auto& point : points_centered)
+    {
+      if (abs(point[0]) > .22 / 2.0 || abs(point[1]) > .22 / 2.0 || abs(point[2]) > .22)
+      {
+        continue;
+      }
+      points.push_back({2.0f * point[0] / .22f, 2.0f * point[1] / .22f, point[2] / .22f});
+    }
+
+    // Run the network
+    std::vector<float> grasps;
+    try
+    {
+      grasps = l2g_->predict(points);
+    }
+    catch (const std::invalid_argument& e)
+    {
+      return tl::make_unexpected(fmt::format("Invalid argument: {}", e.what()));
+    }
+
+    // Copy of grasps into the output vector
+    std::vector<geometry_msgs::msg::PoseStamped> grasps_pose;
+    for (size_t i = 0; i < grasps.size(); i += 8)
+    {
+      geometry_msgs::msg::PoseStamped pose;
+      pose.header = point_cloud_msg.header;
+      pose.header.stamp = shared_resources_->node->now();
+
+      pose.pose.position.x = grasps[i + 0] + centroid.x;
+      pose.pose.position.y = grasps[i + 1] + centroid.y;
+      pose.pose.position.z = grasps[i + 2] + centroid.z;
+
+      pose.pose.orientation.w = grasps[i + 3];
+      pose.pose.orientation.x = grasps[i + 4];
+      pose.pose.orientation.y = grasps[i + 5];
+      pose.pose.orientation.z = grasps[i + 6];
+      grasps_pose.push_back(pose);
+
+      // break after kMaxGrasps grasps have been added
+      if (i >= 8*kMaxGrasps)
+      {
+        break;
+      }
+    }
+
+    setOutput<std::vector<geometry_msgs::msg::PoseStamped>>(kPortGrasps, grasps_pose);
+
+    return true;
+  }
+
+  BT::KeyValueVector L2GBehavior::metadata()
+  {
+    return {
+      {
+        "description",
+        "Generate a vector grasps from a point cloud using the L2G network."
+      }
+    };
+  }
+} // namespace example_behaviors

src/example_behaviors/src/register_behaviors.cpp

@@ -14,6 +14,7 @@
 #include <example_behaviors/example_setup_mtc_place_from_pose.hpp>
 #include <example_behaviors/example_ndt_registration.hpp>
 #include <example_behaviors/example_ransac_registration.hpp>
+#include <example_behaviors/l2g.hpp>
 #include <example_behaviors/example_sam2_segmentation.hpp>
 
 #include <pluginlib/class_list_macros.hpp>
@@ -42,6 +43,7 @@ class ExampleBehaviorsLoader : public moveit_studio::behaviors::SharedResourcesN
     moveit_studio::behaviors::registerBehavior<ExampleNDTRegistration>(factory, "ExampleNDTRegistration", shared_resources);
     moveit_studio::behaviors::registerBehavior<ExampleRANSACRegistration>(factory, "ExampleRANSACRegistration", shared_resources);
     moveit_studio::behaviors::registerBehavior<ExampleSAM2Segmentation>(factory, "ExampleSAM2Segmentation", shared_resources);
+    moveit_studio::behaviors::registerBehavior<L2GBehavior>(factory, "L2GBehavior", shared_resources);
   }
 };
 }  // namespace example_behaviors

src/lab_sim/description/scene.xml

@@ -4,6 +4,10 @@
   <include file="lab_desk/desk_globals.xml" />
   <include file="robotiq_2f85/robotiq2f85_globals.xml" />
 
+<visual>
+    <global offwidth="1920" offheight="1080"/>
+  </visual>
+
   <compiler angle="radian" autolimits="true" />
 
   <option integrator="implicitfast" />
@@ -39,7 +43,7 @@
       pos="-1.0 -2.0 2.966"
       fovy="58"
       mode="fixed"
-      resolution="640 480"
+      resolution="1920 1080"
       euler="0.849 0.000 0.0"
     />
     <!-- lab desk -->
@@ -72,10 +76,11 @@
   </worldbody>
 
   <keyframe>
+      <key qpos='-0.871046 0.729603 0.473715 -0.839072 4.40365e-07 3.86376e-06 -0.544021 -0.5 0.69 0.473487 -0.759682 -7.75899e-06 3.45645e-08 -0.650295 -0.36 0.61 0.473486 -0.759683 -9.25864e-06 5.20597e-07 -0.650294 -0.3 0.699998 0.473469 -0.75968 -2.82938e-05 1.77045e-05 -0.650297 1.07516e-07 0.9 0.474026 -0.839071 -1.39392e-05 -3.69846e-07 0.544021 -0.1 0.849998 0.474023 0.283661 1.33761e-06 1.40874e-05 -0.958925 0.10855 0.55 0.486243 -0.759692 -4.223e-06 1.66946e-05 -0.650283 0.2 0.75 0.499376 1 3.26949e-17 -1.9984e-19 3.18509e-31 1.70843e-05 0.749992 0.499377 1 0.00016562 0.000360588 7.99021e-12 -0.2 0.75 0.499376 1 3.26949e-17 -1.9984e-19 3.18509e-31 1.74013e-21 1.86939e-20 0.00262408 0.000725414 -2.32869e-07 2.82812e-08 -1.55995e-13 0.00260382 0.000146481 0.00285352 -0.00309142 0.00260383 0.000146481 0.00284971 -0.00308363'/>
     <key
       qpos="-0.886764 0.727551 0.472375 -0.834559 -6.95805e-05 -6.1375e-05 -0.550919 -0.499647 0.689928 0.470942 -0.749838 3.91843e-05 -5.08119e-05 -0.661622 -0.3609 0.609537 0.470941 -0.757741 2.96011e-05 4.00136e-05 -0.652556 -0.303591 0.696467 0.470989 -0.7788 -5.32927e-05 2.99347e-06 -0.627272 -0.000373058 0.88828 0.473534 -0.838137 -1.0566e-05 -3.07868e-06 0.54546 -0.0963647 0.840186 0.473528 0.281469 -1.36042e-05 6.60035e-06 -0.95957 0.0484372 0.842182 0.473688 -0.758342 6.07647e-06 2.73022e-05 -0.651856 0.2 0.75 0.499376 1 0 0 0 1.70629e-05 0.749992 0.499377 1 0.00016562 0.000360588 1.95345e-10 -0.2 0.75 0.499376 1 0 0 0 0 0 -1.88491 1.10028 -1.25653 -1.57075 -1.08677e-12 0.0208231 0.000536851 0.0221006 -0.022866 0.020837 0.000547055 0.0221413 -0.0229382"
       qvel="0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0"
       ctrl="0 -1.8849 1.0997 -1.2566 -1.57075 0 0 0"
-    />
+      />
   </keyframe>
 </mujoco>

src/lab_sim/objectives/integration_test.xml

@@ -0,0 +1,32 @@
+<?xml version="1.0" encoding="utf-8" ?>
+<root BTCPP_format="4" main_tree_to_execute="Integration Test Mock">
+  <!--//////////-->
+  <BehaviorTree
+    ID="Integration Test Mock"
+    _description="Mock Graspnet behavior to test taking snapshots and  visualizing pose vector"
+    _favorite="true"
+    _subtreeOnly="false"
+  >
+    <Control ID="Sequence" name="TopLevelSequence">
+      <SubTree
+        ID="Take Snapshots Of Object"
+        _collapsed="false"
+        merged_cloud="{merged_cloud}"
+      />
+      <Action
+        ID="GraspnetExampleBehavior"
+        point_cloud="{merged_cloud}"
+        num_samples="10"
+        stamped_pose="{output_pose}"
+      />
+      <SubTree
+        ID="Visualize Pose Vector"
+        _collapsed="true"
+        pose_vector="{output_pose}"
+      />
+    </Control>
+  </BehaviorTree>
+  <TreeNodesModel>
+    <SubTree ID="Integration Test Mock" />
+  </TreeNodesModel>
+</root>