From 19493735b50646b3a274413c38089c8dc1802f21 Mon Sep 17 00:00:00 2001 From: Griswald Brooks Date: Wed, 5 Aug 2026 03:42:44 -0400 Subject: [PATCH 1/2] refactor(example_behaviors): migrate ExampleSAM2Segmentation to SAM2Segment Co-Authored-By: Claude Opus 5 --- .../example_sam2_segmentation.hpp | 21 +- .../src/example_sam2_segmentation.cpp | 180 ++++++++++++++---- src/moveit_pro_sam2 | 2 +- 3 files changed, 155 insertions(+), 48 deletions(-) diff --git a/src/example_behaviors/include/example_behaviors/example_sam2_segmentation.hpp b/src/example_behaviors/include/example_behaviors/example_sam2_segmentation.hpp index 62e85e645..0435b07a7 100644 --- a/src/example_behaviors/include/example_behaviors/example_sam2_segmentation.hpp +++ b/src/example_behaviors/include/example_behaviors/example_sam2_segmentation.hpp @@ -1,12 +1,13 @@ #pragma once +#include #include #include +#include #include #include -#include -#include +#include #include #include #include @@ -53,10 +54,18 @@ class ExampleSAM2Segmentation : public moveit_pro::behaviors::AsyncBehaviorBase tl::expected doWork() override; private: - std::unique_ptr sam2_; - moveit_pro_ml::ONNXImage onnx_image_; - sensor_msgs::msg::Image mask_image_msg_; - moveit_studio_vision_msgs::msg::Mask2D mask_msg_; + /** + * @brief Load the SAM2 pipeline on first tick and reuse it afterwards. + * @details The model latches to the bundle and runtime it first loaded with. Changing either port + * afterwards fails loudly rather than silently continuing to serve the original model, because the + * pipeline does not support hot reload. + */ + tl::expected ensureLoaded(const std::filesystem::path& bundle_manifest, + const std::string& runtime_id); + + std::optional sam2_; + std::filesystem::path loaded_bundle_manifest_; + std::string loaded_runtime_id_; /** @brief Classes derived from AsyncBehaviorBase must implement getFuture() so that it returns a shared_future class member */ std::shared_future>& getFuture() override diff --git a/src/example_behaviors/src/example_sam2_segmentation.cpp b/src/example_behaviors/src/example_sam2_segmentation.cpp index 05a1888cd..256693de1 100644 --- a/src/example_behaviors/src/example_sam2_segmentation.cpp +++ b/src/example_behaviors/src/example_sam2_segmentation.cpp @@ -1,16 +1,21 @@ #include +#include #include #include #include +#include +#include #include #include #include #include #include -#include +#include +#include #include #include +#include #include namespace @@ -21,54 +26,95 @@ constexpr auto kPortPoint = "pixel_coords"; constexpr auto kPortPointDefault = "{pixel_coords}"; constexpr auto kPortMasks = "masks2d"; constexpr auto kPortMasksDefault = "{masks2d}"; +constexpr auto kPortModelPackage = "model_package"; +constexpr auto kPortModelPackageDefault = "moveit_pro_sam2"; +constexpr auto kPortBundleManifest = "model_bundle_manifest"; +constexpr auto kPortBundleManifestDefault = "models/model.yaml"; +constexpr auto kPortRuntimeId = "runtime_id"; +constexpr auto kPortRuntimeIdDefault = "onnxruntime"; constexpr auto kImageInferenceWidth = 1024; constexpr auto kImageInferenceHeight = 1024; + +/// Number of bytes per pixel each supported ROS encoding stores. +constexpr size_t kRgb8Channels = 3; +constexpr size_t kRgba8Channels = 4; } // 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) +namespace +{ +/** + * @brief Convert a ROS image message to the NHWC tensor the SAM2 pipeline consumes. + * @details The source may carry an alpha channel; the tensor is always three-channel RGB with + * values normalized to [0.0, 1.0]. + */ +tl::expected, std::string> +toImageTensor(const sensor_msgs::msg::Image& image_msg) { - 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) + namespace data = moveit_pro_ml::data; + + const size_t source_channels = image_msg.encoding == "rgb8" ? kRgb8Channels : kRgba8Channels; + const size_t pixel_count = static_cast(image_msg.height) * image_msg.width; + if (image_msg.data.size() < pixel_count * source_channels) + { + return tl::make_unexpected(fmt::format("Image message declares {}x{} {} but carries only {} bytes", image_msg.width, + image_msg.height, image_msg.encoding, image_msg.data.size())); + } + + std::vector values(pixel_count * kRgb8Channels); + for (size_t pixel = 0; pixel < pixel_count; ++pixel) + { + const size_t source = pixel * source_channels; + const size_t destination = pixel * kRgb8Channels; + for (size_t channel = 0; channel < kRgb8Channels; ++channel) + { + values[destination + channel] = static_cast(image_msg.data[source + channel]) / 255.0f; + } + } + + auto tensor = data::Tensor::create( + std::move(values), data::Batch{ 1 }, data::Channels{ static_cast(kRgb8Channels) }, + data::Extent{ data::Height{ static_cast(image_msg.height) }, + data::Width{ static_cast(image_msg.width) } }); + if (!tensor.has_value()) { - onnx_image.data[i] = static_cast(image_msg.data[i]) / 255.0f; - onnx_image.data[i + 1] = static_cast(image_msg.data[i + 1]) / 255.0f; - onnx_image.data[i + 2] = static_cast(image_msg.data[i + 2]) / 255.0f; + return tl::make_unexpected(fmt::format("Failed to build the SAM2 input tensor: {}", tensor.error())); } + return std::move(tensor).value(); } -// 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) +/// @brief Convert a single-channel mask of probabilities to a ROS mask message. +moveit_studio_vision_msgs::msg::Mask2D +toMaskMessage(const moveit_pro_ml::data::Tensor& mask, + const std_msgs::msg::Header& header) { - mask_image_msg.height = static_cast(onnx_image.shape[0]); - mask_image_msg.width = static_cast(onnx_image.shape[1]); + sensor_msgs::msg::Image mask_image_msg; + mask_image_msg.header = header; + mask_image_msg.height = static_cast(mask.height()); + mask_image_msg.width = static_cast(mask.width()); 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.resize(mask.data.size()); + for (size_t i = 0; i < mask.data.size(); ++i) { - mask_image_msg.data[i] = onnx_image.data[i] > 0.5 ? 255 : 0; + mask_image_msg.data[i] = mask.data[i] > 0.5f ? 255 : 0; } - mask_msg.pixels = mask_image_msg; + + moveit_studio_vision_msgs::msg::Mask2D mask_msg; + mask_msg.pixels = std::move(mask_image_msg); mask_msg.x = 0; mask_msg.y = 0; + return mask_msg; } +} // namespace ExampleSAM2Segmentation::ExampleSAM2Segmentation( const std::string& name, const BT::NodeConfiguration& config, const std::shared_ptr& shared_resources) : moveit_pro::behaviors::AsyncBehaviorBase(name, config, shared_resources) { - const std::filesystem::path package_path = ament_index_cpp::get_package_share_directory("moveit_pro_sam2"); - 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" / "sam2_decoder.onnx"; - sam2_ = std::make_unique(encoder_onnx_file, decoder_onnx_file); } BT::PortsList ExampleSAM2Segmentation::providedPorts() @@ -78,17 +124,53 @@ BT::PortsList ExampleSAM2Segmentation::providedPorts() "The input points, as a vector of " "geometry_msgs/PointStamped " "messages to be used for segmentation."), + BT::InputPort(kPortModelPackage, kPortModelPackageDefault, + "ROS package containing the SAM 2 model bundle."), + BT::InputPort(kPortBundleManifest, kPortBundleManifestDefault, + "Path to the SAM 2 model bundle manifest, relative to the model package's share " + "directory. The manifest names every graph the pipeline loads."), + BT::InputPort(kPortRuntimeId, kPortRuntimeIdDefault, + "Which runtimes: section of the bundle manifest to load."), BT::OutputPort>( kPortMasks, kPortMasksDefault, "The masks contained in a vector of moveit_studio_vision_msgs::msg::Mask2D messages.") }; } +tl::expected ExampleSAM2Segmentation::ensureLoaded(const std::filesystem::path& bundle_manifest, + const std::string& runtime_id) +{ + if (sam2_.has_value()) + { + if (bundle_manifest != loaded_bundle_manifest_ || runtime_id != loaded_runtime_id_) + { + return tl::make_unexpected( + fmt::format("The SAM 2 model latched to bundle '{}' (runtime '{}') on its first run and cannot be reloaded " + "with bundle '{}' (runtime '{}'). Stop the Objective, change the ports, and run it again.", + loaded_bundle_manifest_.string(), loaded_runtime_id_, bundle_manifest.string(), runtime_id)); + } + return {}; + } + + auto model = moveit_pro_ml::SAM2Segment::load( + { .bundle_manifest = bundle_manifest, .runtime = moveit_pro_ml::model::RuntimeId{ runtime_id } }); + if (!model.has_value()) + { + return tl::make_unexpected( + fmt::format("Failed to load the SAM 2 model bundle '{}': {}", bundle_manifest.string(), model.error().message)); + } + sam2_ = std::move(model).value(); + loaded_bundle_manifest_ = bundle_manifest; + loaded_runtime_id_ = runtime_id; + return {}; +} + tl::expected ExampleSAM2Segmentation::doWork() { - const auto ports = - moveit_pro::behaviors::getRequiredInputs(getInput(kPortImage), - getInput>(kPortPoint)); + const auto ports = moveit_pro::behaviors::getRequiredInputs( + getInput(kPortImage), + getInput>(kPortPoint), getInput(kPortModelPackage), + getInput(kPortBundleManifest), getInput(kPortRuntimeId)); // Check that all required input data ports were set. if (!ports.has_value()) @@ -96,7 +178,7 @@ tl::expected ExampleSAM2Segmentation::doWork() 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(); + const auto& [image_msg, points_2d, model_package, bundle_manifest, runtime_id] = ports.value(); if (image_msg.encoding != "rgb8" && image_msg.encoding != "rgba8") { @@ -105,10 +187,30 @@ tl::expected ExampleSAM2Segmentation::doWork() 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::filesystem::path manifest_path; + try + { + manifest_path = + std::filesystem::path{ ament_index_cpp::get_package_share_directory(model_package) } / bundle_manifest; + } + catch (const ament_index_cpp::PackageNotFoundError& e) + { + return tl::make_unexpected(fmt::format("Model package '{}' was not found: {}", model_package, e.what())); + } + + if (const auto loaded = ensureLoaded(manifest_path, runtime_id); !loaded.has_value()) + { + return tl::make_unexpected(loaded.error()); + } + + auto image_tensor = toImageTensor(image_msg); + if (!image_tensor.has_value()) + { + return tl::make_unexpected(image_tensor.error()); + } - std::vector point_prompts; + std::vector point_prompts; + point_prompts.reserve(points_2d.size()); for (auto const& point : points_2d) { // Assume all points are the same label @@ -117,20 +219,16 @@ tl::expected ExampleSAM2Segmentation::doWork() { 1.0f } }); } - try + const auto result = + sam2_->predict({ .image = std::move(image_tensor).value(), .point_prompts = std::move(point_prompts) }); + if (!result.has_value()) { - 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>(kPortMasks, { mask_msg_ }); - } - catch (const std::invalid_argument& e) - { - return tl::make_unexpected(fmt::format("Invalid argument: {}", e.what())); + return tl::make_unexpected(fmt::format("SAM 2 segmentation failed: {}", result.error().message)); } + setOutput>(kPortMasks, + { toMaskMessage(result->mask, image_msg.header) }); + return true; } diff --git a/src/moveit_pro_sam2 b/src/moveit_pro_sam2 index 00c69577f..2b6b952b6 160000 --- a/src/moveit_pro_sam2 +++ b/src/moveit_pro_sam2 @@ -1 +1 @@ -Subproject commit 00c69577f97bc36b23b5432725f0cba12ff18065 +Subproject commit 2b6b952b6ed9f6d5e5a1f6221abb31b99a212238 From 9351856e80801c470aa480fe8ef95c4fca854309 Mon Sep 17 00:00:00 2001 From: Griswald Brooks Date: Sat, 8 Aug 2026 04:57:02 -0400 Subject: [PATCH 2/2] fix(example_behaviors): honor image row stride when building the SAM2 tensor Co-Authored-By: Claude Opus 5 --- .../src/example_sam2_segmentation.cpp | 42 ++++++++++++++----- 1 file changed, 32 insertions(+), 10 deletions(-) diff --git a/src/example_behaviors/src/example_sam2_segmentation.cpp b/src/example_behaviors/src/example_sam2_segmentation.cpp index 256693de1..f234ece2f 100644 --- a/src/example_behaviors/src/example_sam2_segmentation.cpp +++ b/src/example_behaviors/src/example_sam2_segmentation.cpp @@ -56,21 +56,43 @@ toImageTensor(const sensor_msgs::msg::Image& image_msg) namespace data = moveit_pro_ml::data; const size_t source_channels = image_msg.encoding == "rgb8" ? kRgb8Channels : kRgba8Channels; - const size_t pixel_count = static_cast(image_msg.height) * image_msg.width; - if (image_msg.data.size() < pixel_count * source_channels) + + // The tensor's strong dimension types reject a non-positive extent by throwing, so catch an empty + // image here and report it through the same error channel as every other malformed message. + if (image_msg.height == 0 || image_msg.width == 0) + { + return tl::make_unexpected( + fmt::format("Image message has a zero extent ({}x{})", image_msg.width, image_msg.height)); + } + + // Rows are `step` bytes apart, not necessarily tightly packed: a publisher may pad each row. Walking + // by pixel would read padding as image data once step exceeds one row of pixels. + const size_t row_bytes = static_cast(image_msg.width) * source_channels; + if (image_msg.step < row_bytes) + { + return tl::make_unexpected(fmt::format("Image message row stride {} is smaller than one {} row of {} pixels ({} " + "bytes)", + image_msg.step, image_msg.encoding, image_msg.width, row_bytes)); + } + if (image_msg.data.size() < static_cast(image_msg.height) * image_msg.step) { - return tl::make_unexpected(fmt::format("Image message declares {}x{} {} but carries only {} bytes", image_msg.width, - image_msg.height, image_msg.encoding, image_msg.data.size())); + return tl::make_unexpected(fmt::format("Image message declares {}x{} {} with row stride {} but carries only {} " + "bytes", + image_msg.width, image_msg.height, image_msg.encoding, image_msg.step, + image_msg.data.size())); } - std::vector values(pixel_count * kRgb8Channels); - for (size_t pixel = 0; pixel < pixel_count; ++pixel) + std::vector values(static_cast(image_msg.height) * image_msg.width * kRgb8Channels); + for (size_t row = 0; row < image_msg.height; ++row) { - const size_t source = pixel * source_channels; - const size_t destination = pixel * kRgb8Channels; - for (size_t channel = 0; channel < kRgb8Channels; ++channel) + for (size_t column = 0; column < image_msg.width; ++column) { - values[destination + channel] = static_cast(image_msg.data[source + channel]) / 255.0f; + const size_t source = row * image_msg.step + column * source_channels; + const size_t destination = (row * image_msg.width + column) * kRgb8Channels; + for (size_t channel = 0; channel < kRgb8Channels; ++channel) + { + values[destination + channel] = static_cast(image_msg.data[source + channel]) / 255.0f; + } } }