24 yolo node for instance segmentation

.devcontainer/devcontainer.json

@@ -1,6 +1,6 @@
 {
-    "name": "vortex-image-segmentation-devcontainer",
-    "image": "vortex-image-segmentation:latest",
+    "name": "vortex-deep-learning-pipelines-devcontainer",
+    "image": "vortex-deep-learning-pipelines:latest",
     "customizations": {
         "vscode": {
             "settings": {

.dockerignore

@@ -27,4 +27,4 @@ venv/
 *.swp
 *.swo
 node_modules/
-bags/
+rosbags/

.gitignore

@@ -26,6 +26,18 @@ share/python-wheels/
 *.egg
 MANIFEST
 
+#VSCode
+.vscode/
+
+# OS Generated files
+.DS_Store
+.DS_Store?
+._*
+.Spotlight-V100
+.Trashes
+ehthumbs.db
+Thumbs.db
+
 # PyInstaller
 #  Usually these files are written by a python script from a template
 #  before PyInstaller builds the exe, so as to inject date/other infos into it.
@@ -173,7 +185,7 @@ build/
 log/
 
 # data
-data/
+rosbags/
 
 # Training outputs and Roboflow datasets
 results/

vortex_image_segmentation/LICENSE

@@ -0,0 +1,17 @@
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in
+all copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL
+THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
+THE SOFTWARE.

vortex_image_segmentation/launch/yolo_segmentation.launch.py

@@ -0,0 +1,19 @@
+from launch import LaunchDescription
+from launch_ros.actions import Node
+from ament_index_python.packages import get_package_share_directory
+import os
+
+
+def generate_launch_description():
+    pkg_share = get_package_share_directory('vortex_image_segmentation')
+    yolo_params = os.path.join(pkg_share, 'params', 'yolo_params.yaml')
+
+    return LaunchDescription([
+        Node(
+            package='vortex_image_segmentation',
+            executable='yolo_seg_node',
+            name='yolo_segmentation_node',
+            output='screen',
+            parameters=[yolo_params]
+        )
+    ])

vortex_image_segmentation/package.xml

@@ -0,0 +1,21 @@
+<?xml version="1.0"?>
+<?xml-model href="http://download.ros.org/schema/package_format3.xsd" schematypens="http://www.w3.org/2001/XMLSchema"?>
+<package format="3">
+  <name>vortex_image_segmentation</name>
+  <version>0.0.0</version>
+  <description>Package for image segmentation nodes</description>
+  <maintainer email="[email protected]">Mads Engesvoll</maintainer>
+  <license>MIT</license>
+
+  <depend>rclpy</depend>
+  <depend>std_msgs</depend>
+
+  <test_depend>ament_copyright</test_depend>
+  <test_depend>ament_flake8</test_depend>
+  <test_depend>ament_pep257</test_depend>
+  <test_depend>python3-pytest</test_depend>
+
+  <export>
+    <build_type>ament_python</build_type>
+  </export>
+</package>

vortex_image_segmentation/params/yolo_params.yaml

@@ -0,0 +1,19 @@
+yolo_segmentation_node:
+  ros__parameters:
+    # Node parameters
+    input_topic: "/gripper_camera/image_raw"
+    output_bbox_topic: "/segmentation/bboxes"
+    output_mask_topic: "/segmentation/mask"
+    debug_topic: "/image_debug"
+    pub_bbox: True
+    pub_mask: True
+    pub_debug: True
+
+    # Implementation parameters
+    model_path: "/ros2_ws/yolo_segmentation_training/results/custom_yolov89/weights/last.pt"
+    device: "cpu"
+    imgsz: 640
+    confidence_threshold: 0.3
+    max_detections: 1
+    compile: True
+    retina_masks: False # We dont need high res masks

vortex_image_segmentation/setup.cfg

@@ -0,0 +1,4 @@
+[develop]
+script_dir=$base/lib/vortex_image_segmentation
+[install]
+install_scripts=$base/lib/vortex_image_segmentation

vortex_image_segmentation/setup.py

@@ -0,0 +1,30 @@
+
+from setuptools import find_packages, setup
+import glob
+
+package_name = 'vortex_image_segmentation'
+
+setup(
+    name=package_name,
+    version='0.0.0',
+    packages=find_packages(exclude=['test']),
+    data_files=[
+        ('share/ament_index/resource_index/packages',
+            ['resource/' + package_name]),
+        ('share/' + package_name, ['package.xml']),
+        ('share/' + package_name + '/launch', glob.glob('launch/*.launch.py')),
+        ('share/' + package_name + '/params', glob.glob('params/*.yaml')),
+    ],
+    install_requires=['setuptools'],
+    zip_safe=True,
+    maintainer='mads',
+    maintainer_email='[email protected]',
+    description='ROS 2 package that provides a YOLO-based instance segmentation node (yolo_seg_node) for real-time segmentation.',
+    license='MIT',
+    tests_require=['pytest'],
+    entry_points={
+        'console_scripts': [
+            'yolo_seg_node = vortex_image_segmentation.yolo_seg_node:main',
+        ],
+    },
+)

vortex_image_segmentation/vortex_image_segmentation/yolo_seg.py

@@ -0,0 +1,78 @@
+"""
+YOLO segmentation model wrapper for inference and visualization using Ultralytics.
+Defines parameter dataclass and main segmentation class.
+"""
+
+from dataclasses import dataclass
+from typing import List
+
+import numpy as np
+import torch
+from ultralytics import YOLO
+from ultralytics.engine.results import Results
+
+
+@dataclass
+class YoloSegmentationParams:
+    """
+    Dataclass for storing YOLO segmentation parameters.
+    """
+    model_path: str
+    device: str
+    confidence_threshold: float
+    max_detections: int
+    imgsz: int
+    compile: bool
+    retina_masks: bool
+
+
+class YoloSegmentation:
+    """
+    Wrapper class for YOLO segmentation model inference and visualization.
+    """
+
+    def __init__(self, params: YoloSegmentationParams) -> None:
+        """
+        Initialize the YOLO segmentation model with given parameters.
+        Args:
+            params (YoloSegmentationParams): Parameters for YOLO segmentation.
+        """
+        self.params: YoloSegmentationParams = params
+        self.model: YOLO = self._load_model()
+
+    def _load_model(self) -> YOLO:
+        """
+        Load the YOLO segmentation model.
+        Returns:
+            YOLO: Ultralytics YOLO model instance.
+        """
+        return YOLO(self.params.model_path, task="segment")
+
+    def predict(self, cv_image: np.ndarray) -> List[Results]:
+        """
+        Run prediction on an input image and return the Results object(s).
+        Args:
+            cv_image (np.ndarray): Input image in OpenCV format.
+        Returns:
+            List[Results]: Ultralytics Results object(s) for the prediction.
+        """
+        results = self.model.predict(
+            source=cv_image,
+            imgsz=self.params.imgsz,
+            conf=self.params.confidence_threshold,
+            device=torch.device(self.params.device),
+            max_det=self.params.max_detections,
+            retina_masks=self.params.retina_masks,
+            # compile=self.params.compile,
+        )
+        return results
+
+    def visualize(self, result: Results) -> np.ndarray:
+        """
+        Generate a visualization image from the Results object.
+        Args:
+            result (Results): Ultralytics Results object from prediction.
+        Returns:
+            np.ndarray: Visualization image with segmentation overlays.
+        """
+        return result.plot()