From 35ab4a69f6405d8ccb6b3755c5dd675103004c83 Mon Sep 17 00:00:00 2001 From: Paul Gesel Date: Thu, 12 Dec 2024 08:31:47 -0700 Subject: [PATCH] working demo Signed-off-by: Paul Gesel --- src/example_behaviors/CMakeLists.txt | 1 - src/example_behaviors/models/l2g.onnx | 4 +- src/example_behaviors/src/l2g.cpp | 112 ++++++++++------ .../src/register_behaviors.cpp | 4 +- src/lab_sim/description/scene.xml | 6 +- src/lab_sim/objectives/integration_test.xml | 5 +- src/lab_sim/objectives/l2g.xml | 121 ++++++++++++++---- .../objectives/segment_waypoint_cloud.xml | 45 +++++++ .../objectives/take_snapshots_of_object.xml | 4 +- src/lab_sim/waypoints/ur_waypoints.yaml | 102 +++++++++++++-- 10 files changed, 321 insertions(+), 83 deletions(-) create mode 100644 src/lab_sim/objectives/segment_waypoint_cloud.xml diff --git a/src/example_behaviors/CMakeLists.txt b/src/example_behaviors/CMakeLists.txt index ec37d6a..1a2649e 100644 --- a/src/example_behaviors/CMakeLists.txt +++ b/src/example_behaviors/CMakeLists.txt @@ -34,7 +34,6 @@ add_library( src/example_ransac_registration.cpp src/example_sam2_segmentation.cpp src/l2g.cpp - src/sam2_segmentation.cpp src/register_behaviors.cpp) target_include_directories( example_behaviors diff --git a/src/example_behaviors/models/l2g.onnx b/src/example_behaviors/models/l2g.onnx index e1eea70..b108aac 100644 --- a/src/example_behaviors/models/l2g.onnx +++ b/src/example_behaviors/models/l2g.onnx @@ -1,3 +1,3 @@ version https://git-lfs.github.com/spec/v1 -oid sha256:ca6f4adf53759848b6697fe9daa81267b6a09cbb098965ccdd438818bf8654ac -size 22982271 +oid sha256:fc1fd4a7cf185225ba29940da16cf052cc9801a4883aecffd82845627e0965b5 +size 15618125 diff --git a/src/example_behaviors/src/l2g.cpp b/src/example_behaviors/src/l2g.cpp index 3f36fc9..edbeeef 100644 --- a/src/example_behaviors/src/l2g.cpp +++ b/src/example_behaviors/src/l2g.cpp @@ -3,23 +3,20 @@ #include #include -#include #include #include +#include #include #include #include -#include -#include #include #include -#include -#include #include +#include +#include - -#include "moveit_pro_ml/onnx_model.hpp" +#include namespace moveit_pro_ml { @@ -34,18 +31,18 @@ namespace moveit_pro_ml [[nodiscard]] std::vector predict(const std::vector>& points) const { // move image into tensor - int num_points = points.size(); + 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 mask out and return reference - auto shape = pred.at("predicted_grasps").onnx_shape; + // 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")); - // const auto grasp_scores_data = model->get_tensor_data(pred.at("grasp_scores")); + // 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)}; } @@ -60,6 +57,8 @@ namespace 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 @@ -77,9 +76,9 @@ namespace example_behaviors { return { BT::InputPort(kPortPointCloud, kPortPointCloudDefault, - "The Image to run segmentation on."), + "The point cloud to grasp. The point cloud should be approximately fit in a 0.22 meter cube."), BT::OutputPort>(kPortGrasps, kPortGraspsDefault, - "The masks contained in a vector of moveit_studio_vision_msgs::msg::Mask2D messages.") + "The grasp poses in a vector of geometry_msgs::msg::PoseStamped messages.") }; } @@ -101,46 +100,75 @@ namespace example_behaviors auto filtered_cloud = std::make_shared>(); pcl::Indices index; pcl::removeNaNFromPointCloud(*cloud, *filtered_cloud, index); - shared_resources_->logger->publishWarnMessage("Pointcloud size: " + filtered_cloud->size()); - double fraction = std::min(10000.0 / filtered_cloud->points.size(), 1.0); + 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); - shared_resources_->logger->publishWarnMessage("Downsampled pointcloud size: " + downsampled_cloud->size()); - try + // 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) { - std::vector> points; - points.reserve(downsampled_cloud->points.size()); - for (auto& point : downsampled_cloud->points) - { - points.push_back({point.x, point.y, point.z}); - } + points_centered.push_back({point.x - centroid.x, point.y - centroid.y, point.z - centroid.z}); + } - const auto grasps = l2g_->predict(points); - std::vector grasps_pose; - for (size_t i = 0; i < grasps.size(); i += 7) + // 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) { - 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] + grasps[i + 3]) / 2.0; - pose.pose.position.y = (grasps[i + 1] + grasps[i + 4]) / 2.0; - pose.pose.position.z = (grasps[i + 2] + grasps[i + 5]) / 2.0; - grasps_pose.push_back(pose); - - if (i >= 70) - { - break; - } + continue; } + points.push_back({2.0f * point[0] / .22f, 2.0f * point[1] / .22f, point[2] / .22f}); + } - setOutput>(kPortGrasps, grasps_pose); + // 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; } @@ -149,8 +177,8 @@ namespace example_behaviors return { { "description", - "Generate a grasp from a point cloud and output as a geometry_msgs/PoseStamped message." + "Generate a vector grasps from a point cloud using the L2G network." } }; } -} // namespace sam2_segmentation +} // namespace example_behaviors diff --git a/src/example_behaviors/src/register_behaviors.cpp b/src/example_behaviors/src/register_behaviors.cpp index 293ff39..20ca08c 100644 --- a/src/example_behaviors/src/register_behaviors.cpp +++ b/src/example_behaviors/src/register_behaviors.cpp @@ -14,9 +14,8 @@ #include #include #include -#include #include -#include +#include #include @@ -45,7 +44,6 @@ class ExampleBehaviorsLoader : public moveit_studio::behaviors::SharedResourcesN moveit_studio::behaviors::registerBehavior(factory, "ExampleRANSACRegistration", shared_resources); moveit_studio::behaviors::registerBehavior(factory, "ExampleSAM2Segmentation", shared_resources); moveit_studio::behaviors::registerBehavior(factory, "L2GBehavior", shared_resources); - moveit_studio::behaviors::registerBehavior(factory, "SAM2Segmentation", shared_resources); } }; } // namespace example_behaviors diff --git a/src/lab_sim/description/scene.xml b/src/lab_sim/description/scene.xml index 12e33db..8820c8e 100644 --- a/src/lab_sim/description/scene.xml +++ b/src/lab_sim/description/scene.xml @@ -4,6 +4,10 @@ + + + +