├── isaac_ros_segformer ├── test │ ├── dummy_model │ │ ├── .gitkeep │ │ └── .gitignore │ └── test_cases │ │ └── segformer_sample │ │ ├── image.json │ │ └── image.jpg ├── CMakeLists.txt └── package.xml ├── isaac_ros_segment_anything2 ├── isaac_ros_segment_anything2 │ └── __init__.py ├── test │ └── model │ │ └── segment_anything2 │ │ ├── 1 │ │ └── model.onnx │ │ └── config.pbtxt ├── package.xml ├── CMakeLists.txt ├── include │ └── isaac_ros_segment_anything2 │ │ ├── segment_anything2_data_encoder_node.hpp │ │ └── segment_anything2_state_manager.hpp └── scripts │ └── add_object.py ├── isaac_ros_unet ├── test │ ├── test_cases │ │ └── unet_sample │ │ │ ├── image.json │ │ │ └── image.jpg │ ├── dummy_model │ │ └── model.dummy.onnx │ └── unet_decoder_node_test.cpp ├── include │ └── isaac_ros_unet │ │ └── unet_decoder_node.hpp ├── package.xml ├── CMakeLists.txt └── config │ └── unet_decoder_node.yaml ├── .gitignore ├── resources ├── rosbags │ └── segment_anything_sample_data │ │ ├── metadata.yaml │ │ └── segment_anything_sample_data_0.db3 └── peoplesemsegnet_shuffleseg_config.pbtxt ├── isaac_ros_segment_anything ├── test │ ├── model │ │ └── segment_anything │ │ │ ├── 1 │ │ │ └── model.onnx │ │ │ └── config.pbtxt │ ├── segment_anything_data_encoder_node_test.cpp │ └── isaac_ros_segment_anything_tensor_to_image_test.py ├── config │ ├── quickstart_interface_specs.json │ ├── segment_anything_spec_file.yaml │ ├── segment_anything_decoder_node.yaml │ └── segment_anything_data_encoder_node.yaml ├── isaac_ros_segment_anything │ └── __init__.py ├── include │ └── isaac_ros_segment_anything │ │ ├── segment_anything_binarize_tensor.hpp │ │ ├── segment_anything_decoder_node.hpp │ │ ├── segment_anything_data_encoder_node.hpp │ │ ├── segment_anything_dummy_mask_publisher_node.hpp │ │ └── segment_anything_tensor_to_image_node.hpp ├── scripts │ ├── colored_mask_converter_node.py │ ├── visualize_mask.py │ └── torch_to_onnx.py ├── package.xml ├── src │ ├── segment_anything_dummy_mask_publisher_node.cpp │ ├── segment_anything_decoder_node.cpp │ ├── segment_anything_binarize_tensor.cu │ └── segment_anything_data_encoder_node.cpp └── CMakeLists.txt ├── isaac_ros_segment_anything2_interfaces ├── srv │ ├── RemoveObject.srv │ └── AddObjects.srv ├── package.xml └── CMakeLists.txt ├── .gitattributes ├── CONTRIBUTING.md ├── isaac_ros_peoplesemseg_models_install ├── package.xml ├── CMakeLists.txt └── asset_scripts │ ├── install_peoplesemsegnet_shuffleseg.sh │ └── install_peoplesemsegnet_vanilla.sh ├── isaac_ros_gxf_extensions ├── gxf_isaac_ros_unet │ ├── package.xml │ ├── gxf │ │ └── image_segmentation │ │ │ ├── segmentation_postprocessor_ext.cpp │ │ │ ├── segmentation_mask_colorizer.cu.hpp │ │ │ ├── segmentation_postprocessor.cu.hpp │ │ │ ├── segmentation_postprocessor.hpp │ │ │ ├── segmentation_mask_colorizer.hpp │ │ │ ├── segmentation_postprocessing_utils.hpp │ │ │ ├── segmentation_mask_colorizer.cu.cpp │ │ │ └── segmentation_postprocessor.cu.cpp │ └── CMakeLists.txt └── gxf_isaac_ros_segment_anything │ ├── gxf │ └── segment_anything │ │ ├── segment_anything_postprocessor.cu.hpp │ │ ├── segment_anything_msg_compositor.hpp │ │ ├── segment_anything_postprocessor.hpp │ │ ├── segment_anything_ext.cpp │ │ ├── segment_anything_postprocessor.cu.cpp │ │ ├── segment_anything_prompt_processor.hpp │ │ ├── segment_anything_msg_compositor.cpp │ │ └── segment_anything_postprocessor.cpp │ ├── package.xml │ └── CMakeLists.txt └── SECURITY.md /isaac_ros_segformer/test/dummy_model/.gitkeep: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /isaac_ros_segformer/test/dummy_model/.gitignore: -------------------------------------------------------------------------------- 1 | *.plan 2 | *.onnx -------------------------------------------------------------------------------- /isaac_ros_segment_anything2/isaac_ros_segment_anything2/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /isaac_ros_unet/test/test_cases/unet_sample/image.json: -------------------------------------------------------------------------------- 1 | { 2 | "image": "image.jpg", 3 | "encoding": "bgr8" 4 | } -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | # Ignore all pycache files 2 | **/__pycache__/** 3 | 4 | # Ignore TensorRT plan files 5 | *.plan 6 | *.engine 7 | -------------------------------------------------------------------------------- /isaac_ros_segformer/test/test_cases/segformer_sample/image.json: -------------------------------------------------------------------------------- 1 | { 2 | "image": "image.jpg", 3 | "encoding": "bgr8" 4 | } -------------------------------------------------------------------------------- /isaac_ros_unet/test/dummy_model/model.dummy.onnx: -------------------------------------------------------------------------------- 1 | version https://git-lfs.github.com/spec/v1 2 | oid sha256:e8f655c2e5ea9ed2bfc6c09c4491baaec40343b6085d92b95f9febd5de15196d 3 | size 124425250 4 | -------------------------------------------------------------------------------- /isaac_ros_unet/test/test_cases/unet_sample/image.jpg: -------------------------------------------------------------------------------- 1 | version https://git-lfs.github.com/spec/v1 2 | oid sha256:05bbfdf09370c27a4f5bbebd92d8060818e7f0dd10ecfabee2442b5b7c522cca 3 | size 163407 4 | -------------------------------------------------------------------------------- /resources/rosbags/segment_anything_sample_data/metadata.yaml: -------------------------------------------------------------------------------- 1 | version https://git-lfs.github.com/spec/v1 2 | oid sha256:809b2e9d1dc3ddcc0402adfe76eeb6ac246ea3bb93faf50c420e8ed4335c3f6f 3 | size 1987 4 | -------------------------------------------------------------------------------- /isaac_ros_segformer/test/test_cases/segformer_sample/image.jpg: -------------------------------------------------------------------------------- 1 | version https://git-lfs.github.com/spec/v1 2 | oid sha256:05bbfdf09370c27a4f5bbebd92d8060818e7f0dd10ecfabee2442b5b7c522cca 3 | size 163407 4 | -------------------------------------------------------------------------------- /isaac_ros_segment_anything/test/model/segment_anything/1/model.onnx: -------------------------------------------------------------------------------- 1 | version https://git-lfs.github.com/spec/v1 2 | oid sha256:d27d793f68688ff3d9440c527e83b6d5ce948308924867552b724d962b0bd4d4 3 | size 44478323 4 | -------------------------------------------------------------------------------- /isaac_ros_segment_anything2/test/model/segment_anything2/1/model.onnx: -------------------------------------------------------------------------------- 1 | version https://git-lfs.github.com/spec/v1 2 | oid sha256:f3aad4b4cf277fd68a872bcad3460dd55ba8b01893a8b3a0ebaccf2aa99edf94 3 | size 8603 4 | -------------------------------------------------------------------------------- /resources/rosbags/segment_anything_sample_data/segment_anything_sample_data_0.db3: -------------------------------------------------------------------------------- 1 | version https://git-lfs.github.com/spec/v1 2 | oid sha256:ddf095720a5d17265cc9ee2d44886dc492b376ed26011b8deefb5dee457f24e3 3 | size 68509696 4 | -------------------------------------------------------------------------------- /isaac_ros_segment_anything/config/quickstart_interface_specs.json: -------------------------------------------------------------------------------- 1 | { 2 | "camera_resolution": { 3 | "width": 1280, 4 | "height": 720 5 | }, 6 | "input_image": { 7 | "width": 1280, 8 | "height": 720 9 | }, 10 | "camera_model" : "zed2", 11 | "subscribed_topics": { 12 | "image": "/image_rect", 13 | "camera_info": "/camera_info_rect", 14 | "prompt": "/detections_output" 15 | } 16 | } -------------------------------------------------------------------------------- /resources/peoplesemsegnet_shuffleseg_config.pbtxt: -------------------------------------------------------------------------------- 1 | name: "peoplesemsegnet_shuffleseg" 2 | platform: "tensorrt_plan" 3 | max_batch_size: 0 4 | input [ 5 | { 6 | name: "input_2:0" 7 | data_type: TYPE_FP32 8 | dims: [ 1, 3, 544, 960 ] 9 | } 10 | ] 11 | output [ 12 | { 13 | name: "argmax_1" 14 | data_type: TYPE_INT64 15 | dims: [ 1, 544, 960, 1 ] 16 | } 17 | ] 18 | version_policy: { 19 | specific { 20 | versions: [ 1 ] 21 | } 22 | } 23 | 24 | -------------------------------------------------------------------------------- /isaac_ros_segment_anything2_interfaces/srv/RemoveObject.srv: -------------------------------------------------------------------------------- 1 | # This service is used to remove object from the SAM2 tracking. 2 | 3 | 4 | # request 5 | std_msgs/Header request_header 6 | 7 | # object id to remove 8 | string object_id 9 | 10 | 11 | --- 12 | 13 | # response 14 | string[] object_ids # All the remaining object ids in the system 15 | int32[] object_indices # their corresponding indices in the segmentation mask tensor 16 | 17 | bool success # True if object was removed successfully, false otherwise 18 | string message # Optional status or error message 19 | 20 | -------------------------------------------------------------------------------- /isaac_ros_segment_anything/isaac_ros_segment_anything/__init__.py: -------------------------------------------------------------------------------- 1 | # SPDX-FileCopyrightText: NVIDIA CORPORATION & AFFILIATES 2 | # Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. 3 | # 4 | # Licensed under the Apache License, Version 2.0 (the "License"); 5 | # you may not use this file except in compliance with the License. 6 | # You may obtain a copy of the License at 7 | # 8 | # http://www.apache.org/licenses/LICENSE-2.0 9 | # 10 | # Unless required by applicable law or agreed to in writing, software 11 | # distributed under the License is distributed on an "AS IS" BASIS, 12 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | # See the License for the specific language governing permissions and 14 | # limitations under the License. 15 | # 16 | # SPDX-License-Identifier: Apache-2.0 17 | -------------------------------------------------------------------------------- /.gitattributes: -------------------------------------------------------------------------------- 1 | # Ignore Python files in linguist 2 | *.py linguist-detectable=false 3 | 4 | # Images 5 | *.gif filter=lfs diff=lfs merge=lfs -text 6 | *.jpg filter=lfs diff=lfs merge=lfs -text 7 | *.png filter=lfs diff=lfs merge=lfs -text 8 | *.psd filter=lfs diff=lfs merge=lfs -text 9 | 10 | # Archives 11 | *.gz filter=lfs diff=lfs merge=lfs -text 12 | *.tar filter=lfs diff=lfs merge=lfs -text 13 | *.zip filter=lfs diff=lfs merge=lfs -text 14 | 15 | # Documents 16 | *.pdf filter=lfs diff=lfs merge=lfs -text 17 | 18 | # Shared libraries 19 | *.so filter=lfs diff=lfs merge=lfs -text 20 | *.so.* filter=lfs diff=lfs merge=lfs -text 21 | 22 | # ROS Bags 23 | **/resources/**/*.zstd filter=lfs diff=lfs merge=lfs -text 24 | **/resources/**/*.db3 filter=lfs diff=lfs merge=lfs -text 25 | **/resources/**/*.yaml filter=lfs diff=lfs merge=lfs -text 26 | 27 | # Model files 28 | *.onnx filter=lfs diff=lfs merge=lfs -text 29 | 30 | -------------------------------------------------------------------------------- /CONTRIBUTING.md: -------------------------------------------------------------------------------- 1 | # Isaac ROS Contribution Rules 2 | 3 | Any contribution that you make to this repository will 4 | be under the Apache 2 License, as dictated by that 5 | [license](http://www.apache.org/licenses/LICENSE-2.0.html): 6 | 7 | > **5. Submission of Contributions.** Unless You explicitly state otherwise, any Contribution intentionally submitted for inclusion in the Work by You to the Licensor shall be under the terms and conditions of this License, without any additional terms or conditions. Notwithstanding the above, nothing herein shall supersede or modify the terms of any separate license agreement you may have executed with Licensor regarding such Contributions. 8 | 9 | Contributors must sign-off each commit by adding a `Signed-off-by: ...` 10 | line to commit messages to certify that they have the right to submit 11 | the code they are contributing to the project according to the 12 | [Developer Certificate of Origin (DCO)](https://developercertificate.org/). 13 | 14 | [//]: # (202201002) 15 | -------------------------------------------------------------------------------- /isaac_ros_segment_anything2_interfaces/srv/AddObjects.srv: -------------------------------------------------------------------------------- 1 | # This service is used to add objects for the SAM2 tracking. 2 | 3 | 4 | # request 5 | std_msgs/Header request_header 6 | 7 | # Each bbox_id should be unique. Should be empty if no bbox to add. 8 | string[] bbox_object_ids 9 | 10 | # Array of bounding boxes. 11 | vision_msgs/BoundingBox2D[] bbox_coords 12 | 13 | # Each unique id would represent a object. Upto 5 points can be added per object. 14 | string[] point_object_ids 15 | 16 | # Array of points. 17 | vision_msgs/Point2D[] point_coords 18 | 19 | # Should be empty or size should be same as number of points. 20 | int32[] point_labels # 1 for foreground, 0 for background 21 | 22 | --- 23 | 24 | # response 25 | string[] object_ids # All the object ids in the system 26 | int32[] object_indices # their corresponding indices in the segmentation mask tensor 27 | 28 | bool success # True if objects were added successfully, false otherwise 29 | string message # Optional status or error message 30 | 31 | -------------------------------------------------------------------------------- /isaac_ros_segment_anything/test/model/segment_anything/config.pbtxt: -------------------------------------------------------------------------------- 1 | name: "segment_anything" 2 | platform: "onnxruntime_onnx" 3 | max_batch_size: 0 4 | input [ 5 | { 6 | name: "images" 7 | data_type: TYPE_FP32 8 | dims: [ 1, 3, 1024, 1024 ] 9 | }, 10 | { 11 | name: "point_coords" 12 | data_type: TYPE_FP32 13 | dims: [ -1, -1, 2 ] 14 | }, 15 | { 16 | name: "point_labels" 17 | data_type: TYPE_FP32 18 | dims: [ -1, -1 ] 19 | }, 20 | { 21 | name: "mask_input" 22 | data_type: TYPE_FP32 23 | dims: [ 1,1, 256,256 ] 24 | }, 25 | { 26 | name: "has_mask_input" 27 | data_type: TYPE_FP32 28 | dims: [ 1 ] 29 | }, 30 | { 31 | name: "orig_im_size" 32 | data_type: TYPE_FP32 33 | dims: [ 2 ] 34 | } 35 | ] 36 | output [ 37 | { 38 | name: "masks" 39 | data_type: TYPE_FP32 40 | dims: [-1, -1, -1, -1] 41 | }, 42 | { 43 | name: "iou_predictions" 44 | data_type: TYPE_FP32 45 | dims: [-1, 1] 46 | }, 47 | { 48 | name: "low_res_masks" 49 | data_type: TYPE_FP32 50 | dims: [-1, 1, -1, -1] 51 | } 52 | ] 53 | version_policy: { 54 | specific { 55 | versions: [ 1 ] 56 | } 57 | } 58 | 59 | -------------------------------------------------------------------------------- /isaac_ros_peoplesemseg_models_install/package.xml: -------------------------------------------------------------------------------- 1 | 2 | 3 | 20 | 21 | 22 | 23 | isaac_ros_peoplesemseg_models_install 24 | 4.0.0 25 | Scripts for installing people segmentation models 26 | 27 | Isaac ROS Maintainers 28 | Apache-2.0 29 | https://developer.nvidia.com/isaac-ros-gems/ 30 | David Tingdahl 31 | 32 | ament_cmake 33 | 34 | isaac_ros_common 35 | isaac_ros_unet 36 | 37 | 38 | ament_cmake 39 | 40 | 41 | -------------------------------------------------------------------------------- /isaac_ros_segment_anything2/test/model/segment_anything2/config.pbtxt: -------------------------------------------------------------------------------- 1 | name: "segment_anything2" 2 | platform: "onnxruntime_onnx" 3 | max_batch_size: 0 4 | input [ 5 | { 6 | name: "image" 7 | data_type: TYPE_FP32 8 | dims: [ 1, 3, 1024, 1024 ] 9 | }, 10 | { 11 | name: "bbox_coords" 12 | data_type: TYPE_FP32 13 | dims: [ -1, 4 ] 14 | }, 15 | { 16 | name: "point_coords" 17 | data_type: TYPE_FP32 18 | dims: [ -1, -1, 2 ] 19 | }, 20 | { 21 | name: "point_labels" 22 | data_type: TYPE_INT32 23 | dims: [ -1, -1 ] 24 | }, 25 | { 26 | name: "mask_memory" 27 | data_type: TYPE_FP32 28 | dims: [ -1, 4, 64, 64, 64 ] 29 | }, 30 | { 31 | name: "obj_ptr_memory" 32 | data_type: TYPE_FP32 33 | dims: [ -1, 2, 256 ] 34 | }, 35 | { 36 | name: "original_size" 37 | data_type: TYPE_INT32 38 | dims: [ 2 ] 39 | }, 40 | { 41 | name: "permutation" 42 | data_type: TYPE_INT64 43 | dims: [ -1 ] 44 | } 45 | ] 46 | output [ 47 | { 48 | name: "high_res_masks" 49 | data_type: TYPE_FP32 50 | dims: [-1, -1, -1, -1] 51 | }, 52 | { 53 | name: "object_score_logits" 54 | data_type: TYPE_FP32 55 | dims: [-1, -1] 56 | }, 57 | { 58 | name: "maskmem_features" 59 | data_type: TYPE_FP32 60 | dims: [-1, -1, -1, -1] 61 | }, 62 | { 63 | name: "maskmem_pos_enc" 64 | data_type: TYPE_FP32 65 | dims: [-1, -1, -1, -1] 66 | }, 67 | { 68 | name: "obj_ptr_features" 69 | data_type: TYPE_FP32 70 | dims: [-1, -1] 71 | } 72 | ] 73 | version_policy: { 74 | specific { 75 | versions: [ 1 ] 76 | } 77 | } 78 | 79 | -------------------------------------------------------------------------------- /isaac_ros_gxf_extensions/gxf_isaac_ros_unet/package.xml: -------------------------------------------------------------------------------- 1 | 2 | 3 | 20 | 21 | 22 | 23 | gxf_isaac_ros_unet 24 | 4.0.0 25 | Segmentation post-processor extension. 26 | 27 | Isaac ROS Maintainers 28 | Apache-2.0 29 | https://developer.nvidia.com/isaac-ros-gems/ 30 | CY Chen 31 | 32 | ament_cmake_auto 33 | 34 | isaac_ros_common 35 | isaac_ros_gxf 36 | 37 | ament_lint_auto 38 | ament_lint_common 39 | 40 | 41 | ament_cmake 42 | 43 | 44 | -------------------------------------------------------------------------------- /isaac_ros_gxf_extensions/gxf_isaac_ros_segment_anything/gxf/segment_anything/segment_anything_postprocessor.cu.hpp: -------------------------------------------------------------------------------- 1 | // SPDX-FileCopyrightText: NVIDIA CORPORATION & AFFILIATES 2 | // Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. 3 | // 4 | // Licensed under the Apache License, Version 2.0 (the "License"); 5 | // you may not use this file except in compliance with the License. 6 | // You may obtain a copy of the License at 7 | // 8 | // http://www.apache.org/licenses/LICENSE-2.0 9 | // 10 | // Unless required by applicable law or agreed to in writing, software 11 | // distributed under the License is distributed on an "AS IS" BASIS, 12 | // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | // See the License for the specific language governing permissions and 14 | // limitations under the License. 15 | // 16 | // SPDX-License-Identifier: Apache-2.0 17 | #ifndef NVIDIA_ISAAC_ROS_GXF_EXTENSIONS_SEGMENT_ANYTHING_POSTPROCESSOR_CU_HPP_ 18 | #define NVIDIA_ISAAC_ROS_GXF_EXTENSIONS_SEGMENT_ANYTHING_POSTPROCESSOR_CU_HPP_ 19 | 20 | #include 21 | #include 22 | 23 | #include "cuda.h" 24 | #include "cuda_runtime.h" 25 | 26 | namespace nvidia { 27 | namespace isaac_ros { 28 | 29 | struct Shape { 30 | int32_t batch_size; 31 | int32_t height; 32 | int32_t width; 33 | int32_t channels; 34 | }; 35 | 36 | typedef uint8_t output_type_t; 37 | 38 | static constexpr int32_t kExpectedChannelCount = 1; 39 | 40 | void cuda_postprocess(Shape shape, const float* input, output_type_t* output, cudaStream_t stream); 41 | } // namespace isaac_ros 42 | } // namespace nvidia 43 | 44 | #endif // NVIDIA_ISAAC_ROS_GXF_EXTENSIONS_SEGMENTATION_POSTPROCESSOR_CU_HPP_ 45 | -------------------------------------------------------------------------------- /isaac_ros_segment_anything/include/isaac_ros_segment_anything/segment_anything_binarize_tensor.hpp: -------------------------------------------------------------------------------- 1 | // SPDX-FileCopyrightText: NVIDIA CORPORATION & AFFILIATES 2 | // Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. 3 | // 4 | // Licensed under the Apache License, Version 2.0 (the "License"); 5 | // you may not use this file except in compliance with the License. 6 | // You may obtain a copy of the License at 7 | // 8 | // http://www.apache.org/licenses/LICENSE-2.0 9 | // 10 | // Unless required by applicable law or agreed to in writing, software 11 | // distributed under the License is distributed on an "AS IS" BASIS, 12 | // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | // See the License for the specific language governing permissions and 14 | // limitations under the License. 15 | // 16 | // SPDX-License-Identifier: Apache-2.0 17 | 18 | #pragma once 19 | 20 | #include 21 | #include 22 | 23 | namespace nvidia 24 | { 25 | namespace isaac_ros 26 | { 27 | namespace segment_anything 28 | { 29 | 30 | // Binarizes tensor data on GPU by setting all non-zero values to 255 31 | void BinarizeTensorOnGPU(uint8_t * tensor, size_t size, cudaStream_t stream); 32 | 33 | // Struct to hold bounding box coordinates 34 | struct BoundingBox 35 | { 36 | int min_x = 0; 37 | int min_y = 0; 38 | int max_x = 0; 39 | int max_y = 0; 40 | }; 41 | 42 | // Find the bounding box of non-zero values in a binary mask 43 | void FindBoundingBoxOnGPU( 44 | const uint8_t * input_data, int width, int height, 45 | BoundingBox * output_bbox, cudaStream_t stream); 46 | 47 | } // namespace segment_anything 48 | } // namespace isaac_ros 49 | } // namespace nvidia 50 | -------------------------------------------------------------------------------- /isaac_ros_gxf_extensions/gxf_isaac_ros_segment_anything/package.xml: -------------------------------------------------------------------------------- 1 | 2 | 3 | 20 | 21 | 22 | 23 | gxf_isaac_ros_segment_anything 24 | 4.0.0 25 | Segmentation Anything extension. 26 | 27 | Isaac ROS Maintainers 28 | Apache-2.0 29 | https://developer.nvidia.com/isaac-ros-gems/ 30 | Shubham Tyagi 31 | 32 | ament_cmake_auto 33 | 34 | isaac_ros_common 35 | isaac_ros_gxf 36 | isaac_ros_nitros_detection2_d_array_type 37 | ament_lint_auto 38 | ament_lint_common 39 | 40 | 41 | ament_cmake 42 | 43 | 44 | -------------------------------------------------------------------------------- /SECURITY.md: -------------------------------------------------------------------------------- 1 | ## Security 2 | 3 | NVIDIA is dedicated to the security and trust of our software products and services, including all source code repositories managed through our organization. 4 | 5 | If you need to report a security issue, please use the appropriate contact points outlined below. **Please do not report security vulnerabilities through GitHub.** 6 | 7 | ## Reporting Potential Security Vulnerability in an NVIDIA Product 8 | 9 | To report a potential security vulnerability in any NVIDIA product: 10 | - Web: [Security Vulnerability Submission Form](https://www.nvidia.com/object/submit-security-vulnerability.html) 11 | - E-Mail: psirt@nvidia.com 12 | - We encourage you to use the following PGP key for secure email communication: [NVIDIA public PGP Key for communication](https://www.nvidia.com/en-us/security/pgp-key) 13 | - Please include the following information: 14 | - Product/Driver name and version/branch that contains the vulnerability 15 | - Type of vulnerability (code execution, denial of service, buffer overflow, etc.) 16 | - Instructions to reproduce the vulnerability 17 | - Proof-of-concept or exploit code 18 | - Potential impact of the vulnerability, including how an attacker could exploit the vulnerability 19 | 20 | While NVIDIA currently does not have a bug bounty program, we do offer acknowledgement when an externally reported security issue is addressed under our coordinated vulnerability disclosure policy. Please visit our [Product Security Incident Response Team (PSIRT)](https://www.nvidia.com/en-us/security/psirt-policies/) policies page for more information. 21 | 22 | ## NVIDIA Product Security 23 | 24 | For all security-related concerns, please visit NVIDIA's Product Security portal at https://www.nvidia.com/en-us/security 25 | -------------------------------------------------------------------------------- /isaac_ros_gxf_extensions/gxf_isaac_ros_unet/gxf/image_segmentation/segmentation_postprocessor_ext.cpp: -------------------------------------------------------------------------------- 1 | // SPDX-FileCopyrightText: NVIDIA CORPORATION & AFFILIATES 2 | // Copyright (c) 2022-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. 3 | // 4 | // Licensed under the Apache License, Version 2.0 (the "License"); 5 | // you may not use this file except in compliance with the License. 6 | // You may obtain a copy of the License at 7 | // 8 | // http://www.apache.org/licenses/LICENSE-2.0 9 | // 10 | // Unless required by applicable law or agreed to in writing, software 11 | // distributed under the License is distributed on an "AS IS" BASIS, 12 | // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | // See the License for the specific language governing permissions and 14 | // limitations under the License. 15 | // 16 | // SPDX-License-Identifier: Apache-2.0 17 | #include "segmentation_mask_colorizer.hpp" 18 | #include "segmentation_postprocessor.hpp" 19 | 20 | #include "gxf/std/extension_factory_helper.hpp" 21 | 22 | GXF_EXT_FACTORY_BEGIN() 23 | GXF_EXT_FACTORY_SET_INFO(0x15ed574714ef4f56, 0xb277090d5b35a4ff, 24 | "SegmentationPostprocessorExtension", 25 | "Isaac ROS Segmentation PostProcessor Extension", "NVIDIA", "0.0.1", 26 | "LICENSE"); 27 | 28 | GXF_EXT_FACTORY_ADD(0xe9681b9e1b864649, 0x8fc86530f45f9d09, 29 | nvidia::isaac_ros::SegmentationPostprocessor, nvidia::gxf::Codelet, 30 | "Generates a raw segmentation mask from a tensor"); 31 | 32 | GXF_EXT_FACTORY_ADD(0xfbdd0d490ccc4df7, 0xba990847538359b9, 33 | nvidia::isaac_ros::SegmentationMaskColorizer, nvidia::gxf::Codelet, 34 | "Mask generation codelet"); 35 | GXF_EXT_FACTORY_END() 36 | -------------------------------------------------------------------------------- /isaac_ros_peoplesemseg_models_install/CMakeLists.txt: -------------------------------------------------------------------------------- 1 | # SPDX-FileCopyrightText: NVIDIA CORPORATION & AFFILIATES 2 | # Copyright (c) 2021-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. 3 | # 4 | # Licensed under the Apache License, Version 2.0 (the "License"); 5 | # you may not use this file except in compliance with the License. 6 | # You may obtain a copy of the License at 7 | # 8 | # http://www.apache.org/licenses/LICENSE-2.0 9 | # 10 | # Unless required by applicable law or agreed to in writing, software 11 | # distributed under the License is distributed on an "AS IS" BASIS, 12 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | # See the License for the specific language governing permissions and 14 | # limitations under the License. 15 | # 16 | # SPDX-License-Identifier: Apache-2.0 17 | 18 | cmake_minimum_required(VERSION 3.22.1) 19 | project(isaac_ros_peoplesemseg_models_install) 20 | 21 | find_package(ament_cmake_auto REQUIRED) 22 | ament_auto_find_build_dependencies() 23 | 24 | # Download and install models to the asset folder (outside the build) 25 | install_isaac_ros_asset(install_peoplesemsegnet_vanilla) 26 | install_isaac_ros_asset(install_peoplesemsegnet_shuffleseg) 27 | 28 | # Install the installation scripts such that the models can be installed when not building from source 29 | install(PROGRAMS asset_scripts/install_peoplesemsegnet_shuffleseg.sh DESTINATION lib/${PROJECT_NAME}) 30 | install(PROGRAMS asset_scripts/install_peoplesemsegnet_vanilla.sh DESTINATION lib/${PROJECT_NAME}) 31 | 32 | 33 | # Embed versioning information into installed files 34 | ament_index_get_resource(ISAAC_ROS_COMMON_CMAKE_PATH isaac_ros_common_cmake_path isaac_ros_common) 35 | include("${ISAAC_ROS_COMMON_CMAKE_PATH}/isaac_ros_common-version-info.cmake") 36 | generate_version_info(${PROJECT_NAME}) 37 | 38 | ament_auto_package(INSTALL_TO_SHARE) 39 | -------------------------------------------------------------------------------- /isaac_ros_segformer/CMakeLists.txt: -------------------------------------------------------------------------------- 1 | # SPDX-FileCopyrightText: NVIDIA CORPORATION & AFFILIATES 2 | # Copyright (c) 2023-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. 3 | # 4 | # Licensed under the Apache License, Version 2.0 (the "License"); 5 | # you may not use this file except in compliance with the License. 6 | # You may obtain a copy of the License at 7 | # 8 | # http://www.apache.org/licenses/LICENSE-2.0 9 | # 10 | # Unless required by applicable law or agreed to in writing, software 11 | # distributed under the License is distributed on an "AS IS" BASIS, 12 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | # See the License for the specific language governing permissions and 14 | # limitations under the License. 15 | # 16 | # SPDX-License-Identifier: Apache-2.0 17 | 18 | cmake_minimum_required(VERSION 3.22.1) 19 | project(isaac_ros_segformer LANGUAGES C CXX) 20 | 21 | if(CMAKE_COMPILER_IS_GNUCXX OR CMAKE_CXX_COMPILER_ID MATCHES "Clang") 22 | add_compile_options(-Wall -Wextra -Wpedantic) 23 | endif() 24 | 25 | find_package(ament_cmake_auto REQUIRED) 26 | ament_auto_find_build_dependencies() 27 | 28 | 29 | if(BUILD_TESTING) 30 | find_package(ament_lint_auto REQUIRED) 31 | ament_lint_auto_find_test_dependencies() 32 | 33 | 34 | # The FindPythonInterp and FindPythonLibs modules are removed 35 | if(POLICY CMP0148) 36 | cmake_policy(SET CMP0148 OLD) 37 | endif() 38 | 39 | find_package(launch_testing_ament_cmake REQUIRED) 40 | add_launch_test(test/isaac_ros_segformer_pol_test.py TIMEOUT "300") 41 | endif() 42 | 43 | 44 | # Embed versioning information into installed files 45 | ament_index_get_resource(ISAAC_ROS_COMMON_CMAKE_PATH isaac_ros_common_cmake_path isaac_ros_common) 46 | include("${ISAAC_ROS_COMMON_CMAKE_PATH}/isaac_ros_common-version-info.cmake") 47 | generate_version_info(${PROJECT_NAME}) 48 | 49 | ament_auto_package(INSTALL_TO_SHARE launch) 50 | -------------------------------------------------------------------------------- /isaac_ros_gxf_extensions/gxf_isaac_ros_unet/gxf/image_segmentation/segmentation_mask_colorizer.cu.hpp: -------------------------------------------------------------------------------- 1 | // SPDX-FileCopyrightText: NVIDIA CORPORATION & AFFILIATES 2 | // Copyright (c) 2022-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. 3 | // 4 | // Licensed under the Apache License, Version 2.0 (the "License"); 5 | // you may not use this file except in compliance with the License. 6 | // You may obtain a copy of the License at 7 | // 8 | // http://www.apache.org/licenses/LICENSE-2.0 9 | // 10 | // Unless required by applicable law or agreed to in writing, software 11 | // distributed under the License is distributed on an "AS IS" BASIS, 12 | // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | // See the License for the specific language governing permissions and 14 | // limitations under the License. 15 | // 16 | // SPDX-License-Identifier: Apache-2.0 17 | #ifndef NVIDIA_ISAAC_ROS_GXF_EXTENSIONS_SEGMENTATION_MASK_COLORIZER_CU_HPP_ 18 | #define NVIDIA_ISAAC_ROS_GXF_EXTENSIONS_SEGMENTATION_MASK_COLORIZER_CU_HPP_ 19 | 20 | #include 21 | #include 22 | 23 | #include "cuda.h" 24 | #include "cuda_runtime.h" 25 | 26 | namespace nvidia { 27 | namespace isaac_ros { 28 | 29 | template 30 | struct ArrayView { 31 | std::unique_ptr data{nullptr, cudaFree}; 32 | std::size_t size; 33 | }; 34 | 35 | enum class ColorImageEncodings { kRGB8, kBGR8 }; 36 | 37 | void ColorizeSegmentationMask(uint8_t* colored_segmentation_mask, uint32_t width, uint32_t height, 38 | ColorImageEncodings image_encoding, 39 | const uint8_t* raw_segmentation_mask, 40 | const ArrayView& color_palette, cudaStream_t stream); 41 | 42 | } // namespace isaac_ros 43 | } // namespace nvidia 44 | 45 | #endif // NVIDIA_ISAAC_ROS_GXF_EXTENSIONS_SEGMENTATION_MASK_COLORIZER_CU_HPP_ 46 | -------------------------------------------------------------------------------- /isaac_ros_segment_anything2_interfaces/package.xml: -------------------------------------------------------------------------------- 1 | 2 | 3 | 20 | 21 | 22 | 23 | isaac_ros_segment_anything2_interfaces 24 | 4.0.0 25 | Interfaces for Isaac ROS Segment Anything 26 | Isaac ROS Maintainers 27 | Apache-2.0 28 | https://developer.nvidia.com/isaac-ros-gems/ 29 | Shubham Tyagi 30 | 31 | ament_cmake 32 | 33 | rosidl_default_generators 34 | rosidl_default_runtime 35 | 36 | isaac_ros_common 37 | std_msgs 38 | vision_msgs 39 | ament_lint_auto 40 | ament_lint_common 41 | 42 | rosidl_interface_packages 43 | 44 | 45 | ament_cmake 46 | 47 | 48 | -------------------------------------------------------------------------------- /isaac_ros_gxf_extensions/gxf_isaac_ros_segment_anything/gxf/segment_anything/segment_anything_msg_compositor.hpp: -------------------------------------------------------------------------------- 1 | // SPDX-FileCopyrightText: NVIDIA CORPORATION & AFFILIATES 2 | // Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. 3 | // 4 | // Licensed under the Apache License, Version 2.0 (the "License"); 5 | // you may not use this file except in compliance with the License. 6 | // You may obtain a copy of the License at 7 | // 8 | // http://www.apache.org/licenses/LICENSE-2.0 9 | // 10 | // Unless required by applicable law or agreed to in writing, software 11 | // distributed under the License is distributed on an "AS IS" BASIS, 12 | // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | // See the License for the specific language governing permissions and 14 | // limitations under the License. 15 | // 16 | // SPDX-License-Identifier: Apache-2.0 17 | #ifndef NVIDIA_ISAAC_ROS_GXF_EXTENSIONS_SEGMENT_ANYTHING_MSG_COMPOSITOR_HPP_ 18 | #define NVIDIA_ISAAC_ROS_GXF_EXTENSIONS_SEGMENT_ANYTHING_MSG_COMPOSITOR_HPP_ 19 | 20 | #include 21 | 22 | #include "gxf/cuda/cuda_stream.hpp" 23 | #include "gxf/cuda/cuda_stream_pool.hpp" 24 | #include "gxf/std/codelet.hpp" 25 | #include "gxf/std/memory_buffer.hpp" 26 | #include "gxf/core/parameter_parser_std.hpp" 27 | #include "gxf/std/receiver.hpp" 28 | #include "gxf/std/transmitter.hpp" 29 | 30 | namespace nvidia { 31 | namespace isaac_ros { 32 | 33 | class SegmentAnythingMsgCompositor : public gxf::Codelet { 34 | public: 35 | gxf_result_t start() override; 36 | gxf_result_t tick() override; 37 | gxf_result_t stop() override; 38 | gxf_result_t registerInterface(gxf::Registrar* registrar) override; 39 | 40 | private: 41 | gxf::Parameter>> inputs_; 42 | gxf::Parameter> output_; 43 | 44 | }; 45 | } // namespace isaac_ros 46 | } // namespace nvidia 47 | 48 | #endif // NVIDIA_ISAAC_ROS_GXF_EXTENSIONS_SEGMENT_ANYTHING_MSG_COMPOSITOR_HPP_ 49 | -------------------------------------------------------------------------------- /isaac_ros_segment_anything/include/isaac_ros_segment_anything/segment_anything_decoder_node.hpp: -------------------------------------------------------------------------------- 1 | // SPDX-FileCopyrightText: NVIDIA CORPORATION & AFFILIATES 2 | // Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. 3 | // 4 | // Licensed under the Apache License, Version 2.0 (the "License"); 5 | // you may not use this file except in compliance with the License. 6 | // You may obtain a copy of the License at 7 | // 8 | // http://www.apache.org/licenses/LICENSE-2.0 9 | // 10 | // Unless required by applicable law or agreed to in writing, software 11 | // distributed under the License is distributed on an "AS IS" BASIS, 12 | // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | // See the License for the specific language governing permissions and 14 | // limitations under the License. 15 | // 16 | // SPDX-License-Identifier: Apache-2.0 17 | 18 | #ifndef ISAAC_ROS_SEGMENT_ANYTHING__SEGMENT_ANYTHING_DECODER_NODE_HPP_ 19 | #define ISAAC_ROS_SEGMENT_ANYTHING__SEGMENT_ANYTHING_DECODER_NODE_HPP_ 20 | 21 | #include 22 | #include 23 | #include 24 | 25 | #include "rclcpp/rclcpp.hpp" 26 | #include "isaac_ros_nitros/nitros_node.hpp" 27 | 28 | namespace nvidia 29 | { 30 | namespace isaac_ros 31 | { 32 | namespace segment_anything 33 | { 34 | 35 | class SegmentAnythingDecoderNode : public nitros::NitrosNode 36 | { 37 | public: 38 | explicit SegmentAnythingDecoderNode(const rclcpp::NodeOptions options = rclcpp::NodeOptions()); 39 | ~SegmentAnythingDecoderNode(); 40 | 41 | void postLoadGraphCallback() override; 42 | 43 | private: 44 | int16_t mask_width_; 45 | 46 | // The height of the segmentation mask 47 | int16_t mask_height_; 48 | 49 | // Needed to calculate block size. It is max batch size for prompt bboxes 50 | int16_t max_batch_size_; 51 | 52 | std::string tensor_name_; 53 | }; 54 | 55 | } // namespace segment_anything 56 | } // namespace isaac_ros 57 | } // namespace nvidia 58 | 59 | #endif // ISAAC_ROS_SEGMENT_ANYTHING__SEGMENT_ANYTHING_DECODER_NODE_HPP_ 60 | -------------------------------------------------------------------------------- /isaac_ros_segment_anything/test/segment_anything_data_encoder_node_test.cpp: -------------------------------------------------------------------------------- 1 | // SPDX-FileCopyrightText: NVIDIA CORPORATION & AFFILIATES 2 | // Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. 3 | // 4 | // Licensed under the Apache License, Version 2.0 (the "License"); 5 | // you may not use this file except in compliance with the License. 6 | // You may obtain a copy of the License at 7 | // 8 | // http://www.apache.org/licenses/LICENSE-2.0 9 | // 10 | // Unless required by applicable law or agreed to in writing, software 11 | // distributed under the License is distributed on an "AS IS" BASIS, 12 | // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | // See the License for the specific language governing permissions and 14 | // limitations under the License. 15 | // 16 | // SPDX-License-Identifier: Apache-2.0 17 | 18 | #include 19 | #include "segment_anything_data_encoder_node.hpp" 20 | #include "rclcpp/rclcpp.hpp" 21 | 22 | // Objective: to cover code lines where exceptions are thrown 23 | // Approach: send Invalid Arguments for node parameters to trigger the exception 24 | 25 | 26 | TEST(segment_anything_data_encoder_node_test, test_invalid_input_prompt_type) 27 | { 28 | rclcpp::init(0, nullptr); 29 | rclcpp::NodeOptions options; 30 | options.append_parameter_override("prompt_input_type", ""); 31 | EXPECT_THROW( 32 | { 33 | try { 34 | nvidia::isaac_ros::segment_anything::SegmentAnythingDataEncoderNode 35 | segment_anything_data_encoder_node(options); 36 | } catch (const std::invalid_argument & e) { 37 | EXPECT_THAT(e.what(), testing::HasSubstr("Received invalid input prompt type")); 38 | throw; 39 | } catch (const rclcpp::exceptions::InvalidParameterValueException & e) { 40 | EXPECT_THAT(e.what(), testing::HasSubstr("No parameter value set")); 41 | throw; 42 | } 43 | }, std::invalid_argument); 44 | rclcpp::shutdown(); 45 | } 46 | 47 | 48 | int main(int argc, char ** argv) 49 | { 50 | testing::InitGoogleTest(&argc, argv); 51 | return RUN_ALL_TESTS(); 52 | } 53 | -------------------------------------------------------------------------------- /isaac_ros_gxf_extensions/gxf_isaac_ros_unet/gxf/image_segmentation/segmentation_postprocessor.cu.hpp: -------------------------------------------------------------------------------- 1 | // SPDX-FileCopyrightText: NVIDIA CORPORATION & AFFILIATES 2 | // Copyright (c) 2022-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. 3 | // 4 | // Licensed under the Apache License, Version 2.0 (the "License"); 5 | // you may not use this file except in compliance with the License. 6 | // You may obtain a copy of the License at 7 | // 8 | // http://www.apache.org/licenses/LICENSE-2.0 9 | // 10 | // Unless required by applicable law or agreed to in writing, software 11 | // distributed under the License is distributed on an "AS IS" BASIS, 12 | // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | // See the License for the specific language governing permissions and 14 | // limitations under the License. 15 | // 16 | // SPDX-License-Identifier: Apache-2.0 17 | #ifndef NVIDIA_ISAAC_ROS_GXF_EXTENSIONS_SEGMENTATION_POSTPROCESSOR_CU_HPP_ 18 | #define NVIDIA_ISAAC_ROS_GXF_EXTENSIONS_SEGMENTATION_POSTPROCESSOR_CU_HPP_ 19 | 20 | #include 21 | #include 22 | 23 | #include "cuda.h" 24 | #include "cuda_runtime.h" 25 | 26 | namespace nvidia { 27 | namespace isaac_ros { 28 | 29 | struct Shape { 30 | int32_t height; 31 | int32_t width; 32 | int32_t channels; 33 | }; 34 | 35 | enum class NetworkOutputType { 36 | kArgmax, 37 | kSigmoid, 38 | kSoftmax, 39 | }; 40 | 41 | enum class DataFormat { 42 | kNCHW, 43 | kHWC, 44 | kNHWC, 45 | }; 46 | 47 | typedef uint8_t output_type_t; 48 | 49 | static constexpr int64_t kMaxChannelCount = std::numeric_limits::max(); 50 | 51 | void cuda_postprocess(NetworkOutputType network_output_type, DataFormat data_format, Shape shape, 52 | const float* input, output_type_t* output, cudaStream_t stream); 53 | 54 | template 55 | void CopyTensorData(NetworkOutputType network_output_type, DataFormat data_format, Shape shape, 56 | const T* input, output_type_t* output, cudaStream_t stream); 57 | 58 | } // namespace isaac_ros 59 | } // namespace nvidia 60 | 61 | #endif // NVIDIA_ISAAC_ROS_GXF_EXTENSIONS_SEGMENTATION_POSTPROCESSOR_CU_HPP_ 62 | -------------------------------------------------------------------------------- /isaac_ros_segment_anything2_interfaces/CMakeLists.txt: -------------------------------------------------------------------------------- 1 | # SPDX-FileCopyrightText: NVIDIA CORPORATION & AFFILIATES 2 | # Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. 3 | # 4 | # Licensed under the Apache License, Version 2.0 (the "License"); 5 | # you may not use this file except in compliance with the License. 6 | # You may obtain a copy of the License at 7 | # 8 | # http://www.apache.org/licenses/LICENSE-2.0 9 | # 10 | # Unless required by applicable law or agreed to in writing, software 11 | # distributed under the License is distributed on an "AS IS" BASIS, 12 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | # See the License for the specific language governing permissions and 14 | # limitations under the License. 15 | # 16 | # SPDX-License-Identifier: Apache-2.0 17 | 18 | cmake_minimum_required(VERSION 3.22.1) 19 | project(isaac_ros_segment_anything2_interfaces LANGUAGES C CXX) 20 | 21 | if(CMAKE_COMPILER_IS_GNUCXX OR CMAKE_CXX_COMPILER_ID MATCHES "Clang") 22 | add_compile_options(-Wall -Wextra -Wpedantic) 23 | endif() 24 | 25 | # The FindPythonInterp and FindPythonLibs modules are removed 26 | if(POLICY CMP0148) 27 | cmake_policy(SET CMP0148 OLD) 28 | endif() 29 | 30 | find_package(ament_cmake_auto REQUIRED) 31 | ament_auto_find_build_dependencies() 32 | 33 | # Prepare custom Visual SLAM interfaces 34 | find_package(rosidl_default_generators REQUIRED) 35 | 36 | set(MSG_FILES 37 | ) 38 | set(SRV_FILES 39 | "srv/AddObjects.srv" 40 | "srv/RemoveObject.srv" 41 | ) 42 | set(ACTION_FILES 43 | ) 44 | 45 | rosidl_generate_interfaces(${PROJECT_NAME} 46 | ${MSG_FILES} 47 | ${SRV_FILES} 48 | ${ACTION_FILES} 49 | DEPENDENCIES std_msgs vision_msgs 50 | ) 51 | ament_export_dependencies(rosidl_default_runtime) 52 | 53 | if(BUILD_TESTING) 54 | find_package(ament_lint_auto REQUIRED) 55 | ament_lint_auto_find_test_dependencies() 56 | endif() 57 | 58 | 59 | # Embed versioning information into installed files 60 | ament_index_get_resource(ISAAC_ROS_COMMON_CMAKE_PATH isaac_ros_common_cmake_path isaac_ros_common) 61 | include("${ISAAC_ROS_COMMON_CMAKE_PATH}/isaac_ros_common-version-info.cmake") 62 | generate_version_info(${PROJECT_NAME}) 63 | 64 | ament_auto_package() 65 | -------------------------------------------------------------------------------- /isaac_ros_segformer/package.xml: -------------------------------------------------------------------------------- 1 | 2 | 3 | 20 | 21 | 22 | 23 | isaac_ros_segformer 24 | 4.0.0 25 | Segformer model processing 26 | 27 | Isaac ROS Maintainers 28 | Apache-2.0 29 | https://developer.nvidia.com/isaac-ros-gems/ 30 | Shubham Tyagi 31 | 32 | ament_cmake 33 | 34 | rclcpp 35 | rclcpp_components 36 | isaac_ros_dnn_image_encoder 37 | isaac_ros_nitros 38 | isaac_ros_nitros_image_type 39 | isaac_ros_nitros_tensor_list_type 40 | 41 | isaac_ros_common 42 | isaac_ros_gxf 43 | 44 | ament_lint_auto 45 | ament_lint_common 46 | isaac_ros_test 47 | isaac_ros_tensor_rt 48 | isaac_ros_unet 49 | 50 | isaac_ros_tensor_rt 51 | isaac_ros_triton 52 | isaac_ros_unet 53 | 54 | 55 | ament_cmake 56 | 57 | 58 | -------------------------------------------------------------------------------- /isaac_ros_unet/include/isaac_ros_unet/unet_decoder_node.hpp: -------------------------------------------------------------------------------- 1 | // SPDX-FileCopyrightText: NVIDIA CORPORATION & AFFILIATES 2 | // Copyright (c) 2021-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. 3 | // 4 | // Licensed under the Apache License, Version 2.0 (the "License"); 5 | // you may not use this file except in compliance with the License. 6 | // You may obtain a copy of the License at 7 | // 8 | // http://www.apache.org/licenses/LICENSE-2.0 9 | // 10 | // Unless required by applicable law or agreed to in writing, software 11 | // distributed under the License is distributed on an "AS IS" BASIS, 12 | // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | // See the License for the specific language governing permissions and 14 | // limitations under the License. 15 | // 16 | // SPDX-License-Identifier: Apache-2.0 17 | 18 | #ifndef ISAAC_ROS_UNET__UNET_DECODER_NODE_HPP_ 19 | #define ISAAC_ROS_UNET__UNET_DECODER_NODE_HPP_ 20 | 21 | #include 22 | #include 23 | #include 24 | 25 | #include "rclcpp/rclcpp.hpp" 26 | #include "isaac_ros_nitros/nitros_node.hpp" 27 | 28 | namespace nvidia 29 | { 30 | namespace isaac_ros 31 | { 32 | namespace unet 33 | { 34 | 35 | class UNetDecoderNode : public nitros::NitrosNode 36 | { 37 | public: 38 | explicit UNetDecoderNode(const rclcpp::NodeOptions options = rclcpp::NodeOptions()); 39 | ~UNetDecoderNode(); 40 | 41 | void postLoadGraphCallback() override; 42 | 43 | private: 44 | // The color encoding that the colored segmentation mask should be in 45 | // This should be either rgb8 or bgr8 46 | std::string color_segmentation_mask_encoding_; 47 | 48 | // The color palette for the color segmentation mask 49 | // There should be an element for each class 50 | // Note: only the first 24 bits are used 51 | std::vector color_palette_; 52 | 53 | // Whether sigmoid or softmax was performed by the network 54 | std::string network_output_type_; 55 | 56 | // The width of the segmentation mask 57 | int16_t mask_width_; 58 | 59 | // The height of the segmentation mask 60 | int16_t mask_height_; 61 | }; 62 | 63 | } // namespace unet 64 | } // namespace isaac_ros 65 | } // namespace nvidia 66 | 67 | #endif // ISAAC_ROS_UNET__UNET_DECODER_NODE_HPP_ 68 | -------------------------------------------------------------------------------- /isaac_ros_gxf_extensions/gxf_isaac_ros_segment_anything/gxf/segment_anything/segment_anything_postprocessor.hpp: -------------------------------------------------------------------------------- 1 | // SPDX-FileCopyrightText: NVIDIA CORPORATION & AFFILIATES 2 | // Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. 3 | // 4 | // Licensed under the Apache License, Version 2.0 (the "License"); 5 | // you may not use this file except in compliance with the License. 6 | // You may obtain a copy of the License at 7 | // 8 | // http://www.apache.org/licenses/LICENSE-2.0 9 | // 10 | // Unless required by applicable law or agreed to in writing, software 11 | // distributed under the License is distributed on an "AS IS" BASIS, 12 | // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | // See the License for the specific language governing permissions and 14 | // limitations under the License. 15 | // 16 | // SPDX-License-Identifier: Apache-2.0 17 | #ifndef NVIDIA_ISAAC_ROS_GXF_EXTENSIONS_SEGMENT_ANYTHING_POSTPROCESSOR_HPP_ 18 | #define NVIDIA_ISAAC_ROS_GXF_EXTENSIONS_SEGMENT_ANYTHING_POSTPROCESSOR_HPP_ 19 | 20 | #include 21 | 22 | #include "gxf/cuda/cuda_stream.hpp" 23 | #include "gxf/cuda/cuda_stream_pool.hpp" 24 | #include "gxf/std/codelet.hpp" 25 | #include "gxf/std/memory_buffer.hpp" 26 | #include "gxf/core/parameter_parser_std.hpp" 27 | #include "gxf/std/receiver.hpp" 28 | #include "gxf/std/transmitter.hpp" 29 | 30 | #include "segment_anything_postprocessor.cu.hpp" 31 | 32 | namespace nvidia { 33 | namespace isaac_ros { 34 | 35 | class SegmentAnythingPostprocessor : public gxf::Codelet { 36 | public: 37 | gxf_result_t start() override; 38 | gxf_result_t tick() override; 39 | gxf_result_t stop() override; 40 | gxf_result_t registerInterface(gxf::Registrar* registrar) override; 41 | 42 | private: 43 | 44 | gxf::Parameter> in_; 45 | gxf::Parameter in_tensor_name_; 46 | gxf::Parameter> out_; 47 | gxf::Parameter> allocator_; 48 | gxf::Parameter> cuda_stream_pool_; 49 | 50 | gxf::Handle stream_; 51 | }; 52 | } // namespace isaac_ros 53 | } // namespace nvidia 54 | 55 | #endif // NVIDIA_ISAAC_ROS_GXF_EXTENSIONS_SEGMENT_ANYTHING_POSTPROCESSOR_HPP_ 56 | -------------------------------------------------------------------------------- /isaac_ros_peoplesemseg_models_install/asset_scripts/install_peoplesemsegnet_shuffleseg.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | # Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved. 3 | # 4 | # NVIDIA CORPORATION and its licensors retain all intellectual property 5 | # and proprietary rights in and to this software, related documentation 6 | # and any modifications thereto. Any use, reproduction, disclosure or 7 | # distribution of this software and related documentation without an express 8 | # license agreement from NVIDIA CORPORATION is strictly prohibited. 9 | 10 | # Download and TRT-compile PeopleSemSeg-ShuffleSeg models. 11 | # * Models will be stored in the isaac_ros_assets dir 12 | # * The script must be called with the --eula argument prior to downloading. 13 | 14 | set -e 15 | 16 | ASSET_NAME="optimized_deployable_shuffleseg_unet_amr_v1.0" 17 | EULA_URL="https://catalog.ngc.nvidia.com/orgs/nvidia/teams/isaac/models/optimized-peoplesemseg-amr" 18 | ASSET_DIR="${ISAAC_ROS_WS}/isaac_ros_assets/models/peoplesemsegnet/${ASSET_NAME}" 19 | ASSET_INSTALL_PATHS="${ASSET_DIR}/1/model.plan" 20 | MODEL_URL="https://api.ngc.nvidia.com/v2/models/org/nvidia/team/isaac/optimized-peoplesemseg-amr/v1.0/files?redirect=true&path=model.onnx" 21 | 22 | source "isaac_ros_asset_eula.sh" 23 | 24 | mkdir -p $(dirname "$ASSET_INSTALL_PATHS") 25 | 26 | wget "${MODEL_URL}" -O "${ASSET_DIR}/model.onnx" 27 | 28 | echo "Converting PeopleSemSegnet shuffleseg amr onnx file to plan file." 29 | /usr/src/tensorrt/bin/trtexec \ 30 | --maxShapes="input_2":1x544x960x3 \ 31 | --minShapes="input_2":1x544x960x3 \ 32 | --optShapes="input_2":1x544x960x3 \ 33 | --fp16 \ 34 | --saveEngine="${ASSET_INSTALL_PATHS}" \ 35 | --onnx="${ASSET_DIR}/model.onnx" 36 | 37 | config_file_text=$( 38 | cat <${ASSET_DIR}/config.pbtxt 65 | 66 | -------------------------------------------------------------------------------- /isaac_ros_segment_anything/include/isaac_ros_segment_anything/segment_anything_data_encoder_node.hpp: -------------------------------------------------------------------------------- 1 | // SPDX-FileCopyrightText: NVIDIA CORPORATION & AFFILIATES 2 | // Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. 3 | // 4 | // Licensed under the Apache License, Version 2.0 (the "License"); 5 | // you may not use this file except in compliance with the License. 6 | // You may obtain a copy of the License at 7 | // 8 | // http://www.apache.org/licenses/LICENSE-2.0 9 | // 10 | // Unless required by applicable law or agreed to in writing, software 11 | // distributed under the License is distributed on an "AS IS" BASIS, 12 | // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | // See the License for the specific language governing permissions and 14 | // limitations under the License. 15 | // 16 | // SPDX-License-Identifier: Apache-2.0 17 | 18 | #ifndef ISAAC_ROS_SEGMENT_ANYTHING__SEGMENT_ANYTHING_DATA_ENCODER_NODE_HPP_ 19 | #define ISAAC_ROS_SEGMENT_ANYTHING__SEGMENT_ANYTHING_DATA_ENCODER_NODE_HPP_ 20 | 21 | #include 22 | #include 23 | #include 24 | 25 | #include "rclcpp/rclcpp.hpp" 26 | #include "isaac_ros_nitros/nitros_node.hpp" 27 | #include "isaac_ros_tensor_list_interfaces/msg/tensor_list.hpp" 28 | #include "vision_msgs/msg/detection2_d_array.hpp" 29 | #include "isaac_ros_nitros_detection2_d_array_type/nitros_detection2_d_array.hpp" 30 | 31 | namespace nvidia 32 | { 33 | namespace isaac_ros 34 | { 35 | namespace segment_anything 36 | { 37 | 38 | class SegmentAnythingDataEncoderNode : public nitros::NitrosNode 39 | { 40 | public: 41 | explicit SegmentAnythingDataEncoderNode( 42 | const rclcpp::NodeOptions options = rclcpp::NodeOptions() 43 | ); 44 | ~SegmentAnythingDataEncoderNode(); 45 | 46 | void postLoadGraphCallback() override; 47 | 48 | private: 49 | // Needed to calculate block size. It is max batch size for prompt bboxes 50 | int32_t max_batch_size_; 51 | 52 | // Prompt type 53 | std::string prompt_input_type_; 54 | bool has_input_mask_; 55 | std::vector orig_img_dims_; 56 | }; 57 | 58 | } // namespace segment_anything 59 | } // namespace isaac_ros 60 | } // namespace nvidia 61 | 62 | #endif // ISAAC_ROS_SEGMENT_ANYTHING__SEGMENT_ANYTHING_DATA_ENCODER_NODE_HPP_ 63 | -------------------------------------------------------------------------------- /isaac_ros_gxf_extensions/gxf_isaac_ros_segment_anything/gxf/segment_anything/segment_anything_ext.cpp: -------------------------------------------------------------------------------- 1 | // SPDX-FileCopyrightText: NVIDIA CORPORATION & AFFILIATES 2 | // Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. 3 | // 4 | // Licensed under the Apache License, Version 2.0 (the "License"); 5 | // you may not use this file except in compliance with the License. 6 | // You may obtain a copy of the License at 7 | // 8 | // http://www.apache.org/licenses/LICENSE-2.0 9 | // 10 | // Unless required by applicable law or agreed to in writing, software 11 | // distributed under the License is distributed on an "AS IS" BASIS, 12 | // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | // See the License for the specific language governing permissions and 14 | // limitations under the License. 15 | // 16 | // SPDX-License-Identifier: Apache-2.0 17 | #include "segment_anything_postprocessor.hpp" 18 | #include "segment_anything_prompt_processor.hpp" 19 | #include "segment_anything_msg_compositor.hpp" 20 | #include "gxf/std/extension_factory_helper.hpp" 21 | 22 | GXF_EXT_FACTORY_BEGIN() 23 | GXF_EXT_FACTORY_SET_INFO(0xa3ed574714ef4f11, 0xc127090d5b35a477, 24 | "SegmentAnythingExtension", 25 | "Isaac ROS Segmentation PostProcessor Extension", "NVIDIA", "0.0.1", 26 | "LICENSE"); 27 | 28 | GXF_EXT_FACTORY_ADD(0xe9681b9e1b864123, 0x8fc86530f45f9ab2, 29 | nvidia::isaac_ros::SegmentAnythingPostprocessor, nvidia::gxf::Codelet, 30 | "Generates a raw segmentation mask from a tensor"); 31 | GXF_EXT_FACTORY_ADD(0xe8211b9e1b864ab1, 0x2de86530f45f9cd3, 32 | nvidia::isaac_ros::SegmentAnythingPromptProcessor, nvidia::gxf::Codelet, 33 | "Transforms the input bboxes/points to SAM format."); 34 | GXF_EXT_FACTORY_ADD(0xe12acb9e1b8642ba, 0x34c86170f32f9abc, 35 | nvidia::isaac_ros::SegmentAnythingMsgCompositor, nvidia::gxf::Codelet, 36 | "Composes a single msg with all the received tensors."); 37 | GXF_EXT_FACTORY_ADD_0( 38 | 0xa321601525594206, 0xbc12d9f22a134452, 39 | std::vector, 40 | "Array of decoded 2D object detections in an image"); 41 | GXF_EXT_FACTORY_END() 42 | -------------------------------------------------------------------------------- /isaac_ros_unet/package.xml: -------------------------------------------------------------------------------- 1 | 2 | 3 | 20 | 21 | 22 | 23 | isaac_ros_unet 24 | 4.0.0 25 | U-Net model processing 26 | 27 | Isaac ROS Maintainers 28 | Apache-2.0 29 | https://developer.nvidia.com/isaac-ros-gems/ 30 | Kajanan Chinniah 31 | Swapnesh Wani 32 | Shubham Tyagi 33 | 34 | ament_cmake 35 | 36 | rclcpp 37 | rclcpp_components 38 | isaac_ros_dnn_image_encoder 39 | isaac_ros_nitros 40 | isaac_ros_nitros_image_type 41 | isaac_ros_nitros_tensor_list_type 42 | 43 | isaac_ros_common 44 | isaac_ros_gxf 45 | 46 | gxf_isaac_ros_unet 47 | 48 | ament_lint_auto 49 | ament_lint_common 50 | isaac_ros_test 51 | isaac_ros_tensor_rt 52 | ament_cmake_gtest 53 | 54 | isaac_ros_tensor_rt 55 | isaac_ros_triton 56 | 57 | 58 | ament_cmake 59 | 60 | 61 | -------------------------------------------------------------------------------- /isaac_ros_peoplesemseg_models_install/asset_scripts/install_peoplesemsegnet_vanilla.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | # Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved. 3 | # 4 | # NVIDIA CORPORATION and its licensors retain all intellectual property 5 | # and proprietary rights in and to this software, related documentation 6 | # and any modifications thereto. Any use, reproduction, disclosure or 7 | # distribution of this software and related documentation without an express 8 | # license agreement from NVIDIA CORPORATION is strictly prohibited. 9 | 10 | # Download and TRT-compile PeopleSemSeg-Vanilla models. 11 | # * Models will be stored in the isaac_ros_assets dir 12 | # * The script must be called with the --eula argument prior to downloading. 13 | 14 | set -e 15 | 16 | ASSET_NAME="deployable_quantized_vanilla_unet_onnx_v2.0" 17 | EULA_URL="https://catalog.ngc.nvidia.com/orgs/nvidia/teams/tao/models/peoplesemsegnet" 18 | ASSET_DIR="${ISAAC_ROS_WS}/isaac_ros_assets/models/peoplesemsegnet/${ASSET_NAME}" 19 | ASSET_INSTALL_PATHS="${ASSET_DIR}/1/model.plan" 20 | MODEL_URL="https://api.ngc.nvidia.com/v2/models/org/nvidia/team/tao/peoplesemsegnet/deployable_quantized_vanilla_unet_onnx_v2.0/files?redirect=true&path=peoplesemsegnet_vanilla_unet_dynamic_etlt_int8_fp16.onnx" 21 | 22 | source "isaac_ros_asset_eula.sh" 23 | 24 | mkdir -p $(dirname "$ASSET_INSTALL_PATHS") 25 | 26 | wget "${MODEL_URL}" -O "${ASSET_DIR}/model.onnx" 27 | 28 | echo "Converting PeopleSemSegnet onnx file to plan file." 29 | /usr/src/tensorrt/bin/trtexec \ 30 | --maxShapes="input_1:0":1x3x544x960 \ 31 | --minShapes="input_1:0":1x3x544x960 \ 32 | --optShapes="input_1:0":1x3x544x960 \ 33 | --onnx="${ASSET_DIR}/model.onnx" \ 34 | --saveEngine="${ASSET_DIR}/1/model.plan" \ 35 | --fp16 \ 36 | --skipInference 37 | 38 | # Create config.pbtxt 39 | config_file_text=$( 40 | cat <${ASSET_DIR}/config.pbtxt 66 | 67 | -------------------------------------------------------------------------------- /isaac_ros_gxf_extensions/gxf_isaac_ros_unet/gxf/image_segmentation/segmentation_postprocessor.hpp: -------------------------------------------------------------------------------- 1 | // SPDX-FileCopyrightText: NVIDIA CORPORATION & AFFILIATES 2 | // Copyright (c) 2022-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. 3 | // 4 | // Licensed under the Apache License, Version 2.0 (the "License"); 5 | // you may not use this file except in compliance with the License. 6 | // You may obtain a copy of the License at 7 | // 8 | // http://www.apache.org/licenses/LICENSE-2.0 9 | // 10 | // Unless required by applicable law or agreed to in writing, software 11 | // distributed under the License is distributed on an "AS IS" BASIS, 12 | // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | // See the License for the specific language governing permissions and 14 | // limitations under the License. 15 | // 16 | // SPDX-License-Identifier: Apache-2.0 17 | #ifndef NVIDIA_ISAAC_ROS_GXF_EXTENSIONS_SEGMENTATION_POSTPROCESSOR_HPP_ 18 | #define NVIDIA_ISAAC_ROS_GXF_EXTENSIONS_SEGMENTATION_POSTPROCESSOR_HPP_ 19 | 20 | #include 21 | 22 | #include "gxf/cuda/cuda_stream.hpp" 23 | #include "gxf/cuda/cuda_stream_pool.hpp" 24 | #include "gxf/std/codelet.hpp" 25 | #include "gxf/std/memory_buffer.hpp" 26 | #include "gxf/core/parameter_parser_std.hpp" 27 | #include "gxf/std/receiver.hpp" 28 | #include "gxf/std/transmitter.hpp" 29 | 30 | #include "segmentation_postprocessor.cu.hpp" 31 | 32 | namespace nvidia { 33 | namespace isaac_ros { 34 | 35 | class SegmentationPostprocessor : public gxf::Codelet { 36 | public: 37 | gxf_result_t start() override; 38 | gxf_result_t tick() override; 39 | gxf_result_t stop() override; 40 | gxf_result_t registerInterface(gxf::Registrar* registrar) override; 41 | 42 | private: 43 | NetworkOutputType network_output_type_value_; 44 | DataFormat data_format_value_; 45 | 46 | gxf::Parameter> in_; 47 | gxf::Parameter in_tensor_name_; 48 | gxf::Parameter network_output_type_; 49 | gxf::Parameter data_format_; 50 | gxf::Parameter> out_; 51 | gxf::Parameter> allocator_; 52 | gxf::Parameter> cuda_stream_pool_; 53 | 54 | gxf::Handle stream_; 55 | }; 56 | } // namespace isaac_ros 57 | } // namespace nvidia 58 | 59 | #endif // NVIDIA_ISAAC_ROS_GXF_EXTENSIONS_SEGMENTATION_POSTPROCESSOR_HPP_ 60 | -------------------------------------------------------------------------------- /isaac_ros_segment_anything2/package.xml: -------------------------------------------------------------------------------- 1 | 2 | 3 | 20 | 21 | 22 | 23 | isaac_ros_segment_anything2 24 | 4.0.0 25 | Segment Anything2 model processing 26 | 27 | Isaac ROS Maintainers 28 | Apache-2.0 29 | https://developer.nvidia.com/isaac-ros-gems/ 30 | Shubham Tyagi 31 | 32 | ament_cmake 33 | 34 | rclcpp 35 | rclcpp_components 36 | isaac_ros_dnn_image_encoder 37 | isaac_ros_nitros 38 | isaac_ros_nitros_image_type 39 | isaac_ros_nitros_tensor_list_type 40 | isaac_ros_managed_nitros 41 | isaac_ros_segment_anything2_interfaces 42 | isaac_ros_tensor_list_interfaces 43 | isaac_ros_segment_anything 44 | isaac_ros_common 45 | isaac_ros_gxf 46 | ament_lint_auto 47 | ament_lint_common 48 | isaac_ros_test 49 | isaac_ros_triton 50 | gxf_isaac_messages 51 | 52 | gxf_isaac_ros_segment_anything 53 | gxf_isaac_messages 54 | isaac_ros_triton 55 | 56 | 57 | ament_cmake 58 | 59 | 60 | -------------------------------------------------------------------------------- /isaac_ros_gxf_extensions/gxf_isaac_ros_segment_anything/gxf/segment_anything/segment_anything_postprocessor.cu.cpp: -------------------------------------------------------------------------------- 1 | // SPDX-FileCopyrightText: NVIDIA CORPORATION & AFFILIATES 2 | // Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. 3 | // 4 | // Licensed under the Apache License, Version 2.0 (the "License"); 5 | // you may not use this file except in compliance with the License. 6 | // You may obtain a copy of the License at 7 | // 8 | // http://www.apache.org/licenses/LICENSE-2.0 9 | // 10 | // Unless required by applicable law or agreed to in writing, software 11 | // distributed under the License is distributed on an "AS IS" BASIS, 12 | // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | // See the License for the specific language governing permissions and 14 | // limitations under the License. 15 | // 16 | // SPDX-License-Identifier: Apache-2.0 17 | #include "segment_anything_postprocessor.cu.hpp" 18 | 19 | namespace nvidia { 20 | namespace isaac_ros { 21 | 22 | 23 | __forceinline__ __device__ uint32_t nchw_to_index(Shape shape, uint32_t y, uint32_t x, uint32_t c, uint32_t n) { 24 | return n * shape.channels * shape.width * shape.height + (c * shape.height + y) * shape.width + x; 25 | } 26 | 27 | __global__ void postprocessing_kernel(Shape shape, const float* input, output_type_t* output) { 28 | const uint32_t x = blockIdx.x * blockDim.x + threadIdx.x; 29 | const uint32_t y = blockIdx.y * blockDim.y + threadIdx.y; 30 | 31 | if ((x >= shape.width) || (y >= shape.height)) { return; } 32 | // For each batch output check if there is a detection. 33 | for (uint32_t n = 0; n < shape.batch_size ; n++) { 34 | for (uint32_t c = 0; c < shape.channels; c++) { 35 | uint32_t idx = nchw_to_index(shape, y, x, c, n); 36 | const float value = input[idx]; 37 | // If value is greater than 0 then that pixel is of interest 38 | if (value > 0.0) { 39 | output[idx] = 1; 40 | } 41 | else { 42 | output[idx] = 0; 43 | } 44 | } 45 | } 46 | } 47 | 48 | uint16_t ceil_div(uint16_t numerator, uint16_t denominator) { 49 | uint32_t accumulator = numerator + denominator - 1; 50 | return accumulator / denominator; 51 | } 52 | 53 | void cuda_postprocess(Shape shape, const float* input, output_type_t* output, cudaStream_t stream) { 54 | dim3 block(32, 32, 1); 55 | dim3 grid(ceil_div(shape.width, block.x), ceil_div(shape.height, block.y), 1); 56 | 57 | postprocessing_kernel 58 | <<>>(shape, input, output); 59 | 60 | } 61 | 62 | } // namespace isaac_ros 63 | } // namespace nvidia 64 | -------------------------------------------------------------------------------- /isaac_ros_segment_anything/config/segment_anything_spec_file.yaml: -------------------------------------------------------------------------------- 1 | %YAML 1.2 2 | # Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved. 3 | # 4 | # NVIDIA CORPORATION and its licensors retain all intellectual property 5 | # and proprietary rights in and to this software, related documentation 6 | # and any modifications thereto. Any use, reproduction, disclosure or 7 | # distribution of this software and related documentation without an express 8 | # license agreement from NVIDIA CORPORATION is strictly prohibited. 9 | --- 10 | id: [0xa3ed574714ef4f11, 0xc127090d5b35a477] 11 | name: SegmentAnythingExtension 12 | version: 0.2.0 13 | components: 14 | - id: [0xe9681b9e1b864123, 0x8fc86530f45f9ab2] 15 | type: nvidia::isaac_ros::SegmentAnythingPostprocessor 16 | input_output_groups: 17 | - input_keys: [in] 18 | output_keys: [out] 19 | input_format_keys: [data_format] 20 | output_format_keys: [] 21 | supported_formats: 22 | - platforms: [any] 23 | details: 24 | - input_formats: [nitros_tensor_list_nchw_rgb_f32] 25 | output_formats: [nitros_tensor_list_nchw] 26 | costs: 27 | throughput: 10bytes/s 28 | latency: 10ms 29 | power: 100J 30 | accuracy: 100% 31 | - id: [0xe8211b9e1b864ab1, 0x2de86530f45f9cd3] 32 | type: nvidia::isaac_ros::SegmentAnythingPromptProcessor 33 | input_output_groups: 34 | - input_keys: [in] 35 | output_keys: [out_points] 36 | input_format_keys: [] 37 | output_format_keys: [] 38 | supported_formats: 39 | - platforms: [any] 40 | details: 41 | - input_formats: [nitros_detection2_d_array] 42 | output_formats: [nitros_tensor_list_nchw] 43 | costs: 44 | throughput: 10bytes/s 45 | latency: 10ms 46 | power: 100J 47 | accuracy: 100% 48 | - id: [0xe12acb9e1b8642ba, 0x34c86170f32f9abc] 49 | type: nvidia::isaac_ros::SegmentAnythingMsgCompositor 50 | input_output_groups: 51 | - input_keys: [] 52 | output_keys: [] 53 | input_format_keys: [] 54 | output_format_keys: [] 55 | supported_formats: 56 | - platforms: [any] 57 | details: 58 | - input_formats: [] 59 | output_formats: [] 60 | costs: 61 | throughput: 10bytes/s 62 | latency: 10ms 63 | power: 100J 64 | accuracy: 100% 65 | 66 | 67 | -------------------------------------------------------------------------------- /isaac_ros_segment_anything/include/isaac_ros_segment_anything/segment_anything_dummy_mask_publisher_node.hpp: -------------------------------------------------------------------------------- 1 | // SPDX-FileCopyrightText: NVIDIA CORPORATION & AFFILIATES 2 | // Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. 3 | // 4 | // Licensed under the Apache License, Version 2.0 (the "License"); 5 | // you may not use this file except in compliance with the License. 6 | // You may obtain a copy of the License at 7 | // 8 | // http://www.apache.org/licenses/LICENSE-2.0 9 | // 10 | // Unless required by applicable law or agreed to in writing, software 11 | // distributed under the License is distributed on an "AS IS" BASIS, 12 | // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | // See the License for the specific language governing permissions and 14 | // limitations under the License. 15 | // 16 | // SPDX-License-Identifier: Apache-2.0 17 | 18 | #ifndef ISAAC_ROS_SEGMENT_ANYTHING__SEGMENT_ANYTHING_DUMMY_MASK_PUBLISHER_NODE_HPP_ 19 | #define ISAAC_ROS_SEGMENT_ANYTHING__SEGMENT_ANYTHING_DUMMY_MASK_PUBLISHER_NODE_HPP_ 20 | 21 | #include 22 | #include 23 | 24 | #include "cuda_runtime.h" // NOLINT - include .h without directory 25 | 26 | #include "rclcpp/rclcpp.hpp" 27 | 28 | #include "isaac_ros_managed_nitros/managed_nitros_subscriber.hpp" 29 | #include "isaac_ros_managed_nitros/managed_nitros_publisher.hpp" 30 | 31 | #include "isaac_ros_nitros_tensor_list_type/nitros_tensor_list_view.hpp" 32 | #include "isaac_ros_nitros_tensor_list_type/nitros_tensor_list.hpp" 33 | namespace nvidia 34 | { 35 | namespace isaac_ros 36 | { 37 | namespace segment_anything 38 | { 39 | 40 | class DummyMaskPublisher : public rclcpp::Node 41 | { 42 | public: 43 | explicit DummyMaskPublisher(const rclcpp::NodeOptions options = rclcpp::NodeOptions()); 44 | 45 | ~DummyMaskPublisher(); 46 | 47 | private: 48 | void InputCallback(const nvidia::isaac_ros::nitros::NitrosTensorListView & msg); 49 | 50 | // Subscription to input NitrosTensorList messages 51 | std::shared_ptr> nitros_sub_; 53 | 54 | // Publisher for output NitrosTensorList messages 55 | std::shared_ptr> nitros_pub_; 57 | 58 | // Name of tensor in NitrosTensorList 59 | std::string tensor_name_{}; 60 | 61 | // CUDA stream for async operations 62 | cudaStream_t stream_; 63 | }; 64 | 65 | } // namespace segment_anything 66 | } // namespace isaac_ros 67 | } // namespace nvidia 68 | 69 | #endif // ISAAC_ROS_SEGMENT_ANYTHING__SEGMENT_ANYTHING_DUMMY_MASK_PUBLISHER_NODE_HPP_ 70 | -------------------------------------------------------------------------------- /isaac_ros_segment_anything/scripts/colored_mask_converter_node.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python3 2 | 3 | # SPDX-FileCopyrightText: NVIDIA CORPORATION & AFFILIATES 4 | # Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. 5 | # 6 | # Licensed under the Apache License, Version 2.0 (the "License"); 7 | # you may not use this file except in compliance with the License. 8 | # You may obtain a copy of the License at 9 | # 10 | # http://www.apache.org/licenses/LICENSE-2.0 11 | # 12 | # Unless required by applicable law or agreed to in writing, software 13 | # distributed under the License is distributed on an "AS IS" BASIS, 14 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 15 | # See the License for the specific language governing permissions and 16 | # limitations under the License. 17 | # 18 | # SPDX-License-Identifier: Apache-2.0 19 | 20 | import cv2 21 | import cv_bridge 22 | import numpy as np 23 | import rclpy 24 | from rclpy.node import Node 25 | from rclpy.qos import DurabilityPolicy, HistoryPolicy, QoSProfile, ReliabilityPolicy 26 | from sensor_msgs.msg import Image 27 | 28 | 29 | class ColoredMaskConverterNode(Node): 30 | 31 | def __init__(self): 32 | super().__init__('colored_mask_converter_node') 33 | 34 | self.declare_parameter('color', [255, 0, 0]) 35 | 36 | self.bridge = cv_bridge.CvBridge() 37 | self.mask = None 38 | 39 | self.publisher = self.create_publisher(Image, 'colored_segmentation_mask', 10) 40 | 41 | self.create_subscription( 42 | Image, 'image', self.image_callback, 43 | QoSProfile( 44 | reliability=ReliabilityPolicy.BEST_EFFORT, 45 | durability=DurabilityPolicy.VOLATILE, 46 | history=HistoryPolicy.KEEP_LAST, 47 | depth=10)) 48 | 49 | self.create_subscription(Image, 'binary_segmentation_mask', self.mask_callback, 10) 50 | 51 | def mask_callback(self, msg): 52 | self.mask = msg 53 | 54 | def image_callback(self, msg): 55 | if self.mask is not None: 56 | src = self.bridge.imgmsg_to_cv2(msg, 'rgb8') 57 | binary_mask = self.bridge.imgmsg_to_cv2(self.mask, 'mono8') 58 | colored_mask = np.zeros(src.shape, dtype=np.uint8) 59 | color = self.get_parameter('color').get_parameter_value().integer_array_value 60 | colored_mask[np.where(binary_mask[:, :] > 0)] = color 61 | dst = cv2.addWeighted(src, 1.0, colored_mask, 1.0, 0) 62 | image = self.bridge.cv2_to_imgmsg(dst, 'rgb8') 63 | self.publisher.publish(image) 64 | else: 65 | self.publisher.publish(msg) 66 | 67 | 68 | def main(): 69 | rclpy.init() 70 | rclpy.spin(ColoredMaskConverterNode()) 71 | rclpy.shutdown() 72 | 73 | 74 | if __name__ == '__main__': 75 | main() 76 | -------------------------------------------------------------------------------- /isaac_ros_segment_anything/package.xml: -------------------------------------------------------------------------------- 1 | 2 | 3 | 20 | 21 | 22 | 23 | isaac_ros_segment_anything 24 | 4.0.0 25 | Segment Anything model processing 26 | 27 | Isaac ROS Maintainers 28 | Apache-2.0 29 | https://developer.nvidia.com/isaac-ros-gems/ 30 | Shubham Tyagi 31 | Karanbir Chahal 32 | 33 | ament_cmake 34 | 35 | rclcpp 36 | rclcpp_components 37 | rclpy 38 | isaac_ros_dnn_image_encoder 39 | isaac_ros_nitros 40 | isaac_ros_nitros_image_type 41 | isaac_ros_nitros_tensor_list_type 42 | isaac_ros_managed_nitros 43 | isaac_ros_nitros_detection2_d_array_type 44 | isaac_ros_nitros_topic_tools 45 | isaac_ros_segment_anything2_interfaces 46 | isaac_ros_tensor_list_interfaces 47 | geometry_msgs 48 | sensor_msgs 49 | vision_msgs 50 | 51 | isaac_ros_common 52 | isaac_ros_gxf 53 | ament_index_cpp 54 | 55 | ament_lint_auto 56 | ament_lint_common 57 | isaac_ros_test 58 | isaac_ros_triton 59 | gxf_isaac_messages 60 | ament_cmake_gtest 61 | 62 | gxf_isaac_ros_segment_anything 63 | gxf_isaac_messages 64 | isaac_ros_triton 65 | 66 | 67 | ament_cmake 68 | 69 | 70 | -------------------------------------------------------------------------------- /isaac_ros_segment_anything2/CMakeLists.txt: -------------------------------------------------------------------------------- 1 | # SPDX-FileCopyrightText: NVIDIA CORPORATION & AFFILIATES 2 | # Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. 3 | # 4 | # Licensed under the Apache License, Version 2.0 (the "License"); 5 | # you may not use this file except in compliance with the License. 6 | # You may obtain a copy of the License at 7 | # 8 | # http://www.apache.org/licenses/LICENSE-2.0 9 | # 10 | # Unless required by applicable law or agreed to in writing, software 11 | # distributed under the License is distributed on an "AS IS" BASIS, 12 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | # See the License for the specific language governing permissions and 14 | # limitations under the License. 15 | # 16 | # SPDX-License-Identifier: Apache-2.0 17 | 18 | cmake_minimum_required(VERSION 3.22.1) 19 | project(isaac_ros_segment_anything2 LANGUAGES C CXX) 20 | 21 | if(CMAKE_COMPILER_IS_GNUCXX OR CMAKE_CXX_COMPILER_ID MATCHES "Clang") 22 | add_compile_options(-Wall -Wextra -Wpedantic) 23 | endif() 24 | 25 | find_package(ament_cmake_auto REQUIRED) 26 | ament_auto_find_build_dependencies() 27 | enable_language(CUDA) 28 | 29 | ament_auto_add_library(segment_anything2_state_manager SHARED src/segment_anything2_state_manager.cpp) 30 | ament_auto_add_library(segment_anything2_data_encoder_node SHARED src/segment_anything2_data_encoder_node.cpp) 31 | rclcpp_components_register_nodes(segment_anything2_data_encoder_node "nvidia::isaac_ros::segment_anything2::SegmentAnything2DataEncoderNode") 32 | set(node_plugins 33 | "${node_plugins}nvidia::isaac_ros::segment_anything2::SegmentAnything2DataEncoderNode;$\n") 34 | set_target_properties(segment_anything2_data_encoder_node PROPERTIES 35 | BUILD_WITH_INSTALL_RPATH TRUE 36 | BUILD_RPATH_USE_ORIGIN TRUE 37 | INSTALL_RPATH_USE_LINK_PATH TRUE) 38 | 39 | ament_python_install_package(${PROJECT_NAME}) 40 | install(PROGRAMS scripts/sam2_wrapper.py DESTINATION lib/${PROJECT_NAME}) 41 | install(PROGRAMS scripts/sam2_onnx_exporter.py DESTINATION lib/${PROJECT_NAME}) 42 | install(PROGRAMS scripts/add_object.py DESTINATION lib/${PROJECT_NAME}) 43 | 44 | if(BUILD_TESTING) 45 | find_package(ament_lint_auto REQUIRED) 46 | # The FindPythonInterp and FindPythonLibs modules are removed 47 | if(POLICY CMP0148) 48 | cmake_policy(SET CMP0148 OLD) 49 | endif() 50 | 51 | find_package(launch_testing_ament_cmake REQUIRED) 52 | add_launch_test(test/isaac_ros_segment_anything2_pol.py TIMEOUT "300") 53 | endif() 54 | 55 | 56 | # Embed versioning information into installed files 57 | ament_index_get_resource(ISAAC_ROS_COMMON_CMAKE_PATH isaac_ros_common_cmake_path isaac_ros_common) 58 | include("${ISAAC_ROS_COMMON_CMAKE_PATH}/isaac_ros_common-version-info.cmake") 59 | generate_version_info(${PROJECT_NAME}) 60 | 61 | ament_auto_package(INSTALL_TO_SHARE launch) 62 | -------------------------------------------------------------------------------- /isaac_ros_gxf_extensions/gxf_isaac_ros_unet/gxf/image_segmentation/segmentation_mask_colorizer.hpp: -------------------------------------------------------------------------------- 1 | // SPDX-FileCopyrightText: NVIDIA CORPORATION & AFFILIATES 2 | // Copyright (c) 2022-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. 3 | // 4 | // Licensed under the Apache License, Version 2.0 (the "License"); 5 | // you may not use this file except in compliance with the License. 6 | // You may obtain a copy of the License at 7 | // 8 | // http://www.apache.org/licenses/LICENSE-2.0 9 | // 10 | // Unless required by applicable law or agreed to in writing, software 11 | // distributed under the License is distributed on an "AS IS" BASIS, 12 | // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | // See the License for the specific language governing permissions and 14 | // limitations under the License. 15 | // 16 | // SPDX-License-Identifier: Apache-2.0 17 | #ifndef NVIDIA_ISAAC_ROS_GXF_EXTENSIONS_SEGMENTATION_MASK_COLORIZER_HPP_ 18 | #define NVIDIA_ISAAC_ROS_GXF_EXTENSIONS_SEGMENTATION_MASK_COLORIZER_HPP_ 19 | 20 | #include 21 | #include 22 | 23 | #include "gxf/cuda/cuda_stream.hpp" 24 | #include "gxf/multimedia/video.hpp" 25 | #include "gxf/std/allocator.hpp" 26 | #include "gxf/std/codelet.hpp" 27 | #include "gxf/std/receiver.hpp" 28 | #include "gxf/std/transmitter.hpp" 29 | #include "segmentation_mask_colorizer.cu.hpp" 30 | #include "gxf/cuda/cuda_stream.hpp" 31 | #include "gxf/cuda/cuda_stream_pool.hpp" 32 | 33 | namespace nvidia { 34 | namespace isaac_ros { 35 | 36 | class SegmentationMaskColorizer : public gxf::Codelet { 37 | public: 38 | gxf_result_t start() override; 39 | gxf_result_t tick() override; 40 | gxf_result_t stop() override; 41 | 42 | gxf_result_t registerInterface(gxf::Registrar* registrar) override; 43 | 44 | private: 45 | gxf::Parameter> input_; 46 | gxf::Parameter> colored_segmentation_output_; 47 | gxf::Parameter> allocator_; 48 | gxf::Parameter> color_palette_vec_; 49 | gxf::Parameter color_segmentation_mask_encoding_str_; 50 | gxf::Parameter> cuda_stream_pool_; 51 | // CUDA stream variables 52 | gxf::Handle cuda_stream_handle_; 53 | cudaStream_t cuda_stream_ = 0; 54 | 55 | ArrayView color_palette_; 56 | ColorImageEncodings color_segmentation_mask_encoding_; 57 | 58 | gxf::Expected generateColorizedSegmentationMask(gxf::Entity& output, gxf::Entity input, 59 | gxf::Handle stream); 60 | 61 | gxf::Expected createColorSegmentationMask( 62 | gxf::Handle colored_segmentation_mask, 63 | gxf::Handle raw_segmentation_mask); 64 | }; 65 | 66 | } // namespace isaac_ros 67 | } // namespace nvidia 68 | 69 | #endif // NVIDIA_ISAAC_ROS_GXF_EXTENSIONS_SEGMENTATION_MASK_COLORIZER_HPP_ 70 | -------------------------------------------------------------------------------- /isaac_ros_gxf_extensions/gxf_isaac_ros_segment_anything/gxf/segment_anything/segment_anything_prompt_processor.hpp: -------------------------------------------------------------------------------- 1 | // SPDX-FileCopyrightText: NVIDIA CORPORATION & AFFILIATES 2 | // Copyright (c) 2024-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. 3 | // 4 | // Licensed under the Apache License, Version 2.0 (the "License"); 5 | // you may not use this file except in compliance with the License. 6 | // You may obtain a copy of the License at 7 | // 8 | // http://www.apache.org/licenses/LICENSE-2.0 9 | // 10 | // Unless required by applicable law or agreed to in writing, software 11 | // distributed under the License is distributed on an "AS IS" BASIS, 12 | // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | // See the License for the specific language governing permissions and 14 | // limitations under the License. 15 | // 16 | // SPDX-License-Identifier: Apache-2.0 17 | #ifndef NVIDIA_ISAAC_ROS_GXF_EXTENSIONS_SEGMENT_ANYTHING_BBOX_PROCESSOR_HPP_ 18 | #define NVIDIA_ISAAC_ROS_GXF_EXTENSIONS_SEGMENT_ANYTHING_BBOX_PROCESSOR_HPP_ 19 | 20 | #include 21 | 22 | #include "gxf/std/codelet.hpp" 23 | #include "gxf/std/memory_buffer.hpp" 24 | #include "gxf/core/parameter_parser_std.hpp" 25 | #include "gxf/std/receiver.hpp" 26 | #include "gxf/std/transmitter.hpp" 27 | #include "detection2_d_array_message.hpp" 28 | #include "isaac_ros_nitros_detection2_d_array_type/nitros_detection2_d_array.hpp" 29 | 30 | #include "gxf/cuda/cuda_stream.hpp" 31 | #include "gxf/cuda/cuda_stream_pool.hpp" 32 | 33 | namespace nvidia { 34 | namespace isaac_ros { 35 | 36 | enum class PromptType { 37 | kBbox, 38 | kPoint 39 | }; 40 | 41 | class SegmentAnythingPromptProcessor : public gxf::Codelet { 42 | public: 43 | gxf_result_t start() override; 44 | gxf_result_t tick() override; 45 | gxf_result_t stop() override; 46 | gxf_result_t registerInterface(gxf::Registrar* registrar) override; 47 | void detectionToSAMPrompt( 48 | std::vector& detections, 49 | std::vector& prompt_vec, 50 | std::vector& label_vec 51 | ); 52 | private: 53 | gxf::Parameter> in_; 54 | gxf::Parameter> out_points_; 55 | gxf::Parameter> allocator_; 56 | gxf::Parameter> orig_img_dim_; 57 | gxf::Parameter max_batch_size_; 58 | gxf::Parameter prompt_type_name_; 59 | gxf::Parameter has_input_mask_; 60 | gxf::Parameter> cuda_stream_pool_; 61 | // CUDA stream variables 62 | gxf::Handle cuda_stream_handle_; 63 | cudaStream_t cuda_stream_ = 0; 64 | 65 | const uint16_t IMAGE_WIDTH_ = 1024; 66 | const uint16_t IMAGE_HEIGHT_ = 1024; 67 | uint16_t resized_width_; 68 | uint16_t resized_height_; 69 | uint32_t num_points_ = 2; 70 | PromptType prompt_type_value_; 71 | }; 72 | } // namespace isaac_ros 73 | } // namespace nvidia 74 | 75 | #endif // NVIDIA_ISAAC_ROS_GXF_EXTENSIONS_SEGMENT_ANYTHING_BBOX_PROCESSOR_HPP_ 76 | -------------------------------------------------------------------------------- /isaac_ros_gxf_extensions/gxf_isaac_ros_segment_anything/CMakeLists.txt: -------------------------------------------------------------------------------- 1 | # SPDX-FileCopyrightText: NVIDIA CORPORATION & AFFILIATES 2 | # Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. 3 | # 4 | # Licensed under the Apache License, Version 2.0 (the "License"); 5 | # you may not use this file except in compliance with the License. 6 | # You may obtain a copy of the License at 7 | # 8 | # http://www.apache.org/licenses/LICENSE-2.0 9 | # 10 | # Unless required by applicable law or agreed to in writing, software 11 | # distributed under the License is distributed on an "AS IS" BASIS, 12 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | # See the License for the specific language governing permissions and 14 | # limitations under the License. 15 | # 16 | # SPDX-License-Identifier: Apache-2.0 17 | cmake_minimum_required(VERSION 3.22.1) 18 | project(gxf_isaac_ros_segment_anything LANGUAGES C CXX) 19 | 20 | if(CMAKE_COMPILER_IS_GNUCXX OR CMAKE_CXX_COMPILER_ID MATCHES "Clang") 21 | add_compile_options(-fPIC -w) 22 | endif() 23 | 24 | # Dependencies 25 | find_package(ament_cmake_auto REQUIRED) 26 | ament_auto_find_build_dependencies() 27 | find_package(CUDAToolkit) 28 | find_package(yaml-cpp) 29 | find_package(isaac_ros_nitros_detection2_d_array_type REQUIRED) 30 | enable_language(CUDA) 31 | # Create extension 32 | ament_auto_add_library(${PROJECT_NAME} SHARED 33 | gxf/segment_anything/segment_anything_ext.cpp 34 | gxf/segment_anything/segment_anything_postprocessor.cpp 35 | gxf/segment_anything/segment_anything_postprocessor.hpp 36 | gxf/segment_anything/segment_anything_postprocessor.cu.cpp 37 | gxf/segment_anything/segment_anything_postprocessor.cu.hpp 38 | gxf/segment_anything/segment_anything_prompt_processor.cpp 39 | gxf/segment_anything/segment_anything_prompt_processor.hpp 40 | gxf/segment_anything/segment_anything_msg_compositor.cpp 41 | gxf/segment_anything/segment_anything_msg_compositor.hpp 42 | ) 43 | 44 | # Mark as CUDA files with non-standard extensions 45 | set_source_files_properties( 46 | gxf/segment_anything/segment_anything_postprocessor.cu.cpp 47 | gxf/segment_anything/segment_anything_postprocessor.cu.hpp 48 | PROPERTIES LANGUAGE CUDA 49 | ) 50 | target_link_libraries(${PROJECT_NAME} 51 | CUDA::cudart 52 | yaml-cpp 53 | ) 54 | target_include_directories(${PROJECT_NAME} PRIVATE "${CMAKE_CURRENT_SOURCE_DIR}/gxf" 55 | ${isaac_ros_nitros_detection2_d_array_type_INCLUDE_DIRS}) 56 | 57 | set_target_properties(${PROJECT_NAME} PROPERTIES 58 | BUILD_WITH_INSTALL_RPATH TRUE 59 | BUILD_RPATH_USE_ORIGIN TRUE 60 | INSTALL_RPATH_USE_LINK_PATH TRUE 61 | ) 62 | 63 | # Install the binary file 64 | install(TARGETS ${PROJECT_NAME} DESTINATION share/${PROJECT_NAME}/gxf/lib) 65 | 66 | 67 | # Embed versioning information into installed files 68 | ament_index_get_resource(ISAAC_ROS_COMMON_CMAKE_PATH isaac_ros_common_cmake_path isaac_ros_common) 69 | include("${ISAAC_ROS_COMMON_CMAKE_PATH}/isaac_ros_common-version-info.cmake") 70 | generate_version_info(${PROJECT_NAME}) 71 | 72 | ament_auto_package(INSTALL_TO_SHARE) -------------------------------------------------------------------------------- /isaac_ros_gxf_extensions/gxf_isaac_ros_unet/CMakeLists.txt: -------------------------------------------------------------------------------- 1 | # SPDX-FileCopyrightText: NVIDIA CORPORATION & AFFILIATES 2 | # Copyright (c) 2022-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. 3 | # 4 | # Licensed under the Apache License, Version 2.0 (the "License"); 5 | # you may not use this file except in compliance with the License. 6 | # You may obtain a copy of the License at 7 | # 8 | # http://www.apache.org/licenses/LICENSE-2.0 9 | # 10 | # Unless required by applicable law or agreed to in writing, software 11 | # distributed under the License is distributed on an "AS IS" BASIS, 12 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | # See the License for the specific language governing permissions and 14 | # limitations under the License. 15 | # 16 | # SPDX-License-Identifier: Apache-2.0 17 | 18 | cmake_minimum_required(VERSION 3.22.1) 19 | project(gxf_isaac_ros_unet LANGUAGES C CXX) 20 | 21 | if(CMAKE_COMPILER_IS_GNUCXX OR CMAKE_CXX_COMPILER_ID MATCHES "Clang") 22 | add_compile_options(-fPIC -w) 23 | endif() 24 | 25 | find_package(ament_cmake_auto REQUIRED) 26 | ament_auto_find_build_dependencies() 27 | 28 | # Dependencies 29 | find_package(CUDAToolkit) 30 | find_package(yaml-cpp) 31 | 32 | enable_language(CUDA) 33 | 34 | # Create extension 35 | ament_auto_add_library(${PROJECT_NAME} SHARED 36 | gxf/image_segmentation/segmentation_mask_colorizer.cpp 37 | gxf/image_segmentation/segmentation_mask_colorizer.hpp 38 | gxf/image_segmentation/segmentation_postprocessing_utils.hpp 39 | gxf/image_segmentation/segmentation_postprocessor_ext.cpp 40 | gxf/image_segmentation/segmentation_postprocessor.cpp 41 | gxf/image_segmentation/segmentation_postprocessor.hpp 42 | gxf/image_segmentation/segmentation_postprocessor.cu.cpp 43 | gxf/image_segmentation/segmentation_postprocessor.cu.hpp 44 | gxf/image_segmentation/segmentation_mask_colorizer.cu.cpp 45 | gxf/image_segmentation/segmentation_mask_colorizer.cu.hpp 46 | ) 47 | 48 | # Mark as CUDA files with non-standard extensions 49 | set_source_files_properties( 50 | gxf/image_segmentation/segmentation_postprocessor.cu.cpp 51 | gxf/image_segmentation/segmentation_postprocessor.cu.hpp 52 | gxf/image_segmentation/segmentation_mask_colorizer.cu.cpp 53 | gxf/image_segmentation/segmentation_mask_colorizer.cu.hpp 54 | PROPERTIES LANGUAGE CUDA 55 | ) 56 | 57 | target_link_libraries(${PROJECT_NAME} 58 | CUDA::cudart 59 | yaml-cpp 60 | ) 61 | 62 | target_include_directories(${PROJECT_NAME} PRIVATE "${CMAKE_CURRENT_SOURCE_DIR}/gxf") 63 | 64 | set_target_properties(${PROJECT_NAME} PROPERTIES 65 | BUILD_WITH_INSTALL_RPATH TRUE 66 | BUILD_RPATH_USE_ORIGIN TRUE 67 | INSTALL_RPATH_USE_LINK_PATH TRUE 68 | ) 69 | 70 | # Install the binary file 71 | install(TARGETS ${PROJECT_NAME} DESTINATION share/${PROJECT_NAME}/gxf/lib) 72 | 73 | 74 | # Embed versioning information into installed files 75 | ament_index_get_resource(ISAAC_ROS_COMMON_CMAKE_PATH isaac_ros_common_cmake_path isaac_ros_common) 76 | include("${ISAAC_ROS_COMMON_CMAKE_PATH}/isaac_ros_common-version-info.cmake") 77 | generate_version_info(${PROJECT_NAME}) 78 | 79 | ament_auto_package(INSTALL_TO_SHARE) 80 | 81 | -------------------------------------------------------------------------------- /isaac_ros_segment_anything/config/segment_anything_decoder_node.yaml: -------------------------------------------------------------------------------- 1 | %YAML 1.2 2 | # SPDX-FileCopyrightText: NVIDIA CORPORATION & AFFILIATES 3 | # Copyright (c) 2024-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. 4 | # 5 | # Licensed under the Apache License, Version 2.0 (the "License"); 6 | # you may not use this file except in compliance with the License. 7 | # You may obtain a copy of the License at 8 | # 9 | # http://www.apache.org/licenses/LICENSE-2.0 10 | # 11 | # Unless required by applicable law or agreed to in writing, software 12 | # distributed under the License is distributed on an "AS IS" BASIS, 13 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 14 | # See the License for the specific language governing permissions and 15 | # limitations under the License. 16 | # 17 | # SPDX-License-Identifier: Apache-2.0 18 | --- 19 | name: segmentation_postprocessor 20 | components: 21 | - name: input_tensor 22 | type: nvidia::gxf::DoubleBufferReceiver 23 | parameters: 24 | capacity: 1 25 | - type: nvidia::gxf::MessageAvailableSchedulingTerm 26 | parameters: 27 | receiver: input_tensor 28 | min_size: 1 29 | - name: output_buffer 30 | type: nvidia::gxf::DoubleBufferTransmitter 31 | parameters: 32 | capacity: 1 33 | - type: nvidia::gxf::DownstreamReceptiveSchedulingTerm 34 | parameters: 35 | transmitter: output_buffer 36 | min_size: 1 37 | - name: allocator 38 | type: nvidia::gxf::BlockMemoryPool 39 | parameters: 40 | storage_type: 1 41 | block_size: 522240 42 | num_blocks: 40 43 | - name: stream 44 | type: nvidia::gxf::CudaStreamPool 45 | parameters: 46 | dev_id: 0 47 | stream_flags: 1 48 | stream_priority: 0 49 | reserved_size: 1 50 | max_size: 1 51 | nvtx_identifier: "segment_anything_decoder_node_stream" 52 | - type: nvidia::isaac_ros::SegmentAnythingPostprocessor 53 | parameters: 54 | in: input_tensor 55 | out: output_buffer 56 | allocator: allocator 57 | cuda_stream_pool: stream 58 | - type: nvidia::gxf::MemoryAvailableSchedulingTerm 59 | parameters: 60 | allocator: allocator 61 | min_blocks: 1 62 | --- 63 | name: raw_segmentation_mask_sink 64 | components: 65 | - name: input_raw_segmentation_mask 66 | type: nvidia::gxf::DoubleBufferReceiver 67 | parameters: 68 | capacity: 1 69 | - type: nvidia::gxf::MessageAvailableSchedulingTerm 70 | parameters: 71 | receiver: input_raw_segmentation_mask 72 | min_size: 1 73 | - name: sink 74 | type: nvidia::isaac_ros::MessageRelay 75 | parameters: 76 | source: input_raw_segmentation_mask 77 | --- 78 | components: 79 | - type: nvidia::gxf::Connection 80 | parameters: 81 | source: segmentation_postprocessor/output_buffer 82 | target: raw_segmentation_mask_sink/input_raw_segmentation_mask 83 | --- 84 | components: 85 | - name: clock 86 | type: nvidia::gxf::RealtimeClock 87 | - type: nvidia::gxf::MultiThreadScheduler 88 | parameters: 89 | clock: clock 90 | stop_on_deadlock: false 91 | check_recession_period_ms: 1 92 | worker_thread_number: 2 93 | worker_thread_name_id: "segment_anything_decoder_node" 94 | - type: nvidia::gxf::JobStatistics 95 | parameters: 96 | clock: clock 97 | -------------------------------------------------------------------------------- /isaac_ros_segment_anything/include/isaac_ros_segment_anything/segment_anything_tensor_to_image_node.hpp: -------------------------------------------------------------------------------- 1 | // SPDX-FileCopyrightText: NVIDIA CORPORATION & AFFILIATES 2 | // Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. 3 | // 4 | // Licensed under the Apache License, Version 2.0 (the "License"); 5 | // you may not use this file except in compliance with the License. 6 | // You may obtain a copy of the License at 7 | // 8 | // http://www.apache.org/licenses/LICENSE-2.0 9 | // 10 | // Unless required by applicable law or agreed to in writing, software 11 | // distributed under the License is distributed on an "AS IS" BASIS, 12 | // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | // See the License for the specific language governing permissions and 14 | // limitations under the License. 15 | // 16 | // SPDX-License-Identifier: Apache-2.0 17 | 18 | #pragma once 19 | 20 | #include 21 | #include 22 | 23 | #include "isaac_ros_common/qos.hpp" 24 | #include "isaac_ros_managed_nitros/managed_nitros_publisher.hpp" 25 | #include "isaac_ros_managed_nitros/managed_nitros_subscriber.hpp" 26 | #include "isaac_ros_nitros_image_type/nitros_image.hpp" 27 | #include "isaac_ros_nitros_tensor_list_type/nitros_tensor_list_view.hpp" 28 | #include "rclcpp/rclcpp.hpp" 29 | #include "vision_msgs/msg/detection2_d_array.hpp" 30 | 31 | namespace nvidia 32 | { 33 | namespace isaac_ros 34 | { 35 | namespace segment_anything 36 | { 37 | 38 | /** 39 | * @brief Node that converts a TensorList containing segmentation masks to a binary mask image. 40 | * 41 | * This node subscribes to a NITROS TensorList topic and converts the segmentation masks to a 42 | * binary mask image using NITROS. It only supports batch size 1 tensors. 43 | */ 44 | class TensorToImageNode : public rclcpp::Node 45 | { 46 | public: 47 | /** 48 | * @brief Constructor for TensorListToBinaryMaskNode. 49 | * @param options The node options. 50 | */ 51 | explicit TensorToImageNode(const rclcpp::NodeOptions & options); 52 | 53 | /** 54 | * @brief Destructor for TensorToImageNode. 55 | */ 56 | ~TensorToImageNode() override = default; 57 | 58 | private: 59 | /** 60 | * @brief Callback for processing incoming TensorList messages. 61 | * @param tensor_list The received TensorList message. 62 | */ 63 | void TensorListCallback(const nvidia::isaac_ros::nitros::NitrosTensorListView & tensor_list); 64 | 65 | // QoS settings 66 | rclcpp::QoS input_qos_; 67 | rclcpp::QoS output_qos_; 68 | 69 | // NITROS subscribers and publishers 70 | std::shared_ptr< 71 | nvidia::isaac_ros::nitros::ManagedNitrosSubscriber< 72 | nvidia::isaac_ros::nitros::NitrosTensorListView>> tensor_list_sub_; 73 | 74 | std::shared_ptr< 75 | nvidia::isaac_ros::nitros::ManagedNitrosPublisher< 76 | nvidia::isaac_ros::nitros::NitrosImage>> binary_mask_pub_; 77 | 78 | // Standard ROS publisher for bounding boxes 79 | rclcpp::Publisher::SharedPtr detection_pub_; 80 | 81 | // CUDA stream for GPU operations 82 | cudaStream_t stream_; 83 | }; 84 | 85 | } // namespace segment_anything 86 | } // namespace isaac_ros 87 | } // namespace nvidia 88 | -------------------------------------------------------------------------------- /isaac_ros_unet/CMakeLists.txt: -------------------------------------------------------------------------------- 1 | # SPDX-FileCopyrightText: NVIDIA CORPORATION & AFFILIATES 2 | # Copyright (c) 2021-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. 3 | # 4 | # Licensed under the Apache License, Version 2.0 (the "License"); 5 | # you may not use this file except in compliance with the License. 6 | # You may obtain a copy of the License at 7 | # 8 | # http://www.apache.org/licenses/LICENSE-2.0 9 | # 10 | # Unless required by applicable law or agreed to in writing, software 11 | # distributed under the License is distributed on an "AS IS" BASIS, 12 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | # See the License for the specific language governing permissions and 14 | # limitations under the License. 15 | # 16 | # SPDX-License-Identifier: Apache-2.0 17 | 18 | cmake_minimum_required(VERSION 3.22.1) 19 | project(isaac_ros_unet LANGUAGES C CXX) 20 | 21 | if(CMAKE_COMPILER_IS_GNUCXX OR CMAKE_CXX_COMPILER_ID MATCHES "Clang") 22 | add_compile_options(-Wall -Wextra -Wpedantic) 23 | endif() 24 | 25 | find_package(ament_cmake_auto REQUIRED) 26 | ament_auto_find_build_dependencies() 27 | 28 | # Decoder node 29 | ament_auto_add_library(unet_decoder_node SHARED src/unet_decoder_node.cpp) 30 | rclcpp_components_register_nodes(unet_decoder_node "nvidia::isaac_ros::unet::UNetDecoderNode") 31 | set(node_plugins "${node_plugins}nvidia::isaac_ros::unet::UNetDecoderNode;$\n") 32 | set_target_properties(unet_decoder_node PROPERTIES 33 | BUILD_WITH_INSTALL_RPATH TRUE 34 | BUILD_RPATH_USE_ORIGIN TRUE 35 | INSTALL_RPATH_USE_LINK_PATH TRUE 36 | ) 37 | 38 | if(BUILD_TESTING) 39 | find_package(ament_lint_auto REQUIRED) 40 | find_package(ament_cmake_gtest REQUIRED) 41 | ament_lint_auto_find_test_dependencies() 42 | 43 | # Force use of ROS2 vendor googletest to avoid version conflicts 44 | set(CMAKE_IGNORE_PATH "/usr/src/googletest" ${CMAKE_IGNORE_PATH}) 45 | set(CMAKE_IGNORE_PATH "/usr/include/gtest" ${CMAKE_IGNORE_PATH}) 46 | set(CMAKE_IGNORE_PATH "/usr/include/gmock" ${CMAKE_IGNORE_PATH}) 47 | 48 | # Explicitly use ROS2 vendor googletest/gmock 49 | include_directories(BEFORE /opt/ros/$ENV{ROS_DISTRO}/src/gtest_vendor/include) 50 | include_directories(BEFORE /opt/ros/$ENV{ROS_DISTRO}/src/gmock_vendor/include) 51 | 52 | # Gtest for unet decoder node 53 | ament_add_gtest(unet_decoder_node_test test/unet_decoder_node_test.cpp) 54 | target_link_libraries(unet_decoder_node_test unet_decoder_node) 55 | target_include_directories(unet_decoder_node_test PUBLIC include/isaac_ros_unet/) 56 | ament_target_dependencies(unet_decoder_node_test rclcpp) 57 | ament_target_dependencies(unet_decoder_node_test isaac_ros_nitros) 58 | 59 | # The FindPythonInterp and FindPythonLibs modules are removed 60 | if(POLICY CMP0148) 61 | cmake_policy(SET CMP0148 OLD) 62 | endif() 63 | 64 | find_package(launch_testing_ament_cmake REQUIRED) 65 | add_launch_test(test/isaac_ros_unet_pol_test.py TIMEOUT "600") 66 | endif() 67 | 68 | 69 | # Embed versioning information into installed files 70 | ament_index_get_resource(ISAAC_ROS_COMMON_CMAKE_PATH isaac_ros_common_cmake_path isaac_ros_common) 71 | include("${ISAAC_ROS_COMMON_CMAKE_PATH}/isaac_ros_common-version-info.cmake") 72 | generate_version_info(${PROJECT_NAME}) 73 | 74 | ament_auto_package(INSTALL_TO_SHARE config launch) 75 | -------------------------------------------------------------------------------- /isaac_ros_gxf_extensions/gxf_isaac_ros_unet/gxf/image_segmentation/segmentation_postprocessing_utils.hpp: -------------------------------------------------------------------------------- 1 | // SPDX-FileCopyrightText: NVIDIA CORPORATION & AFFILIATES 2 | // Copyright (c) 2022-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. 3 | // 4 | // Licensed under the Apache License, Version 2.0 (the "License"); 5 | // you may not use this file except in compliance with the License. 6 | // You may obtain a copy of the License at 7 | // 8 | // http://www.apache.org/licenses/LICENSE-2.0 9 | // 10 | // Unless required by applicable law or agreed to in writing, software 11 | // distributed under the License is distributed on an "AS IS" BASIS, 12 | // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | // See the License for the specific language governing permissions and 14 | // limitations under the License. 15 | // 16 | // SPDX-License-Identifier: Apache-2.0 17 | #ifndef NVIDIA_ISAAC_ROS_GXF_EXTENSIONS_SEGMENTATION_POSTPROCESSING_UTILS_HPP_ 18 | #define NVIDIA_ISAAC_ROS_GXF_EXTENSIONS_SEGMENTATION_POSTPROCESSING_UTILS_HPP_ 19 | 20 | #include 21 | 22 | #include "gxf/multimedia/video.hpp" 23 | #include "gxf/std/allocator.hpp" 24 | 25 | namespace nvidia { 26 | namespace isaac_ros { 27 | 28 | template 29 | inline constexpr uint8_t GetChannelSize(); 30 | 31 | template <> 32 | inline constexpr uint8_t GetChannelSize() { 33 | return 3; 34 | } 35 | 36 | template <> 37 | inline constexpr uint8_t GetChannelSize() { 38 | return 3; 39 | } 40 | 41 | template <> 42 | inline constexpr uint8_t GetChannelSize() { 43 | return 1; 44 | } 45 | 46 | template 47 | inline const char* GetColorName(); 48 | 49 | template <> 50 | inline const char* GetColorName() { 51 | return "RGB"; 52 | } 53 | 54 | template <> 55 | inline const char* GetColorName() { 56 | return "BGR"; 57 | } 58 | 59 | template <> 60 | inline const char* GetColorName() { 61 | return "gray"; 62 | } 63 | 64 | template 65 | gxf::Expected AllocateVideoBuffer(gxf::Handle video_buffer, size_t width, 66 | size_t height, gxf::MemoryStorageType memory_storage_type, 67 | gxf::Handle allocator) { 68 | if (width % 2 != 0 || height % 2 != 0) { 69 | GXF_LOG_ERROR("Image width and height must be even, but received width: %zu, height: %zu", 70 | width, height); 71 | return gxf::Unexpected{GXF_FAILURE}; 72 | } 73 | 74 | std::array planes{ 75 | gxf::ColorPlane(GetColorName(), GetChannelSize(), GetChannelSize() * width)}; 76 | gxf::VideoFormatSize video_format_size; 77 | const uint64_t size = video_format_size.size(width, height, planes); 78 | const std::vector planes_filled{planes.begin(), planes.end()}; 79 | const gxf::VideoBufferInfo buffer_info{static_cast(width), 80 | static_cast(height), T, planes_filled, 81 | gxf::SurfaceLayout::GXF_SURFACE_LAYOUT_PITCH_LINEAR}; 82 | return video_buffer->resizeCustom(buffer_info, size, memory_storage_type, allocator); 83 | } 84 | 85 | } // namespace isaac_ros 86 | } // namespace nvidia 87 | 88 | #endif // NVIDIA_ISAAC_ROS_GXF_EXTENSIONS_SEGMENTATION_POSTPROCESSING_UTILS_HPP_ 89 | -------------------------------------------------------------------------------- /isaac_ros_segment_anything/scripts/visualize_mask.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python3 2 | 3 | # SPDX-FileCopyrightText: NVIDIA CORPORATION & AFFILIATES 4 | # Copyright (c) 2024-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. 5 | # 6 | # Licensed under the Apache License, Version 2.0 (the "License"); 7 | # you may not use this file except in compliance with the License. 8 | # You may obtain a copy of the License at 9 | # 10 | # http://www.apache.org/licenses/LICENSE-2.0 11 | # 12 | # Unless required by applicable law or agreed to in writing, software 13 | # distributed under the License is distributed on an "AS IS" BASIS, 14 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 15 | # See the License for the specific language governing permissions and 16 | # limitations under the License. 17 | # 18 | # SPDX-License-Identifier: Apache-2.0 19 | import cv2 20 | import cv_bridge 21 | from isaac_ros_tensor_list_interfaces.msg import TensorList 22 | from message_filters import Subscriber, TimeSynchronizer 23 | import numpy as np 24 | import rclpy 25 | from rclpy.node import Node 26 | from sensor_msgs.msg import Image 27 | 28 | 29 | class SegmentAnythingVisualization(Node): 30 | 31 | def __init__(self): 32 | super().__init__('segment_anything_visualizer') 33 | 34 | self._bridge = cv_bridge.CvBridge() 35 | self._mask_subscriber = Subscriber( 36 | self, TensorList, '/segment_anything/raw_segmentation_mask') 37 | 38 | self._image_subscriber = Subscriber( 39 | self, Image, '/yolov8_encoder/resize/image') 40 | 41 | self._processed_image_pub = self.create_publisher( 42 | Image, '/segment_anything/colored_segmentation_mask', 10) 43 | 44 | # Pre-convert color palette from hex to RGB for better performance 45 | hex_colors = ['800000', '556B2F', '008080', '000080', '9ACD32', 'FF0000', 46 | 'FF8C00', 'FFD700', '00FF00', 'BA55D3', '00FA9A', '00FFFF', 47 | '0000FF', 'F08080', 'FF00FF', '1E90FF', 'DDA0DD', 'FF1493', 48 | '87CEFA', 'FFDEAD'] 49 | 50 | self._color_palette = np.array([ 51 | [int(color[i:i+2], 16) for i in (0, 2, 4)] 52 | for color in hex_colors 53 | ], dtype=np.uint8) 54 | 55 | self.sync = TimeSynchronizer([self._image_subscriber, self._mask_subscriber], 50) 56 | self.sync.registerCallback(self.callback) 57 | 58 | def callback(self, img, masks): 59 | input_image = self._bridge.imgmsg_to_cv2(img, 'rgb8') 60 | tensor = masks.tensors[0] 61 | 62 | # Extract dimensions directly from shape 63 | dimensions = tensor.shape.dims.tolist() 64 | num_masks = dimensions[0] 65 | 66 | # Reshape data more efficiently 67 | data = np.array(tensor.data.tolist(), dtype=np.uint8).reshape(dimensions) 68 | 69 | # Start with the original image instead of an empty array 70 | result_image = input_image.copy() 71 | # Apply mask colors directly to the result image 72 | for n in range(num_masks): 73 | pallet_idx = min(n, len(self._color_palette) - 1) 74 | mask = data[n, 0, :, :] > 0 75 | if np.any(mask): # Only process if mask has positive values 76 | # For pixels where mask exists, blend with 50% transparency 77 | result_image[mask] = cv2.addWeighted( 78 | input_image[mask], 0.4, 79 | np.full((np.count_nonzero(mask), 3), 80 | self._color_palette[pallet_idx], dtype=np.uint8), 81 | 0.6, 0 82 | ) 83 | 84 | # Publish the processed image 85 | processed_img = self._bridge.cv2_to_imgmsg(result_image, encoding='rgb8') 86 | processed_img.header = img.header # Maintain timestamp and frame_id 87 | self._processed_image_pub.publish(processed_img) 88 | 89 | 90 | def main(args=None): 91 | rclpy.init(args=args) 92 | node = SegmentAnythingVisualization() 93 | rclpy.spin(node) 94 | node.destroy_node() 95 | rclpy.shutdown() 96 | 97 | 98 | if __name__ == '__main__': 99 | main() 100 | -------------------------------------------------------------------------------- /isaac_ros_segment_anything/src/segment_anything_dummy_mask_publisher_node.cpp: -------------------------------------------------------------------------------- 1 | // SPDX-FileCopyrightText: NVIDIA CORPORATION & AFFILIATES 2 | // Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. 3 | // 4 | // Licensed under the Apache License, Version 2.0 (the "License"); 5 | // you may not use this file except in compliance with the License. 6 | // You may obtain a copy of the License at 7 | // 8 | // http://www.apache.org/licenses/LICENSE-2.0 9 | // 10 | // Unless required by applicable law or agreed to in writing, software 11 | // distributed under the License is distributed on an "AS IS" BASIS, 12 | // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | // See the License for the specific language governing permissions and 14 | // limitations under the License. 15 | // 16 | // SPDX-License-Identifier: Apache-2.0 17 | 18 | #include 19 | #include 20 | 21 | #include "isaac_ros_common/cuda_stream.hpp" 22 | #include "isaac_ros_nitros_tensor_list_type/nitros_tensor_builder.hpp" 23 | #include "isaac_ros_nitros_tensor_list_type/nitros_tensor_list_builder.hpp" 24 | #include "isaac_ros_segment_anything/segment_anything_dummy_mask_publisher_node.hpp" 25 | namespace nvidia 26 | { 27 | namespace isaac_ros 28 | { 29 | namespace segment_anything 30 | { 31 | 32 | DummyMaskPublisher::DummyMaskPublisher(const rclcpp::NodeOptions options) 33 | : rclcpp::Node("dummy_mask_publisher", options), 34 | nitros_sub_{std::make_shared>( 36 | this, 37 | "tensor_pub", 38 | nvidia::isaac_ros::nitros::nitros_tensor_list_nchw_rgb_f32_t::supported_type_name, 39 | std::bind(&DummyMaskPublisher::InputCallback, this, 40 | std::placeholders::_1))}, 41 | nitros_pub_{std::make_shared>( 43 | this, "mask", 44 | nvidia::isaac_ros::nitros::nitros_tensor_list_nchw_rgb_f32_t::supported_type_name)}, 45 | tensor_name_{declare_parameter("tensor_name", "input_mask")} 46 | { 47 | // Initialize CUDA stream 48 | CHECK_CUDA_ERROR( 49 | ::nvidia::isaac_ros::common::initNamedCudaStream( 50 | stream_, "isaac_ros_dummy_mask_publisher_node"), 51 | "Error initializing CUDA stream"); 52 | } 53 | 54 | DummyMaskPublisher::~DummyMaskPublisher() = default; 55 | 56 | void DummyMaskPublisher::InputCallback(const nvidia::isaac_ros::nitros::NitrosTensorListView & msg) 57 | { 58 | // Buffer size for mask 59 | size_t buffer_size{256 * 256 * 4}; 60 | 61 | // Allocate CUDA buffer in the stream to not block the main default stream and other work. 62 | void * buffer; 63 | CHECK_CUDA_ERROR( 64 | cudaMallocAsync(&buffer, buffer_size, stream_), 65 | "Failed to allocate GPU memory"); 66 | CHECK_CUDA_ERROR( 67 | cudaMemsetAsync(buffer, 0, buffer_size, stream_), 68 | "Failed to zero GPU memory"); 69 | 70 | // Adding header data 71 | std_msgs::msg::Header header; 72 | header.stamp.sec = msg.GetTimestampSeconds(); 73 | header.stamp.nanosec = msg.GetTimestampNanoseconds(); 74 | header.frame_id = msg.GetFrameId(); 75 | 76 | // Sync the stream 77 | CHECK_CUDA_ERROR(cudaStreamSynchronize(stream_), "Failed to synchronize CUDA stream"); 78 | 79 | // Create tensor list with tensor wrapping CUDA buffer 80 | nvidia::isaac_ros::nitros::NitrosTensorList tensor_list = 81 | nvidia::isaac_ros::nitros::NitrosTensorListBuilder() 82 | .WithHeader(header) 83 | .AddTensor( 84 | tensor_name_, 85 | ( 86 | nvidia::isaac_ros::nitros::NitrosTensorBuilder() 87 | .WithShape({1, 1, 256, 256}) 88 | .WithDataType(nvidia::isaac_ros::nitros::NitrosDataType::kFloat32) 89 | .WithData(buffer) 90 | .Build() 91 | ) 92 | ) 93 | .Build(); 94 | nitros_pub_->publish(tensor_list); 95 | } 96 | 97 | } // namespace segment_anything 98 | } // namespace isaac_ros 99 | } // namespace nvidia 100 | 101 | // Register as component 102 | #include "rclcpp_components/register_node_macro.hpp" 103 | RCLCPP_COMPONENTS_REGISTER_NODE(nvidia::isaac_ros::segment_anything::DummyMaskPublisher) 104 | -------------------------------------------------------------------------------- /isaac_ros_gxf_extensions/gxf_isaac_ros_segment_anything/gxf/segment_anything/segment_anything_msg_compositor.cpp: -------------------------------------------------------------------------------- 1 | // SPDX-FileCopyrightText: NVIDIA CORPORATION & AFFILIATES 2 | // Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. 3 | // 4 | // Licensed under the Apache License, Version 2.0 (the "License"); 5 | // you may not use this file except in compliance with the License. 6 | // You may obtain a copy of the License at 7 | // 8 | // http://www.apache.org/licenses/LICENSE-2.0 9 | // 10 | // Unless required by applicable law or agreed to in writing, software 11 | // distributed under the License is distributed on an "AS IS" BASIS, 12 | // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | // See the License for the specific language governing permissions and 14 | // limitations under the License. 15 | // 16 | // SPDX-License-Identifier: Apache-2.0 17 | #include 18 | #include 19 | 20 | #include "cuda.h" 21 | #include "cuda_runtime.h" 22 | #include "gxf/cuda/cuda_stream_id.hpp" 23 | #include "gxf/multimedia/video.hpp" 24 | #include "gxf/std/tensor.hpp" 25 | #include "gxf/std/timestamp.hpp" 26 | #include "segment_anything_msg_compositor.hpp" 27 | 28 | namespace nvidia { 29 | namespace isaac_ros { 30 | 31 | namespace { 32 | 33 | gxf::Expected AddInputTimestampToOutput(gxf::Entity& output, gxf::Entity input) { 34 | std::string timestamp_name{"timestamp"}; 35 | auto maybe_timestamp = input.get(timestamp_name.c_str()); 36 | 37 | // Default to unnamed 38 | if (!maybe_timestamp) { 39 | timestamp_name = std::string{""}; 40 | maybe_timestamp = input.get(timestamp_name.c_str()); 41 | } 42 | 43 | if (!maybe_timestamp) { 44 | GXF_LOG_ERROR("Failed to get input timestamp!"); 45 | return gxf::ForwardError(maybe_timestamp); 46 | } 47 | 48 | auto maybe_out_timestamp = output.add(timestamp_name.c_str()); 49 | if (!maybe_out_timestamp) { 50 | GXF_LOG_ERROR("Failed to add timestamp to output message!"); 51 | return gxf::ForwardError(maybe_out_timestamp); 52 | } 53 | 54 | *maybe_out_timestamp.value() = *maybe_timestamp.value(); 55 | return gxf::Success; 56 | } 57 | } // namespace 58 | 59 | gxf_result_t SegmentAnythingMsgCompositor::registerInterface(gxf::Registrar* registrar) { 60 | gxf::Expected result; 61 | result &= registrar->parameter(inputs_, "inputs", "Input", "List of input receivers"); 62 | result &= registrar->parameter(output_, "output", "Output", "Output Transmitter"); 63 | return gxf::ToResultCode(result); 64 | } 65 | 66 | gxf_result_t SegmentAnythingMsgCompositor::start() { 67 | return GXF_SUCCESS; 68 | } 69 | 70 | gxf_result_t SegmentAnythingMsgCompositor::tick() { 71 | // Process input message 72 | auto out_message = gxf::Entity::New(context()); 73 | bool assign_ts = false; 74 | bool is_data_valid = true; 75 | for (const auto & rx : inputs_.get()) { 76 | const auto in_message = rx->receive(); 77 | 78 | if (!in_message || in_message.value().is_null()) { return GXF_CONTRACT_MESSAGE_NOT_AVAILABLE; } 79 | // Look for invalid_frame tensor, it means the current input has no corresponding prompt. 80 | auto maybe_invalid_data = in_message.value().get("invalid_frame"); 81 | if (maybe_invalid_data) { 82 | is_data_valid = false; 83 | } 84 | auto maybe_tensor_list = in_message.value().findAll(); 85 | if (!maybe_tensor_list) { return gxf::ToResultCode(maybe_tensor_list); } 86 | auto tensor_list = maybe_tensor_list.value(); 87 | 88 | for (int idx = 0; idx(in_tensor.name()); 93 | if (!out_tensor) { return gxf::ToResultCode(out_tensor); } 94 | 95 | out_tensor.value()->wrapMemoryBuffer(in_tensor->shape(), 96 | in_tensor->element_type(),gxf::PrimitiveTypeSize(in_tensor->element_type()), 97 | gxf::Unexpected{GXF_UNINITIALIZED_VALUE},in_tensor->move_buffer()); 98 | } 99 | if(!assign_ts){ 100 | auto maybe_added_timestamp = AddInputTimestampToOutput(out_message.value(), in_message.value()); 101 | if(maybe_added_timestamp){ 102 | assign_ts = true; 103 | } 104 | } 105 | } 106 | // Only publish if the data is valid. 107 | if (is_data_valid) { 108 | output_->publish(std::move(out_message.value())); 109 | } 110 | return GXF_SUCCESS; 111 | } 112 | 113 | gxf_result_t SegmentAnythingMsgCompositor::stop() { 114 | return GXF_SUCCESS; 115 | } 116 | 117 | } // namespace isaac_ros 118 | } // namespace nvidia 119 | -------------------------------------------------------------------------------- /isaac_ros_segment_anything2/include/isaac_ros_segment_anything2/segment_anything2_data_encoder_node.hpp: -------------------------------------------------------------------------------- 1 | // SPDX-FileCopyrightText: NVIDIA CORPORATION & AFFILIATES 2 | // Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. 3 | // 4 | // Licensed under the Apache License, Version 2.0 (the "License"); 5 | // you may not use this file except in compliance with the License. 6 | // You may obtain a copy of the License at 7 | // 8 | // http://www.apache.org/licenses/LICENSE-2.0 9 | // 10 | // Unless required by applicable law or agreed to in writing, software 11 | // distributed under the License is distributed on an "AS IS" BASIS, 12 | // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | // See the License for the specific language governing permissions and 14 | // limitations under the License. 15 | // 16 | // SPDX-License-Identifier: Apache-2.0 17 | 18 | #ifndef ISAAC_ROS_SEGMENT_ANYTHING2__SEGMENT_ANYTHING2_DATA_ENCODER_NODE_HPP_ 19 | #define ISAAC_ROS_SEGMENT_ANYTHING2__SEGMENT_ANYTHING2_DATA_ENCODER_NODE_HPP_ 20 | 21 | #include 22 | #include 23 | #include 24 | #include "isaac_ros_common/qos.hpp" 25 | #include "rclcpp/rclcpp.hpp" 26 | #include "sensor_msgs/msg/image.hpp" 27 | #include "isaac_ros_tensor_list_interfaces/msg/tensor_list.hpp" 28 | #include "vision_msgs/msg/detection2_d_array.hpp" 29 | #include "isaac_ros_nitros_detection2_d_array_type/nitros_detection2_d_array.hpp" 30 | #include "isaac_ros_managed_nitros/managed_nitros_publisher.hpp" 31 | #include "isaac_ros_managed_nitros/managed_nitros_subscriber.hpp" 32 | #include "isaac_ros_segment_anything2/segment_anything2_state_manager.hpp" 33 | #include "isaac_ros_nitros_tensor_list_type/nitros_tensor_list_view.hpp" 34 | #include "isaac_ros_segment_anything2_interfaces/srv/add_objects.hpp" 35 | #include "isaac_ros_segment_anything2_interfaces/srv/remove_object.hpp" 36 | #include "isaac_ros_common/cuda_stream.hpp" 37 | 38 | namespace nvidia 39 | { 40 | namespace isaac_ros 41 | { 42 | namespace segment_anything2 43 | { 44 | 45 | class SegmentAnything2DataEncoderNode : public rclcpp::Node 46 | { 47 | public: 48 | explicit SegmentAnything2DataEncoderNode( 49 | const rclcpp::NodeOptions options = rclcpp::NodeOptions() 50 | ); 51 | ~SegmentAnything2DataEncoderNode(); 52 | 53 | private: 54 | // QoS settings 55 | rclcpp::QoS image_qos_; 56 | rclcpp::QoS memory_qos_; 57 | rclcpp::QoS encoded_data_qos_; 58 | 59 | // Callbacks for subscribers 60 | void ImageCallback(const nvidia::isaac_ros::nitros::NitrosTensorListView & msg); 61 | void MemoryCallback(const nvidia::isaac_ros::nitros::NitrosTensorListView & msg); 62 | 63 | // Service callback for adding objects 64 | void AddObjectsCallback( 65 | const std::shared_ptr< 66 | isaac_ros_segment_anything2_interfaces::srv::AddObjects::Request> request, 67 | std::shared_ptr< 68 | isaac_ros_segment_anything2_interfaces::srv::AddObjects::Response> response); 69 | 70 | // Service callback to remove object 71 | void RemoveObjectCallback( 72 | const std::shared_ptr< 73 | isaac_ros_segment_anything2_interfaces::srv::RemoveObject::Request> request, 74 | std::shared_ptr< 75 | isaac_ros_segment_anything2_interfaces::srv::RemoveObject::Response> response); 76 | 77 | // Publisher for output NitrosTensorList messages 78 | std::shared_ptr< 79 | nvidia::isaac_ros::nitros::ManagedNitrosPublisher< 80 | nvidia::isaac_ros::nitros::NitrosTensorList>> encoded_data_pub_; 81 | // Subscribers 82 | std::shared_ptr> image_sub_; 84 | std::shared_ptr> memory_sub_; 86 | 87 | // Service for adding objects to the segmentation model 88 | rclcpp::Service< 89 | isaac_ros_segment_anything2_interfaces::srv::AddObjects>::SharedPtr add_objects_srv_; 90 | 91 | // Service for adding objects to the segmentation model 92 | rclcpp::Service< 93 | isaac_ros_segment_anything2_interfaces::srv::RemoveObject>::SharedPtr remove_object_srv_; 94 | // Maximum number of objects that can be added to the state manager. 95 | // Required because post processor always caps the number of objects. 96 | int32_t max_num_objects_; 97 | 98 | std::vector orig_img_dims_param_; 99 | int32_t * original_size_buffer_; 100 | std::unique_ptr sam2_state_manager_; 101 | cudaStream_t stream_; 102 | }; 103 | 104 | } // namespace segment_anything2 105 | } // namespace isaac_ros 106 | } // namespace nvidia 107 | 108 | #endif // ISAAC_ROS_SEGMENT_ANYTHING2__SEGMENT_ANYTHING2_DATA_ENCODER_NODE_HPP_ 109 | -------------------------------------------------------------------------------- /isaac_ros_unet/test/unet_decoder_node_test.cpp: -------------------------------------------------------------------------------- 1 | // SPDX-FileCopyrightText: NVIDIA CORPORATION & AFFILIATES 2 | // Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. 3 | // 4 | // Licensed under the Apache License, Version 2.0 (the "License"); 5 | // you may not use this file except in compliance with the License. 6 | // You may obtain a copy of the License at 7 | // 8 | // http://www.apache.org/licenses/LICENSE-2.0 9 | // 10 | // Unless required by applicable law or agreed to in writing, software 11 | // distributed under the License is distributed on an "AS IS" BASIS, 12 | // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | // See the License for the specific language governing permissions and 14 | // limitations under the License. 15 | // 16 | // SPDX-License-Identifier: Apache-2.0 17 | 18 | #include 19 | #include "unet_decoder_node.hpp" 20 | #include "rclcpp/rclcpp.hpp" 21 | 22 | // Objective: to cover code lines where exceptions are thrown 23 | // Approach: send Invalid Arguments for node parameters to trigger the exception 24 | 25 | 26 | TEST(unet_decoder_node_test, test_empty_color_segmentation_mask_encoding) 27 | { 28 | rclcpp::init(0, nullptr); 29 | rclcpp::NodeOptions options; 30 | options.append_parameter_override("color_segmentation_mask_encoding", ""); 31 | EXPECT_THROW( 32 | { 33 | try { 34 | nvidia::isaac_ros::unet::UNetDecoderNode unet_decoder_node(options); 35 | } catch (const std::invalid_argument & e) { 36 | EXPECT_THAT(e.what(), testing::HasSubstr("Received empty color segmentation mask encoding!")); 37 | throw; 38 | } catch (const rclcpp::exceptions::InvalidParameterValueException & e) { 39 | EXPECT_THAT(e.what(), testing::HasSubstr("No parameter value set")); 40 | throw; 41 | } 42 | }, std::invalid_argument); 43 | rclcpp::shutdown(); 44 | } 45 | 46 | TEST(unet_decoder_node_test, test_invalid_color_segmentation_mask_encoding) 47 | { 48 | rclcpp::init(0, nullptr); 49 | rclcpp::NodeOptions options; 50 | options.append_parameter_override("color_segmentation_mask_encoding", "gbr8"); 51 | options.append_parameter_override("color_palette", std::vector{1}); 52 | EXPECT_THROW( 53 | { 54 | try { 55 | nvidia::isaac_ros::unet::UNetDecoderNode unet_decoder_node(options); 56 | } catch (const std::invalid_argument & e) { 57 | EXPECT_THAT( 58 | e.what(), 59 | testing::HasSubstr("Received invalid color segmentation mask encoding")); 60 | throw; 61 | } catch (const rclcpp::exceptions::InvalidParameterValueException & e) { 62 | EXPECT_THAT(e.what(), testing::HasSubstr("No parameter value set")); 63 | throw; 64 | } 65 | }, std::invalid_argument); 66 | rclcpp::shutdown(); 67 | } 68 | 69 | TEST(unet_decoder_node_test, test_empty_color_palette) 70 | { 71 | rclcpp::init(0, nullptr); 72 | rclcpp::NodeOptions options; 73 | options.append_parameter_override("color_segmentation_mask_encoding", "rgb8"); 74 | EXPECT_THROW( 75 | { 76 | try { 77 | nvidia::isaac_ros::unet::UNetDecoderNode unet_decoder_node(options); 78 | } catch (const std::invalid_argument & e) { 79 | EXPECT_THAT( 80 | e.what(), 81 | testing::HasSubstr( 82 | "Received empty color palette! Fill this with a 24-bit hex color for each class!")); 83 | throw; 84 | } catch (const rclcpp::exceptions::InvalidParameterValueException & e) { 85 | EXPECT_THAT(e.what(), testing::HasSubstr("No parameter value set")); 86 | throw; 87 | } 88 | }, std::invalid_argument); 89 | rclcpp::shutdown(); 90 | } 91 | 92 | TEST(unet_decoder_node_test, test_invalid_network_output_type) 93 | { 94 | rclcpp::init(0, nullptr); 95 | rclcpp::NodeOptions options; 96 | // options.arguments( 97 | // { 98 | // "--ros-args", 99 | // "-p", "color_segmentation_mask_encoding:='rgb8'", 100 | // "-p", "color_palette:=[1]", 101 | // "-p", "network_output_type:='invalid'", 102 | // }); 103 | options.append_parameter_override("color_segmentation_mask_encoding", "rgb8"); 104 | options.append_parameter_override("color_palette", std::vector(1)); 105 | options.append_parameter_override("network_output_type", "invalid"); 106 | EXPECT_THROW( 107 | { 108 | try { 109 | nvidia::isaac_ros::unet::UNetDecoderNode unet_decoder_node(options); 110 | } catch (const std::invalid_argument & e) { 111 | EXPECT_THAT(e.what(), testing::HasSubstr("Received invalid network output type: ")); 112 | throw; 113 | } catch (const rclcpp::exceptions::InvalidParameterValueException & e) { 114 | EXPECT_THAT(e.what(), testing::HasSubstr("No parameter value set")); 115 | throw; 116 | } 117 | }, std::invalid_argument); 118 | rclcpp::shutdown(); 119 | } 120 | 121 | 122 | int main(int argc, char ** argv) 123 | { 124 | testing::InitGoogleTest(&argc, argv); 125 | return RUN_ALL_TESTS(); 126 | } 127 | -------------------------------------------------------------------------------- /isaac_ros_segment_anything/src/segment_anything_decoder_node.cpp: -------------------------------------------------------------------------------- 1 | // SPDX-FileCopyrightText: NVIDIA CORPORATION & AFFILIATES 2 | // Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. 3 | // 4 | // Licensed under the Apache License, Version 2.0 (the "License"); 5 | // you may not use this file except in compliance with the License. 6 | // You may obtain a copy of the License at 7 | // 8 | // http://www.apache.org/licenses/LICENSE-2.0 9 | // 10 | // Unless required by applicable law or agreed to in writing, software 11 | // distributed under the License is distributed on an "AS IS" BASIS, 12 | // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | // See the License for the specific language governing permissions and 14 | // limitations under the License. 15 | // 16 | // SPDX-License-Identifier: Apache-2.0 17 | 18 | #include "isaac_ros_segment_anything/segment_anything_decoder_node.hpp" 19 | 20 | #include 21 | #include 22 | #include 23 | #include 24 | #include 25 | 26 | #include "rclcpp/rclcpp.hpp" 27 | #include "sensor_msgs/image_encodings.hpp" 28 | #include "isaac_ros_nitros_image_type/nitros_image.hpp" 29 | #include "isaac_ros_nitros_tensor_list_type/nitros_tensor_list.hpp" 30 | 31 | namespace nvidia 32 | { 33 | namespace isaac_ros 34 | { 35 | namespace segment_anything 36 | { 37 | namespace 38 | { 39 | 40 | using nvidia::gxf::optimizer::GraphIOGroupSupportedDataTypesInfoList; 41 | 42 | constexpr char INPUT_COMPONENT_KEY[] = "segmentation_postprocessor/input_tensor"; 43 | constexpr char INPUT_DEFAULT_TENSOR_FORMAT[] = "nitros_tensor_list_nchw_rgb_f32"; 44 | constexpr char INPUT_TOPIC_NAME[] = "tensor_sub"; 45 | 46 | constexpr char RAW_OUTPUT_COMPONENT_KEY[] = "raw_segmentation_mask_sink/sink"; 47 | constexpr char RAW_OUTPUT_DEFAULT_TENSOR_FORMAT[] = "nitros_tensor_list_nchw"; 48 | constexpr char RAW_OUTPUT_TOPIC_NAME[] = "segment_anything/raw_segmentation_mask"; 49 | 50 | constexpr char APP_YAML_FILENAME[] = "config/segment_anything_decoder_node.yaml"; 51 | constexpr char PACKAGE_NAME[] = "isaac_ros_segment_anything"; 52 | 53 | const std::vector> EXTENSIONS = { 54 | {"isaac_ros_gxf", "gxf/lib/std/libgxf_std.so"}, 55 | {"isaac_ros_gxf", "gxf/lib/cuda/libgxf_cuda.so"}, 56 | {"gxf_isaac_ros_segment_anything", "gxf/lib/libgxf_isaac_ros_segment_anything.so"}}; 57 | 58 | const std::vector PRESET_EXTENSION_SPEC_NAMES = {}; 59 | const std::vector EXTENSION_SPEC_FILENAMES = { 60 | "config/segment_anything_spec_file.yaml" 61 | }; 62 | const std::vector GENERATOR_RULE_FILENAMES = {}; 63 | 64 | #pragma GCC diagnostic push 65 | #pragma GCC diagnostic ignored "-Wpedantic" 66 | const nitros::NitrosPublisherSubscriberConfigMap CONFIG_MAP = { 67 | {INPUT_COMPONENT_KEY, 68 | { 69 | .type = nitros::NitrosPublisherSubscriberType::NEGOTIATED, 70 | .qos = rclcpp::QoS(1), 71 | .compatible_data_format = INPUT_DEFAULT_TENSOR_FORMAT, 72 | .topic_name = INPUT_TOPIC_NAME, 73 | }}, 74 | {RAW_OUTPUT_COMPONENT_KEY, 75 | { 76 | .type = nitros::NitrosPublisherSubscriberType::NEGOTIATED, 77 | .qos = rclcpp::QoS(1), 78 | .compatible_data_format = RAW_OUTPUT_DEFAULT_TENSOR_FORMAT, 79 | .topic_name = RAW_OUTPUT_TOPIC_NAME, 80 | .frame_id_source_key = INPUT_COMPONENT_KEY, 81 | }} 82 | }; 83 | #pragma GCC diagnostic pop 84 | } // namespace 85 | 86 | SegmentAnythingDecoderNode::SegmentAnythingDecoderNode(const rclcpp::NodeOptions options) 87 | : nitros::NitrosNode( 88 | options, 89 | APP_YAML_FILENAME, 90 | CONFIG_MAP, 91 | PRESET_EXTENSION_SPEC_NAMES, 92 | EXTENSION_SPEC_FILENAMES, 93 | GENERATOR_RULE_FILENAMES, 94 | EXTENSIONS, 95 | PACKAGE_NAME), 96 | mask_width_(declare_parameter("mask_width", 960)), 97 | mask_height_(declare_parameter("mask_height", 544)), 98 | max_batch_size_(declare_parameter("max_batch_size", 20)), 99 | tensor_name_(declare_parameter("tensor_name", "")) 100 | { 101 | registerSupportedType(); 102 | registerSupportedType(); 103 | startNitrosNode(); 104 | } 105 | 106 | void SegmentAnythingDecoderNode::postLoadGraphCallback() 107 | { 108 | const uint64_t block_size(mask_width_ * mask_height_ * max_batch_size_); 109 | 110 | getNitrosContext().setParameterUInt64( 111 | "segmentation_postprocessor", 112 | "nvidia::gxf::BlockMemoryPool", "block_size", block_size 113 | ); 114 | getNitrosContext().setParameterStr( 115 | "segmentation_postprocessor", 116 | "nvidia::isaac_ros::SegmentAnythingPostprocessor", 117 | "tensor_name", tensor_name_ 118 | ); 119 | } 120 | 121 | SegmentAnythingDecoderNode::~SegmentAnythingDecoderNode() = default; 122 | 123 | } // namespace segment_anything 124 | } // namespace isaac_ros 125 | } // namespace nvidia 126 | 127 | // Register as component 128 | #include "rclcpp_components/register_node_macro.hpp" 129 | RCLCPP_COMPONENTS_REGISTER_NODE(nvidia::isaac_ros::segment_anything::SegmentAnythingDecoderNode) 130 | -------------------------------------------------------------------------------- /isaac_ros_segment_anything/test/isaac_ros_segment_anything_tensor_to_image_test.py: -------------------------------------------------------------------------------- 1 | # SPDX-FileCopyrightText: NVIDIA CORPORATION & AFFILIATES 2 | # Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. 3 | # 4 | # Licensed under the Apache License, Version 2.0 (the "License"); 5 | # you may not use this file except in compliance with the License. 6 | # You may obtain a copy of the License at 7 | # 8 | # http://www.apache.org/licenses/LICENSE-2.0 9 | # 10 | # Unless required by applicable law or agreed to in writing, software 11 | # distributed under the License is distributed on an "AS IS" BASIS, 12 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | # See the License for the specific language governing permissions and 14 | # limitations under the License. 15 | # 16 | # SPDX-License-Identifier: Apache-2.0 17 | 18 | import time 19 | 20 | from isaac_ros_tensor_list_interfaces.msg import Tensor, TensorList, TensorShape 21 | from isaac_ros_test import IsaacROSBaseTest 22 | import launch 23 | from launch_ros.actions import ComposableNodeContainer 24 | from launch_ros.descriptions import ComposableNode 25 | import launch_testing 26 | import numpy as np 27 | import pytest 28 | import rclpy 29 | from sensor_msgs.msg import Image 30 | 31 | 32 | @pytest.mark.rostest 33 | def generate_test_description(): 34 | """Generate test description for testing tensor to image node.""" 35 | tensor_to_image_node = ComposableNode( 36 | name='tensor_to_image', 37 | package='isaac_ros_segment_anything', 38 | plugin='nvidia::isaac_ros::segment_anything::TensorToImageNode', 39 | ) 40 | 41 | container = ComposableNodeContainer( 42 | name='segment_anything_container', 43 | namespace='', 44 | package='rclcpp_components', 45 | executable='component_container_mt', 46 | composable_node_descriptions=[ 47 | tensor_to_image_node 48 | ], 49 | output='screen' 50 | ) 51 | 52 | return launch.LaunchDescription([ 53 | container, 54 | launch_testing.actions.ReadyToTest() 55 | ]) 56 | 57 | 58 | class TensorToImageTest(IsaacROSBaseTest): 59 | """Test for converting tensor to binary mask image.""" 60 | 61 | def test_tensor_to_image_conversion(self): 62 | """Test tensor to image conversion with a sample tensor.""" 63 | self.node = rclpy.create_node('tensor_to_image_test') 64 | 65 | # Create publishers and subscribers 66 | self.tensor_pub = self.node.create_publisher( 67 | TensorList, 'segmentation_tensor', 10) 68 | 69 | received_messages = [] 70 | 71 | def image_callback(msg): 72 | self.node.get_logger().info('Received image message') 73 | received_messages.append(msg) 74 | 75 | self.image_sub = self.node.create_subscription( 76 | Image, 'binary_mask', image_callback, 10) 77 | 78 | # Create a sample tensor (simulating a segmentation mask) 79 | height, width = 480, 640 80 | test_tensor = np.zeros((1, 1, height, width), dtype=np.uint8) 81 | # Create a simple pattern (e.g., a rectangle) in the middle 82 | test_tensor[0, 0, height//4:3*height//4, width//4:3*width//4] = 255.0 83 | 84 | # Create TensorList message 85 | tensor_msg = TensorList() 86 | tensor_msg.tensors = [] 87 | 88 | tensor = Tensor() 89 | tensor_shape = TensorShape() 90 | 91 | tensor_shape.rank = 4 # [batch, channel, height, width] 92 | tensor_shape.dims = [1, 1, height, width] 93 | 94 | tensor.shape = tensor_shape 95 | tensor.name = 'segmentation' 96 | tensor.data_type = 2 # UINT8 97 | tensor.strides = [] # Let pynitros handle strides 98 | tensor.data = test_tensor.tobytes() 99 | 100 | tensor_msg.tensors.append(tensor) 101 | 102 | # Wait for subscriptions to be ready 103 | time.sleep(2) 104 | 105 | # Publish tensor 106 | self.tensor_pub.publish(tensor_msg) 107 | 108 | # Wait for results (timeout after 10 seconds) 109 | end_time = time.time() + 10 110 | while time.time() < end_time and not received_messages: 111 | rclpy.spin_once(self.node, timeout_sec=0.1) 112 | 113 | # Verify results 114 | self.assertTrue(len(received_messages) > 0, 'No image message received') 115 | 116 | received_image = received_messages[0] 117 | self.assertEqual(received_image.height, height) 118 | self.assertEqual(received_image.width, width) 119 | self.assertEqual(received_image.encoding, 'mono8') 120 | 121 | # Convert received image data to numpy array for verification 122 | received_data = np.frombuffer(received_image.data, dtype=np.uint8).reshape(height, width) 123 | 124 | # Verify the pattern is preserved (allowing for some tolerance due to thresholding) 125 | expected_mask = ((test_tensor[0] > 0.5) * 255).reshape(height, width) 126 | 127 | np.testing.assert_array_equal(received_data, expected_mask, 128 | 'Received image does not match expected binary mask') 129 | 130 | self.node.get_logger().info('Test completed successfully') 131 | -------------------------------------------------------------------------------- /isaac_ros_segment_anything/src/segment_anything_binarize_tensor.cu: -------------------------------------------------------------------------------- 1 | // SPDX-FileCopyrightText: NVIDIA CORPORATION & AFFILIATES 2 | // Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. 3 | // 4 | // Licensed under the Apache License, Version 2.0 (the "License"); 5 | // you may not use this file except in compliance with the License. 6 | // You may obtain a copy of the License at 7 | // 8 | // http://www.apache.org/licenses/LICENSE-2.0 9 | // 10 | // Unless required by applicable law or agreed to in writing, software 11 | // distributed under the License is distributed on an "AS IS" BASIS, 12 | // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | // See the License for the specific language governing permissions and 14 | // limitations under the License. 15 | // 16 | // SPDX-License-Identifier: Apache-2.0 17 | 18 | #include "isaac_ros_segment_anything/segment_anything_binarize_tensor.hpp" 19 | 20 | #include 21 | #include 22 | 23 | namespace nvidia 24 | { 25 | namespace isaac_ros 26 | { 27 | namespace segment_anything 28 | { 29 | 30 | namespace 31 | { 32 | 33 | constexpr size_t kBlockSize = 256; 34 | 35 | } 36 | 37 | __global__ void BinarizeTensorKernel(uint8_t * tensor, const size_t size) 38 | { 39 | const int idx = blockIdx.x * blockDim.x + threadIdx.x; 40 | if (idx < size) { 41 | tensor[idx] = (tensor[idx] > 0) ? 255 : 0; 42 | } 43 | } 44 | 45 | void BinarizeTensorOnGPU(uint8_t * tensor, const size_t size, cudaStream_t stream) 46 | { 47 | const int num_blocks = (size + kBlockSize - 1) / kBlockSize; 48 | BinarizeTensorKernel <<< num_blocks, kBlockSize, 0, stream >>> (tensor, size); 49 | } 50 | 51 | // CUDA kernel to find min/max coordinates of non-zero values 52 | __global__ void FindBoundingBoxKernel( 53 | const uint8_t * data, int width, int height, 54 | int * min_x, int * min_y, int * max_x, int * max_y) 55 | { 56 | extern __shared__ int shared_data[]; 57 | // Set up shared memory for reduction 58 | int * s_min_x = &shared_data[0]; 59 | int * s_min_y = &shared_data[blockDim.x]; 60 | int * s_max_x = &shared_data[2 * blockDim.x]; 61 | int * s_max_y = &shared_data[3 * blockDim.x]; 62 | 63 | const int tid = threadIdx.x; 64 | const int gid = blockIdx.x * blockDim.x + threadIdx.x; 65 | 66 | // Initialize shared memory with extreme values 67 | s_min_x[tid] = width; 68 | s_min_y[tid] = height; 69 | s_max_x[tid] = -1; 70 | s_max_y[tid] = -1; 71 | 72 | // Process multiple pixels per thread if needed 73 | for (int i = gid; i < width * height; i += blockDim.x * gridDim.x) { 74 | if (data[i] > 0) { // Non-zero pixel found 75 | const int x = i % width; 76 | const int y = i / width; 77 | 78 | s_min_x[tid] = min(s_min_x[tid], x); 79 | s_min_y[tid] = min(s_min_y[tid], y); 80 | s_max_x[tid] = max(s_max_x[tid], x); 81 | s_max_y[tid] = max(s_max_y[tid], y); 82 | } 83 | } 84 | 85 | __syncthreads(); 86 | 87 | // Perform reduction in shared memory 88 | for (int s = blockDim.x / 2; s > 0; s >>= 1) { 89 | if (tid < s) { 90 | s_min_x[tid] = min(s_min_x[tid], s_min_x[tid + s]); 91 | s_min_y[tid] = min(s_min_y[tid], s_min_y[tid + s]); 92 | s_max_x[tid] = max(s_max_x[tid], s_max_x[tid + s]); 93 | s_max_y[tid] = max(s_max_y[tid], s_max_y[tid + s]); 94 | } 95 | __syncthreads(); 96 | } 97 | 98 | // Write results to global memory 99 | if (tid == 0) { 100 | atomicMin(min_x, s_min_x[0]); 101 | atomicMin(min_y, s_min_y[0]); 102 | atomicMax(max_x, s_max_x[0]); 103 | atomicMax(max_y, s_max_y[0]); 104 | } 105 | } 106 | 107 | void FindBoundingBoxOnGPU( 108 | const uint8_t * input_data, int width, int height, BoundingBox * output_bbox, cudaStream_t stream) 109 | { 110 | // Initialize device memory for results 111 | int * d_min_x; 112 | int * d_min_y; 113 | int * d_max_x; 114 | int * d_max_y; 115 | 116 | // Allocate device memory 117 | cudaMallocAsync(&d_min_x, sizeof(int), stream); 118 | cudaMallocAsync(&d_min_y, sizeof(int), stream); 119 | cudaMallocAsync(&d_max_x, sizeof(int), stream); 120 | cudaMallocAsync(&d_max_y, sizeof(int), stream); 121 | 122 | // Initialize with extreme values 123 | int h_min_x = width; 124 | int h_min_y = height; 125 | int h_max_x = -1; 126 | int h_max_y = -1; 127 | 128 | cudaMemcpyAsync(d_min_x, &h_min_x, sizeof(int), cudaMemcpyHostToDevice, stream); 129 | cudaMemcpyAsync(d_min_y, &h_min_y, sizeof(int), cudaMemcpyHostToDevice, stream); 130 | cudaMemcpyAsync(d_max_x, &h_max_x, sizeof(int), cudaMemcpyHostToDevice, stream); 131 | cudaMemcpyAsync(d_max_y, &h_max_y, sizeof(int), cudaMemcpyHostToDevice, stream); 132 | 133 | // Launch kernel 134 | const int num_blocks = (width * height + kBlockSize - 1) / kBlockSize; 135 | const int shared_mem_size = 4 * kBlockSize * sizeof(int); 136 | 137 | FindBoundingBoxKernel <<< num_blocks, kBlockSize, shared_mem_size, stream >>> ( 138 | input_data, width, height, d_min_x, d_min_y, d_max_x, d_max_y); 139 | 140 | // Copy results back to host 141 | cudaMemcpyAsync(&output_bbox->min_x, d_min_x, sizeof(int), cudaMemcpyDeviceToHost, stream); 142 | cudaMemcpyAsync(&output_bbox->min_y, d_min_y, sizeof(int), cudaMemcpyDeviceToHost, stream); 143 | cudaMemcpyAsync(&output_bbox->max_x, d_max_x, sizeof(int), cudaMemcpyDeviceToHost, stream); 144 | cudaMemcpyAsync(&output_bbox->max_y, d_max_y, sizeof(int), cudaMemcpyDeviceToHost, stream); 145 | 146 | // Free device memory 147 | cudaFreeAsync(d_min_x, stream); 148 | cudaFreeAsync(d_min_y, stream); 149 | cudaFreeAsync(d_max_x, stream); 150 | cudaFreeAsync(d_max_y, stream); 151 | } 152 | 153 | } // namespace segment_anything 154 | } // namespace isaac_ros 155 | } // namespace nvidia 156 | -------------------------------------------------------------------------------- /isaac_ros_segment_anything/scripts/torch_to_onnx.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python3 2 | 3 | # SPDX-FileCopyrightText: NVIDIA CORPORATION & AFFILIATES 4 | # Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. 5 | # 6 | # Licensed under the Apache License, Version 2.0 (the "License"); 7 | # you may not use this file except in compliance with the License. 8 | # You may obtain a copy of the License at 9 | # 10 | # http://www.apache.org/licenses/LICENSE-2.0 11 | # 12 | # Unless required by applicable law or agreed to in writing, software 13 | # distributed under the License is distributed on an "AS IS" BASIS, 14 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 15 | # See the License for the specific language governing permissions and 16 | # limitations under the License. 17 | # 18 | # SPDX-License-Identifier: Apache-2.0 19 | 20 | import argparse 21 | import warnings 22 | 23 | import torch 24 | import torch.nn as nn 25 | 26 | try: 27 | import onnxruntime 28 | onnxruntime_exists = True 29 | except ImportError: 30 | onnxruntime_exists = False 31 | 32 | 33 | class Model(nn.Module): 34 | 35 | def __init__( 36 | self, 37 | sam_onnx 38 | ) -> None: 39 | super().__init__() 40 | self.sam_onnx = sam_onnx 41 | 42 | @torch.no_grad() 43 | def forward( 44 | self, 45 | images: torch.Tensor, 46 | point_coords: torch.Tensor, 47 | point_labels: torch.Tensor, 48 | mask_input: torch.Tensor, 49 | has_mask_input: torch.Tensor, 50 | orig_im_size: torch.Tensor, 51 | ): 52 | image_embeddings = self.sam_onnx.model.image_encoder.forward(images) 53 | return self.sam_onnx.forward( 54 | image_embeddings, point_coords, 55 | point_labels, mask_input, 56 | has_mask_input, orig_im_size 57 | ) 58 | 59 | 60 | def get_parser(): 61 | parser = argparse.ArgumentParser( 62 | description='Export the Image Encoder, SAM prompt encoder and mask decoder \ 63 | to an ONNX model.' 64 | ) 65 | 66 | parser.add_argument( 67 | '--checkpoint', type=str, required=True, help='The path to the SAM model checkpoint.' 68 | ) 69 | 70 | parser.add_argument( 71 | '--output', type=str, required=True, help='The filename to save the ONNX model to.' 72 | ) 73 | 74 | parser.add_argument( 75 | '--model-type', 76 | type=str, 77 | required=True, 78 | help='In ["default", "vit_h", "vit_l", "vit_b", "vit_t"]. Type of SAM model to export.', 79 | ) 80 | 81 | parser.add_argument( 82 | '--sam-type', 83 | type=str, 84 | required=True, 85 | help='In ["SAM", "MobileSAM"]. Which type of SAM model to export.', 86 | ) 87 | 88 | parser.add_argument( 89 | '--opset', 90 | type=int, 91 | default=17, 92 | help='The ONNX opset version to use. Must be >=11', 93 | ) 94 | 95 | return parser 96 | 97 | 98 | def run_export( 99 | model_type: str, 100 | checkpoint: str, 101 | output: str, 102 | opset: int, 103 | sam_type: str 104 | ): 105 | if (sam_type == 'SAM'): 106 | from segment_anything import sam_model_registry 107 | from segment_anything.utils.onnx import SamOnnxModel 108 | else: 109 | from mobile_sam import sam_model_registry 110 | from mobile_sam.utils.onnx import SamOnnxModel 111 | 112 | print('Loading model...') 113 | sam = sam_model_registry[model_type](checkpoint=checkpoint) 114 | 115 | decoder_onnx_model = SamOnnxModel( 116 | model=sam, 117 | return_single_mask=True 118 | ) 119 | 120 | dynamic_axes = { 121 | 'point_coords': {0: 'batch_size', 1: 'num_points'}, 122 | 'point_labels': {0: 'batch_size', 1: 'num_points'}, 123 | } 124 | 125 | embed_size = sam.prompt_encoder.image_embedding_size 126 | mask_input_size = [4 * x for x in embed_size] 127 | 128 | dummy_inputs = { 129 | 'images': torch.randn(1, 3, 1024, 1024), 130 | 'point_coords': torch.randint(low=0, high=1024, size=(1, 5, 2), dtype=torch.float), 131 | 'point_labels': torch.randint(low=0, high=4, size=(1, 5), dtype=torch.float), 132 | 'mask_input': torch.randn(1, 1, *mask_input_size, dtype=torch.float), 133 | 'has_mask_input': torch.tensor([1], dtype=torch.float), 134 | 'orig_im_size': torch.tensor([1500, 2250], dtype=torch.float), 135 | } 136 | 137 | onnx_model = Model(decoder_onnx_model) 138 | _ = onnx_model(**dummy_inputs) 139 | 140 | output_names = ['masks', 'iou_predictions', 'low_res_masks'] 141 | 142 | with warnings.catch_warnings(): 143 | warnings.filterwarnings('ignore', category=torch.jit.TracerWarning) 144 | warnings.filterwarnings('ignore', category=UserWarning) 145 | with open(output, 'wb') as f: 146 | print(f'Exporting onnx model to {output}...') 147 | torch.onnx.export( 148 | onnx_model, 149 | tuple(dummy_inputs.values()), 150 | f, 151 | export_params=True, 152 | verbose=False, 153 | opset_version=opset, 154 | do_constant_folding=True, 155 | input_names=list(dummy_inputs.keys()), 156 | output_names=output_names, 157 | dynamic_axes=dynamic_axes, 158 | dynamo=False # Use legacy exporter for dynamic shapes 159 | ) 160 | 161 | if onnxruntime_exists: 162 | ort_inputs = {k: to_numpy(v) for k, v in dummy_inputs.items()} 163 | # set cpu provider default 164 | providers = ['CPUExecutionProvider'] 165 | ort_session = onnxruntime.InferenceSession(output, providers=providers) 166 | _ = ort_session.run(None, ort_inputs) 167 | print('Model has successfully been run with ONNXRuntime.') 168 | 169 | 170 | def to_numpy(tensor): 171 | return tensor.cpu().numpy() 172 | 173 | 174 | if __name__ == '__main__': 175 | parser = get_parser() 176 | args = parser.parse_args() 177 | run_export( 178 | model_type=args.model_type, 179 | checkpoint=args.checkpoint, 180 | output=args.output, 181 | opset=args.opset, 182 | sam_type=args.sam_type 183 | ) 184 | -------------------------------------------------------------------------------- /isaac_ros_gxf_extensions/gxf_isaac_ros_unet/gxf/image_segmentation/segmentation_mask_colorizer.cu.cpp: -------------------------------------------------------------------------------- 1 | // SPDX-FileCopyrightText: NVIDIA CORPORATION & AFFILIATES 2 | // Copyright (c) 2022-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. 3 | // 4 | // Licensed under the Apache License, Version 2.0 (the "License"); 5 | // you may not use this file except in compliance with the License. 6 | // You may obtain a copy of the License at 7 | // 8 | // http://www.apache.org/licenses/LICENSE-2.0 9 | // 10 | // Unless required by applicable law or agreed to in writing, software 11 | // distributed under the License is distributed on an "AS IS" BASIS, 12 | // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | // See the License for the specific language governing permissions and 14 | // limitations under the License. 15 | // 16 | // SPDX-License-Identifier: Apache-2.0 17 | #include "segmentation_mask_colorizer.cu.hpp" 18 | 19 | #include "cuda.h" 20 | #include "cuda_runtime.h" 21 | 22 | namespace nvidia { 23 | namespace isaac_ros { 24 | 25 | namespace { 26 | 27 | __device__ inline uint32_t GetFlattenedIndex(const uint32_t v, const uint32_t u, const uint32_t c, 28 | const uint32_t width, const uint32_t channels) { 29 | return v * width * channels + u * channels + c; 30 | } 31 | 32 | __device__ inline uint32_t GetMonoIndex(const uint32_t v, const uint32_t u, const uint32_t width) { 33 | return GetFlattenedIndex(v, u, 0, width, 1); 34 | } 35 | 36 | template 37 | __device__ inline uint32_t GetRedIndex(const uint32_t v, const uint32_t u, const uint32_t width); 38 | 39 | template <> 40 | __device__ inline uint32_t GetRedIndex(const uint32_t v, 41 | const uint32_t u, 42 | const uint32_t width) { 43 | return GetFlattenedIndex(v, u, 0, width, 3); 44 | } 45 | 46 | template <> 47 | __device__ inline uint32_t GetRedIndex(const uint32_t v, 48 | const uint32_t u, 49 | const uint32_t width) { 50 | return GetFlattenedIndex(v, u, 2, width, 3); 51 | } 52 | 53 | template 54 | __device__ inline uint32_t GetGreenIndex(const uint32_t v, const uint32_t u, const uint32_t width) { 55 | static_assert(e == ColorImageEncodings::kBGR8 || e == ColorImageEncodings::kRGB8, 56 | "Error: received invalid encoding!"); 57 | return GetFlattenedIndex(v, u, 1, width, 3); 58 | } 59 | 60 | template 61 | __device__ inline uint32_t GetBlueIndex(const uint32_t v, const uint32_t u, const uint32_t width); 62 | 63 | template <> 64 | __device__ inline uint32_t GetBlueIndex(const uint32_t v, 65 | const uint32_t u, 66 | const uint32_t width) { 67 | return GetFlattenedIndex(v, u, 2, width, 3); 68 | } 69 | 70 | template <> 71 | __device__ inline uint32_t GetBlueIndex(const uint32_t v, 72 | const uint32_t u, 73 | const uint32_t width) { 74 | return GetFlattenedIndex(v, u, 0, width, 3); 75 | } 76 | 77 | template 78 | __global__ void ColorizeSegmentationMaskImpl(uint8_t* colored_segmentation_mask, 79 | const uint32_t width, const uint32_t height, 80 | const uint8_t* raw_segmentation_mask, 81 | const int64_t* class_to_color, 82 | const size_t class_to_color_size) { 83 | uint32_t u_idx{blockIdx.x * blockDim.x + threadIdx.x}; 84 | uint32_t u_stride{gridDim.x * blockDim.x}; 85 | 86 | uint32_t v_idx{blockIdx.y * blockDim.y + threadIdx.y}; 87 | uint32_t v_stride{gridDim.y * blockDim.y}; 88 | 89 | for (uint32_t v = v_idx; v < height; v += v_stride) { 90 | for (uint32_t u = u_idx; u < width; u += u_stride) { 91 | const uint8_t class_id{raw_segmentation_mask[GetMonoIndex(v, u, width)]}; 92 | const int64_t selected_color{class_to_color[class_id % class_to_color_size]}; 93 | colored_segmentation_mask[GetRedIndex(v, u, width)] = (selected_color >> 16) & 0xFF; 94 | colored_segmentation_mask[GetGreenIndex(v, u, width)] = (selected_color >> 8) & 0xFF; 95 | colored_segmentation_mask[GetBlueIndex(v, u, width)] = selected_color & 0xFF; 96 | } 97 | } 98 | } 99 | 100 | } // namespace 101 | 102 | void ColorizeSegmentationMask(uint8_t* colored_segmentation_mask, const uint32_t width, 103 | const uint32_t height, const ColorImageEncodings image_encoding, 104 | const uint8_t* raw_segmentation_mask, 105 | const ArrayView& color_palette, cudaStream_t stream) { 106 | dim3 threads_per_block{32, 32, 1}; 107 | dim3 blocks{(width + threads_per_block.x - 1) / threads_per_block.x, 108 | (height + threads_per_block.y - 1) / threads_per_block.y, 1}; 109 | const int64_t* class_to_color = color_palette.data.get(); 110 | const size_t class_to_color_size = color_palette.size; 111 | 112 | switch (image_encoding) { 113 | case ColorImageEncodings::kRGB8: 114 | ColorizeSegmentationMaskImpl 115 | <<>>(colored_segmentation_mask, width, height, 116 | raw_segmentation_mask, class_to_color, 117 | class_to_color_size); 118 | break; 119 | case ColorImageEncodings::kBGR8: 120 | ColorizeSegmentationMaskImpl 121 | <<>>(colored_segmentation_mask, width, height, 122 | raw_segmentation_mask, class_to_color, 123 | class_to_color_size); 124 | break; 125 | } 126 | } 127 | 128 | } // namespace isaac_ros 129 | } // namespace nvidia 130 | -------------------------------------------------------------------------------- /isaac_ros_unet/config/unet_decoder_node.yaml: -------------------------------------------------------------------------------- 1 | %YAML 1.2 2 | # SPDX-FileCopyrightText: NVIDIA CORPORATION & AFFILIATES 3 | # Copyright (c) 2022-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. 4 | # 5 | # Licensed under the Apache License, Version 2.0 (the "License"); 6 | # you may not use this file except in compliance with the License. 7 | # You may obtain a copy of the License at 8 | # 9 | # http://www.apache.org/licenses/LICENSE-2.0 10 | # 11 | # Unless required by applicable law or agreed to in writing, software 12 | # distributed under the License is distributed on an "AS IS" BASIS, 13 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 14 | # See the License for the specific language governing permissions and 15 | # limitations under the License. 16 | # 17 | # SPDX-License-Identifier: Apache-2.0 18 | --- 19 | name: segmentation_postprocessor 20 | components: 21 | - name: input_tensor 22 | type: nvidia::gxf::DoubleBufferReceiver 23 | parameters: 24 | capacity: 12 25 | - type: nvidia::gxf::MessageAvailableSchedulingTerm 26 | parameters: 27 | receiver: input_tensor 28 | min_size: 1 29 | - name: output_video_buffer 30 | type: nvidia::gxf::DoubleBufferTransmitter 31 | parameters: 32 | capacity: 12 33 | - type: nvidia::gxf::DownstreamReceptiveSchedulingTerm 34 | parameters: 35 | transmitter: output_video_buffer 36 | min_size: 1 37 | - name: allocator 38 | type: nvidia::gxf::BlockMemoryPool 39 | parameters: 40 | storage_type: 1 41 | block_size: 522240 42 | num_blocks: 40 43 | - name: stream 44 | type: nvidia::gxf::CudaStreamPool 45 | parameters: 46 | dev_id: 0 47 | stream_flags: 1 48 | stream_priority: 0 49 | reserved_size: 1 50 | max_size: 1 51 | nvtx_identifier: "unet_decoder_node_stream" 52 | - type: nvidia::isaac_ros::SegmentationPostprocessor 53 | parameters: 54 | in: input_tensor 55 | out: output_video_buffer 56 | allocator: allocator 57 | network_output_type: softmax 58 | data_format: NHWC 59 | cuda_stream_pool: stream 60 | - type: nvidia::gxf::MemoryAvailableSchedulingTerm 61 | parameters: 62 | allocator: allocator 63 | min_blocks: 1 64 | --- 65 | name: segmentation_broadcaster 66 | components: 67 | - name: input_video_buffer 68 | type: nvidia::gxf::DoubleBufferReceiver 69 | parameters: 70 | capacity: 1 71 | - type: nvidia::gxf::MessageAvailableSchedulingTerm 72 | parameters: 73 | receiver: input_video_buffer 74 | min_size: 1 75 | - type: nvidia::gxf::Broadcast 76 | parameters: 77 | source: input_video_buffer 78 | - name: output_video_buffer 79 | type: nvidia::gxf::DoubleBufferTransmitter 80 | parameters: 81 | capacity: 1 82 | - type: nvidia::gxf::DownstreamReceptiveSchedulingTerm 83 | parameters: 84 | transmitter: output_video_buffer 85 | min_size: 1 86 | - name: output_raw_segmentation_mask 87 | type: nvidia::gxf::DoubleBufferTransmitter 88 | parameters: 89 | capacity: 1 90 | - type: nvidia::gxf::DownstreamReceptiveSchedulingTerm 91 | parameters: 92 | transmitter: output_raw_segmentation_mask 93 | min_size: 1 94 | --- 95 | name: segmentation_mask_generator 96 | components: 97 | - name: input_video_buffer 98 | type: nvidia::gxf::DoubleBufferReceiver 99 | parameters: 100 | capacity: 12 101 | - type: nvidia::gxf::MessageAvailableSchedulingTerm 102 | parameters: 103 | receiver: input_video_buffer 104 | min_size: 1 105 | - name: output_colored_segmentation_mask 106 | type: nvidia::gxf::DoubleBufferTransmitter 107 | parameters: 108 | capacity: 12 109 | - type: nvidia::gxf::DownstreamReceptiveSchedulingTerm 110 | parameters: 111 | transmitter: output_colored_segmentation_mask 112 | min_size: 1 113 | - name: allocator 114 | type: nvidia::gxf::BlockMemoryPool 115 | parameters: 116 | storage_type: 1 117 | block_size: 1566720 118 | num_blocks: 40 119 | - name: stream_pool 120 | type: nvidia::gxf::CudaStreamPool 121 | parameters: 122 | dev_id: 0 123 | stream_flags: 1 124 | stream_priority: 0 125 | reserved_size: 1 126 | max_size: 1 127 | nvtx_identifier: "segmentation_mask_generator_stream_pool" 128 | - type: nvidia::isaac_ros::SegmentationMaskColorizer 129 | parameters: 130 | raw_segmentation_mask_input: input_video_buffer 131 | colored_segmentation_mask_output: output_colored_segmentation_mask 132 | allocator: allocator 133 | color_palette: [] 134 | color_segmentation_mask_encoding: "rgb8" 135 | stream_pool: stream_pool 136 | --- 137 | name: raw_segmentation_mask_sink 138 | components: 139 | - name: input_raw_segmentation_mask 140 | type: nvidia::gxf::DoubleBufferReceiver 141 | parameters: 142 | capacity: 2 143 | - type: nvidia::gxf::MessageAvailableSchedulingTerm 144 | parameters: 145 | receiver: input_raw_segmentation_mask 146 | min_size: 1 147 | - name: sink 148 | type: nvidia::isaac_ros::MessageRelay 149 | parameters: 150 | source: input_raw_segmentation_mask 151 | --- 152 | name: colored_segmentation_mask_sink 153 | components: 154 | - name: input_colored_segmentation_mask 155 | type: nvidia::gxf::DoubleBufferReceiver 156 | parameters: 157 | capacity: 2 158 | - type: nvidia::gxf::MessageAvailableSchedulingTerm 159 | parameters: 160 | receiver: input_colored_segmentation_mask 161 | min_size: 1 162 | - name: sink 163 | type: nvidia::isaac_ros::MessageRelay 164 | parameters: 165 | source: input_colored_segmentation_mask 166 | --- 167 | components: 168 | - type: nvidia::gxf::Connection 169 | parameters: 170 | source: segmentation_postprocessor/output_video_buffer 171 | target: segmentation_broadcaster/input_video_buffer 172 | - type: nvidia::gxf::Connection 173 | parameters: 174 | source: segmentation_broadcaster/output_raw_segmentation_mask 175 | target: raw_segmentation_mask_sink/input_raw_segmentation_mask 176 | - type: nvidia::gxf::Connection 177 | parameters: 178 | source: segmentation_broadcaster/output_video_buffer 179 | target: segmentation_mask_generator/input_video_buffer 180 | - type: nvidia::gxf::Connection 181 | parameters: 182 | source: segmentation_mask_generator/output_colored_segmentation_mask 183 | target: colored_segmentation_mask_sink/input_colored_segmentation_mask 184 | --- 185 | components: 186 | - name: clock 187 | type: nvidia::gxf::RealtimeClock 188 | - type: nvidia::gxf::EventBasedScheduler 189 | parameters: 190 | clock: clock 191 | stop_on_deadlock: false 192 | worker_thread_number: 2 193 | worker_thread_name_id: "unet_decoder_node" 194 | - type: nvidia::gxf::JobStatistics 195 | parameters: 196 | clock: clock 197 | -------------------------------------------------------------------------------- /isaac_ros_segment_anything/config/segment_anything_data_encoder_node.yaml: -------------------------------------------------------------------------------- 1 | %YAML 1.2 2 | # SPDX-FileCopyrightText: NVIDIA CORPORATION & AFFILIATES 3 | # Copyright (c) 2024-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. 4 | # 5 | # Licensed under the Apache License, Version 2.0 (the "License"); 6 | # you may not use this file except in compliance with the License. 7 | # You may obtain a copy of the License at 8 | # 9 | # http://www.apache.org/licenses/LICENSE-2.0 10 | # 11 | # Unless required by applicable law or agreed to in writing, software 12 | # distributed under the License is distributed on an "AS IS" BASIS, 13 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 14 | # See the License for the specific language governing permissions and 15 | # limitations under the License. 16 | # 17 | # SPDX-License-Identifier: Apache-2.0 18 | --- 19 | name: sync 20 | components: 21 | - name: image_receiver 22 | type: nvidia::gxf::DoubleBufferReceiver 23 | parameters: 24 | capacity: 2 25 | - type: nvidia::gxf::MessageAvailableSchedulingTerm 26 | parameters: 27 | receiver: image_receiver 28 | min_size: 1 29 | - name: prompt_receiver 30 | type: nvidia::gxf::DoubleBufferReceiver 31 | parameters: 32 | capacity: 2 33 | - type: nvidia::gxf::MessageAvailableSchedulingTerm 34 | parameters: 35 | receiver: prompt_receiver 36 | min_size: 1 37 | - name: mask_receiver 38 | type: nvidia::gxf::DoubleBufferReceiver 39 | parameters: 40 | capacity: 1 41 | - type: nvidia::gxf::MessageAvailableSchedulingTerm 42 | parameters: 43 | receiver: mask_receiver 44 | min_size: 1 45 | - name: image_transmitter 46 | type: nvidia::gxf::DoubleBufferTransmitter 47 | parameters: 48 | capacity: 1 49 | - type: nvidia::gxf::DownstreamReceptiveSchedulingTerm 50 | parameters: 51 | transmitter: image_transmitter 52 | min_size: 1 53 | - name: prompt_transmitter 54 | type: nvidia::gxf::DoubleBufferTransmitter 55 | parameters: 56 | capacity: 1 57 | - type: nvidia::gxf::DownstreamReceptiveSchedulingTerm 58 | parameters: 59 | transmitter: prompt_transmitter 60 | min_size: 1 61 | - name: mask_transmitter 62 | type: nvidia::gxf::DoubleBufferTransmitter 63 | parameters: 64 | capacity: 1 65 | - type: nvidia::gxf::DownstreamReceptiveSchedulingTerm 66 | parameters: 67 | transmitter: mask_transmitter 68 | min_size: 1 69 | - type: nvidia::gxf::Synchronization 70 | parameters: 71 | inputs: [image_receiver, prompt_receiver, mask_receiver] 72 | outputs: [image_transmitter, prompt_transmitter, mask_transmitter] 73 | --- 74 | name: prompt_processor 75 | components: 76 | - name: input_prompt 77 | type: nvidia::gxf::DoubleBufferReceiver 78 | parameters: 79 | capacity: 1 80 | - type: nvidia::gxf::MessageAvailableSchedulingTerm 81 | parameters: 82 | receiver: input_prompt 83 | min_size: 1 84 | - name: output_buffer 85 | type: nvidia::gxf::DoubleBufferTransmitter 86 | parameters: 87 | capacity: 1 88 | - type: nvidia::gxf::DownstreamReceptiveSchedulingTerm 89 | parameters: 90 | transmitter: output_buffer 91 | min_size: 1 92 | - name: allocator 93 | type: nvidia::gxf::BlockMemoryPool 94 | parameters: 95 | storage_type: 1 96 | block_size: 2048 97 | num_blocks: 200 98 | - name: stream_pool 99 | type: nvidia::gxf::CudaStreamPool 100 | parameters: 101 | dev_id: 0 102 | stream_flags: 1 103 | stream_priority: 0 104 | reserved_size: 1 105 | max_size: 1 106 | nvtx_identifier: "segment_anything_prompt_processor_stream_pool" 107 | - type: nvidia::isaac_ros::SegmentAnythingPromptProcessor 108 | parameters: 109 | in: input_prompt 110 | out_points: output_buffer 111 | allocator: allocator 112 | stream_pool: stream_pool 113 | - type: nvidia::gxf::MemoryAvailableSchedulingTerm 114 | parameters: 115 | allocator: allocator 116 | min_blocks: 1 117 | --- 118 | name: message_compositor 119 | components: 120 | - name: input_prompt 121 | type: nvidia::gxf::DoubleBufferReceiver 122 | parameters: 123 | capacity: 1 124 | - type: nvidia::gxf::MessageAvailableSchedulingTerm 125 | parameters: 126 | receiver: input_prompt 127 | min_size: 1 128 | - name: input_image 129 | type: nvidia::gxf::DoubleBufferReceiver 130 | parameters: 131 | capacity: 1 132 | - type: nvidia::gxf::MessageAvailableSchedulingTerm 133 | parameters: 134 | receiver: input_image 135 | min_size: 1 136 | - name: input_mask 137 | type: nvidia::gxf::DoubleBufferReceiver 138 | parameters: 139 | capacity: 1 140 | - type: nvidia::gxf::MessageAvailableSchedulingTerm 141 | parameters: 142 | receiver: input_mask 143 | min_size: 1 144 | - name: output_buffer 145 | type: nvidia::gxf::DoubleBufferTransmitter 146 | parameters: 147 | capacity: 1 148 | - type: nvidia::gxf::DownstreamReceptiveSchedulingTerm 149 | parameters: 150 | transmitter: output_buffer 151 | min_size: 1 152 | - type: nvidia::isaac_ros::SegmentAnythingMsgCompositor 153 | parameters: 154 | inputs: [input_prompt, input_image, input_mask] 155 | output: output_buffer 156 | --- 157 | name: sink 158 | components: 159 | - name: input 160 | type: nvidia::gxf::DoubleBufferReceiver 161 | parameters: 162 | capacity: 2 163 | - type: nvidia::gxf::MessageAvailableSchedulingTerm 164 | parameters: 165 | receiver: input 166 | min_size: 1 167 | - name: output 168 | type: nvidia::isaac_ros::MessageRelay 169 | parameters: 170 | source: input 171 | --- 172 | components: 173 | - type: nvidia::gxf::Connection 174 | parameters: 175 | source: sync/prompt_transmitter 176 | target: prompt_processor/input_prompt 177 | --- 178 | components: 179 | - type: nvidia::gxf::Connection 180 | parameters: 181 | source: prompt_processor/output_buffer 182 | target: message_compositor/input_prompt 183 | --- 184 | components: 185 | - type: nvidia::gxf::Connection 186 | parameters: 187 | source: sync/image_transmitter 188 | target: message_compositor/input_image 189 | --- 190 | components: 191 | - type: nvidia::gxf::Connection 192 | parameters: 193 | source: sync/mask_transmitter 194 | target: message_compositor/input_mask 195 | 196 | --- 197 | components: 198 | - type: nvidia::gxf::Connection 199 | parameters: 200 | source: message_compositor/output_buffer 201 | target: sink/input 202 | --- 203 | components: 204 | - name: clock 205 | type: nvidia::gxf::RealtimeClock 206 | - type: nvidia::gxf::MultiThreadScheduler 207 | parameters: 208 | clock: clock 209 | stop_on_deadlock: false 210 | check_recession_period_ms: 1 211 | worker_thread_number: 2 212 | worker_thread_name_id: "segment_anything_data_encoder_node" 213 | - type: nvidia::gxf::JobStatistics 214 | parameters: 215 | clock: clock 216 | -------------------------------------------------------------------------------- /isaac_ros_segment_anything/CMakeLists.txt: -------------------------------------------------------------------------------- 1 | # SPDX-FileCopyrightText: NVIDIA CORPORATION & AFFILIATES 2 | # Copyright (c) 2024-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. 3 | # 4 | # Licensed under the Apache License, Version 2.0 (the "License"); 5 | # you may not use this file except in compliance with the License. 6 | # You may obtain a copy of the License at 7 | # 8 | # http://www.apache.org/licenses/LICENSE-2.0 9 | # 10 | # Unless required by applicable law or agreed to in writing, software 11 | # distributed under the License is distributed on an "AS IS" BASIS, 12 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | # See the License for the specific language governing permissions and 14 | # limitations under the License. 15 | # 16 | # SPDX-License-Identifier: Apache-2.0 17 | 18 | cmake_minimum_required(VERSION 3.22.1) 19 | project(isaac_ros_segment_anything LANGUAGES C CXX) 20 | 21 | if(CMAKE_COMPILER_IS_GNUCXX OR CMAKE_CXX_COMPILER_ID MATCHES "Clang") 22 | add_compile_options(-Wall -Wextra -Wpedantic) 23 | endif() 24 | 25 | find_package(ament_cmake_auto REQUIRED) 26 | ament_auto_find_build_dependencies() 27 | enable_language(CUDA) 28 | # Decoder node 29 | 30 | ament_auto_add_library(segment_anything_decoder_node SHARED src/segment_anything_decoder_node.cpp) 31 | rclcpp_components_register_nodes(segment_anything_decoder_node "nvidia::isaac_ros::segment_anything::SegmentAnythingDecoderNode") 32 | set(node_plugins 33 | "${node_plugins}nvidia::isaac_ros::segment_anything::SegmentAnythingDecoderNode;$\n") 34 | set_target_properties(segment_anything_decoder_node PROPERTIES 35 | BUILD_WITH_INSTALL_RPATH TRUE 36 | BUILD_RPATH_USE_ORIGIN TRUE 37 | INSTALL_RPATH_USE_LINK_PATH TRUE) 38 | 39 | ament_auto_add_library(segment_anything_data_encoder_node SHARED src/segment_anything_data_encoder_node.cpp) 40 | rclcpp_components_register_nodes(segment_anything_data_encoder_node "nvidia::isaac_ros::segment_anything::SegmentAnythingDataEncoderNode") 41 | set(node_plugins 42 | "${node_plugins}nvidia::isaac_ros::segment_anything::SegmentAnythingDataEncoderNode;$\n") 43 | set_target_properties(segment_anything_data_encoder_node PROPERTIES 44 | BUILD_WITH_INSTALL_RPATH TRUE 45 | BUILD_RPATH_USE_ORIGIN TRUE 46 | INSTALL_RPATH_USE_LINK_PATH TRUE) 47 | 48 | ament_auto_add_library(segment_anything_dummy_mask_publisher_node SHARED src/segment_anything_dummy_mask_publisher_node.cpp) 49 | rclcpp_components_register_nodes(segment_anything_dummy_mask_publisher_node "nvidia::isaac_ros::segment_anything::DummyMaskPublisher") 50 | set(node_plugins 51 | "${node_plugins}nvidia::isaac_ros::segment_anything::DummyMaskPublisher;$\n") 52 | set_target_properties(segment_anything_dummy_mask_publisher_node PROPERTIES 53 | BUILD_WITH_INSTALL_RPATH TRUE 54 | BUILD_RPATH_USE_ORIGIN TRUE 55 | INSTALL_RPATH_USE_LINK_PATH TRUE) 56 | 57 | ament_auto_add_library(segment_anything_point_triggered_node SHARED src/segment_anything_point_triggered_node.cpp) 58 | rclcpp_components_register_nodes( 59 | segment_anything_point_triggered_node "nvidia::isaac_ros::segment_anything::SegmentAnythingPointTriggeredNode") 60 | set(node_plugins 61 | "${node_plugins}nvidia::isaac_ros::segment_anything::SegmentAnythingPointTriggeredNode;$\n") 62 | set_target_properties(segment_anything_point_triggered_node PROPERTIES 63 | BUILD_WITH_INSTALL_RPATH TRUE 64 | BUILD_RPATH_USE_ORIGIN TRUE 65 | INSTALL_RPATH_USE_LINK_PATH TRUE) 66 | 67 | # Segment Anything Tensor to Image 68 | ament_auto_add_library(segment_anything_tensor_to_image_node SHARED 69 | src/segment_anything_tensor_to_image_node.cpp 70 | src/segment_anything_binarize_tensor.cu) 71 | rclcpp_components_register_nodes(segment_anything_tensor_to_image_node 72 | "nvidia::isaac_ros::segment_anything::TensorToImageNode") 73 | set(node_plugins 74 | "${node_plugins}nvidia::isaac_ros::segment_anything::TensorToImageNode;$\n") 75 | set_target_properties(segment_anything_tensor_to_image_node PROPERTIES 76 | BUILD_WITH_INSTALL_RPATH TRUE 77 | BUILD_RPATH_USE_ORIGIN TRUE 78 | INSTALL_RPATH_USE_LINK_PATH TRUE) 79 | # Mark as CUDA files with non-standard extensions 80 | set_source_files_properties( 81 | src/segment_anything_binarize_tensor.cu 82 | include/isaac_ros_segment_anything/segment_anything_binarize_tensor.hpp 83 | PROPERTIES LANGUAGE CUDA 84 | ) 85 | 86 | ament_python_install_package(${PROJECT_NAME}) 87 | install(PROGRAMS scripts/colored_mask_converter_node.py DESTINATION lib/${PROJECT_NAME}) 88 | install(PROGRAMS scripts/torch_to_onnx.py DESTINATION lib/${PROJECT_NAME}) 89 | install(PROGRAMS scripts/visualize_mask.py DESTINATION lib/${PROJECT_NAME}) 90 | 91 | if(BUILD_TESTING) 92 | find_package(ament_lint_auto REQUIRED) 93 | find_package(ament_cmake_gtest REQUIRED) 94 | ament_lint_auto_find_test_dependencies() 95 | 96 | # Force use of ROS2 vendor googletest to avoid version conflicts 97 | set(CMAKE_IGNORE_PATH "/usr/src/googletest" ${CMAKE_IGNORE_PATH}) 98 | set(CMAKE_IGNORE_PATH "/usr/include/gtest" ${CMAKE_IGNORE_PATH}) 99 | set(CMAKE_IGNORE_PATH "/usr/include/gmock" ${CMAKE_IGNORE_PATH}) 100 | 101 | # Explicitly use ROS2 vendor googletest/gmock 102 | include_directories(BEFORE /opt/ros/$ENV{ROS_DISTRO}/src/gtest_vendor/include) 103 | include_directories(BEFORE /opt/ros/$ENV{ROS_DISTRO}/src/gmock_vendor/include) 104 | 105 | # Gtests for segment anything data encoder node 106 | ament_add_gtest(segment_anything_data_encoder_node_test test/segment_anything_data_encoder_node_test.cpp) 107 | target_link_libraries(segment_anything_data_encoder_node_test segment_anything_data_encoder_node) 108 | target_include_directories(segment_anything_data_encoder_node_test PUBLIC include/isaac_ros_segment_anything/) 109 | ament_target_dependencies(segment_anything_data_encoder_node_test rclcpp) 110 | ament_target_dependencies(segment_anything_data_encoder_node_test isaac_ros_nitros) 111 | 112 | # The FindPythonInterp and FindPythonLibs modules are removed 113 | if(POLICY CMP0148) 114 | cmake_policy(SET CMP0148 OLD) 115 | endif() 116 | 117 | find_package(launch_testing_ament_cmake REQUIRED) 118 | add_launch_test(test/isaac_ros_segment_anything_onnx_pol.py TIMEOUT "300") 119 | add_launch_test(test/isaac_ros_segment_anything_point_prompt_pol.py TIMEOUT "300") 120 | add_launch_test(test/isaac_ros_segment_anything_tensor_to_image_test.py) 121 | endif() 122 | 123 | 124 | # Embed versioning information into installed files 125 | ament_index_get_resource(ISAAC_ROS_COMMON_CMAKE_PATH isaac_ros_common_cmake_path isaac_ros_common) 126 | include("${ISAAC_ROS_COMMON_CMAKE_PATH}/isaac_ros_common-version-info.cmake") 127 | generate_version_info(${PROJECT_NAME}) 128 | 129 | ament_auto_package(INSTALL_TO_SHARE config launch) 130 | -------------------------------------------------------------------------------- /isaac_ros_gxf_extensions/gxf_isaac_ros_unet/gxf/image_segmentation/segmentation_postprocessor.cu.cpp: -------------------------------------------------------------------------------- 1 | // SPDX-FileCopyrightText: NVIDIA CORPORATION & AFFILIATES 2 | // Copyright (c) 2022-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. 3 | // 4 | // Licensed under the Apache License, Version 2.0 (the "License"); 5 | // you may not use this file except in compliance with the License. 6 | // You may obtain a copy of the License at 7 | // 8 | // http://www.apache.org/licenses/LICENSE-2.0 9 | // 10 | // Unless required by applicable law or agreed to in writing, software 11 | // distributed under the License is distributed on an "AS IS" BASIS, 12 | // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | // See the License for the specific language governing permissions and 14 | // limitations under the License. 15 | // 16 | // SPDX-License-Identifier: Apache-2.0 17 | #include "segmentation_postprocessor.cu.hpp" 18 | 19 | namespace nvidia { 20 | namespace isaac_ros { 21 | 22 | __forceinline__ __device__ uint32_t hwc_to_index(Shape shape, uint32_t y, uint32_t x, uint32_t c) { 23 | return (y * shape.width + x) * shape.channels + c; 24 | } 25 | 26 | __forceinline__ __device__ uint32_t nchw_to_index(Shape shape, uint32_t y, uint32_t x, uint32_t c) { 27 | return (c * shape.height + y) * shape.width + x; 28 | } 29 | 30 | template 31 | __forceinline__ __device__ uint32_t data_format_to_index(Shape shape, uint32_t y, uint32_t x, 32 | uint32_t c) {} 33 | 34 | template <> 35 | __forceinline__ __device__ uint32_t data_format_to_index(Shape shape, uint32_t y, 36 | uint32_t x, uint32_t c) { 37 | return hwc_to_index(shape, y, x, c); 38 | } 39 | 40 | template <> 41 | __forceinline__ __device__ uint32_t data_format_to_index(Shape shape, uint32_t y, 42 | uint32_t x, 43 | uint32_t c) { 44 | return hwc_to_index(shape, y, x, c); 45 | } 46 | 47 | template <> 48 | __forceinline__ __device__ uint32_t data_format_to_index(Shape shape, uint32_t y, 49 | uint32_t x, 50 | uint32_t c) { 51 | return nchw_to_index(shape, y, x, c); 52 | } 53 | 54 | __forceinline__ __device__ uint32_t hw1_to_index(Shape shape, uint32_t y, uint32_t x) { 55 | return y * shape.width + x; 56 | } 57 | 58 | template 59 | __global__ void postprocessing_kernel(Shape shape, const float* input, output_type_t* output, 60 | NetworkOutputType network_output_type) { 61 | const uint32_t x = blockIdx.x * blockDim.x + threadIdx.x; 62 | const uint32_t y = blockIdx.y * blockDim.y + threadIdx.y; 63 | 64 | if ((x >= shape.width) || (y >= shape.height)) { return; } 65 | 66 | float max_value = 0.0f; 67 | uint8_t max_index = 0; 68 | 69 | switch (network_output_type) { 70 | case NetworkOutputType::kSigmoid: { 71 | const float value = input[data_format_to_index(shape, y, x, 0)]; 72 | max_index = value >= 0.5f ? 1 : 0; 73 | } break; 74 | case NetworkOutputType::kSoftmax: { 75 | for (uint32_t c = 0; c < shape.channels; c++) { 76 | const float value = input[data_format_to_index(shape, y, x, c)]; 77 | if (value > max_value) { 78 | max_value = value; 79 | max_index = c; 80 | } 81 | } 82 | } break; 83 | } 84 | 85 | output[hw1_to_index(shape, y, x)] = max_index; 86 | } 87 | 88 | template 89 | __global__ void CopyTensorDataCuda(Shape shape, const T* input, output_type_t* output, 90 | NetworkOutputType network_output_type) { 91 | const uint32_t x = blockIdx.x * blockDim.x + threadIdx.x; 92 | const uint32_t y = blockIdx.y * blockDim.y + threadIdx.y; 93 | 94 | if ((x >= shape.width) || (y >= shape.height)) { return; } 95 | 96 | const int32_t class_id{input[data_format_to_index(shape, y, x, 0)]}; 97 | output[hw1_to_index(shape, y, x)] = static_cast(class_id); 98 | } 99 | 100 | uint16_t ceil_div(uint16_t numerator, uint16_t denominator) { 101 | uint32_t accumulator = numerator + denominator - 1; 102 | return accumulator / denominator; 103 | } 104 | 105 | void cuda_postprocess(NetworkOutputType network_output_type, DataFormat data_format, Shape shape, 106 | const float* input, output_type_t* output, cudaStream_t stream) { 107 | dim3 block(32, 32, 1); 108 | dim3 grid(ceil_div(shape.width, block.x), ceil_div(shape.height, block.y), 1); 109 | switch (data_format) { 110 | case DataFormat::kNCHW: 111 | postprocessing_kernel 112 | <<>>(shape, input, output, network_output_type); 113 | break; 114 | case DataFormat::kHWC: 115 | postprocessing_kernel 116 | <<>>(shape, input, output, network_output_type); 117 | break; 118 | case DataFormat::kNHWC: 119 | postprocessing_kernel 120 | <<>>(shape, input, output, network_output_type); 121 | break; 122 | } 123 | } 124 | 125 | template 126 | void CopyTensorData(NetworkOutputType network_output_type, DataFormat data_format, Shape shape, 127 | const T* input, output_type_t* output, cudaStream_t stream) { 128 | if (network_output_type != NetworkOutputType::kArgmax) { return; } 129 | dim3 block(32, 32, 1); 130 | dim3 grid(ceil_div(shape.width, block.x), ceil_div(shape.height, block.y), 1); 131 | switch (data_format) { 132 | case DataFormat::kNCHW: 133 | CopyTensorDataCuda 134 | <<>>(shape, input, output, network_output_type); 135 | break; 136 | case DataFormat::kHWC: 137 | CopyTensorDataCuda 138 | <<>>(shape, input, output, network_output_type); 139 | break; 140 | case DataFormat::kNHWC: 141 | CopyTensorDataCuda 142 | <<>>(shape, input, output, network_output_type); 143 | break; 144 | } 145 | } 146 | template void CopyTensorData(NetworkOutputType network_output_type, DataFormat data_format, Shape shape, 147 | const int32_t* input, output_type_t* output, cudaStream_t stream); 148 | template void CopyTensorData(NetworkOutputType network_output_type, DataFormat data_format, Shape shape, 149 | const int64_t* input, output_type_t* output, cudaStream_t stream); 150 | 151 | 152 | } // namespace isaac_ros 153 | } // namespace nvidia 154 | -------------------------------------------------------------------------------- /isaac_ros_segment_anything/src/segment_anything_data_encoder_node.cpp: -------------------------------------------------------------------------------- 1 | // SPDX-FileCopyrightText: NVIDIA CORPORATION & AFFILIATES 2 | // Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. 3 | // 4 | // Licensed under the Apache License, Version 2.0 (the "License"); 5 | // you may not use this file except in compliance with the License. 6 | // You may obtain a copy of the License at 7 | // 8 | // http://www.apache.org/licenses/LICENSE-2.0 9 | // 10 | // Unless required by applicable law or agreed to in writing, software 11 | // distributed under the License is distributed on an "AS IS" BASIS, 12 | // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | // See the License for the specific language governing permissions and 14 | // limitations under the License. 15 | // 16 | // SPDX-License-Identifier: Apache-2.0 17 | 18 | #include "isaac_ros_segment_anything/segment_anything_data_encoder_node.hpp" 19 | 20 | #include 21 | #include 22 | #include 23 | #include 24 | #include 25 | 26 | #include "rclcpp/rclcpp.hpp" 27 | #include "sensor_msgs/image_encodings.hpp" 28 | #include "isaac_ros_nitros_image_type/nitros_image.hpp" 29 | #include "isaac_ros_nitros_tensor_list_type/nitros_tensor_list.hpp" 30 | #include "isaac_ros_nitros_detection2_d_array_type/nitros_detection2_d_array.hpp" 31 | 32 | namespace nvidia 33 | { 34 | namespace isaac_ros 35 | { 36 | namespace segment_anything 37 | { 38 | namespace 39 | { 40 | 41 | using nvidia::gxf::optimizer::GraphIOGroupSupportedDataTypesInfoList; 42 | 43 | constexpr char INPUT_COMPONENT_KEY[] = "sync/prompt_receiver"; 44 | constexpr char INPUT_DEFAULT_TENSOR_FORMAT[] = "nitros_detection2_d_array"; 45 | constexpr char INPUT_TOPIC_NAME[] = "prompts"; 46 | 47 | constexpr char INPUT_IMAGE_COMPONENT_KEY[] = "sync/image_receiver"; 48 | constexpr char INPUT_IMAGE_DEFAULT_TENSOR_FORMAT[] = "nitros_tensor_list_nchw_rgb_f32"; 49 | constexpr char INPUT_IMAGE_TOPIC_NAME[] = "tensor_pub"; 50 | 51 | constexpr char INPUT_MASK_COMPONENT_KEY[] = "sync/mask_receiver"; 52 | constexpr char INPUT_MASK_DEFAULT_TENSOR_FORMAT[] = "nitros_tensor_list_nchw_rgb_f32"; 53 | constexpr char INPUT_MASK_TOPIC_NAME[] = "mask"; 54 | 55 | constexpr char RAW_OUTPUT_COMPONENT_KEY[] = "sink/output"; 56 | constexpr char RAW_OUTPUT_DEFAULT_TENSOR_FORMAT[] = "nitros_tensor_list_nchw"; 57 | constexpr char RAW_OUTPUT_TOPIC_NAME[] = "tensor"; 58 | 59 | constexpr char APP_YAML_FILENAME[] = "config/segment_anything_data_encoder_node.yaml"; 60 | constexpr char PACKAGE_NAME[] = "isaac_ros_segment_anything"; 61 | 62 | const std::vector> EXTENSIONS = { 63 | {"isaac_ros_gxf", "gxf/lib/std/libgxf_std.so"}, 64 | {"isaac_ros_gxf", "gxf/lib/cuda/libgxf_cuda.so"}, 65 | {"gxf_isaac_ros_segment_anything", "gxf/lib/libgxf_isaac_ros_segment_anything.so"}}; 66 | 67 | const std::vector PRESET_EXTENSION_SPEC_NAMES = {}; 68 | const std::vector EXTENSION_SPEC_FILENAMES = { 69 | "config/segment_anything_spec_file.yaml" 70 | }; 71 | const std::vector GENERATOR_RULE_FILENAMES = {}; 72 | 73 | #pragma GCC diagnostic push 74 | #pragma GCC diagnostic ignored "-Wpedantic" 75 | const nitros::NitrosPublisherSubscriberConfigMap CONFIG_MAP = { 76 | {INPUT_COMPONENT_KEY, 77 | { 78 | .type = nitros::NitrosPublisherSubscriberType::NEGOTIATED, 79 | .qos = rclcpp::QoS(1), 80 | .compatible_data_format = INPUT_DEFAULT_TENSOR_FORMAT, 81 | .topic_name = INPUT_TOPIC_NAME, 82 | }}, 83 | {RAW_OUTPUT_COMPONENT_KEY, 84 | { 85 | .type = nitros::NitrosPublisherSubscriberType::NEGOTIATED, 86 | .qos = rclcpp::QoS(1), 87 | .compatible_data_format = RAW_OUTPUT_DEFAULT_TENSOR_FORMAT, 88 | .topic_name = RAW_OUTPUT_TOPIC_NAME, 89 | .frame_id_source_key = INPUT_COMPONENT_KEY, 90 | }}, 91 | {INPUT_IMAGE_COMPONENT_KEY, 92 | { 93 | .type = nitros::NitrosPublisherSubscriberType::NEGOTIATED, 94 | .qos = rclcpp::QoS(1), 95 | .compatible_data_format = INPUT_IMAGE_DEFAULT_TENSOR_FORMAT, 96 | .topic_name = INPUT_IMAGE_TOPIC_NAME, 97 | }}, 98 | {INPUT_MASK_COMPONENT_KEY, 99 | { 100 | .type = nitros::NitrosPublisherSubscriberType::NEGOTIATED, 101 | .qos = rclcpp::QoS(1), 102 | .compatible_data_format = INPUT_MASK_DEFAULT_TENSOR_FORMAT, 103 | .topic_name = INPUT_MASK_TOPIC_NAME, 104 | }} 105 | }; 106 | #pragma GCC diagnostic pop 107 | } // namespace 108 | 109 | bool IsSupportedInputPromptType(const std::string & network_output_type) 110 | { 111 | return network_output_type == std::string{"bbox"} || 112 | network_output_type == std::string{"point"}; 113 | } 114 | 115 | SegmentAnythingDataEncoderNode::SegmentAnythingDataEncoderNode(const rclcpp::NodeOptions options) 116 | : nitros::NitrosNode( 117 | options, 118 | APP_YAML_FILENAME, 119 | CONFIG_MAP, 120 | PRESET_EXTENSION_SPEC_NAMES, 121 | EXTENSION_SPEC_FILENAMES, 122 | GENERATOR_RULE_FILENAMES, 123 | EXTENSIONS, 124 | PACKAGE_NAME), 125 | max_batch_size_(declare_parameter("max_batch_size", 20)), 126 | prompt_input_type_(declare_parameter("prompt_input_type", "bbox")), 127 | has_input_mask_(declare_parameter("has_input_mask", false)), 128 | orig_img_dims_(declare_parameter>("orig_img_dims", {632, 1200})) 129 | { 130 | if (!IsSupportedInputPromptType(prompt_input_type_)) { 131 | RCLCPP_ERROR( 132 | get_logger(), 133 | "Received invalid input prompt type: %s!", 134 | prompt_input_type_.c_str()); 135 | throw std::invalid_argument("Received invalid input prompt type." + prompt_input_type_); 136 | } 137 | 138 | registerSupportedType(); 139 | registerSupportedType(); 140 | startNitrosNode(); 141 | } 142 | 143 | void SegmentAnythingDataEncoderNode::postLoadGraphCallback() 144 | { 145 | getNitrosContext().setParameterStr( 146 | "prompt_processor", 147 | "nvidia::isaac_ros::SegmentAnythingPromptProcessor", 148 | "prompt_type_name", prompt_input_type_ 149 | ); 150 | 151 | getNitrosContext().setParameterInt32( 152 | "prompt_processor", 153 | "nvidia::isaac_ros::SegmentAnythingPromptProcessor", 154 | "max_batch_size", max_batch_size_ 155 | ); 156 | 157 | getNitrosContext().setParameterBool( 158 | "prompt_processor", 159 | "nvidia::isaac_ros::SegmentAnythingPromptProcessor", 160 | "has_input_mask", has_input_mask_ 161 | ); 162 | 163 | std::vector orig_img_dims(orig_img_dims_.begin(), orig_img_dims_.end()); 164 | getNitrosContext().setParameter1DInt32Vector( 165 | "prompt_processor", 166 | "nvidia::isaac_ros::SegmentAnythingPromptProcessor", 167 | "orig_img_dim", orig_img_dims 168 | ); 169 | } 170 | 171 | SegmentAnythingDataEncoderNode::~SegmentAnythingDataEncoderNode() = default; 172 | 173 | } // namespace segment_anything 174 | } // namespace isaac_ros 175 | } // namespace nvidia 176 | 177 | // Register as component 178 | #include "rclcpp_components/register_node_macro.hpp" 179 | RCLCPP_COMPONENTS_REGISTER_NODE(nvidia::isaac_ros::segment_anything::SegmentAnythingDataEncoderNode) 180 | -------------------------------------------------------------------------------- /isaac_ros_gxf_extensions/gxf_isaac_ros_segment_anything/gxf/segment_anything/segment_anything_postprocessor.cpp: -------------------------------------------------------------------------------- 1 | // SPDX-FileCopyrightText: NVIDIA CORPORATION & AFFILIATES 2 | // Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. 3 | // 4 | // Licensed under the Apache License, Version 2.0 (the "License"); 5 | // you may not use this file except in compliance with the License. 6 | // You may obtain a copy of the License at 7 | // 8 | // http://www.apache.org/licenses/LICENSE-2.0 9 | // 10 | // Unless required by applicable law or agreed to in writing, software 11 | // distributed under the License is distributed on an "AS IS" BASIS, 12 | // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | // See the License for the specific language governing permissions and 14 | // limitations under the License. 15 | // 16 | // SPDX-License-Identifier: Apache-2.0 17 | #include "segment_anything_postprocessor.hpp" 18 | 19 | #include 20 | #include 21 | 22 | #include "cuda.h" 23 | #include "cuda_runtime.h" 24 | #include "gxf/cuda/cuda_stream_id.hpp" 25 | #include "gxf/std/tensor.hpp" 26 | #include "gxf/std/timestamp.hpp" 27 | #include "segment_anything_postprocessor.cu.hpp" 28 | 29 | namespace nvidia { 30 | namespace isaac_ros { 31 | 32 | namespace { 33 | 34 | gxf::Expected> GetTensor(gxf::Entity entity, const char* tensor_name) { 35 | // Try getting the named tensor first 36 | auto maybe_tensor = entity.get(tensor_name); 37 | 38 | // Default to getting any tensor 39 | if (!maybe_tensor) { maybe_tensor = entity.get(); } 40 | 41 | if (!maybe_tensor) { GXF_LOG_ERROR("Failed to get any tensor from input message!"); } 42 | 43 | return maybe_tensor; 44 | } 45 | 46 | gxf::Expected GetShape(gxf::Handle in_tensor) { 47 | Shape shape{}; 48 | 49 | shape.batch_size = in_tensor->shape().dimension(0); 50 | shape.channels = in_tensor->shape().dimension(1); 51 | shape.height = in_tensor->shape().dimension(2); 52 | shape.width = in_tensor->shape().dimension(3); 53 | 54 | // Only single channel 55 | if (shape.channels != kExpectedChannelCount) { 56 | GXF_LOG_ERROR("Received %d input channels, which is larger than the maximum allowable %d", 57 | shape.channels, kExpectedChannelCount); 58 | return gxf::Unexpected{GXF_FAILURE}; 59 | } 60 | 61 | return shape; 62 | } 63 | 64 | gxf::Expected GenerateSegmentationMask(Shape shape, 65 | gxf::Handle in_tensor, 66 | uint8_t* output_raw_buffer, 67 | cudaStream_t stream) { 68 | gxf::Expected in_tensor_data = in_tensor->data(); 69 | if (!in_tensor_data) { 70 | GXF_LOG_ERROR("Failed to get input tensor data!"); 71 | return gxf::ForwardError(in_tensor_data); 72 | } 73 | 74 | cuda_postprocess(shape, in_tensor_data.value(), output_raw_buffer,stream); 75 | cudaError_t kernel_result = cudaGetLastError(); 76 | if (kernel_result != cudaSuccess) { 77 | GXF_LOG_ERROR("Error while executing kernel: %s", cudaGetErrorString(kernel_result)); 78 | return gxf::Unexpected{GXF_FAILURE}; 79 | } 80 | 81 | return gxf::Success; 82 | } 83 | 84 | gxf::Expected AddInputTimestampToOutput(gxf::Entity& output, gxf::Entity input) { 85 | std::string timestamp_name{"timestamp"}; 86 | auto maybe_timestamp = input.get(timestamp_name.c_str()); 87 | 88 | // Default to unnamed 89 | if (!maybe_timestamp) { 90 | timestamp_name = std::string{""}; 91 | maybe_timestamp = input.get(timestamp_name.c_str()); 92 | } 93 | 94 | if (!maybe_timestamp) { 95 | GXF_LOG_ERROR("Failed to get input timestamp!"); 96 | return gxf::ForwardError(maybe_timestamp); 97 | } 98 | 99 | auto maybe_out_timestamp = output.add(timestamp_name.c_str()); 100 | if (!maybe_out_timestamp) { 101 | GXF_LOG_ERROR("Failed to add timestamp to output message!"); 102 | return gxf::ForwardError(maybe_out_timestamp); 103 | } 104 | 105 | *maybe_out_timestamp.value() = *maybe_timestamp.value(); 106 | return gxf::Success; 107 | } 108 | 109 | gxf::Expected AddStreamToOutput(gxf::Entity& output, gxf::Handle stream) { 110 | auto maybe_stream_id = output.add(); 111 | if (!maybe_stream_id) { return gxf::ForwardError(maybe_stream_id); } 112 | maybe_stream_id.value()->stream_cid = stream.cid(); 113 | if (maybe_stream_id.value()->stream_cid == kNullUid) { 114 | GXF_LOG_ERROR("Error: cuda stream handle is null!"); 115 | return gxf::Unexpected{GXF_FAILURE}; 116 | } 117 | return gxf::Success; 118 | } 119 | 120 | } // namespace 121 | 122 | gxf_result_t SegmentAnythingPostprocessor::registerInterface(gxf::Registrar* registrar) { 123 | gxf::Expected result; 124 | result &= registrar->parameter(in_, "in", "Input", "Input channel."); 125 | result &= registrar->parameter(in_tensor_name_, "in_tensor_name", "InputTensorName", 126 | "Name of the input tensor.", std::string("")); 127 | result &= registrar->parameter(out_, "out", "Output", "Output channel."); 128 | result &= registrar->parameter(allocator_, "allocator", "Allocator", "Output Allocator"); 129 | result &= registrar->parameter(cuda_stream_pool_, "cuda_stream_pool"); 130 | return gxf::ToResultCode(result); 131 | } 132 | 133 | gxf_result_t SegmentAnythingPostprocessor::start() { 134 | auto maybe_stream = cuda_stream_pool_.get()->allocateStream(); 135 | if (!maybe_stream) { return gxf::ToResultCode(maybe_stream); } 136 | 137 | stream_ = std::move(maybe_stream.value()); 138 | if (!stream_->stream()) { 139 | GXF_LOG_ERROR("Error: allocated stream is not initialized!"); 140 | return GXF_FAILURE; 141 | } 142 | 143 | return GXF_SUCCESS; 144 | } 145 | 146 | gxf_result_t SegmentAnythingPostprocessor::tick() { 147 | // Process input message 148 | const auto in_message = in_->receive(); 149 | 150 | if (!in_message || in_message.value().is_null()) { return GXF_CONTRACT_MESSAGE_NOT_AVAILABLE; } 151 | 152 | auto maybe_tensor = GetTensor(in_message.value(), in_tensor_name_.get().c_str()); 153 | if (!maybe_tensor) { return gxf::ToResultCode(maybe_tensor); } 154 | 155 | gxf::Handle in_tensor = maybe_tensor.value(); 156 | 157 | auto maybe_shape = GetShape(in_tensor); 158 | if (!maybe_shape) { return gxf::ToResultCode(maybe_shape); } 159 | 160 | Shape shape = maybe_shape.value(); 161 | 162 | auto out_message = gxf::Entity::New(context()); 163 | 164 | if (!out_message) { 165 | GXF_LOG_ERROR("Failed to allocate message"); 166 | return gxf::ToResultCode(out_message); 167 | } 168 | 169 | // Update the CUDA device to run the kernel on 170 | if (stream_->dev_id() >= 0) { 171 | if (cudaSetDevice(stream_->dev_id()) != cudaSuccess) { 172 | GXF_LOG_ERROR("Failed to set device_id: %d", stream_->dev_id()); 173 | return GXF_FAILURE; 174 | } 175 | } 176 | 177 | auto maybe_cuda_stream = stream_->stream(); 178 | if (!maybe_cuda_stream) { return gxf::ToResultCode(maybe_cuda_stream); } 179 | 180 | cudaStream_t cuda_stream = maybe_cuda_stream.value(); 181 | 182 | gxf::Expected maybe_kernel_result; 183 | 184 | auto raw_tensor = out_message.value().add(); 185 | 186 | if (!raw_tensor) { 187 | GXF_LOG_ERROR("Failed to allocate output raw tensor"); 188 | return gxf::ToResultCode(raw_tensor); 189 | } 190 | 191 | auto result = raw_tensor.value()->reshapeCustom( 192 | in_tensor->shape(), gxf::PrimitiveType::kUnsigned8, 193 | gxf::PrimitiveTypeSize(gxf::PrimitiveType::kUnsigned8), 194 | gxf::Unexpected{GXF_UNINITIALIZED_VALUE}, in_tensor->storage_type(), allocator_.get()); 195 | 196 | uint8_t* output_video_buffer_data = static_cast(raw_tensor.value()->pointer()); 197 | if (!output_video_buffer_data) { 198 | GXF_LOG_ERROR("Failed to get raw tensor buffer data!"); 199 | return GXF_FAILURE; 200 | } 201 | 202 | maybe_kernel_result = 203 | GenerateSegmentationMask(shape, in_tensor, output_video_buffer_data, cuda_stream); 204 | 205 | auto maybe_added_timestamp = AddInputTimestampToOutput(out_message.value(), in_message.value()); 206 | if (!maybe_added_timestamp) { return gxf::ToResultCode(maybe_added_timestamp); } 207 | 208 | 209 | auto maybe_added_stream = AddStreamToOutput(out_message.value(), stream_); 210 | if (!maybe_added_stream) { return gxf::ToResultCode(maybe_added_stream); } 211 | 212 | // Sync the stream, before sending it across the framework. 213 | // This is important to do since Isaac does not support non synced work across 214 | // codelets and multiple nodes. 215 | cudaError_t sync_result = cudaStreamSynchronize(cuda_stream); 216 | if (sync_result != cudaSuccess) { 217 | GXF_LOG_ERROR("Error while synchronizing CUDA stream: %s", cudaGetErrorString(sync_result)); 218 | return GXF_FAILURE; 219 | } 220 | 221 | return gxf::ToResultCode(out_->publish(std::move(out_message.value()))); 222 | } 223 | 224 | gxf_result_t SegmentAnythingPostprocessor::stop() { 225 | return GXF_SUCCESS; 226 | } 227 | 228 | } // namespace isaac_ros 229 | } // namespace nvidia 230 | -------------------------------------------------------------------------------- /isaac_ros_segment_anything2/include/isaac_ros_segment_anything2/segment_anything2_state_manager.hpp: -------------------------------------------------------------------------------- 1 | // SPDX-FileCopyrightText: NVIDIA CORPORATION & AFFILIATES 2 | // Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. 3 | // 4 | // Licensed under the Apache License, Version 2.0 (the "License"); 5 | // you may not use this file except in compliance with the License. 6 | // You may obtain a copy of the License at 7 | // 8 | // http://www.apache.org/licenses/LICENSE-2.0 9 | // 10 | // Unless required by applicable law or agreed to in writing, software 11 | // distributed under the License is distributed on an "AS IS" BASIS, 12 | // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | // See the License for the specific language governing permissions and 14 | // limitations under the License. 15 | // 16 | // SPDX-License-Identifier: Apache-2.0 17 | 18 | #ifndef ISAAC_ROS_SEGMENT_ANYTHING2__SEGMENT_ANYTHING2_STATE_MANAGER_HPP_ 19 | #define ISAAC_ROS_SEGMENT_ANYTHING2__SEGMENT_ANYTHING2_STATE_MANAGER_HPP_ 20 | 21 | #include 22 | #include 23 | #include 24 | #include 25 | #include 26 | #include 27 | #include 28 | #include 29 | #include 30 | #include 31 | #include 32 | #include 33 | #include 34 | #include "isaac_ros_common/cuda_stream.hpp" 35 | 36 | namespace nvidia 37 | { 38 | namespace isaac_ros 39 | { 40 | namespace segment_anything2 41 | { 42 | 43 | // Enum to track how an object was initially defined 44 | enum class SAM2PromptType 45 | { 46 | BBOX, // Object defined by bounding box 47 | POINTS, // Object defined by points 48 | UNKNOWN // Not set or unknown 49 | }; 50 | 51 | 52 | enum class TimestampType 53 | { 54 | MEMORY_UPDATE, 55 | OBJECT_UPDATE, 56 | LAST_FRAME_WITH_PROMPT 57 | }; 58 | 59 | struct BBox 60 | { 61 | float top_left_x; 62 | float top_left_y; 63 | float bottom_right_x; 64 | float bottom_right_y; 65 | BBox(float top_left_x, float top_left_y, float bottom_right_x, float bottom_right_y) 66 | : top_left_x(top_left_x), top_left_y(top_left_y), 67 | bottom_right_x(bottom_right_x), bottom_right_y(bottom_right_y) 68 | {} 69 | std::array toArray() const 70 | { 71 | return {top_left_x, top_left_y, bottom_right_x, bottom_right_y}; 72 | } 73 | }; 74 | 75 | struct Point 76 | { 77 | float x; 78 | float y; 79 | int label; 80 | Point(float x, float y, int label) 81 | : x(x), y(y), label(label) 82 | {} 83 | }; 84 | 85 | struct SAM2BufferData 86 | { 87 | float * mask_mem = nullptr; 88 | float * obj_ptr_mem = nullptr; 89 | float * bbox_coords = nullptr; 90 | float * point_coords = nullptr; 91 | int * point_labels = nullptr; 92 | int64_t * permutation = nullptr; 93 | int batch_size = 0; 94 | int num_points = 0; 95 | int num_bboxes = 0; 96 | }; 97 | 98 | // SAM2ObjectMemory encapsulates the tensor memories for a single tracked object 99 | class SAM2Object 100 | { 101 | public: 102 | SAM2Object(const std::string & object_id, SAM2PromptType prompt_type); 103 | ~SAM2Object(); 104 | 105 | // Initialize memory with zeros 106 | void allocateMemory(cudaStream_t stream); 107 | void deallocateMemory(); 108 | // Access memory buffers 109 | float * getMaskMemory() {return mask_memory_;} 110 | float * getObjPtrMemory() {return obj_ptr_memory_;} 111 | int32_t getMaskIdx() {return mask_idx_;} 112 | int64_t getObjectUpdateTimestamp() {return object_update_timestamp_;} 113 | int64_t getMemoryUpdateTimestamp() {return memory_update_timestamp_;} 114 | int64_t getLastFrameWithPromptTimestamp() {return last_frame_with_prompt_timestamp_;} 115 | bool updateTimestamp(TimestampType type, int64_t timestamp); 116 | // Update memory for this object 117 | void updateMemory( 118 | const float * mask_memory, 119 | const float * mask_pos_enc, 120 | const float * obj_ptr, 121 | cudaStream_t stream, 122 | bool update_condition_memory = false); 123 | 124 | // Store bounding box used to define this object 125 | void setInitialBoundingBox(const BBox & bbox); 126 | 127 | // Store points used to define this object 128 | bool setPoints( 129 | const std::vector & points, 130 | const std::vector & labels); 131 | 132 | // Get prompt type 133 | SAM2PromptType getPromptType() const {return prompt_type_;} 134 | // Get prompt data (returns either BBox or vector based on prompt type) 135 | std::variant> getPromptData() const; 136 | 137 | private: 138 | float * mask_memory_; // Shape [4, mem_dim, mem_dim, mem_dim] 139 | float * obj_ptr_memory_; // Shape [2, 256] 140 | 141 | std::string object_id_; 142 | 143 | // Prompt type for this object 144 | SAM2PromptType prompt_type_ = SAM2PromptType::UNKNOWN; 145 | 146 | // Bbox data if defined by bbox 147 | std::optional bbox_; 148 | // Points data if defined by points 149 | std::vector points_; 150 | 151 | // Timestamp of the last update to the object 152 | // Used to determine when was the last time the object bbox/points/or mask index was updated 153 | // We only update the memories if the memory timestamp is greater than the object update 154 | // timestamp. Let's say the object for which mask index was 2 was deleted, 155 | // i.e. object whose index was expected at index 3 is now updated to index 2 in state manager 156 | // but if an old memory msg comes where this information wasn't passed, we can use timestamp to 157 | // determine if we should update the memory or not. 158 | int64_t object_update_timestamp_ = 0; 159 | int64_t memory_update_timestamp_ = 0; 160 | 161 | // Last timestamp when we sent a prompt for this object 162 | // We continue to send objectprompts with an image unless we have received a memory for 163 | // that object but for initial frames we get late response from the model. Once we have 164 | // received a memory for that object we can ignore other frames up until 165 | // last_frame_with_prompt_timestamp_ 166 | int64_t last_frame_with_prompt_timestamp_ = 0; 167 | int32_t mask_idx_; 168 | }; 169 | 170 | class SegmentAnything2DataEncoderNode; 171 | // SAM2StateManager manages the tracking state of all objects 172 | class SAM2StateManager 173 | { 174 | public: 175 | explicit SAM2StateManager(SegmentAnything2DataEncoderNode * node); 176 | 177 | static constexpr int kMaxPointsPerObject = 5; 178 | // Add a new object to track with bbox 179 | // pair of object id and mask index 180 | std::vector addObjects( 181 | const std::vector & object_ids_bbox, 182 | const std::vector & object_ids_points, 183 | const std::vector & bbox_coords, 184 | const std::vector> & points, 185 | const std::vector> & labels, 186 | int64_t timestamp, cudaStream_t stream); 187 | // Update memories for all objects after inference 188 | void updateAllMemories( 189 | const float * mask_memories, 190 | const float * mask_pos_enc, 191 | const float * obj_ptr_memories, 192 | const float * object_score_logits, 193 | cudaStream_t stream, 194 | int64_t batch_size, 195 | int64_t memory_timestamp); 196 | 197 | int64_t getNumberOfObjects() const {return objects_.size();} 198 | 199 | SAM2BufferData getBuffers( 200 | cudaStream_t stream, 201 | const int64_t timestamp); 202 | 203 | bool removeObject(const std::string & obj_id, const int64_t timestamp); 204 | 205 | void getObjectIdsToOutputMaskIdx( 206 | std::vector & object_ids, 207 | std::vector & output_mask_idx); 208 | 209 | std::vector getAllObjectIds(); 210 | 211 | private: 212 | SegmentAnything2DataEncoderNode * node_; 213 | // Map object IDs to objects 214 | std::unordered_map> objects_; 215 | 216 | // Mapping from object IDs to output mask indices 217 | std::unordered_map obj_ids_to_output_mask_idx_; 218 | 219 | // vector of object ids in the order they will be processed by the model 220 | std::vector output_mask_idx_to_obj_id_; 221 | std::vector new_bbox_ids_; 222 | std::vector new_point_ids_; 223 | 224 | int64_t last_image_received_timestamp_ = 0; 225 | 226 | std::tuple getPointBuffer( 227 | cudaStream_t stream, int64_t timestamp); 228 | float * getBboxBuffer( 229 | cudaStream_t stream, int64_t timestamp); 230 | int64_t * getPermutationBuffer(cudaStream_t stream); 231 | // Add a new object to track (generic version) 232 | bool addObject( 233 | const std::string & obj_id, 234 | const SAM2PromptType prompt_type, 235 | int64_t timestamp, 236 | cudaStream_t stream, 237 | const std::optional & bbox = std::nullopt, 238 | const std::optional> & points = std::nullopt, 239 | const std::optional> & labels = std::nullopt); 240 | 241 | void removeObjectMaskMapping(const std::string & obj_id, const int64_t timestamp); 242 | }; 243 | 244 | } // namespace segment_anything2 245 | } // namespace isaac_ros 246 | } // namespace nvidia 247 | #endif // ISAAC_ROS_SEGMENT_ANYTHING2__SEGMENT_ANYTHING2_STATE_MANAGER_HPP_ 248 | -------------------------------------------------------------------------------- /isaac_ros_segment_anything2/scripts/add_object.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python3 2 | 3 | # SPDX-FileCopyrightText: NVIDIA CORPORATION & AFFILIATES 4 | # Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. 5 | # 6 | # Licensed under the Apache License, Version 2.0 (the "License"); 7 | # you may not use this file except in compliance with the License. 8 | # You may obtain a copy of the License at 9 | # 10 | # http://www.apache.org/licenses/LICENSE-2.0 11 | # 12 | # Unless required by applicable law or agreed to in writing, software 13 | # distributed under the License is distributed on an "AS IS" BASIS, 14 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 15 | # See the License for the specific language governing permissions and 16 | # limitations under the License. 17 | # 18 | # SPDX-License-Identifier: Apache-2.0 19 | 20 | import argparse 21 | import rclpy 22 | from rclpy.node import Node 23 | from vision_msgs.msg import BoundingBox2D, Point2D 24 | from isaac_ros_segment_anything2_interfaces.srv import AddObjects 25 | from std_msgs.msg import Header 26 | import sys 27 | 28 | 29 | class SingleObjectAdder(Node): 30 | def __init__(self): 31 | super().__init__('single_object_adder') 32 | 33 | # Create service client for AddObjects 34 | self.add_objects_client = self.create_client(AddObjects, 'add_objects') 35 | 36 | # Wait for service to be available 37 | timeout_count = 0 38 | while not self.add_objects_client.wait_for_service(timeout_sec=1.0): 39 | self.get_logger().info('Waiting for add_objects service...') 40 | timeout_count += 1 41 | if timeout_count > 10: 42 | self.get_logger().error('Service not available after 10 seconds. Exiting.') 43 | return 44 | 45 | self.get_logger().info('Connected to add_objects service!') 46 | 47 | def add_bbox_object(self, object_id, x_center, y_center, width, height, 48 | frame_id='camera_frame'): 49 | """Add a single bounding box object""" 50 | 51 | # Create bounding box 52 | bbox = BoundingBox2D() 53 | bbox.center.position.x = float(x_center) 54 | bbox.center.position.y = float(y_center) 55 | bbox.size_x = float(width) 56 | bbox.size_y = float(height) 57 | 58 | # Create service request 59 | request = AddObjects.Request() 60 | request.request_header = Header() 61 | request.request_header.frame_id = frame_id 62 | 63 | request.bbox_object_ids = [object_id] 64 | request.bbox_coords = [bbox] 65 | request.point_object_ids = [] 66 | request.point_coords = [] 67 | request.point_labels = [] 68 | 69 | return self._call_service(request) 70 | 71 | def add_point_object(self, object_id, x, y, point_label=1, 72 | frame_id='camera_frame'): 73 | """Add a single point object""" 74 | 75 | # Create point 76 | point = Point2D() 77 | point.x = float(x) 78 | point.y = float(y) 79 | 80 | # Create service request 81 | request = AddObjects.Request() 82 | request.request_header = Header() 83 | request.request_header.frame_id = frame_id 84 | 85 | request.bbox_object_ids = [] 86 | request.bbox_coords = [] 87 | request.point_object_ids = [object_id] 88 | request.point_coords = [point] 89 | request.point_labels = [int(point_label)] 90 | 91 | return self._call_service(request) 92 | 93 | def _call_service(self, request): 94 | """Call the AddObjects service""" 95 | try: 96 | self.get_logger().info('Calling add_objects service...') 97 | future = self.add_objects_client.call_async(request) 98 | 99 | # Wait for response with 5-second timeout 100 | self.get_logger().info('Waiting for service response (5 second timeout)...') 101 | rclpy.spin_until_future_complete(self, future, timeout_sec=5.0) 102 | 103 | if future.done(): 104 | self.get_logger().info('Service call completed') 105 | try: 106 | response = future.result() 107 | 108 | if response.success: 109 | self.get_logger().info(f'Service call succeeded: {response.message}') 110 | self.get_logger().info(f'Object IDs: {response.object_ids}') 111 | self.get_logger().info(f'Object Indices: {response.object_indices}') 112 | return True 113 | else: 114 | self.get_logger().error(f'Service call failed: {response.message}') 115 | if (len(response.not_added_object_ids) > 0): 116 | self.get_logger().error( 117 | f'Not added object ids: {response.not_added_object_ids}' 118 | ) 119 | return False 120 | except Exception as result_error: 121 | self.get_logger().error(f'Error getting service result: {str(result_error)}') 122 | return False 123 | else: 124 | # Timeout occurred 125 | self.get_logger().error( 126 | 'Service call timed out after 5 seconds - no response received' 127 | ) 128 | # Cancel the future to clean up 129 | future.cancel() 130 | return False 131 | 132 | except Exception as e: 133 | self.get_logger().error(f'Service call error: {str(e)}') 134 | return False 135 | 136 | 137 | def main(): 138 | parser = argparse.ArgumentParser(description='Utility to add a single object.') 139 | 140 | # Common arguments 141 | parser.add_argument('--object-id', '-i', required=True, type=str, 142 | help='Object ID/name (e.g., "bottle", "person1", "box")') 143 | parser.add_argument('--frame-id', '-f', default='camera_frame', type=str, 144 | help='Frame ID for the object (default: camera_frame)') 145 | 146 | # Create subparsers for bbox and point 147 | subparsers = parser.add_subparsers(dest='type', help='Object type') 148 | 149 | # Bounding box subparser 150 | bbox_parser = subparsers.add_parser('bbox', help='Add bounding box object') 151 | bbox_parser.add_argument('--x-center', '-x', required=True, type=float, 152 | help='X coordinate of bounding box center') 153 | bbox_parser.add_argument('--y-center', '-y', required=True, type=float, 154 | help='Y coordinate of bounding box center') 155 | bbox_parser.add_argument('--width', '-w', required=True, type=float, 156 | help='Width of bounding box') 157 | bbox_parser.add_argument('--height', '-H', required=True, type=float, 158 | help='Height of bounding box') 159 | 160 | # Point subparser 161 | point_parser = subparsers.add_parser('point', help='Add point object') 162 | point_parser.add_argument('--x', '-x', required=True, type=float, 163 | help='X coordinate of point') 164 | point_parser.add_argument('--y', '-y', required=True, type=float, 165 | help='Y coordinate of point') 166 | point_parser.add_argument('--label', '-l', default=1, type=int, 167 | help='Point label (1=foreground, 0=background, default: 1)') 168 | 169 | # Parse arguments 170 | args = parser.parse_args() 171 | 172 | if args.type is None: 173 | parser.print_help() 174 | print('\nError: Must specify either bbox or point subcommand') 175 | return 1 176 | 177 | # Initialize ROS2 178 | rclpy.init() 179 | 180 | try: 181 | # Create node 182 | adder = SingleObjectAdder() 183 | 184 | success = False 185 | if args.type == 'bbox': 186 | print(f'Adding bounding box object {args.object_id} at ' 187 | f'({args.x_center}, {args.y_center}) with size ' 188 | f'{args.width}x{args.height}') 189 | success = adder.add_bbox_object( 190 | args.object_id, 191 | args.x_center, 192 | args.y_center, 193 | args.width, 194 | args.height, 195 | args.frame_id 196 | ) 197 | elif args.type == 'point': 198 | print(f'Adding point object {args.object_id} at ({args.x}, {args.y}) ' 199 | f'with label {args.label}') 200 | success = adder.add_point_object( 201 | args.object_id, 202 | args.x, 203 | args.y, 204 | args.label, 205 | args.frame_id 206 | ) 207 | 208 | if success: 209 | print('Object added successfully!') 210 | return 0 211 | else: 212 | print('Failed to add object (check logs for details - may be timeout ' 213 | 'or service error)') 214 | return 1 215 | 216 | except KeyboardInterrupt: 217 | print('\nInterrupted by user') 218 | return 1 219 | except Exception as e: 220 | print(f'Error: {e}') 221 | return 1 222 | finally: 223 | rclpy.shutdown() 224 | 225 | 226 | if __name__ == '__main__': 227 | sys.exit(main()) 228 | --------------------------------------------------------------------------------