add recognition publisher and message (#200)

* add recognition publisher and message

* fix checking both subscription count and publish

Author: renan028

webots_ros2_core/package.xml

@@ -13,6 +13,7 @@
 
   <exec_depend>rclpy</exec_depend>
   <exec_depend>std_msgs</exec_depend>
+  <exec_depend>vision_msgs</exec_depend>
   <exec_depend>builtin_interfaces</exec_depend>
   <exec_depend>webots_ros2_msgs</exec_depend>
   <exec_depend>python3-tk</exec_depend>

webots_ros2_core/webots_ros2_core/devices/camera_device.py

@@ -15,9 +15,14 @@
 """Camera device."""
 
 from sensor_msgs.msg import Image, CameraInfo
+from vision_msgs.msg import Detection2D, Detection2DArray, ObjectHypothesisWithPose
+from geometry_msgs.msg import Point, Quaternion
+from std_msgs.msg import ColorRGBA
+from webots_ros2_msgs.msg import WbCameraRecognitionObject, WbCameraRecognitionObjects
 from rclpy.time import Time
 from rclpy.qos import DurabilityPolicy, HistoryPolicy, QoSProfile, QoSReliabilityPolicy, qos_profile_sensor_data
 from .sensor_device import SensorDevice
+from webots_ros2_core.math.quaternions import axangle2quat
 
 
 class CameraDevice(SensorDevice):
@@ -43,6 +48,8 @@ class CameraDevice(SensorDevice):
     def __init__(self, node, device_key, wb_device, params=None):
         super().__init__(node, device_key, wb_device, params)
         self._camera_info_publisher = None
+        self._recognition_publisher = None
+        self._recognition_webots_publisher = None
         self._image_publisher = None
 
         # Create topics
@@ -62,18 +69,31 @@ def __init__(self, node, device_key, wb_device, params=None):
                     history=HistoryPolicy.KEEP_LAST,
                 )
             )
+            if self._wb_device.hasRecognition():
+                self._recognition_publisher = self._node.create_publisher(
+                    Detection2DArray,
+                    self._topic_name + '/recognitions',
+                    qos_profile_sensor_data
+                )
+                self._recognition_webots_publisher = self._node.create_publisher(
+                    WbCameraRecognitionObjects,
+                    self._topic_name + '/recognitions/webots',
+                    qos_profile_sensor_data
+                )
 
             # CameraInfo data
             self.__message_info = CameraInfo()
-            self.__message_info.header.stamp = Time(seconds=self._node.robot.getTime()).to_msg()
+            self.__message_info.header.stamp = Time(
+                seconds=self._node.robot.getTime()).to_msg()
             self.__message_info.height = self._wb_device.getHeight()
             self.__message_info.width = self._wb_device.getWidth()
             self.__message_info.distortion_model = 'plumb_bob'
             focal_length = self._wb_device.getFocalLength()
             if focal_length == 0:
                 focal_length = 570.34  # Identical to Orbbec Astra
             self.__message_info.d = [0.0, 0.0, 0.0, 0.0, 0.0]
-            self.__message_info.r = [1.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 1.0]
+            self.__message_info.r = [
+                1.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 1.0]
             self.__message_info.k = [
                 focal_length, 0.0, self._wb_device.getWidth() / 2,
                 0.0, focal_length, self._wb_device.getHeight() / 2,
@@ -87,7 +107,8 @@ def __init__(self, node, device_key, wb_device, params=None):
             self._camera_info_publisher.publish(self.__message_info)
 
             # Load parameters
-            camera_period_param = node.declare_parameter(wb_device.getName() + '_period', self._timestep)
+            camera_period_param = node.declare_parameter(
+                wb_device.getName() + '_period', self._timestep)
             self._camera_period = camera_period_param.value
 
     def step(self):
@@ -119,7 +140,78 @@ def step(self):
             msg.encoding = 'bgra8'
             self._image_publisher.publish(msg)
 
-            self.__message_info.header.stamp = Time(seconds=self._node.robot.getTime()).to_msg()
+            self.__message_info.header.stamp = Time(
+                seconds=self._node.robot.getTime()).to_msg()
             self._camera_info_publisher.publish(self.__message_info)
+
+            if self._wb_device.hasRecognition() and (
+                    self._recognition_publisher.get_subscription_count() > 0 or
+                    self._recognition_webots_publisher.get_subscription_count() > 0):
+                self._wb_device.recognitionEnable(self._timestep)
+                objects = self._wb_device.getRecognitionObjects()
+
+                if objects is None:
+                    return
+
+                # Recognition data
+                reco_msg = Detection2DArray()
+                reco_msg_webots = WbCameraRecognitionObjects()
+                reco_msg.header.stamp = stamp
+                reco_msg_webots.header.stamp = stamp
+                reco_msg.header.frame_id = self._frame_id
+                reco_msg_webots.header.frame_id = self._frame_id
+                for obj in objects:
+                    # Getting Object Info
+                    position = Point()
+                    orientation = Quaternion()
+                    position.x = obj.get_position()[0]
+                    position.y = obj.get_position()[1]
+                    position.z = obj.get_position()[2]
+                    axangle = obj.get_orientation()
+                    quat = axangle2quat(axangle[:-1], axangle[-1])
+                    orientation.w = quat[0]
+                    orientation.x = quat[1]
+                    orientation.y = quat[2]
+                    orientation.z = quat[3]
+                    obj_model = obj.get_model().decode('UTF-8')
+                    obj_center = [float(i)
+                                  for i in obj.get_position_on_image()]
+                    obj_size = [float(i) for i in obj.get_size_on_image()]
+                    obj_id = obj.get_id()
+                    obj_colors = obj.get_colors()
+                    # Object Info -> Detection2D
+                    reco_obj = Detection2D()
+                    hyp = ObjectHypothesisWithPose()
+                    hyp.id = obj_model
+                    hyp.pose.pose.position = position
+                    hyp.pose.pose.orientation = orientation
+                    reco_obj.results.append(hyp)
+                    reco_obj.bbox.center.x = obj_center[0]
+                    reco_obj.bbox.center.y = obj_center[1]
+                    reco_obj.bbox.size_x = obj_size[0]
+                    reco_obj.bbox.size_y = obj_size[1]
+                    reco_msg.detections.append(reco_obj)
+
+                    # Object Info -> WbCameraRecognitionObject
+                    reco_webots_obj = WbCameraRecognitionObject()
+                    reco_webots_obj.id = obj_id
+                    reco_webots_obj.model = obj_model
+                    reco_webots_obj.pose.pose.position = position
+                    reco_webots_obj.pose.pose.orientation = orientation
+                    reco_webots_obj.bbox.center.x = obj_center[0]
+                    reco_webots_obj.bbox.center.y = obj_center[1]
+                    reco_webots_obj.bbox.size_x = obj_size[0]
+                    reco_webots_obj.bbox.size_y = obj_size[1]
+                    for i in range(0, obj.get_number_of_colors()):
+                        color = ColorRGBA()
+                        color.r = obj_colors[3 * i]
+                        color.g = obj_colors[3 * i + 1]
+                        color.b = obj_colors[3 * i + 2]
+                        reco_webots_obj.colors.append(color)
+                    reco_msg_webots.objects.append(reco_webots_obj)
+                self._recognition_webots_publisher.publish(reco_msg_webots)
+                self._recognition_publisher.publish(reco_msg)
+            else:
+                self._wb_device.recognitionDisable()
         else:
             self._wb_device.disable()

webots_ros2_msgs/CMakeLists.txt

@@ -18,23 +18,34 @@ endif()
 find_package(ament_cmake REQUIRED)
 find_package(builtin_interfaces REQUIRED)
 find_package(rosidl_default_generators REQUIRED)
+find_package(geometry_msgs REQUIRED)
+find_package(std_msgs REQUIRED)
+find_package(vision_msgs REQUIRED)
 
+set(msg_files
+  "msg/WbCameraRecognitionObject.msg"
+  "msg/WbCameraRecognitionObjects.msg"
+)
 set(srv_files
   "srv/SetDifferentialWheelSpeed.srv"
   "srv/SetInt.srv"
 )
 
+rosidl_generate_interfaces(${PROJECT_NAME}
+  ${msg_files} ${srv_files}
+  DEPENDENCIES builtin_interfaces std_msgs geometry_msgs vision_msgs
+  ADD_LINTER_TESTS
+)
+
 if(BUILD_TESTING)
   find_package(ament_lint_auto REQUIRED)
   set(ament_cmake_copyright_FOUND TRUE)
   set(ament_cmake_cpplint_FOUND TRUE)
   ament_lint_auto_find_test_dependencies()
 endif()
 
-rosidl_generate_interfaces(${PROJECT_NAME}
-  ${srv_files}
-  DEPENDENCIES builtin_interfaces
-  ADD_LINTER_TESTS
-)
+
 ament_export_dependencies(rosidl_default_runtime)
+ament_export_include_directories(include)
+
 ament_package()

webots_ros2_msgs/msg/WbCameraRecognitionObject.msg

@@ -0,0 +1,5 @@
+int32 id
+geometry_msgs/PoseStamped pose
+vision_msgs/BoundingBox2D bbox 
+std_msgs/ColorRGBA[] colors
+string model

webots_ros2_msgs/msg/WbCameraRecognitionObjects.msg

@@ -0,0 +1,2 @@
+std_msgs/Header header
+WbCameraRecognitionObject[] objects

webots_ros2_msgs/package.xml

@@ -15,9 +15,15 @@
   <buildtool_depend>rosidl_default_generators</buildtool_depend>
 
   <build_depend>builtin_interfaces</build_depend>
+  <build_depend>geometry_msgs</build_depend>
+  <build_depend>std_msgs</build_depend>
+  <build_depend>vision_msgs</build_depend>
 
   <exec_depend>builtin_interfaces</exec_depend>
   <exec_depend>rosidl_default_runtime</exec_depend>
+  <exec_depend>geometry_msgs</exec_depend>
+  <exec_depend>std_msgs</exec_depend>
+  <exec_depend>vision_msgs</exec_depend>
 
   <test_depend>ament_lint_auto</test_depend>
   <test_depend>ament_lint_common</test_depend>