Add functions to publish segmentation masks, and update parameter

README.md

@@ -1,11 +1,11 @@
 # ultralytics_ros [![ROS2-humble Industrial CI](https://github.com/Alpaca-zip/ultralytics_ros/actions/workflows/humble-ci.yml/badge.svg)](https://github.com/Alpaca-zip/ultralytics_ros/actions/workflows/humble-ci.yml) [![ROS2-humble Docker Build Check](https://github.com/Alpaca-zip/ultralytics_ros/actions/workflows/humble-docker-build-check.yml/badge.svg)](https://github.com/Alpaca-zip/ultralytics_ros/actions/workflows/humble-docker-build-check.yml)
-ROS2 package for real-time object detection using the Ultralytics YOLO, enabling flexible integration with various robotics applications.
+ROS 2 package for real-time object detection using the Ultralytics YOLO, enabling flexible integration with various robotics applications.
 
 |  `tracker_node`  |  `tracker_with_cloud_node`  |
 | :------------: | :-----------------------: |
 | <img src="https://github.com/Alpaca-zip/ultralytics_ros/assets/84959376/9da7dbbf-5cc0-41bc-be82-d481abbf552a" width="450px"> | **TODO** |
 
-- The `tracker_node` provides real-time object detection on incoming ROS image messages using the Ultralytics YOLO model.
+- The `tracker_node` provides real-time object detection on incoming ROS 2 image messages using the Ultralytics YOLO model.
 
 ## Setup ⚙
 ```
@@ -24,37 +24,51 @@ $ ros2 launch ultralytics_ros tracker.launch.xml debug:=true
 ```
 ## `tracker_node`
 ### Params
+### Params
 - `yolo_model`: Pre-trained Weights.  
-For yolov8, you can choose `yolov8*.pt`, `yolov8*-seg.pt`, `yolov8*-pose.pt`.
+For yolov8, you can choose `yolov8*.pt`, `yolov8*-seg.pt`.
 
   |  YOLOv8  |  <img src="https://github.com/Alpaca-zip/ultralytics_ros/assets/84959376/08770080-bf20-470b-8269-eee7a7c41acc" width="350px">  |
   | :-------------: | :-------------: |
   |  **YOLOv8-seg**  |  <img src="https://github.com/Alpaca-zip/ultralytics_ros/assets/84959376/7bb6650c-769d-41c1-86f7-39fcbf01bc7c" width="350px">  |
-  |  **YOLOv8-pose**  |  <img src="https://github.com/Alpaca-zip/ultralytics_ros/assets/84959376/46d2a5ef-193b-4f83-a0b3-6cc0d5a3756c" width="350px">  |
 
   See also: https://docs.ultralytics.com/models/
-- `publish_rate`: Publish rate for `image_topic`.
-- `image_topic`: Topic name for image.
-- `detection_topic`: Topic name for 2D bounding box.
+- `input_topic`: Topic name for input image.
+- `result_topic`: Topic name of the custom message containing the 2D bounding box and the mask image.
+- `result_image_topic`: Topic name of the image on which the detection and segmentation results are plotted.
 - `conf_thres`: Confidence threshold below which boxes will be filtered out.
 - `iou_thres`: IoU threshold below which boxes will be filtered out during NMS.
 - `max_det`: Maximum number of boxes to keep after NMS.
 - `tracker`: Tracking algorithms.
-- `classes`: List of class indices to consider.  
-See also: https://github.com/ultralytics/ultralytics/blob/main/ultralytics/datasets/coco128.yaml 
-- `debug`:  If true, run simple viewer.
-- `debug_conf`:  Whether to plot the detection confidence score.
-- `debug_line_width`: Line width of the bounding boxes.
-- `debug_font_size`: Font size of the text.
-- `debug_labels`: Font to use for the text.
-- `debug_font`: Whether to plot the label of bounding boxes.
-- `debug_boxes`: Whether to plot the bounding boxes.
+- `device`: Device to run the model on(e.g. cpu or cuda:0).
+  ```xml
+  <arg name="device" default="cpu"/>
+  ```
+  ```xml
+  <arg name="device" default="cuda:0"/>
+  ```
+- `classes`: List of class indices to consider.
+  ```xml
+  <param name="classes" value="0, 1" value-sep=", "/> <!-- person, bicycle -->
+  ```
+  See also: https://github.com/ultralytics/ultralytics/blob/main/ultralytics/cfg/datasets/coco128.yaml 
+- `result_conf`:  Whether to plot the detection confidence score.
+- `result_line_width`: Line width of the bounding boxes.
+- `result_font_size`: Font size of the text.
+- `result_labels`: Font to use for the text.
+- `result_font`: Whether to plot the label of bounding boxes.
+- `result_boxes`: Whether to plot the bounding boxes.
 ### Topics
 - Subscribed Topics:
-  - Image data from `image_topic` parameter. ([sensor_msgs/Image](https://docs.ros.org/en/api/sensor_msgs/html/msg/Image.html))
+  - Image data from `input_topic` parameter. ([sensor_msgs/Image](https://github.com/ros2/common_interfaces/blob/humble/sensor_msgs/msg/Image.msg))
 - Published Topics:
-  - Debug images to `/debug_image` topic. ([sensor_msgs/Image](https://docs.ros.org/en/api/sensor_msgs/html/msg/Image.html))
-  - Detected objects(2D bounding box) to `detection_topic` parameter. (vision_msgs/Detection2DArray)
+  - Plotted images to `result_image_topic` parameter. ([sensor_msgs/Image](https://github.com/ros2/common_interfaces/blob/humble/sensor_msgs/msg/Image.msg))
+  - Detected objects(2D bounding box, mask image) to `result_topic` parameter. (ultralytics_ros/YoloResult)
+    ```
+    std_msgs/Header header
+    vision_msgs/Detection2DArray detections
+    sensor_msgs/Image[] masks
+    ```
 ## `tracker_with_cloud_node`
 **TODO**
 ## Docker with KITTI datasets 🐳
@@ -65,7 +79,7 @@ See also: https://github.com/ultralytics/ultralytics/blob/main/ultralytics/datas
 $ docker pull alpacazip/ultralytics_ros:humble
 $ docker run -p 6080:80 --shm-size=512m alpacazip/ultralytics_ros:humble
 ```
-### Run tracker_node & tracker_with_cloud_node
+### Run tracker_node
 ```
 $ ros2 launch ultralytics_ros kitti_tracker.launch.xml
 $ cd ~/colcon_ws/src/ultralytics_ros/ros2bag && ros2 bag play kitti_2011_09_26_drive_0106_synced --clock --loop

launch/kitti_tracker.launch.xml

@@ -1,41 +1,42 @@
 <?xml version="1.0" encoding="UTF-8"?>
 <launch>
-  <arg name="yolo_model" default="yolov8n.pt"/>
-  <arg name="publish_rate" default="10"/>
-  <arg name="detection_topic" default="detection_result"/>
-  <arg name="image_topic" default="kitti/camera_color_left/image_raw"/>
+  <arg name="debug" default="true"/>
+  <arg name="yolo_model" default="yolov8m-seg.pt"/>
+  <arg name="input_topic" default="/kitti/camera_color_left/image_raw"/>
+  <arg name="result_topic" default="/yolo_result"/>
+  <arg name="result_image_topic" default="/yolo_image"/>
   <arg name="conf_thres" default="0.25"/>
   <arg name="iou_thres" default="0.45"/>
   <arg name="max_det" default="300"/>
   <arg name="tracker" default="bytetrack.yaml"/>
-  <arg name="debug" default="true"/>
-  <arg name="debug_conf" default="true"/>
-  <arg name="debug_line_width" default="1"/>
-  <arg name="debug_font_size" default="1"/>
-  <arg name="debug_font" default="Arial.ttf"/>
-  <arg name="debug_labels" default="true"/>
-  <arg name="debug_boxes" default="true"/>
-  <!-- Object detection node -->
+  <arg name="device" default="cuda:0"/>
+  <arg name="result_conf" default="true"/>
+  <arg name="result_line_width" default="1"/>
+  <arg name="result_font_size" default="1"/>
+  <arg name="result_font" default="Arial.ttf"/>
+  <arg name="result_labels" default="true"/>
+  <arg name="result_boxes" default="true"/>
+
   <node pkg="ultralytics_ros" exec="tracker_node.py" output="screen">
     <param name="yolo_model" value="$(var yolo_model)"/>
-    <param name="publish_rate" value="$(var publish_rate)"/>
-    <param name="detection_topic" value="$(var detection_topic)"/>
-    <param name="image_topic" value="$(var image_topic)"/>
+    <param name="input_topic" value="$(var input_topic)"/>
+    <param name="result_topic" value="$(var result_topic)"/>
+    <param name="result_image_topic" value="$(var result_image_topic)"/>
     <param name="conf_thres" value="$(var conf_thres)"/>
     <param name="iou_thres" value="$(var iou_thres)"/>
     <param name="max_det" value="$(var max_det)"/>
     <param name="tracker" value="$(var tracker)"/>
-    <param name="debug" value="$(var debug)"/>
-    <param name="debug_conf" value="$(var debug_conf)"/>
-    <param name="debug_line_width" value="$(var debug_line_width)"/>
-    <param name="debug_font_size" value="$(var debug_font_size)"/>
-    <param name="debug_font" value="$(var debug_font)"/>
-    <param name="debug_labels" value="$(var debug_labels)"/>
-    <param name="debug_boxes" value="$(var debug_boxes)"/>
+    <param name="result_conf" value="$(var result_conf)"/>
+    <param name="result_line_width" value="$(var result_line_width)"/>
+    <param name="result_font_size" value="$(var result_font_size)"/>
+    <param name="result_font" value="$(var result_font)"/>
+    <param name="result_labels" value="$(var result_labels)"/>
+    <param name="result_boxes" value="$(var result_boxes)"/>
     <!--<param name="classes" value="0, 1, 2" value-sep=", "/>-->
+    <param name="device" value="$(var device)"/>
   </node>
-  <!-- Image view -->
+
   <node if="$(var debug)" pkg="image_view" exec="image_view" output="screen">
-    <remap from="image" to="debug_image"/>
+    <remap from="image" to="$(var result_image_topic)"/>
   </node>
 </launch>

launch/tracker.launch.xml

@@ -1,41 +1,42 @@
 <?xml version="1.0" encoding="UTF-8"?>
 <launch>
-  <arg name="yolo_model" default="yolov8n.pt"/>
-  <arg name="publish_rate" default="10"/>
-  <arg name="detection_topic" default="detection_result"/>
-  <arg name="image_topic" default="image_raw"/>
+  <arg name="debug" default="false"/>
+  <arg name="yolo_model" default="yolov8m-seg.pt"/>
+  <arg name="input_topic" default="/image_raw"/>
+  <arg name="result_topic" default="/yolo_result"/>
+  <arg name="result_image_topic" default="/yolo_image"/>
   <arg name="conf_thres" default="0.25"/>
   <arg name="iou_thres" default="0.45"/>
   <arg name="max_det" default="300"/>
   <arg name="tracker" default="bytetrack.yaml"/>
-  <arg name="debug" default="false"/>
-  <arg name="debug_conf" default="true"/>
-  <arg name="debug_line_width" default="1"/>
-  <arg name="debug_font_size" default="1"/>
-  <arg name="debug_font" default="Arial.ttf"/>
-  <arg name="debug_labels" default="true"/>
-  <arg name="debug_boxes" default="true"/>
-  <!-- Object detection node -->
+  <arg name="device" default="cuda:0"/>
+  <arg name="result_conf" default="true"/>
+  <arg name="result_line_width" default="1"/>
+  <arg name="result_font_size" default="1"/>
+  <arg name="result_font" default="Arial.ttf"/>
+  <arg name="result_labels" default="true"/>
+  <arg name="result_boxes" default="true"/>
+
   <node pkg="ultralytics_ros" exec="tracker_node.py" output="screen">
     <param name="yolo_model" value="$(var yolo_model)"/>
-    <param name="publish_rate" value="$(var publish_rate)"/>
-    <param name="detection_topic" value="$(var detection_topic)"/>
-    <param name="image_topic" value="$(var image_topic)"/>
+    <param name="input_topic" value="$(var input_topic)"/>
+    <param name="result_topic" value="$(var result_topic)"/>
+    <param name="result_image_topic" value="$(var result_image_topic)"/>
     <param name="conf_thres" value="$(var conf_thres)"/>
     <param name="iou_thres" value="$(var iou_thres)"/>
     <param name="max_det" value="$(var max_det)"/>
     <param name="tracker" value="$(var tracker)"/>
-    <param name="debug" value="$(var debug)"/>
-    <param name="debug_conf" value="$(var debug_conf)"/>
-    <param name="debug_line_width" value="$(var debug_line_width)"/>
-    <param name="debug_font_size" value="$(var debug_font_size)"/>
-    <param name="debug_font" value="$(var debug_font)"/>
-    <param name="debug_labels" value="$(var debug_labels)"/>
-    <param name="debug_boxes" value="$(var debug_boxes)"/>
+    <param name="result_conf" value="$(var result_conf)"/>
+    <param name="result_line_width" value="$(var result_line_width)"/>
+    <param name="result_font_size" value="$(var result_font_size)"/>
+    <param name="result_font" value="$(var result_font)"/>
+    <param name="result_labels" value="$(var result_labels)"/>
+    <param name="result_boxes" value="$(var result_boxes)"/>
     <!--<param name="classes" value="0, 1, 2" value-sep=", "/>-->
+    <param name="device" value="$(var device)"/>
   </node>
-  <!-- Image view -->
+
   <node if="$(var debug)" pkg="image_view" exec="image_view" output="screen">
-    <remap from="image" to="debug_image"/>
+    <remap from="image" to="$(var result_image_topic)"/>
   </node>
 </launch>