diff --git a/isaac_ros_rtdetr/CMakeLists.txt b/isaac_ros_rtdetr/CMakeLists.txt new file mode 100644 index 0000000..eaf30c8 --- /dev/null +++ b/isaac_ros_rtdetr/CMakeLists.txt @@ -0,0 +1,61 @@ +# SPDX-FileCopyrightText: NVIDIA CORPORATION & AFFILIATES +# Copyright (c) 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +# SPDX-License-Identifier: Apache-2.0 + +cmake_minimum_required(VERSION 3.22.1) +project(isaac_ros_rtdetr LANGUAGES C CXX) + +if(CMAKE_COMPILER_IS_GNUCXX OR CMAKE_CXX_COMPILER_ID MATCHES "Clang") + add_compile_options(-Wall -Wextra -Wpedantic) +endif() + +find_package(ament_cmake_auto REQUIRED) +find_package(rclcpp REQUIRED) +find_package(sensor_msgs REQUIRED) +find_package(vision_msgs REQUIRED) +find_package(geometry_msgs REQUIRED) +find_package(cv_bridge REQUIRED) +find_package(OpenCV REQUIRED) +include_directories(${OpenCV_INCLUDE_DIRS}) + +ament_auto_find_build_dependencies() +find_package(ament_cmake_python REQUIRED) + +# rtdetr Decoder node +ament_auto_add_library(rtdetr_decoder_node SHARED src/rtdetr_decoder_node.cpp) +rclcpp_components_register_nodes(rtdetr_decoder_node "nvidia::isaac_ros::rtdetr::RTDETRDecoderNode") +set(node_plugins "${node_plugins}nvidia::isaac_ros::rtdetr::RTDETRDecoderNode;$\n") + +target_link_libraries(rtdetr_decoder_node ${OpenCV_LIBRARIES}) + +if(BUILD_TESTING) + find_package(ament_lint_auto REQUIRED) + ament_lint_auto_find_test_dependencies() + + + # The FindPythonInterp and FindPythonLibs modules are removed + if(POLICY CMP0148) + cmake_policy(SET CMP0148 OLD) + endif() + + find_package(launch_testing_ament_cmake REQUIRED) +endif() + +# Visualizer python scripts + + + +ament_auto_package(INSTALL_TO_SHARE launch) diff --git a/isaac_ros_rtdetr/include/isaac_ros_rtdetr/rtdetr_decoder_node.hpp b/isaac_ros_rtdetr/include/isaac_ros_rtdetr/rtdetr_decoder_node.hpp new file mode 100644 index 0000000..f800e29 --- /dev/null +++ b/isaac_ros_rtdetr/include/isaac_ros_rtdetr/rtdetr_decoder_node.hpp @@ -0,0 +1,69 @@ +// SPDX-FileCopyrightText: NVIDIA CORPORATION & AFFILIATES +// Copyright (c) 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +// +// SPDX-License-Identifier: Apache-2.0 + +#ifndef ISAAC_ROS_YOLOV8__YOLOV8_DECODER_NODE_HPP_ +#define ISAAC_ROS_YOLOV8__YOLOV8_DECODER_NODE_HPP_ + +#include +#include + +#include "rclcpp/rclcpp.hpp" + +#include "isaac_ros_managed_nitros/managed_nitros_subscriber.hpp" + +#include "std_msgs/msg/string.hpp" +#include "vision_msgs/msg/detection2_d_array.hpp" +#include "isaac_ros_nitros_tensor_list_type/nitros_tensor_list_view.hpp" + +namespace nvidia +{ +namespace isaac_ros +{ +namespace rtdetr +{ + +class RTDETRDecoderNode : public rclcpp::Node +{ +public: + explicit RTDETRDecoderNode(const rclcpp::NodeOptions options = rclcpp::NodeOptions()); + + ~RTDETRDecoderNode(); + +private: + void InputCallback(const nvidia::isaac_ros::nitros::NitrosTensorListView & msg); + + // Subscription to input NitrosTensorList messages + std::shared_ptr> nitros_sub_; + + // Publisher for output Detection2DArray messages + rclcpp::Publisher::SharedPtr pub_; + + // Name of tensor in NitrosTensorList + std::string tensor_conf_name_{}; + std::string tensor_bboxes_name_{}; + std::string tensor_labels_name_{}; + + double confidence_threshold_; + +}; + +} // namespace rtdetr +} // namespace isaac_ros +} // namespace nvidia + +#endif // ISAAC_ROS_YOLOV8__YOLOV8_DECODER_NODE_HPP_ diff --git a/isaac_ros_rtdetr/isaac_ros_yolonas/__init__.py b/isaac_ros_rtdetr/isaac_ros_yolonas/__init__.py new file mode 100644 index 0000000..6bbcb13 --- /dev/null +++ b/isaac_ros_rtdetr/isaac_ros_yolonas/__init__.py @@ -0,0 +1,16 @@ +# SPDX-FileCopyrightText: NVIDIA CORPORATION & AFFILIATES +# Copyright (c) 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +# SPDX-License-Identifier: Apache-2.0 diff --git a/isaac_ros_rtdetr/launch/isaac_ros_yolov8_visualize.launch.py b/isaac_ros_rtdetr/launch/isaac_ros_yolov8_visualize.launch.py new file mode 100644 index 0000000..17f5c2a --- /dev/null +++ b/isaac_ros_rtdetr/launch/isaac_ros_yolov8_visualize.launch.py @@ -0,0 +1,46 @@ +# SPDX-FileCopyrightText: NVIDIA CORPORATION & AFFILIATES +# Copyright (c) 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +# SPDX-License-Identifier: Apache-2.0 + +import os + +from ament_index_python.packages import get_package_share_directory +from launch import LaunchDescription +from launch.actions import IncludeLaunchDescription +from launch.launch_description_sources import PythonLaunchDescriptionSource +from launch_ros.actions import Node + + +def generate_launch_description(): + my_package_dir = get_package_share_directory('isaac_ros_rtdetr') + return LaunchDescription([ + IncludeLaunchDescription( + PythonLaunchDescriptionSource([os.path.join( + my_package_dir, 'launch'), + '/rtdetr_tensor_rt.launch.py']) + ), + Node( + package='isaac_ros_rtdetr', + executable='isaac_ros_rtdetr_visualizer.py', + name='rtdetr_visualizer' + ), + Node( + package='rqt_image_view', + executable='rqt_image_view', + name='image_view', + arguments=['/rtdetr_processed_image'] + ) + ]) diff --git a/isaac_ros_rtdetr/launch/yolov8_tensor_rt.launch.py b/isaac_ros_rtdetr/launch/yolov8_tensor_rt.launch.py new file mode 100644 index 0000000..9a80a45 --- /dev/null +++ b/isaac_ros_rtdetr/launch/yolov8_tensor_rt.launch.py @@ -0,0 +1,137 @@ +# SPDX-FileCopyrightText: NVIDIA CORPORATION & AFFILIATES +# Copyright (c) 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +# SPDX-License-Identifier: Apache-2.0 + +import launch +from launch.actions import DeclareLaunchArgument +from launch.substitutions import LaunchConfiguration +from launch_ros.actions import ComposableNodeContainer +from launch_ros.descriptions import ComposableNode + + +def generate_launch_description(): + """Generate launch description for TensorRT ROS 2 node.""" + # By default loads and runs mobilenetv2-1.0 included in isaac_ros_dnn_inference/models + launch_args = [ + DeclareLaunchArgument( + 'model_file_path', + default_value='', + description='The absolute file path to the ONNX file'), + DeclareLaunchArgument( + 'engine_file_path', + default_value='', + description='The absolute file path to the TensorRT engine file'), + DeclareLaunchArgument( + 'input_tensor_names', + default_value='["input_tensor"]', + description='A list of tensor names to bound to the specified input binding names'), + DeclareLaunchArgument( + 'input_binding_names', + default_value='[""]', + description='A list of input tensor binding names (specified by model)'), + DeclareLaunchArgument( + 'output_tensor_names', + default_value='["output_tensor"]', + description='A list of tensor names to bound to the specified output binding names'), + DeclareLaunchArgument( + 'output_binding_names', + default_value='[""]', + description='A list of output tensor binding names (specified by model)'), + DeclareLaunchArgument( + 'verbose', + default_value='False', + description='Whether TensorRT should verbosely log or not'), + DeclareLaunchArgument( + 'force_engine_update', + default_value='False', + description='Whether TensorRT should update the TensorRT engine file or not'), + ] + + # DNN Image Encoder parameters + input_image_width = LaunchConfiguration('input_image_width') + input_image_height = LaunchConfiguration('input_image_height') + network_image_width = LaunchConfiguration('network_image_width') + network_image_height = LaunchConfiguration('network_image_height') + image_mean = LaunchConfiguration('image_mean') + image_stddev = LaunchConfiguration('image_stddev') + + # TensorRT parameters + model_file_path = LaunchConfiguration('model_file_path') + engine_file_path = LaunchConfiguration('engine_file_path') + input_tensor_names = LaunchConfiguration('input_tensor_names') + input_binding_names = LaunchConfiguration('input_binding_names') + output_tensor_names = LaunchConfiguration('output_tensor_names') + output_binding_names = LaunchConfiguration('output_binding_names') + verbose = LaunchConfiguration('verbose') + force_engine_update = LaunchConfiguration('force_engine_update') + + # YOLOv8 Decoder parameters + confidence_threshold = LaunchConfiguration('confidence_threshold') + nms_threshold = LaunchConfiguration('nms_threshold') + + encoder_node = ComposableNode( + name='dnn_image_encoder', + package='isaac_ros_dnn_image_encoder', + plugin='nvidia::isaac_ros::dnn_inference::DnnImageEncoderNode', + remappings=[('encoded_tensor', 'tensor_pub')], + parameters=[{ + 'input_image_width': input_image_width, + 'input_image_height': input_image_height, + 'network_image_width': network_image_width, + 'network_image_height': network_image_height, + 'image_mean': image_mean, + 'image_stddev': image_stddev, + }] + ) + + tensor_rt_node = ComposableNode( + name='tensor_rt', + package='isaac_ros_tensor_rt', + plugin='nvidia::isaac_ros::dnn_inference::TensorRTNode', + parameters=[{ + 'model_file_path': model_file_path, + 'engine_file_path': engine_file_path, + 'output_binding_names': output_binding_names, + 'output_tensor_names': output_tensor_names, + 'input_tensor_names': input_tensor_names, + 'input_binding_names': input_binding_names, + 'verbose': verbose, + 'force_engine_update': force_engine_update + }] + ) + + rtdetr_decoder_node = ComposableNode( + name='rtdetr_decoder_node', + package='isaac_ros_rtdetr', + plugin='nvidia::isaac_ros::rtdetr::RTDETRDecoderNode', + parameters=[{ + 'confidence_threshold': confidence_threshold, + 'nms_threshold': nms_threshold, + }] + ) + + tensor_rt_container = ComposableNodeContainer( + name='tensor_rt_container', + package='rclcpp_components', + executable='component_container_mt', + composable_node_descriptions=[encoder_node, tensor_rt_node, rtdetr_decoder_node], + output='screen', + arguments=['--ros-args', '--log-level', 'INFO'], + namespace='' + ) + + final_launch_description = launch_args + [tensor_rt_container] + return launch.LaunchDescription(final_launch_description) diff --git a/isaac_ros_rtdetr/launch_commands.md b/isaac_ros_rtdetr/launch_commands.md new file mode 100644 index 0000000..006f2bb --- /dev/null +++ b/isaac_ros_rtdetr/launch_commands.md @@ -0,0 +1,27 @@ +# Custom Nitros YOLOv8 + +This sample shows how to use your custom model decoder with Isaac ROS Managed Nitros. We consider the task of Object Detection using YOLOv8 with [Isaac ROS DNN Inference](https://github.com/NVIDIA-ISAAC-ROS/isaac_ros_dnn_inference). + +### Launch the image publisher in Terminal 1: +``` +ros2 run image_publisher image_publisher_node --ros-args --remap /image_raw:=/image` +``` + +For example (sample image people_cycles.jpg provided in this repo): +``` +ros2 run image_publisher image_publisher_node people_cycles.jpg --ros-args --remap /image_raw:=/image +``` + +### Launch the inference graph in Terminal 2: +``` +ros2 launch isaac_ros_rtdetr rtdetr_tensor_rt.launch.py engine_file_path:=rtdetrs_fp16.plan input_binding_names:=['images'] output_binding_names:=['output0'] network_image_width:=640 network_image_height:=640 model_file_path:=rtdetrs_fp16.onnx force_engine_update:=False image_mean:=[0.485,0.456,0.406] image_stddev:=[0.229,0.224,0.225] input_image_width:=640 input_image_height:=640 +``` + +Results will be published to `/detections_output` as Detection2DArray messages. + +## Visualizing results: +``` +ros2 launch isaac_ros_rtdetr isaac_ros_rtdetr_visualize.launch.py engine_file_path:=rtdetrs_fp16.plan input_binding_names:=['images'] output_binding_names:=['output0'] network_image_width:=640 network_image_height:=640 model_file_path:=rtdetrs_fp16.onnx force_engine_update:=False image_mean:=[0.485,0.456,0.406] image_stddev:=[0.229,0.224,0.225] input_image_width:=640 input_image_height:=640 +``` + +An RQT image window will pop up to display resulting bounding boxes on the input image. These output images are published on the `/rtdetr_processed_image` topic. diff --git a/isaac_ros_rtdetr/package.xml b/isaac_ros_rtdetr/package.xml new file mode 100644 index 0000000..db89bcd --- /dev/null +++ b/isaac_ros_rtdetr/package.xml @@ -0,0 +1,61 @@ + + + + + + + isaac_ros_rtdetr + 2.0.0 + Isaac ROS YOLOv8 decoding + + Hemal Shah + Apache-2.0 + https://developer.nvidia.com/isaac-ros-gems/ + Asawaree Bhide + + ament_cmake_auto + + + rclcpp + rclcpp_components + std_msgs + sensor_msgs + vision_msgs + geometry_msgs + tf2_geometry_msgs + tf2_ros + tf2_msgs + isaac_ros_nitros + isaac_ros_managed_nitros + isaac_ros_nitros_interfaces + isaac_ros_tensor_list_interfaces + python3-opencv + libopencv-dev + cv_bridge + + isaac_ros_common + + ament_lint_auto + ament_lint_common + isaac_ros_test + + + ament_cmake + + diff --git a/isaac_ros_rtdetr/scripts/isaac_ros_rtdetr_visualizer.py b/isaac_ros_rtdetr/scripts/isaac_ros_rtdetr_visualizer.py new file mode 100755 index 0000000..4c9bc45 --- /dev/null +++ b/isaac_ros_rtdetr/scripts/isaac_ros_rtdetr_visualizer.py @@ -0,0 +1,104 @@ +#!/usr/bin/env python3 + +# SPDX-FileCopyrightText: NVIDIA CORPORATION & AFFILIATES +# Copyright (c) 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +# SPDX-License-Identifier: Apache-2.0 + +# This script listens for images and object detections on the image, +# then renders the output boxes on top of the image and publishes +# the result as an image message + +import cv2 +import cv_bridge +import message_filters +import rclpy +from rclpy.node import Node +from sensor_msgs.msg import Image +from vision_msgs.msg import Detection2DArray + +names = { + 0: 'seed', +} + + +class RTDETRVisualizer(Node): + QUEUE_SIZE = 10 + color = (0, 255, 0) + bbox_thickness = 2 + + def __init__(self): + super().__init__('rtdetr_visualizer') + self._bridge = cv_bridge.CvBridge() + self._processed_image_pub = self.create_publisher( + Image, 'rtdetr_processed_image', self.QUEUE_SIZE) + + self._detections_subscription = message_filters.Subscriber( + self, + Detection2DArray, + 'detections_output') + self._image_subscription = message_filters.Subscriber( + self, + Image, + 'image') + + self.time_synchronizer = message_filters.TimeSynchronizer( + [self._detections_subscription, self._image_subscription], + self.QUEUE_SIZE) + + self.time_synchronizer.registerCallback(self.detections_callback) + + def detections_callback(self, detections_msg, img_msg): + txt_color = (255, 0, 255) + cv2_img = self._bridge.imgmsg_to_cv2(img_msg) + for detection in detections_msg.detections: + center_x = detection.bbox.center.position.x + center_y = detection.bbox.center.position.y + width = detection.bbox.size_x + height = detection.bbox.size_y + + label = names[int(detection.results[0].hypothesis.class_id)] + conf_score = detection.results[0].hypothesis.score + label = f'{label} {conf_score:.2f}' + + min_pt = (round(center_x - (width / 2.0)), + round(center_y - (height / 2.0))) + max_pt = (round(center_x + (width / 2.0)), + round(center_y + (height / 2.0))) + + lw = max(round((img_msg.height + img_msg.width) / 2 * 0.003), 2) # line width + tf = max(lw - 1, 1) # font thickness + # text width, height + w, h = cv2.getTextSize(label, 0, fontScale=lw / 3, thickness=tf)[0] + outside = min_pt[1] - h >= 3 + + cv2.rectangle(cv2_img, min_pt, max_pt, + self.color, self.bbox_thickness) + cv2.putText(cv2_img, label, (min_pt[0], min_pt[1]-2 if outside else min_pt[1]+h+2), + 0, lw / 3, txt_color, thickness=tf, lineType=cv2.LINE_AA) + + processed_img = self._bridge.cv2_to_imgmsg( + cv2_img, encoding=img_msg.encoding) + self._processed_image_pub.publish(processed_img) + + +def main(): + rclpy.init() + rclpy.spin(RTDETRVisualizer()) + rclpy.shutdown() + + +if __name__ == '__main__': + main() diff --git a/isaac_ros_rtdetr/src/rtdetr_decoder_node.cpp b/isaac_ros_rtdetr/src/rtdetr_decoder_node.cpp new file mode 100644 index 0000000..b1ac5bd --- /dev/null +++ b/isaac_ros_rtdetr/src/rtdetr_decoder_node.cpp @@ -0,0 +1,135 @@ +// SPDX-FileCopyrightText: NVIDIA CORPORATION & AFFILIATES +// Copyright (c) 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +// +// SPDX-License-Identifier: Apache-2.0 + +#include "isaac_ros_rtdetr/rtdetr_decoder_node.hpp" + +#include +#include +#include +#include +#include + +#include "isaac_ros_nitros_tensor_list_type/nitros_tensor_list_view.hpp" +#include "isaac_ros_nitros_tensor_list_type/nitros_tensor_list.hpp" + + +#include +#include +#include + +#include "vision_msgs/msg/detection2_d_array.hpp" +#include "tf2_geometry_msgs/tf2_geometry_msgs.hpp" + + +namespace nvidia +{ +namespace isaac_ros +{ +namespace rtdetr +{ +RTDETRDecoderNode::RTDETRDecoderNode(const rclcpp::NodeOptions options) +: rclcpp::Node("rtdetr_decoder_node", options), + nitros_sub_{std::make_shared>( + this, + "tensor_sub", + nvidia::isaac_ros::nitros::nitros_tensor_list_nchw_rgb_f32_t::supported_type_name, + std::bind(&RTDETRDecoderNode::InputCallback, this, + std::placeholders::_1))}, + pub_{create_publisher( + "detections_output", 50)}, + tensor_conf_name_{declare_parameter("tensor_conf_name", "scores")}, + tensor_bboxes_name_{declare_parameter("tensor_bboxes_name", "boxes")}, + tensor_labels_name_{declare_parameter("tensor_conf_name", "labels")}, + confidence_threshold_{declare_parameter("confidence_threshold", 0.25)} +{} + +RTDETRDecoderNode::~RTDETRDecoderNode() = default; + +void RTDETRDecoderNode::InputCallback(const nvidia::isaac_ros::nitros::NitrosTensorListView & msg) +{ + auto bbox_tensor = msg.GetNamedTensor(tensor_bboxes_name_); + auto conf_tensor = msg.GetNamedTensor(tensor_conf_name_); + auto labels_tensor = msg.GetNamedTensor(tensor_labels_name_); + size_t buffer_size{bbox_tensor.GetTensorSize()}; + size_t buffer_size_conf{conf_tensor.GetTensorSize()}; + size_t buffer_size_labels{labels_tensor.GetTensorSize()}; + std::vector results_vector{}; + std::vector labels_vec{}; + std::vector conf_results_vector{}; + results_vector.resize(buffer_size); + conf_results_vector.resize(buffer_size_conf); + labels_vec.resize(buffer_size_labels); + + cudaMemcpy(results_vector.data(), bbox_tensor.GetBuffer(), buffer_size, cudaMemcpyDefault); + cudaMemcpy(conf_results_vector.data(), conf_tensor.GetBuffer(), buffer_size_conf, cudaMemcpyDefault); + cudaMemcpy(labels_vec.data(), labels_tensor.GetBuffer(), buffer_size, cudaMemcpyDefault); + + + + + float * results_data = reinterpret_cast(results_vector.data()); + float * conf_results_data = reinterpret_cast(conf_results_vector.data()); + long long* labels_data = reinterpret_cast(labels_vec.data()); + + vision_msgs::msg::Detection2DArray final_detections_arr; + + for (int i = 0; i < labels_vec.size(); i++) { + if(conf_results_vector[i]>confidence_threshold_){ + float x_min = results_vector[4*i]; + float y_min = results_vector[4*i+1]; + float x_max = results_vector[4*i+2]; + float y_max = results_vector[4*i+3]; + + vision_msgs::msg::Detection2D detection; + vision_msgs::msg::BoundingBox2D bbox; + float w = x_max - x_min; + float h = y_max - y_min; + float x_center = (x_max+x_min)/2; + float y_center = (y_max + y_min)/2; + detection.bbox.center.position.x = x_center; + detection.bbox.center.position.y = y_center; + detection.bbox.size_x = w; + detection.bbox.size_y = h; + + + // Class probabilities + vision_msgs::msg::ObjectHypothesisWithPose hyp; + hyp.hypothesis.class_id = std::to_string(labels_data[i]); + hyp.hypothesis.score = conf_results_vector[i]; + detection.results.push_back(hyp); + + detection.header.stamp.sec = msg.GetTimestampSeconds(); + detection.header.stamp.nanosec = msg.GetTimestampNanoseconds(); + + + + } + } + + final_detections_arr.header.stamp.sec = msg.GetTimestampSeconds(); + final_detections_arr.header.stamp.nanosec = msg.GetTimestampNanoseconds(); + pub_->publish(final_detections_arr); +} + +} // namespace rtdetr +} // namespace isaac_ros +} // namespace nvidia + +// Register as component +#include "rclcpp_components/register_node_macro.hpp" +RCLCPP_COMPONENTS_REGISTER_NODE(nvidia::isaac_ros::rtdetr::RTDETRDecoderNode) diff --git a/isaac_ros_yolonas/CMakeLists.txt b/isaac_ros_yolonas/CMakeLists.txt new file mode 100644 index 0000000..e4abb39 --- /dev/null +++ b/isaac_ros_yolonas/CMakeLists.txt @@ -0,0 +1,65 @@ +# SPDX-FileCopyrightText: NVIDIA CORPORATION & AFFILIATES +# Copyright (c) 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +# SPDX-License-Identifier: Apache-2.0 + +cmake_minimum_required(VERSION 3.22.1) +project(isaac_ros_yolonas LANGUAGES C CXX) + +if(CMAKE_COMPILER_IS_GNUCXX OR CMAKE_CXX_COMPILER_ID MATCHES "Clang") + add_compile_options(-Wall -Wextra -Wpedantic) +endif() + +find_package(ament_cmake_auto REQUIRED) +find_package(rclcpp REQUIRED) +find_package(sensor_msgs REQUIRED) +find_package(vision_msgs REQUIRED) +find_package(geometry_msgs REQUIRED) +find_package(cv_bridge REQUIRED) +find_package(OpenCV REQUIRED) +include_directories(${OpenCV_INCLUDE_DIRS}) + +ament_auto_find_build_dependencies() +find_package(ament_cmake_python REQUIRED) + +# yolonas Decoder node +ament_auto_add_library(yolonas_decoder_node SHARED src/yolonas_decoder_node.cpp) +rclcpp_components_register_nodes(yolonas_decoder_node "nvidia::isaac_ros::yolonas::YoloNasDecoderNode") +set(node_plugins "${node_plugins}nvidia::isaac_ros::yolonas::YoloNasDecoderNode;$\n") + +target_link_libraries(yolonas_decoder_node ${OpenCV_LIBRARIES}) + +if(BUILD_TESTING) + find_package(ament_lint_auto REQUIRED) + ament_lint_auto_find_test_dependencies() + + + # The FindPythonInterp and FindPythonLibs modules are removed + if(POLICY CMP0148) + cmake_policy(SET CMP0148 OLD) + endif() + + find_package(launch_testing_ament_cmake REQUIRED) +endif() + +# Visualizer python scripts +ament_python_install_package(${PROJECT_NAME}) + +install(PROGRAMS + scripts/isaac_ros_yolonas_visualizer.py + DESTINATION lib/${PROJECT_NAME} +) + +ament_auto_package(INSTALL_TO_SHARE launch) diff --git a/isaac_ros_yolonas/include/isaac_ros_yolonas/yolonas_decoder_node.hpp b/isaac_ros_yolonas/include/isaac_ros_yolonas/yolonas_decoder_node.hpp new file mode 100644 index 0000000..015113d --- /dev/null +++ b/isaac_ros_yolonas/include/isaac_ros_yolonas/yolonas_decoder_node.hpp @@ -0,0 +1,71 @@ +// SPDX-FileCopyrightText: NVIDIA CORPORATION & AFFILIATES +// Copyright (c) 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +// +// SPDX-License-Identifier: Apache-2.0 + +#ifndef ISAAC_ROS_YOLOV8__YOLOV8_DECODER_NODE_HPP_ +#define ISAAC_ROS_YOLOV8__YOLOV8_DECODER_NODE_HPP_ + +#include +#include + +#include "rclcpp/rclcpp.hpp" + +#include "isaac_ros_managed_nitros/managed_nitros_subscriber.hpp" + +#include "std_msgs/msg/string.hpp" +#include "vision_msgs/msg/detection2_d_array.hpp" +#include "isaac_ros_nitros_tensor_list_type/nitros_tensor_list_view.hpp" + +namespace nvidia +{ +namespace isaac_ros +{ +namespace yolonas +{ + +class YoloNasDecoderNode : public rclcpp::Node +{ +public: + explicit YoloNasDecoderNode(const rclcpp::NodeOptions options = rclcpp::NodeOptions()); + + ~YoloNasDecoderNode(); + +private: + void InputCallback(const nvidia::isaac_ros::nitros::NitrosTensorListView & msg); + + // Subscription to input NitrosTensorList messages + std::shared_ptr> nitros_sub_; + + // Publisher for output Detection2DArray messages + rclcpp::Publisher::SharedPtr pub_; + + // Name of tensor in NitrosTensorList + std::string tensor_conf_name_{}; + std::string tensor_bboxes_name_{}; + + + // YOLOv8 Decoder Parameters + double confidence_threshold_{}; + double nms_threshold_{}; + int num_classes_{}; +}; + +} // namespace yolonas +} // namespace isaac_ros +} // namespace nvidia + +#endif // ISAAC_ROS_YOLOV8__YOLOV8_DECODER_NODE_HPP_ diff --git a/isaac_ros_yolonas/isaac_ros_yolonas/__init__.py b/isaac_ros_yolonas/isaac_ros_yolonas/__init__.py new file mode 100644 index 0000000..6bbcb13 --- /dev/null +++ b/isaac_ros_yolonas/isaac_ros_yolonas/__init__.py @@ -0,0 +1,16 @@ +# SPDX-FileCopyrightText: NVIDIA CORPORATION & AFFILIATES +# Copyright (c) 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +# SPDX-License-Identifier: Apache-2.0 diff --git a/isaac_ros_yolonas/launch/isaac_ros_yolov8_visualize.launch.py b/isaac_ros_yolonas/launch/isaac_ros_yolov8_visualize.launch.py new file mode 100644 index 0000000..3bc29a9 --- /dev/null +++ b/isaac_ros_yolonas/launch/isaac_ros_yolov8_visualize.launch.py @@ -0,0 +1,46 @@ +# SPDX-FileCopyrightText: NVIDIA CORPORATION & AFFILIATES +# Copyright (c) 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +# SPDX-License-Identifier: Apache-2.0 + +import os + +from ament_index_python.packages import get_package_share_directory +from launch import LaunchDescription +from launch.actions import IncludeLaunchDescription +from launch.launch_description_sources import PythonLaunchDescriptionSource +from launch_ros.actions import Node + + +def generate_launch_description(): + my_package_dir = get_package_share_directory('isaac_ros_yolonas') + return LaunchDescription([ + IncludeLaunchDescription( + PythonLaunchDescriptionSource([os.path.join( + my_package_dir, 'launch'), + '/yolonas_tensor_rt.launch.py']) + ), + Node( + package='isaac_ros_yolonas', + executable='isaac_ros_yolonas_visualizer.py', + name='yolonas_visualizer' + ), + Node( + package='rqt_image_view', + executable='rqt_image_view', + name='image_view', + arguments=['/yolonas_processed_image'] + ) + ]) diff --git a/isaac_ros_yolonas/launch/yolov8_tensor_rt.launch.py b/isaac_ros_yolonas/launch/yolov8_tensor_rt.launch.py new file mode 100644 index 0000000..bcccc5c --- /dev/null +++ b/isaac_ros_yolonas/launch/yolov8_tensor_rt.launch.py @@ -0,0 +1,137 @@ +# SPDX-FileCopyrightText: NVIDIA CORPORATION & AFFILIATES +# Copyright (c) 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +# SPDX-License-Identifier: Apache-2.0 + +import launch +from launch.actions import DeclareLaunchArgument +from launch.substitutions import LaunchConfiguration +from launch_ros.actions import ComposableNodeContainer +from launch_ros.descriptions import ComposableNode + + +def generate_launch_description(): + """Generate launch description for TensorRT ROS 2 node.""" + # By default loads and runs mobilenetv2-1.0 included in isaac_ros_dnn_inference/models + launch_args = [ + DeclareLaunchArgument( + 'model_file_path', + default_value='', + description='The absolute file path to the ONNX file'), + DeclareLaunchArgument( + 'engine_file_path', + default_value='', + description='The absolute file path to the TensorRT engine file'), + DeclareLaunchArgument( + 'input_tensor_names', + default_value='["input_tensor"]', + description='A list of tensor names to bound to the specified input binding names'), + DeclareLaunchArgument( + 'input_binding_names', + default_value='[""]', + description='A list of input tensor binding names (specified by model)'), + DeclareLaunchArgument( + 'output_tensor_names', + default_value='["output_tensor"]', + description='A list of tensor names to bound to the specified output binding names'), + DeclareLaunchArgument( + 'output_binding_names', + default_value='[""]', + description='A list of output tensor binding names (specified by model)'), + DeclareLaunchArgument( + 'verbose', + default_value='False', + description='Whether TensorRT should verbosely log or not'), + DeclareLaunchArgument( + 'force_engine_update', + default_value='False', + description='Whether TensorRT should update the TensorRT engine file or not'), + ] + + # DNN Image Encoder parameters + input_image_width = LaunchConfiguration('input_image_width') + input_image_height = LaunchConfiguration('input_image_height') + network_image_width = LaunchConfiguration('network_image_width') + network_image_height = LaunchConfiguration('network_image_height') + image_mean = LaunchConfiguration('image_mean') + image_stddev = LaunchConfiguration('image_stddev') + + # TensorRT parameters + model_file_path = LaunchConfiguration('model_file_path') + engine_file_path = LaunchConfiguration('engine_file_path') + input_tensor_names = LaunchConfiguration('input_tensor_names') + input_binding_names = LaunchConfiguration('input_binding_names') + output_tensor_names = LaunchConfiguration('output_tensor_names') + output_binding_names = LaunchConfiguration('output_binding_names') + verbose = LaunchConfiguration('verbose') + force_engine_update = LaunchConfiguration('force_engine_update') + + # YOLOv8 Decoder parameters + confidence_threshold = LaunchConfiguration('confidence_threshold') + nms_threshold = LaunchConfiguration('nms_threshold') + + encoder_node = ComposableNode( + name='dnn_image_encoder', + package='isaac_ros_dnn_image_encoder', + plugin='nvidia::isaac_ros::dnn_inference::DnnImageEncoderNode', + remappings=[('encoded_tensor', 'tensor_pub')], + parameters=[{ + 'input_image_width': input_image_width, + 'input_image_height': input_image_height, + 'network_image_width': network_image_width, + 'network_image_height': network_image_height, + 'image_mean': image_mean, + 'image_stddev': image_stddev, + }] + ) + + tensor_rt_node = ComposableNode( + name='tensor_rt', + package='isaac_ros_tensor_rt', + plugin='nvidia::isaac_ros::dnn_inference::TensorRTNode', + parameters=[{ + 'model_file_path': model_file_path, + 'engine_file_path': engine_file_path, + 'output_binding_names': output_binding_names, + 'output_tensor_names': output_tensor_names, + 'input_tensor_names': input_tensor_names, + 'input_binding_names': input_binding_names, + 'verbose': verbose, + 'force_engine_update': force_engine_update + }] + ) + + yolonas_decoder_node = ComposableNode( + name='yolonas_decoder_node', + package='isaac_ros_yolonas', + plugin='nvidia::isaac_ros::yolonas::YoloNasDecoderNode', + parameters=[{ + 'confidence_threshold': confidence_threshold, + 'nms_threshold': nms_threshold, + }] + ) + + tensor_rt_container = ComposableNodeContainer( + name='tensor_rt_container', + package='rclcpp_components', + executable='component_container_mt', + composable_node_descriptions=[encoder_node, tensor_rt_node, yolonas_decoder_node], + output='screen', + arguments=['--ros-args', '--log-level', 'INFO'], + namespace='' + ) + + final_launch_description = launch_args + [tensor_rt_container] + return launch.LaunchDescription(final_launch_description) diff --git a/isaac_ros_yolonas/launch_commands.md b/isaac_ros_yolonas/launch_commands.md new file mode 100644 index 0000000..e531c55 --- /dev/null +++ b/isaac_ros_yolonas/launch_commands.md @@ -0,0 +1,27 @@ +# Custom Nitros YOLOv8 + +This sample shows how to use your custom model decoder with Isaac ROS Managed Nitros. We consider the task of Object Detection using YOLOv8 with [Isaac ROS DNN Inference](https://github.com/NVIDIA-ISAAC-ROS/isaac_ros_dnn_inference). + +### Launch the image publisher in Terminal 1: +``` +ros2 run image_publisher image_publisher_node --ros-args --remap /image_raw:=/image` +``` + +For example (sample image people_cycles.jpg provided in this repo): +``` +ros2 run image_publisher image_publisher_node people_cycles.jpg --ros-args --remap /image_raw:=/image +``` + +### Launch the inference graph in Terminal 2: +``` +ros2 launch isaac_ros_yolonas yolonas_tensor_rt.launch.py engine_file_path:=yolonass_fp16.plan input_binding_names:=['images'] output_binding_names:=['output0'] network_image_width:=640 network_image_height:=640 model_file_path:=yolonass_fp16.onnx force_engine_update:=False image_mean:=[0.485,0.456,0.406] image_stddev:=[0.229,0.224,0.225] input_image_width:=640 input_image_height:=640 +``` + +Results will be published to `/detections_output` as Detection2DArray messages. + +## Visualizing results: +``` +ros2 launch isaac_ros_yolonas isaac_ros_yolonas_visualize.launch.py engine_file_path:=yolonass_fp16.plan input_binding_names:=['images'] output_binding_names:=['output0'] network_image_width:=640 network_image_height:=640 model_file_path:=yolonass_fp16.onnx force_engine_update:=False image_mean:=[0.485,0.456,0.406] image_stddev:=[0.229,0.224,0.225] input_image_width:=640 input_image_height:=640 +``` + +An RQT image window will pop up to display resulting bounding boxes on the input image. These output images are published on the `/yolonas_processed_image` topic. diff --git a/isaac_ros_yolonas/package.xml b/isaac_ros_yolonas/package.xml new file mode 100644 index 0000000..c52d445 --- /dev/null +++ b/isaac_ros_yolonas/package.xml @@ -0,0 +1,61 @@ + + + + + + + isaac_ros_yolonas + 2.0.0 + Isaac ROS YOLOv8 decoding + + Hemal Shah + Apache-2.0 + https://developer.nvidia.com/isaac-ros-gems/ + Asawaree Bhide + + ament_cmake_auto + + + rclcpp + rclcpp_components + std_msgs + sensor_msgs + vision_msgs + geometry_msgs + tf2_geometry_msgs + tf2_ros + tf2_msgs + isaac_ros_nitros + isaac_ros_managed_nitros + isaac_ros_nitros_interfaces + isaac_ros_tensor_list_interfaces + python3-opencv + libopencv-dev + cv_bridge + + isaac_ros_common + + ament_lint_auto + ament_lint_common + isaac_ros_test + + + ament_cmake + + diff --git a/isaac_ros_yolonas/scripts/isaac_ros_yolonas_visualizer.py b/isaac_ros_yolonas/scripts/isaac_ros_yolonas_visualizer.py new file mode 100755 index 0000000..75fab61 --- /dev/null +++ b/isaac_ros_yolonas/scripts/isaac_ros_yolonas_visualizer.py @@ -0,0 +1,104 @@ +#!/usr/bin/env python3 + +# SPDX-FileCopyrightText: NVIDIA CORPORATION & AFFILIATES +# Copyright (c) 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +# SPDX-License-Identifier: Apache-2.0 + +# This script listens for images and object detections on the image, +# then renders the output boxes on top of the image and publishes +# the result as an image message + +import cv2 +import cv_bridge +import message_filters +import rclpy +from rclpy.node import Node +from sensor_msgs.msg import Image +from vision_msgs.msg import Detection2DArray + +names = { + 0: 'seed', +} + + +class YoloNasVisualizer(Node): + QUEUE_SIZE = 10 + color = (0, 255, 0) + bbox_thickness = 2 + + def __init__(self): + super().__init__('yolonas_visualizer') + self._bridge = cv_bridge.CvBridge() + self._processed_image_pub = self.create_publisher( + Image, 'yolonas_processed_image', self.QUEUE_SIZE) + + self._detections_subscription = message_filters.Subscriber( + self, + Detection2DArray, + 'detections_output') + self._image_subscription = message_filters.Subscriber( + self, + Image, + 'image') + + self.time_synchronizer = message_filters.TimeSynchronizer( + [self._detections_subscription, self._image_subscription], + self.QUEUE_SIZE) + + self.time_synchronizer.registerCallback(self.detections_callback) + + def detections_callback(self, detections_msg, img_msg): + txt_color = (255, 0, 255) + cv2_img = self._bridge.imgmsg_to_cv2(img_msg) + for detection in detections_msg.detections: + center_x = detection.bbox.center.position.x + center_y = detection.bbox.center.position.y + width = detection.bbox.size_x + height = detection.bbox.size_y + + label = names[int(detection.results[0].hypothesis.class_id)] + conf_score = detection.results[0].hypothesis.score + label = f'{label} {conf_score:.2f}' + + min_pt = (round(center_x - (width / 2.0)), + round(center_y - (height / 2.0))) + max_pt = (round(center_x + (width / 2.0)), + round(center_y + (height / 2.0))) + + lw = max(round((img_msg.height + img_msg.width) / 2 * 0.003), 2) # line width + tf = max(lw - 1, 1) # font thickness + # text width, height + w, h = cv2.getTextSize(label, 0, fontScale=lw / 3, thickness=tf)[0] + outside = min_pt[1] - h >= 3 + + cv2.rectangle(cv2_img, min_pt, max_pt, + self.color, self.bbox_thickness) + cv2.putText(cv2_img, label, (min_pt[0], min_pt[1]-2 if outside else min_pt[1]+h+2), + 0, lw / 3, txt_color, thickness=tf, lineType=cv2.LINE_AA) + + processed_img = self._bridge.cv2_to_imgmsg( + cv2_img, encoding=img_msg.encoding) + self._processed_image_pub.publish(processed_img) + + +def main(): + rclpy.init() + rclpy.spin(YoloNasVisualizer()) + rclpy.shutdown() + + +if __name__ == '__main__': + main() diff --git a/isaac_ros_yolonas/src/yolonas_decoder_node.cpp b/isaac_ros_yolonas/src/yolonas_decoder_node.cpp new file mode 100644 index 0000000..6b1a82b --- /dev/null +++ b/isaac_ros_yolonas/src/yolonas_decoder_node.cpp @@ -0,0 +1,166 @@ +// SPDX-FileCopyrightText: NVIDIA CORPORATION & AFFILIATES +// Copyright (c) 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +// +// SPDX-License-Identifier: Apache-2.0 + +#include "isaac_ros_yolonas/yolonas_decoder_node.hpp" + +#include +#include +#include +#include +#include + +#include "isaac_ros_nitros_tensor_list_type/nitros_tensor_list_view.hpp" +#include "isaac_ros_nitros_tensor_list_type/nitros_tensor_list.hpp" + + +#include +#include +#include + +#include "vision_msgs/msg/detection2_d_array.hpp" +#include "tf2_geometry_msgs/tf2_geometry_msgs.hpp" + + +namespace nvidia +{ +namespace isaac_ros +{ +namespace yolonas +{ +YoloNasDecoderNode::YoloNasDecoderNode(const rclcpp::NodeOptions options) +: rclcpp::Node("yolonas_decoder_node", options), + nitros_sub_{std::make_shared>( + this, + "tensor_sub", + nvidia::isaac_ros::nitros::nitros_tensor_list_nchw_rgb_f32_t::supported_type_name, + std::bind(&YoloNasDecoderNode::InputCallback, this, + std::placeholders::_1))}, + pub_{create_publisher( + "detections_output", 50)}, + tensor_conf_name_{declare_parameter("tensor_conf_name", "conf")}, + tensor_bboxes_name_{declare_parameter("tensor_bboxes_name", "bboxes")}, + confidence_threshold_{declare_parameter("confidence_threshold", 0.25)}, + nms_threshold_{declare_parameter("nms_threshold", 0.45)}, + num_classes_{declare_parameter("num_classes", 1)} +{} + +YoloNasDecoderNode::~YoloNasDecoderNode() = default; + +void YoloNasDecoderNode::InputCallback(const nvidia::isaac_ros::nitros::NitrosTensorListView & msg) +{ + auto bbox_tensor = msg.GetNamedTensor(tensor_bboxes_name_); + auto conf_tensor = msg.GetNamedTensor(tensor_conf_name_); + size_t buffer_size{bbox_tensor.GetTensorSize()}; + size_t buffer_size_conf{conf_tensor.GetTensorSize()}; + std::vector results_vector{}; + std::vector conf_results_vector{}; + results_vector.resize(buffer_size); + conf_results_vector.resize(buffer_size_conf); + + cudaMemcpy(results_vector.data(), bbox_tensor.GetBuffer(), buffer_size, cudaMemcpyDefault); + cudaMemcpy(conf_results_vector.data(), conf_tensor.GetBuffer(), buffer_size_conf, cudaMemcpyDefault); + + std::vector bboxes; + std::vector scores; + std::vector indices; + std::vector classes; + + + int out_dim = conf_tensor.GetElementCount(); + float * results_data = reinterpret_cast(results_vector.data()); + float * conf_results_data = reinterpret_cast(conf_results_vector.data()); + RCLCPP_INFO(this->get_logger(), "Output tensor size: %d", out_dim); + RCLCPP_INFO(this->get_logger(), "Confidence tensor size: %d", conf_tensor.GetTensorSize()); + int num_objs = 0; + for (int i = 0; i < out_dim; i++) { + float x0 = *(results_data + i*4); + float y0 = *(results_data + i*4 + 1); + float x1 = *(results_data + i*4 + 2); + float y1 = *(results_data + i*4 + 3); + + float width = x1 - x0; + float height = y1 - y0; + + + std::vector conf; + for (int j = 0; j < num_classes_; j++) { + conf.push_back(*(conf_results_data + (i * num_classes_) + j)); + if(conf[j]>=confidence_threshold_) num_objs++; + } + + std::vector::iterator ind_max_conf; + ind_max_conf = std::max_element(std::begin(conf), std::end(conf)); + int max_index = distance(std::begin(conf), ind_max_conf); + float val_max_conf = *max_element(std::begin(conf), std::end(conf)); + + bboxes.push_back(cv::Rect(x0, y0, width, height)); + indices.push_back(i); + scores.push_back(val_max_conf); + classes.push_back(max_index); + } + + RCLCPP_INFO(this->get_logger(), "Count of bboxes: %lu", bboxes.size()); + cv::dnn::NMSBoxes(bboxes, scores, confidence_threshold_, nms_threshold_, indices, 5); + RCLCPP_INFO(this->get_logger(), "# boxes after NMS: %lu", indices.size()); + RCLCPP_INFO(this->get_logger(), "# objects after NMS: %d", num_objs); + vision_msgs::msg::Detection2DArray final_detections_arr; + + for (size_t i = 0; i < indices.size(); i++) { + int ind = indices[i]; + vision_msgs::msg::Detection2D detection; + + geometry_msgs::msg::Pose center; + geometry_msgs::msg::Point position; + geometry_msgs::msg::Quaternion orientation; + + // 2D object Bbox + vision_msgs::msg::BoundingBox2D bbox; + float w = bboxes[ind].width; + float h = bboxes[ind].height; + float x_center = bboxes[ind].x + 0.5*w; + float y_center = bboxes[ind].y + 0.5*h; + detection.bbox.center.position.x = x_center; + detection.bbox.center.position.y = y_center; + detection.bbox.size_x = w; + detection.bbox.size_y = h; + + + // Class probabilities + vision_msgs::msg::ObjectHypothesisWithPose hyp; + hyp.hypothesis.class_id = std::to_string(classes.at(ind)); + hyp.hypothesis.score = scores.at(ind); + detection.results.push_back(hyp); + + detection.header.stamp.sec = msg.GetTimestampSeconds(); + detection.header.stamp.nanosec = msg.GetTimestampNanoseconds(); + + final_detections_arr.detections.push_back(detection); + } + + final_detections_arr.header.stamp.sec = msg.GetTimestampSeconds(); + final_detections_arr.header.stamp.nanosec = msg.GetTimestampNanoseconds(); + pub_->publish(final_detections_arr); +} + +} // namespace yolonas +} // namespace isaac_ros +} // namespace nvidia + +// Register as component +#include "rclcpp_components/register_node_macro.hpp" +RCLCPP_COMPONENTS_REGISTER_NODE(nvidia::isaac_ros::yolonas::YoloNasDecoderNode)