refactor: gdino and gsam as agents (#501)

Author: rachwalk

src/rai_bringup/launch/openset.launch.py

@@ -14,36 +14,16 @@
 
 
 from launch import LaunchDescription
-from launch.actions import IncludeLaunchDescription
-from launch.launch_description_sources import AnyLaunchDescriptionSource
-from launch.substitutions import PathJoinSubstitution
-from launch_ros.substitutions import FindPackageShare
+from launch.actions import ExecuteProcess
 
 
 def generate_launch_description():
     return LaunchDescription(
         [
-            IncludeLaunchDescription(
-                AnyLaunchDescriptionSource(
-                    PathJoinSubstitution(
-                        [
-                            FindPackageShare("rai_open_set_vision"),
-                            "launch",
-                            "gdino_launch.xml",
-                        ]
-                    )
-                )
-            ),
-            IncludeLaunchDescription(
-                AnyLaunchDescriptionSource(
-                    PathJoinSubstitution(
-                        [
-                            FindPackageShare("rai_open_set_vision"),
-                            "launch",
-                            "gsam_launch.xml",
-                        ]
-                    )
-                )
+            ExecuteProcess(
+                cmd=["python", "run_vision_agents.py"],
+                cwd="src/rai_extensions/rai_open_set_vision/scripts",
+                output="screen",
             ),
         ]
     )

src/rai_core/rai/agents/__init__.py

@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from rai.agents.base import BaseAgent
 from rai.agents.conversational_agent import create_conversational_agent
 from rai.agents.react_agent import ReActAgent
 from rai.agents.runner import AgentRunner, wait_for_shutdown
@@ -20,6 +21,7 @@
 
 __all__ = [
     "AgentRunner",
+    "BaseAgent",
     "ReActAgent",
     "ToolRunner",
     "create_conversational_agent",

src/rai_core/rai/tools/ros2/generic/topics.py

@@ -135,11 +135,11 @@ def _run(
         msg_type = type(message.payload)
         if msg_type == Image:
             image = CvBridge().imgmsg_to_cv2(  # type: ignore
-                message.payload, desired_encoding="rgb8"
+                message.payload, desired_encoding="bgr8"
             )
         elif msg_type == CompressedImage:
             image = CvBridge().compressed_imgmsg_to_cv2(  # type: ignore
-                message.payload, desired_encoding="rgb8"
+                message.payload, desired_encoding="bgr8"
             )
         else:
             raise ValueError(

src/rai_extensions/rai_open_set_vision/rai_open_set_vision/__init__.py

@@ -13,12 +13,21 @@
 # limitations under the License.
 
 
-from .services.grounding_dino import GDINO_NODE_NAME, GDINO_SERVICE_NAME
+from .agents.grounded_sam import GSAM_NODE_NAME, GSAM_SERVICE_NAME, GroundedSamAgent
+from .agents.grounding_dino import (
+    GDINO_NODE_NAME,
+    GDINO_SERVICE_NAME,
+    GroundingDinoAgent,
+)
 from .tools import GetDetectionTool, GetDistanceToObjectsTool
 
 __all__ = [
     "GDINO_NODE_NAME",
     "GDINO_SERVICE_NAME",
+    "GSAM_NODE_NAME",
+    "GSAM_SERVICE_NAME",
     "GetDetectionTool",
     "GetDistanceToObjectsTool",
+    "GroundedSamAgent",
+    "GroundingDinoAgent",
 ]

src/rai_extensions/rai_open_set_vision/rai_open_set_vision/services/__init__.py renamed to src/rai_extensions/rai_open_set_vision/rai_open_set_vision/agents/__init__.py

@@ -11,3 +11,13 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+
+from .base_vision_agent import BaseVisionAgent
+from .grounded_sam import GroundedSamAgent
+from .grounding_dino import GroundingDinoAgent
+
+__all__ = [
+    "BaseVisionAgent",
+    "GroundedSamAgent",
+    "GroundingDinoAgent",
+]

src/rai_extensions/rai_open_set_vision/rai_open_set_vision/agents/base_vision_agent.py

@@ -0,0 +1,80 @@
+# Copyright (C) 2025 Robotec.AI
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#         http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+import os
+import subprocess
+from pathlib import Path
+
+from rai.agents import BaseAgent
+from rai.communication.ros2 import ROS2Connector
+
+
+class BaseVisionAgent(BaseAgent):
+    WEIGHTS_URL: str = ""
+    WEIGHTS_FILENAME: str = ""
+
+    def __init__(
+        self,
+        weights_path: str | Path = Path.home() / Path(".cache/rai/"),
+        ros2_name: str = "",
+    ):
+        super().__init__()
+        self._weights_path = Path(weights_path)
+        os.makedirs(self._weights_path, exist_ok=True)
+        self._init_weight_path()
+        self.weight_path = self._weights_path
+        self.ros2_connector = ROS2Connector(ros2_name)
+
+    def _init_weight_path(self):
+        try:
+            if self.WEIGHTS_FILENAME == "":
+                raise ValueError("WEIGHTS_FILENAME is not set")
+
+            install_path = (
+                self._weights_path / "vision" / "weights" / self.WEIGHTS_FILENAME
+            )
+            # make sure the file exists
+            if install_path.exists():
+                self._weights_path = install_path
+            else:
+                self._download_weights(install_path)
+                self._weights_path = install_path
+
+        except Exception:
+            self.logger.error("Could not find package path")
+            raise Exception("Could not find package path")
+
+    def _download_weights(self, path: Path):
+        try:
+            os.makedirs(path.parent, exist_ok=True)
+            subprocess.run(
+                [
+                    "wget",
+                    self.WEIGHTS_URL,
+                    "-O",
+                    path,
+                    "--progress=dot:giga",
+                ]
+            )
+        except Exception:
+            self.logger.error("Could not download weights")
+            raise Exception("Could not download weights")
+
+    def _remove_weights(self, path: Path):
+        if path.exists():
+            os.remove(path)
+
+    def stop(self):
+        self.ros2_connector.shutdown()

src/rai_extensions/rai_open_set_vision/rai_open_set_vision/agents/grounded_sam.py

@@ -0,0 +1,76 @@
+# Copyright (C) 2024 Robotec.AI
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#         http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+from pathlib import Path
+
+import numpy as np
+from cv_bridge import CvBridge
+
+from rai_interfaces.srv import RAIGroundedSam
+from rai_open_set_vision.agents.base_vision_agent import BaseVisionAgent
+from rai_open_set_vision.vision_markup.segmenter import GDSegmenter
+
+GSAM_NODE_NAME = "grounded_sam"
+GSAM_SERVICE_NAME = "grounded_sam_segment"
+
+
+class GroundedSamAgent(BaseVisionAgent):
+    WEIGHTS_URL = (
+        "https://dl.fbaipublicfiles.com/segment_anything_2/072824/sam2_hiera_large.pt"
+    )
+    WEIGHTS_FILENAME = "sam2_hiera_large.pt"
+
+    def __init__(
+        self,
+        weights_path: str | Path = Path.home() / Path(".cache/rai"),
+        ros2_name: str = GSAM_NODE_NAME,
+    ):
+        super().__init__(weights_path, ros2_name)
+        try:
+            self._segmenter = GDSegmenter(self._weights_path)
+        except Exception:
+            self.logger.error(
+                "Could not load model. The weights might be corrupted. Redownloading..."
+            )
+            self._remove_weights(self.weight_path)
+            self._init_weight_path()
+            self.segmenter = GDSegmenter(self.weight_path)
+
+    def run(self):
+        self.ros2_connector.create_service(
+            service_name=GSAM_SERVICE_NAME,
+            on_request=self._segment_callback,
+            service_type="rai_interfaces/srv/RAIGroundedSam",
+        )
+
+    def _segment_callback(self, request, response: RAIGroundedSam.Response):
+        received_boxes = []
+        for detection in request.detections.detections:
+            received_boxes.append(detection.bbox)
+
+        image = request.source_img
+
+        assert self._segmenter is not None
+        masks = self._segmenter.get_segmentation(image, received_boxes)
+        bridge = CvBridge()
+        img_arr = []
+        for mask in masks:
+            if len(mask.shape) > 2:  # Check if the mask has multiple channels
+                mask = np.squeeze(mask)
+            arr = (mask * 255).astype(np.uint8)  # Convert binary 0/1 to 0/255
+            img_arr.append(bridge.cv2_to_imgmsg(arr, encoding="mono8"))
+
+        response.masks = img_arr
+        return response

src/rai_extensions/rai_open_set_vision/rai_open_set_vision/agents/grounding_dino.py

@@ -0,0 +1,77 @@
+# Copyright (C) 2024 Robotec.AI
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#         http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+from pathlib import Path
+
+from rai_interfaces.msg import RAIDetectionArray
+from rai_open_set_vision.agents.base_vision_agent import BaseVisionAgent
+from rai_open_set_vision.vision_markup.boxer import GDBoxer
+
+GDINO_NODE_NAME = "grounding_dino"
+GDINO_SERVICE_NAME = "grounding_dino_classify"
+
+
+class GroundingDinoAgent(BaseVisionAgent):
+    WEIGHTS_URL = "https://github.com/IDEA-Research/GroundingDINO/releases/download/v0.1.0-alpha/groundingdino_swint_ogc.pth"
+    WEIGHTS_FILENAME = "groundingdino_swint_ogc.pth"
+
+    def __init__(
+        self,
+        weights_path: str | Path = Path.home() / Path(".cache/rai"),
+        ros2_name: str = GDINO_NODE_NAME,
+    ):
+        super().__init__(weights_path, ros2_name)
+        try:
+            self._boxer = GDBoxer(self._weights_path)
+        except Exception:
+            self.logger.error(
+                "Could not load model. The weights might be corrupted. Redownloading..."
+            )
+            self._remove_weights(self.weight_path)
+            self._init_weight_path()
+            self.segmenter = GDBoxer(self.weight_path)
+
+    def run(self):
+        self.ros2_connector.create_service(
+            GDINO_SERVICE_NAME,
+            self._classify_callback,
+            service_type="rai_interfaces/srv/RAIGroundingDino",
+        )
+
+    def _classify_callback(self, request, response: RAIDetectionArray):
+        self.logger.info(
+            f"Request received: {request.classes}, {request.box_threshold}, {request.text_threshold}"
+        )
+
+        class_array = request.classes.split(",")
+        class_array = [class_name.strip() for class_name in class_array]
+        class_dict = {class_name: i for i, class_name in enumerate(class_array)}
+
+        boxes = self._boxer.get_boxes(
+            request.source_img,
+            class_array,
+            request.box_threshold,
+            request.text_threshold,
+        )
+
+        ts = self.ros2_connector._node.get_clock().now().to_msg()
+        response.detections.detections = [  # type: ignore
+            box.to_detection_msg(class_dict, ts)  # type: ignore
+            for box in boxes
+        ]
+        response.detections.header.stamp = ts  # type: ignore
+        response.detections.detection_classes = class_array  # type: ignore
+
+        return response