This ROS2 package provides an integration of YOLONAS with TensorRT for high-performance object detection. The package allows flexibility in specifying input topics and model names through launch parameters.
Check out the following gifs demonstrating the YOLONAS TensorRT ROS2 package in action:
To quickly get started with the YOLO TensorRT ROS2 package, follow these steps:
-
Install prerequisites:
Before using this package, ensure you have the following dependencies installed:
- ROS2 Humble
- TensorRT
tensorrt==8.6.1.post1
- Pycuda
pycuda==2024.1
- OpenCV
- PIL (Python Imaging Library)
vision_msgs
cv_bridge
pip install tensorrt==8.6.1.post1 --extra-index-url https://pypi.nvidia.com
pip install pycuda==2024.1
pip install opencv-python
pip install pillow
sudo apt-get install ros-humble-vision-msgs
sudo apt-get install ros-humble-cv-bridge
-
Clone the repository and install dependencies:
cd ~/ros2_ws/src git clone https://github.com/djetshu/yolonas_trt_ros2.git cd ~/ros2_ws rosdep install --from-paths src --ignore-src -r -y
-
Build the package:
colcon build --packages-select yolo_py
-
Source the workspace:
source ~/ros2_ws/install/setup.bash
-
Download the TensorRT model and ROSBAG files from this link.
-
Move the YOLO NAS TRT model to the following directory:
mv /path/to/downloaded/yolo_nas_m_flower.trt ~/ros2_ws/src/yolonas_trt_ros2/models/
-
Modify the class names to your custom classes' names in
ros2_ws/src/yolonas_trt_ros2/yolo_py/dataset.py
. -
Launch the YOLO node with the custom model:
ros2 launch yolo_py yolo_trt.launch.py input_topic:=/oak/rgb/image_raw model_name:=yolo_nas_m_flower.trt
To use this package, you need to launch the yolo_trt
node. The package includes a launch file that allows you to specify the input topic and the model name.
rviz_config
: Path to the RViz config file. Default is set to the RViz config provided in the package.rviz_en
: Enable or disable RViz. Default istrue
.input_topic
: The input topic for YOLO. Default is/camera/image
.model_name
: The name of the YOLO model to use. Default isyolo_nas_m.trt
.
This node subscribes to an image topic, processes the image with YOLO using TensorRT, and publishes the detection results.
-
Subscribed Topics:
/image
(sensor_msgs/Image): Input image topic.
-
Published Topics:
/detection
(vision_msgs/Detection2DArray): Detection results./detection_image
(sensor_msgs/Image): Image with detection results.
To launch the YOLO node:
ros2 launch yolo_py yolo_trt.launch.py input_topic:=/your_custom_topic model_name:=your_custom_model.trt
For inquiries, collaboration opportunities, or questions feel free to contact:
- Email: [email protected]
Consult this previous repository to train YoloNAS and convert to TensorRT model.