To build an end-to-end, hardware-accelerated computer vision and tracking pipeline on NVIDIA Jetson using ROS 2, you need to avoid CPU bottlenecking. Passing images back and forth between host RAM and GPU VRAM using naive cv_bridge creates massive latency.
This architecture leverages CUDA zero-copy pinned memory, TensorRT FP16 execution, and ByteTrack / SORT tracking, wrapped cleanly in a ROS 2 node.
Prerequisites & Exporting the Engine
Run these steps directly on your Jetson (or export via trtexec with matching JetPack CUDA versions).
1. Export YOLO to TensorRT FP16 Engine
Bash
# Export using Ultralytics CLI (run on your Jetson)
yolo export model=yolov8n.pt format=engine half=True workspace=2
This outputs yolov8n.engine, optimized specifically for your Jetson’s Ampere/Orin GPU kernels.
2. Dependencies
Bash
sudo apt install ros-humble-cv-bridge ros-humble-vision-msgs
pip install pycuda ultralytics opencv-python
High-Performance ROS 2 TensorRT Tracking Node
Save the following code as trt_tracker_node.py in your ROS 2 package. It handles CUDA context binding, asynchronous memory streaming, TensorRT inference, bounding box extraction, and tracking output.
Python
#!/usr/bin/env python3
import rclpy
from rclpy.node import Node
from sensor_msgs.msg import Image
from vision_msgs.msg import Detection2DArray, Detection2D, ObjectHypothesisWithPose
from cv_bridge import CvBridge
import cv2
import numpy as np
import tensorrt as trt
import pycuda.driver as cuda
import pycuda.autoinit # Automatically initializes CUDA driver context
from ultralytics.trackers.byte_tracker import BYTETracker # Lightweight byte-tracker
from argparse import Namespace
class TensorRTROS2Tracker(Node):
def __init__(self):
super().__init__('tensorrt_tracker_node')
# Parameters
self.declare_parameter('engine_path', 'yolov8n.engine')
self.declare_parameter('input_topic', '/image_raw')
self.declare_parameter('conf_thresh', 0.4)
engine_path = self.get_parameter('engine_path').get_parameter_value().string_value
input_topic = self.get_parameter('input_topic').get_parameter_value().string_value
self.conf_thresh = self.get_parameter('conf_thresh').get_parameter_value().double_value
self.bridge = CvBridge()
# 1. Initialize TensorRT Engine
self.logger = trt.Logger(trt.Logger.WARNING)
with open(engine_path, "rb") as f, trt.Runtime(self.logger) as runtime:
self.engine = runtime.deserialize_cuda_engine(f.read())
self.context = self.engine.create_execution_context()
# 2. Allocate Host & Device Buffers
self.inputs, self.outputs, self.bindings, self.stream = self._allocate_buffers()
# 3. Initialize ByteTracker
tracker_args = Namespace(
track_high_thresh=0.5, track_low_thresh=0.1, new_track_thresh=0.6,
track_buffer=30, match_thresh=0.8, fuse_score=True
)
self.tracker = BYTETracker(tracker_args, frame_rate=30)
# 4. ROS 2 Subscriptions and Publishers
self.sub_img = self.create_subscription(Image, input_topic, self.image_callback, 10)
self.pub_detections = self.create_publisher(Detection2DArray, '/tracked_objects', 10)
self.pub_debug_img = self.create_publisher(Image, '/debug_image', 10)
self.get_logger().info("TensorRT Tracker Node Active.")
def _allocate_buffers(self):
inputs, outputs, bindings = [], [], []
stream = cuda.Stream()
for binding in self.engine:
size = trt.volume(self.engine.get_binding_shape(binding)) * self.engine.has_implicit_batch_dimension
dtype = trt.nptype(self.engine.get_binding_dtype(binding))
# Allocate pinned host and device memory
host_mem = cuda.pagelocked_empty(abs(size), dtype)
device_mem = cuda.mem_alloc(host_mem.nbytes)
bindings.append(int(device_mem))
if self.engine.binding_is_input(binding):
inputs.append({'host': host_mem, 'device': device_mem})
self.input_shape = self.engine.get_binding_shape(binding)
else:
outputs.append({'host': host_mem, 'device': device_mem})
return inputs, outputs, bindings, stream
def preprocess(self, img):
""" Resize and normalize image into NCHW format """
h, w, c = img.shape
# Input size expected by YOLO (e.g., 640x640)
input_h, input_w = self.input_shape[2], self.input_shape[3]
resized = cv2.resize(img, (input_w, input_h))
rgb = cv2.cvtColor(resized, cv2.COLOR_BGR2RGB)
# HWC to CHW -> Normalize 0-1
blob = rgb.transpose((2, 0, 1)).astype(np.float32) / 255.0
return np.ascontiguousarray(blob), h, w
def infer(self, blob):
""" Run TensorRT Inference Asynchronously """
np.copyto(self.inputs[0]['host'], blob.ravel())
# Transfer input data to GPU
cuda.memcpy_htod_async(self.inputs[0]['device'], self.inputs[0]['host'], self.stream)
# Execute Engine
self.context.execute_async_v2(bindings=self.bindings, stream_handle=self.stream.handle)
# Transfer predictions back to Host
cuda.memcpy_dtoh_async(self.outputs[0]['host'], self.outputs[0]['device'], self.stream)
self.stream.synchronize()
return self.outputs[0]['host']
def image_callback(self, msg):
# Convert ROS Image to OpenCV Mat
frame = self.bridge.imgmsg_to_cv2(msg, desired_encoding='bgr8')
blob, orig_h, orig_w = self.preprocess(frame)
# Run Engine
raw_output = self.infer(blob)
# Post-Processing & Tracking Logic
# (YOLOv8 tensor shape typically [1, 84, 8400] -> Reshape to [8400, 84])
output = raw_output.reshape((84, 8400)).T
bboxes = output[:, :4]
scores = np.max(output[:, 4:], axis=1)
class_ids = np.argmax(output[:, 4:], axis=1)
# Filter low confidence
mask = scores > self.conf_thresh
bboxes = bboxes[mask]
scores = scores[mask]
class_ids = class_ids[mask]
if len(bboxes) == 0:
return
# Scale bboxes back to original image size
scale_x = orig_w / self.input_shape[3]
scale_y = orig_h / self.input_shape[2]
detections_for_tracker = []
for box, score, cls_id in zip(bboxes, scores, class_ids):
xc, yc, w, h = box
x1 = (xc - w / 2) * scale_x
y1 = (yc - h / 2) * scale_y
x2 = (xc + w / 2) * scale_x
y2 = (yc + h / 2) * scale_y
detections_for_tracker.append([x1, y1, x2, y2, score, cls_id])
# Feed to ByteTracker
online_targets = self.tracker.update(
np.array(detections_for_tracker), [orig_h, orig_w], (orig_h, orig_w)
)
# Build ROS2 Detection2DArray Message
det_array = Detection2DArray()
det_array.header = msg.header
for t in online_targets:
tlwh = t.tlwh
track_id = t.track_id
det = Detection2D()
det.header = msg.header
det.bbox.center.position.x = float(tlwh[0] + tlwh[2] / 2)
det.bbox.center.position.y = float(tlwh[1] + tlwh[3] / 2)
det.bbox.size_x = float(tlwh[2])
det.bbox.size_y = float(tlwh[3])
hyp = ObjectHypothesisWithPose()
hyp.hypothesis.class_id = str(t.cls)
hyp.hypothesis.score = float(t.score)
det.results.append(hyp)
det_array.detections.append(det)
# Draw visual debug boxes
cv2.rectangle(frame, (int(tlwh[0]), int(tlwh[1])),
(int(tlwh[0] + tlwh[2]), int(tlwh[1] + tlwh[3])), (0, 255, 0), 2)
cv2.putText(frame, f"ID: {track_id}", (int(tlwh[0]), int(tlwh[1] - 10)),
cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0, 255, 0), 2)
# Publish Outputs
self.pub_detections.publish(det_array)
self.pub_debug_img.publish(self.bridge.cv2_to_imgmsg(frame, encoding='bgr8'))
def main(args=None):
rclpy.init(args=args)
node = TensorRTROS2Tracker()
try:
rclpy.spin(node)
except KeyboardInterrupt:
pass
finally:
node.destroy_node()
rclpy.shutdown()
if __name__ == '__main__':
main()
Jetson-Specific Optimization Checklist
To squeeze the lowest latency and highest throughput out of this pipeline on NVIDIA Jetson hardware:
- Max Out Hardware Clocks: Run the Jetson power management script before executing ROS 2 nodes: Bash
sudo nvpmodel -m 0 # Set board to max performance mode sudo jetson_clocks # Lock CPU/GPU clock frequencies to maximum - Zero-Copy Memory Pipeline: For zero-copy streaming without conversion CPU usage, consume frame memory using
gstreamerhardware decoders (nvv4l2decoder/nvvidconv) directly from your camera node. - Multi-Threading: If tracking frames at $>60\text{ FPS}$, separate the
image_callbackimage reception and TensorRT inference loop using ROS 2ReentrantCallbackGroupand aMultiThreadedExecutor.