Skip to content

Commit 2d81e6e

Browse files
Support third party RT-DETR models
1 parent ed86aac commit 2d81e6e

File tree

4 files changed

+230
-4
lines changed

4 files changed

+230
-4
lines changed

isaac_ros_rtdetr/include/isaac_ros_rtdetr/rtdetr_preprocessor_node.hpp

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -64,6 +64,7 @@ class RtDetrPreprocessorNode : public rclcpp::Node
6464
std::string output_size_tensor_name_{};
6565
int64_t image_height_{};
6666
int64_t image_width_{};
67+
bool use_max_dim_for_orig_size_{};
6768
cudaStream_t stream_;
6869
};
6970

Lines changed: 219 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,219 @@
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 launch
19+
from launch.actions import DeclareLaunchArgument
20+
from launch.substitutions import LaunchConfiguration
21+
from launch_ros.actions import ComposableNodeContainer
22+
from launch_ros.descriptions import ComposableNode
23+
24+
MODEL_INPUT_SIZE = 640 # RT-DETR models expect 640x640 encoded image size
25+
MODEL_NUM_CHANNELS = 3 # RT-DETR models expect 3 image channels
26+
27+
28+
def generate_launch_description():
29+
"""Generate launch description for testing relevant nodes."""
30+
launch_args = [
31+
DeclareLaunchArgument(
32+
'model_file_path',
33+
default_value='',
34+
description='The absolute file path to the ONNX file'),
35+
DeclareLaunchArgument(
36+
'engine_file_path',
37+
default_value='',
38+
description='The absolute file path to the TensorRT engine file'),
39+
DeclareLaunchArgument(
40+
'input_image_width',
41+
default_value='640',
42+
description='The input image width'),
43+
DeclareLaunchArgument(
44+
'input_image_height',
45+
default_value='480',
46+
description='The input image height'),
47+
DeclareLaunchArgument(
48+
'input_tensor_names',
49+
default_value='["images", "orig_target_sizes"]',
50+
description='A list of tensor names to bound to the specified input binding names'),
51+
DeclareLaunchArgument(
52+
'input_binding_names',
53+
default_value='["images", "orig_target_sizes"]',
54+
description='A list of input tensor binding names (specified by model)'),
55+
DeclareLaunchArgument(
56+
'output_tensor_names',
57+
default_value='["labels", "boxes", "scores"]',
58+
description='A list of tensor names to bound to the specified output binding names'),
59+
DeclareLaunchArgument(
60+
'output_binding_names',
61+
default_value='["labels", "boxes", "scores"]',
62+
description='A list of output tensor binding names (specified by model)'),
63+
DeclareLaunchArgument(
64+
'verbose',
65+
default_value='False',
66+
description='Whether TensorRT should verbosely log or not'),
67+
DeclareLaunchArgument(
68+
'force_engine_update',
69+
default_value='False',
70+
description='Whether TensorRT should update the TensorRT engine file or not'),
71+
DeclareLaunchArgument(
72+
'confidence_threshold',
73+
default_value='0.6',
74+
description='The minimum score for a bounding box to be published.',
75+
),
76+
]
77+
78+
# Image Encoding parameters
79+
input_image_width = LaunchConfiguration('input_image_width')
80+
input_image_height = LaunchConfiguration('input_image_height')
81+
82+
# TensorRT parameters
83+
model_file_path = LaunchConfiguration('model_file_path')
84+
engine_file_path = LaunchConfiguration('engine_file_path')
85+
input_tensor_names = LaunchConfiguration('input_tensor_names')
86+
input_binding_names = LaunchConfiguration('input_binding_names')
87+
output_tensor_names = LaunchConfiguration('output_tensor_names')
88+
output_binding_names = LaunchConfiguration('output_binding_names')
89+
verbose = LaunchConfiguration('verbose')
90+
force_engine_update = LaunchConfiguration('force_engine_update')
91+
92+
confidence_threshold = LaunchConfiguration('confidence_threshold')
93+
94+
resize_node = ComposableNode(
95+
name='resize_node',
96+
package='isaac_ros_image_proc',
97+
plugin='nvidia::isaac_ros::image_proc::ResizeNode',
98+
parameters=[{
99+
'input_width': input_image_width,
100+
'input_height': input_image_height,
101+
'output_width': MODEL_INPUT_SIZE,
102+
'output_height': MODEL_INPUT_SIZE,
103+
'keep_aspect_ratio': False,
104+
'encoding_desired': 'rgb8',
105+
'disable_padding': False
106+
}],
107+
remappings=[
108+
('image', 'image_rect'),
109+
('camera_info', 'camera_info_rect')],
110+
)
111+
112+
image_format_node = ComposableNode(
113+
name='image_format_node',
114+
package='isaac_ros_image_proc',
115+
plugin='nvidia::isaac_ros::image_proc::ImageFormatConverterNode',
116+
parameters=[{
117+
'encoding_desired': 'rgb8',
118+
'image_width': MODEL_INPUT_SIZE,
119+
'image_height': MODEL_INPUT_SIZE
120+
}],
121+
remappings=[
122+
('image_raw', 'resize/image'),
123+
('image', 'image_rgb')]
124+
)
125+
126+
image_to_tensor_node = ComposableNode(
127+
name='image_to_tensor_node',
128+
package='isaac_ros_tensor_proc',
129+
plugin='nvidia::isaac_ros::dnn_inference::ImageToTensorNode',
130+
parameters=[{
131+
'scale': True,
132+
'tensor_name': 'image',
133+
}],
134+
remappings=[
135+
('image', 'image_rgb'),
136+
('tensor', 'normalized_tensor'),
137+
]
138+
)
139+
140+
interleave_to_planar_node = ComposableNode(
141+
name='interleaved_to_planar_node',
142+
package='isaac_ros_tensor_proc',
143+
plugin='nvidia::isaac_ros::dnn_inference::InterleavedToPlanarNode',
144+
parameters=[{
145+
'input_tensor_shape': [MODEL_INPUT_SIZE, MODEL_INPUT_SIZE, MODEL_NUM_CHANNELS]
146+
}],
147+
remappings=[
148+
('interleaved_tensor', 'normalized_tensor')
149+
]
150+
)
151+
152+
reshape_node = ComposableNode(
153+
name='reshape_node',
154+
package='isaac_ros_tensor_proc',
155+
plugin='nvidia::isaac_ros::dnn_inference::ReshapeNode',
156+
parameters=[{
157+
'output_tensor_name': 'input_tensor',
158+
'input_tensor_shape': [MODEL_NUM_CHANNELS, MODEL_INPUT_SIZE, MODEL_INPUT_SIZE],
159+
'output_tensor_shape': [1, MODEL_NUM_CHANNELS, MODEL_INPUT_SIZE, MODEL_INPUT_SIZE]
160+
}],
161+
remappings=[
162+
('tensor', 'planar_tensor')
163+
],
164+
)
165+
166+
rtdetr_preprocessor_node = ComposableNode(
167+
name='rtdetr_preprocessor',
168+
package='isaac_ros_rtdetr',
169+
plugin='nvidia::isaac_ros::rtdetr::RtDetrPreprocessorNode',
170+
parameters=[{
171+
'image_width': input_image_width,
172+
'image_height': input_image_height,
173+
'use_max_dim_for_orig_size': False,
174+
}],
175+
remappings=[
176+
('encoded_tensor', 'reshaped_tensor')
177+
]
178+
)
179+
180+
tensor_rt_node = ComposableNode(
181+
name='tensor_rt',
182+
package='isaac_ros_tensor_rt',
183+
plugin='nvidia::isaac_ros::dnn_inference::TensorRTNode',
184+
parameters=[{
185+
'model_file_path': model_file_path,
186+
'engine_file_path': engine_file_path,
187+
'output_binding_names': output_binding_names,
188+
'output_tensor_names': output_tensor_names,
189+
'input_tensor_names': input_tensor_names,
190+
'input_binding_names': input_binding_names,
191+
'verbose': verbose,
192+
'force_engine_update': force_engine_update
193+
}]
194+
)
195+
196+
rtdetr_decoder_node = ComposableNode(
197+
name='rtdetr_decoder',
198+
package='isaac_ros_rtdetr',
199+
plugin='nvidia::isaac_ros::rtdetr::RtDetrDecoderNode',
200+
parameters=[{
201+
'confidence_threshold': confidence_threshold,
202+
}],
203+
)
204+
205+
container = ComposableNodeContainer(
206+
name='rtdetr_container',
207+
namespace='rtdetr_container',
208+
package='rclcpp_components',
209+
executable='component_container_mt',
210+
composable_node_descriptions=[
211+
resize_node, image_format_node,
212+
image_to_tensor_node, interleave_to_planar_node, reshape_node,
213+
rtdetr_preprocessor_node, tensor_rt_node, rtdetr_decoder_node
214+
],
215+
output='screen'
216+
)
217+
218+
final_launch_description = launch_args + [container]
219+
return launch.LaunchDescription(final_launch_description)

isaac_ros_rtdetr/package.xml

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -21,7 +21,7 @@ SPDX-License-Identifier: Apache-2.0
2121
<?xml-model href="http://download.ros.org/schema/package_format3.xsd" schematypens="http://www.w3.org/2001/XMLSchema"?>
2222
<package format="3">
2323
<name>isaac_ros_rtdetr</name>
24-
<version>4.0.0</version>
24+
<version>4.0.1</version>
2525
<description>RT-DETR model processing</description>
2626

2727
<maintainer email="[email protected]">Isaac ROS Maintainers</maintainer>

isaac_ros_rtdetr/src/rtdetr_preprocessor_node.cpp

Lines changed: 9 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,8 @@
1717

1818
#include "isaac_ros_rtdetr/rtdetr_preprocessor_node.hpp"
1919

20+
#include <algorithm>
21+
2022
#include "isaac_ros_nitros_tensor_list_type/nitros_tensor_builder.hpp"
2123
#include "isaac_ros_nitros_tensor_list_type/nitros_tensor_list.hpp"
2224
#include "isaac_ros_nitros_tensor_list_type/nitros_tensor_list_builder.hpp"
@@ -59,7 +61,8 @@ RtDetrPreprocessorNode::RtDetrPreprocessorNode(const rclcpp::NodeOptions options
5961
"output_size_tensor_name",
6062
"orig_target_sizes")},
6163
image_height_{declare_parameter<int64_t>("image_height", 480)},
62-
image_width_{declare_parameter<int64_t>("image_width", 640)}
64+
image_width_{declare_parameter<int64_t>("image_width", 640)},
65+
use_max_dim_for_orig_size_{declare_parameter<bool>("use_max_dim_for_orig_size", true)}
6366
{
6467
CHECK_CUDA_ERROR(
6568
::nvidia::isaac_ros::common::initNamedCudaStream(
@@ -89,9 +92,12 @@ void RtDetrPreprocessorNode::InputCallback(
8992
output_image_buffer, input_image_tensor.GetBuffer(),
9093
input_image_tensor.GetTensorSize(), cudaMemcpyDefault, stream_);
9194

92-
int64_t image_size = std::max(image_height_, image_width_);
95+
const int64_t orig_width = use_max_dim_for_orig_size_ ?
96+
std::max(image_height_, image_width_) : image_width_;
97+
const int64_t orig_height = use_max_dim_for_orig_size_ ?
98+
std::max(image_height_, image_width_) : image_height_;
9399

94-
int64_t output_size[2]{image_size, image_size};
100+
int64_t output_size[2]{orig_width, orig_height};
95101
void * output_size_buffer;
96102
cudaMallocAsync(&output_size_buffer, sizeof(output_size), stream_);
97103
cudaMemcpyAsync(output_size_buffer, output_size, sizeof(output_size), cudaMemcpyDefault, stream_);

0 commit comments

Comments
 (0)