diff --git a/src/example_behaviors/CMakeLists.txt b/src/example_behaviors/CMakeLists.txt index 2c7cd6a0..1a2649e6 100644 --- a/src/example_behaviors/CMakeLists.txt +++ b/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( diff --git a/src/example_behaviors/COLCON_IGNORE b/src/example_behaviors/COLCON_IGNORE deleted file mode 100644 index e69de29b..00000000 diff --git a/src/example_behaviors/include/example_behaviors/l2g.hpp b/src/example_behaviors/include/example_behaviors/l2g.hpp new file mode 100644 index 00000000..0ed93191 --- /dev/null +++ b/src/example_behaviors/include/example_behaviors/l2g.hpp @@ -0,0 +1,69 @@ +#pragma once + +#include +#include +#include + +#include +#include +#include +#include +#include +#include + + +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& 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 doWork() override; + + +private: + std::shared_ptr l2g_; + + /** @brief Classes derived from AsyncBehaviorBase must implement getFuture() so that it returns a shared_future class member */ + std::shared_future>& getFuture() override + { + return future_; + } + + /** @brief Classes derived from AsyncBehaviorBase must have this shared_future as a class member */ + std::shared_future> future_; + +}; +} // namespace sam2_segmentation diff --git a/src/example_behaviors/models/l2g.onnx b/src/example_behaviors/models/l2g.onnx new file mode 100644 index 00000000..b108aac6 --- /dev/null +++ b/src/example_behaviors/models/l2g.onnx @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:fc1fd4a7cf185225ba29940da16cf052cc9801a4883aecffd82845627e0965b5 +size 15618125 diff --git a/src/example_behaviors/src/l2g.cpp b/src/example_behaviors/src/l2g.cpp new file mode 100644 index 00000000..edbeeef2 --- /dev/null +++ b/src/example_behaviors/src/l2g.cpp @@ -0,0 +1,184 @@ +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include +#include +#include + +#include + +namespace moveit_pro_ml +{ + class L2GModel + { + public: + L2GModel(const std::filesystem::path& onnx_file) + : model{std::make_shared(onnx_file)} + { + } + + [[nodiscard]] std::vector predict(const std::vector>& points) const + { + // move image into tensor + const int num_points = points.size(); + const auto point_data = model->create_dynamic_tensor(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(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 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& 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(onnx_file); + } + + BT::PortsList L2GBehavior::providedPorts() + { + return { + BT::InputPort(kPortPointCloud, kPortPointCloudDefault, + "The point cloud to grasp. The point cloud should be approximately fit in a 0.22 meter cube."), + BT::OutputPort>(kPortGrasps, kPortGraspsDefault, + "The grasp poses in a vector of geometry_msgs::msg::PoseStamped messages.") + }; + } + + tl::expected L2GBehavior::doWork() + { + const auto ports = moveit_studio::behaviors::getRequiredInputs( + getInput(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::fromROSMsg(point_cloud_msg, *cloud); + auto filtered_cloud = std::make_shared>(); + 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(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> 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> 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 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 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>(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 diff --git a/src/example_behaviors/src/register_behaviors.cpp b/src/example_behaviors/src/register_behaviors.cpp index e95a4c03..20ca08ca 100644 --- a/src/example_behaviors/src/register_behaviors.cpp +++ b/src/example_behaviors/src/register_behaviors.cpp @@ -14,6 +14,7 @@ #include #include #include +#include #include #include @@ -42,6 +43,7 @@ class ExampleBehaviorsLoader : public moveit_studio::behaviors::SharedResourcesN moveit_studio::behaviors::registerBehavior(factory, "ExampleNDTRegistration", shared_resources); moveit_studio::behaviors::registerBehavior(factory, "ExampleRANSACRegistration", shared_resources); moveit_studio::behaviors::registerBehavior(factory, "ExampleSAM2Segmentation", shared_resources); + moveit_studio::behaviors::registerBehavior(factory, "L2GBehavior", shared_resources); } }; } // namespace example_behaviors diff --git a/src/lab_sim/description/scene.xml b/src/lab_sim/description/scene.xml index 9e7726e7..8820c8e2 100644 --- a/src/lab_sim/description/scene.xml +++ b/src/lab_sim/description/scene.xml @@ -4,6 +4,10 @@ + + + +