feat: unified ROS2 hri message

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,16 +15,21 @@
 import threading
 import time
 import uuid
-from typing import Any, Callable, Dict, List, Literal, Optional, Tuple
+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
+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
 
 from rai.communication import (
@@ -41,6 +46,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 +209,100 @@ class ROS2HRIMessage(HRIMessage):
     def __init__(self, payload: HRIPayload, message_author: Literal["ai", "human"]):
         super().__init__(payload, message_author)
 
+    @classmethod
+    def from_ros2(cls, msg: Any, message_author: Literal["ai", "human"]):
+        from rai_interfaces.msg import HRIMessage as ROS2HRIMessage_
+
+        if not isinstance(msg, ROS2HRIMessage_):
+            raise ValueError(
+                "ROS2HRIMessage can only be created from rai_interfaces/msg/HRIMessage"
+            )
+
+        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):
+        # TODO: This import causes circular import. Fix it.
+        from rai.tools.ros2.utils import ros2_message_to_dict
+
+        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 ros2_message_to_dict(
+            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 +319,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 +332,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 +363,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
+int16 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

@@ -71,7 +71,7 @@ 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(msg_type="std_msgs/msg/String")
         topic_api.configure_publisher(topic_name, cfg)
         assert topic_api._publishers[topic_name] is not None
     finally:
@@ -86,8 +86,7 @@ def test_ros2_configre_subscriber(ros_setup, request: pytest.FixtureRequest):
     try:
         topic_api = ConfigurableROS2TopicAPI(node)
         cfg = TopicConfig(
-            topic_name,
-            "std_msgs/msg/String",
+            msg_type="std_msgs/msg/String",
             is_subscriber=True,
             subscriber_callback=lambda _: None,
         )
@@ -109,8 +108,7 @@ def test_ros2_single_message_publish_configured(
     try:
         topic_api = ConfigurableROS2TopicAPI(node)
         cfg = TopicConfig(
-            topic_name,
-            "std_msgs/msg/String",
+            msg_type="std_msgs/msg/String",
             is_subscriber=False,
         )
         topic_api.configure_publisher(topic_name, cfg)

tests/communication/ros2/test_connectors.py

@@ -16,12 +16,21 @@
 from typing import Any, List
 
 import pytest
-from rai.communication.ros2.connectors import ROS2ARIConnector, ROS2ARIMessage
+from PIL import Image
+from pydub import AudioSegment
+from rai.communication.ros2.connectors import (
+    HRIPayload,
+    ROS2ARIConnector,
+    ROS2ARIMessage,
+    ROS2HRIConnector,
+    ROS2HRIMessage,
+)
 from std_msgs.msg import String
 from std_srvs.srv import SetBool
 
 from .helpers import ActionServer_ as ActionServer
 from .helpers import (
+    HRIMessageSubscriber,
     MessagePublisher,
     MessageReceiver,
     ServiceServer,
@@ -162,3 +171,46 @@ def test_ros2ari_connector_send_goal_erronous_callback(
     finally:
         connector.shutdown()
         shutdown_executors_and_threads(executors, threads)
+
+
+def test_ros2hri_default_message_publish(
+    ros_setup: None, request: pytest.FixtureRequest
+):
+    topic_name = f"{request.node.originalname}_topic"  # type: ignore
+    connector = ROS2HRIConnector(targets=[topic_name])
+    hri_message_receiver = HRIMessageSubscriber(topic_name)
+    executors, threads = multi_threaded_spinner([hri_message_receiver])
+
+    try:
+        images = [Image.new("RGB", (100, 100), color="red")]
+        audios = [AudioSegment.silent(duration=1000)]
+        text = "Hello, HRI!"
+        payload = HRIPayload(images=images, audios=audios, text=text)
+        message = ROS2HRIMessage(payload=payload, message_author="ai")
+        connector.send_message(message, target=topic_name)
+        time.sleep(1)  # wait for the message to be received
+
+        assert len(hri_message_receiver.received_messages) > 0
+        recreated_message = ROS2HRIMessage.from_ros2(
+            hri_message_receiver.received_messages[0], message_author="ai"
+        )
+
+        assert message.text == recreated_message.text
+        assert message.message_author == recreated_message.message_author
+        assert len(message.images) == len(recreated_message.images)
+        assert len(message.audios) == len(recreated_message.audios)
+        assert all(
+            [
+                image_a == image_b
+                for image_a, image_b in zip(message.images, recreated_message.images)
+            ]
+        )
+        assert all(
+            [
+                audio_a == audio_b
+                for audio_a, audio_b in zip(message.audios, recreated_message.audios)
+            ]
+        )
+    finally:
+        connector.shutdown()
+        shutdown_executors_and_threads(executors, threads)