feat: unified ROS2 hri message (#416)

Further work: #418

Author: maciejmajek

src/rai_core/rai/communication/ros2/api.py

@@ -319,8 +319,7 @@ def shutdown(self) -> None:
 
 @dataclass
 class TopicConfig:
-    name: str
-    msg_type: str
+    msg_type: str = "rai_interfaces/msg/HRIMessage"
     auto_qos_matching: bool = True
     qos_profile: Optional[QoSProfile] = None
     is_subscriber: bool = False

src/rai_core/rai/communication/ros2/connectors.py

@@ -15,18 +15,26 @@
 import threading
 import time
 import uuid
-from typing import Any, Callable, Dict, List, Literal, Optional, Tuple
+from collections import OrderedDict
+from typing import Any, Callable, Dict, List, Literal, Optional, Tuple, Union, cast
 
+import numpy as np
 import rclpy
 import rclpy.executors
 import rclpy.node
 import rclpy.time
+import rosidl_runtime_py.convert
+from cv_bridge import CvBridge
+from PIL import Image
+from pydub import AudioSegment
 from rclpy.duration import Duration
 from rclpy.executors import MultiThreadedExecutor
 from rclpy.node import Node
 from rclpy.qos import QoSProfile
+from sensor_msgs.msg import Image as ROS2Image
 from tf2_ros import Buffer, LookupException, TransformListener, TransformStamped
 
+import rai_interfaces.msg
 from rai.communication import (
     ARIConnector,
     ARIMessage,
@@ -41,6 +49,10 @@
     ROS2TopicAPI,
     TopicConfig,
 )
+from rai_interfaces.msg import HRIMessage as ROS2HRIMessage_
+from rai_interfaces.msg._audio_message import (
+    AudioMessage as ROS2HRIMessage__Audio,
+)
 
 
 class ROS2ARIMessage(ARIMessage):
@@ -200,26 +212,95 @@ class ROS2HRIMessage(HRIMessage):
     def __init__(self, payload: HRIPayload, message_author: Literal["ai", "human"]):
         super().__init__(payload, message_author)
 
+    @classmethod
+    def from_ros2(
+        cls, msg: rai_interfaces.msg.HRIMessage, message_author: Literal["ai", "human"]
+    ):
+        cv_bridge = CvBridge()
+        images = [
+            cv_bridge.imgmsg_to_cv2(img_msg, "rgb8")
+            for img_msg in cast(List[ROS2Image], msg.images)
+        ]
+        pil_images = [Image.fromarray(img) for img in images]
+        audio_segments = [
+            AudioSegment(
+                data=audio_msg.audio,
+                frame_rate=audio_msg.sample_rate,
+                sample_width=2,  # bytes, int16
+                channels=audio_msg.channels,
+            )
+            for audio_msg in msg.audios
+        ]
+        return ROS2HRIMessage(
+            payload=HRIPayload(text=msg.text, images=pil_images, audios=audio_segments),
+            message_author=message_author,
+        )
+
+    def to_ros2_dict(self) -> OrderedDict[str, Any]:
+        cv_bridge = CvBridge()
+        assert isinstance(self.payload, HRIPayload)
+        img_msgs = [
+            cv_bridge.cv2_to_imgmsg(np.array(img), "rgb8")
+            for img in self.payload.images
+        ]
+        audio_msgs = [
+            ROS2HRIMessage__Audio(
+                audio=audio.raw_data,
+                sample_rate=audio.frame_rate,
+                channels=audio.channels,
+            )
+            for audio in self.payload.audios
+        ]
+
+        return cast(
+            OrderedDict[str, Any],
+            rosidl_runtime_py.convert.message_to_ordereddict(
+                ROS2HRIMessage_(
+                    text=self.payload.text,
+                    images=img_msgs,
+                    audios=audio_msgs,
+                )
+            ),
+        )
+
 
 class ROS2HRIConnector(HRIConnector[ROS2HRIMessage]):
     def __init__(
         self,
         node_name: str = f"rai_ros2_hri_connector_{str(uuid.uuid4())[-12:]}",
-        targets: List[Tuple[str, TopicConfig]] = [],
-        sources: List[Tuple[str, TopicConfig]] = [],
+        targets: List[Union[str, Tuple[str, TopicConfig]]] = [],
+        sources: List[Union[str, Tuple[str, TopicConfig]]] = [],
     ):
-        configured_targets = [target[0] for target in targets]
-        configured_sources = [source[0] for source in sources]
+        configured_targets = [
+            target[0] if isinstance(target, tuple) else target for target in targets
+        ]
+        configured_sources = [
+            source[0] if isinstance(source, tuple) else source for source in sources
+        ]
 
-        self._configure_publishers(targets)
-        self._configure_subscribers(sources)
+        _targets = [
+            target
+            if isinstance(target, tuple)
+            else (target, TopicConfig(is_subscriber=False))
+            for target in targets
+        ]
+        _sources = [
+            source
+            if isinstance(source, tuple)
+            else (source, TopicConfig(is_subscriber=True))
+            for source in sources
+        ]
 
-        super().__init__(configured_targets, configured_sources)
         self._node = Node(node_name)
         self._topic_api = ConfigurableROS2TopicAPI(self._node)
         self._service_api = ROS2ServiceAPI(self._node)
         self._actions_api = ROS2ActionAPI(self._node)
 
+        self._configure_publishers(_targets)
+        self._configure_subscribers(_sources)
+
+        super().__init__(configured_targets, configured_sources)
+
         self._executor = MultiThreadedExecutor()
         self._executor.add_node(self._node)
         self._thread = threading.Thread(target=self._executor.spin)
@@ -236,7 +317,7 @@ def _configure_subscribers(self, sources: List[Tuple[str, TopicConfig]]):
     def send_message(self, message: ROS2HRIMessage, target: str, **kwargs):
         self._topic_api.publish_configured(
             topic=target,
-            msg_content=message.payload,
+            msg_content=message.to_ros2_dict(),
         )
 
     def receive_message(
@@ -249,16 +330,12 @@ def receive_message(
         auto_topic_type: bool = True,
         **kwargs: Any,
     ) -> ROS2HRIMessage:
-        if msg_type != "std_msgs/msg/String":
-            raise ValueError("ROS2HRIConnector only supports receiving sting messages")
         msg = self._topic_api.receive(
             topic=source,
             timeout_sec=timeout_sec,
-            msg_type=msg_type,
             auto_topic_type=auto_topic_type,
         )
-        payload = HRIPayload(msg.data)
-        return ROS2HRIMessage(payload=payload, message_author=message_author)
+        return ROS2HRIMessage.from_ros2(msg, message_author)
 
     def service_call(
         self, message: ROS2HRIMessage, target: str, timeout_sec: float, **kwargs: Any
@@ -284,3 +361,10 @@ def terminate_action(self, action_handle: str, **kwargs: Any):
         raise NotImplementedError(
             f"{self.__class__.__name__} doesn't support action calls"
         )
+
+    def shutdown(self):
+        self._executor.shutdown()
+        self._thread.join()
+        self._actions_api.shutdown()
+        self._topic_api.shutdown()
+        self._node.destroy_node()

src/rai_interfaces/CMakeLists.txt

@@ -21,6 +21,8 @@ rosidl_generate_interfaces(${PROJECT_NAME}
   "srv/VectorStoreRetrieval.srv"
   "srv/StringList.srv"
   "msg/RAIDetectionArray.msg"
+  "msg/AudioMessage.msg"
+  "msg/HRIMessage.msg"
   "srv/RAIGroundingDino.srv"
   "srv/RAIGroundedSam.srv"
   "action/Task.action"

src/rai_interfaces/msg/AudioMessage.msg

@@ -0,0 +1,19 @@
+#
+# 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.
+#
+
+int16[] audio
+uint16 sample_rate
+uint16 channels

src/rai_interfaces/msg/HRIMessage.msg

@@ -0,0 +1,20 @@
+#
+# 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.
+#
+
+std_msgs/Header header
+string text
+sensor_msgs/Image[] images
+rai_interfaces/AudioMessage[] audios

tests/communication/ros2/helpers.py

@@ -19,7 +19,9 @@
 import numpy as np
 import pytest
 import rclpy
+from cv_bridge import CvBridge
 from nav2_msgs.action import NavigateToPose
+from pydub import AudioSegment
 from rclpy.action import ActionServer, CancelResponse, GoalResponse
 from rclpy.action.server import ServerGoalHandle
 from rclpy.callback_groups import ReentrantCallbackGroup
@@ -30,6 +32,36 @@
 from std_srvs.srv import SetBool
 from tf2_ros import TransformBroadcaster, TransformStamped
 
+from rai_interfaces.msg import HRIMessage
+
+
+class HRIMessagePublisher(Node):
+    def __init__(self, topic: str):
+        super().__init__("test_hri_message_publisher")
+        self.publisher = self.create_publisher(HRIMessage, topic, 10)
+        self.timer = self.create_timer(0.1, self.publish_message)
+        self.cv_bridge = CvBridge()
+
+    def publish_message(self) -> None:
+        msg = HRIMessage()
+        image = self.cv_bridge.cv2_to_imgmsg(np.zeros((100, 100, 3), dtype=np.uint8))
+        msg.images = [image]
+        msg.audios = [AudioSegment.silent(duration=1000)]
+        msg.text = "Hello, HRI!"
+        self.publisher.publish(msg)
+
+
+class HRIMessageSubscriber(Node):
+    def __init__(self, topic: str):
+        super().__init__("test_hri_message_subscriber")
+        self.subscription = self.create_subscription(
+            HRIMessage, topic, self.handle_test_message, 10
+        )
+        self.received_messages: List[HRIMessage] = []
+
+    def handle_test_message(self, msg: HRIMessage) -> None:
+        self.received_messages.append(msg)
+
 
 class ServiceServer(Node):
     def __init__(self, service_name: str):

tests/communication/ros2/test_api.py

@@ -30,6 +30,7 @@
 
 from .helpers import ActionServer_ as ActionServer
 from .helpers import (
+    HRIMessageSubscriber,
     MessagePublisher,
     MessageReceiver,
     ServiceServer,
@@ -71,23 +72,21 @@ def test_ros2_configure_publisher(ros_setup: None, request: pytest.FixtureReques
     executors, threads = multi_threaded_spinner([node])
     try:
         topic_api = ConfigurableROS2TopicAPI(node)
-        cfg = TopicConfig(topic_name, "std_msgs/msg/String")
+        cfg = TopicConfig()
         topic_api.configure_publisher(topic_name, cfg)
         assert topic_api._publishers[topic_name] is not None
     finally:
         shutdown_executors_and_threads(executors, threads)
 
 
-def test_ros2_configre_subscriber(ros_setup, request: pytest.FixtureRequest):
+def test_ros2_configure_subscriber(ros_setup, request: pytest.FixtureRequest):
     topic_name = f"{request.node.originalname}_topic"  # type: ignore
     node_name = f"{request.node.originalname}_node"  # type: ignore
     node = Node(node_name)
     executors, threads = multi_threaded_spinner([node])
     try:
         topic_api = ConfigurableROS2TopicAPI(node)
         cfg = TopicConfig(
-            topic_name,
-            "std_msgs/msg/String",
             is_subscriber=True,
             subscriber_callback=lambda _: None,
         )
@@ -102,25 +101,23 @@ def test_ros2_single_message_publish_configured(
 ) -> None:
     topic_name = f"{request.node.originalname}_topic"  # type: ignore
     node_name = f"{request.node.originalname}_node"  # type: ignore
-    message_receiver = MessageReceiver(topic_name)
+    message_receiver = HRIMessageSubscriber(topic_name)
     node = Node(node_name)
     executors, threads = multi_threaded_spinner([message_receiver, node])
 
     try:
         topic_api = ConfigurableROS2TopicAPI(node)
         cfg = TopicConfig(
-            topic_name,
-            "std_msgs/msg/String",
             is_subscriber=False,
         )
         topic_api.configure_publisher(topic_name, cfg)
         topic_api.publish_configured(
             topic_name,
-            {"data": "Hello, ROS2!"},
+            {"text": "Hello, ROS2!"},
         )
         time.sleep(1)
         assert len(message_receiver.received_messages) == 1
-        assert message_receiver.received_messages[0].data == "Hello, ROS2!"
+        assert message_receiver.received_messages[0].text == "Hello, ROS2!"
     finally:
         shutdown_executors_and_threads(executors, threads)