Add segmentation data publisher to camera plugin

webots_ros2_driver/include/webots_ros2_driver/plugins/static/Ros2Camera.hpp

@@ -47,6 +47,7 @@ namespace webots_ros2_driver {
   private:
     void publishImage();
     void publishRecognition();
+    void publishSegmentation();
 
     WbDeviceTag mCamera;
 
@@ -59,13 +60,17 @@ namespace webots_ros2_driver {
     rclcpp::Publisher<sensor_msgs::msg::CameraInfo>::SharedPtr mCameraInfoPublisher;
     sensor_msgs::msg::CameraInfo mCameraInfoMessage;
 
+    rclcpp::Publisher<sensor_msgs::msg::Image>::SharedPtr mSegmentationPublisher;
+    sensor_msgs::msg::Image mSegmentationMessage;
+
     rclcpp::Publisher<vision_msgs::msg::Detection2DArray>::SharedPtr mRecognitionPublisher;
     rclcpp::Publisher<webots_ros2_msgs::msg::CameraRecognitionObjects>::SharedPtr mWebotsRecognitionPublisher;
     vision_msgs::msg::Detection2DArray mRecognitionMessage;
     webots_ros2_msgs::msg::CameraRecognitionObjects mWebotsRecognitionMessage;
 
     bool mIsEnabled;
     bool mRecognitionIsEnabled;
+    bool mSegmentationIsEnabled;
   };
 
 }  // namespace webots_ros2_driver

webots_ros2_driver/src/plugins/static/Ros2Camera.cpp

@@ -21,6 +21,7 @@ namespace webots_ros2_driver {
     Ros2SensorPlugin::init(node, parameters);
     mIsEnabled = false;
     mRecognitionIsEnabled = false;
+    mSegmentationIsEnabled = false;
     mCamera = wb_robot_get_device(parameters["name"].c_str());
 
     mCameraInfoSuffix = parameters.count("cameraInfoSuffix") ? parameters["cameraInfoSuffix"] : "/camera_info";
@@ -78,6 +79,13 @@ namespace webots_ros2_driver {
       mRecognitionMessage.header.frame_id = mFrameName;
       mWebotsRecognitionMessage.header.frame_id = mFrameName;
     }
+
+    // Segmentation publisher
+    if (wb_camera_recognition_has_segmentation(mCamera)) {
+      mSegmentationPublisher = mNode->create_publisher<sensor_msgs::msg::Image>(mTopicName + "/segmentation", 
+                                                                                rclcpp::SensorDataQoS().reliable());
+      mSegmentationMessage = mImageMessage;
+    }
   }
 
   void Ros2Camera::step() {
@@ -89,7 +97,10 @@ namespace webots_ros2_driver {
     const bool recognitionSubscriptionsExist =
       (mRecognitionPublisher != nullptr && mRecognitionPublisher->get_subscription_count() > 0) ||
       (mWebotsRecognitionPublisher != nullptr && mWebotsRecognitionPublisher->get_subscription_count() > 0);
-    const bool shouldBeEnabled = mAlwaysOn || imageSubscriptionsExist || recognitionSubscriptionsExist;
+    const bool segmentationSubscriptionsExist = 
+      mSegmentationPublisher != nullptr && mSegmentationPublisher->get_subscription_count() > 0;
+    const bool shouldBeEnabled = mAlwaysOn || imageSubscriptionsExist || 
+      recognitionSubscriptionsExist || segmentationSubscriptionsExist;
 
     if (shouldBeEnabled != mIsEnabled) {
       if (shouldBeEnabled)
@@ -99,19 +110,29 @@ namespace webots_ros2_driver {
       mIsEnabled = shouldBeEnabled;
     }
 
-    if (recognitionSubscriptionsExist != mRecognitionIsEnabled) {
-      if (recognitionSubscriptionsExist)
+    if ( (recognitionSubscriptionsExist || segmentationSubscriptionsExist) != mRecognitionIsEnabled) {
+      if (recognitionSubscriptionsExist || segmentationSubscriptionsExist)
         wb_camera_recognition_enable(mCamera, mPublishTimestepSyncedMs);
       else
         wb_camera_recognition_disable(mCamera);
-      mRecognitionIsEnabled = recognitionSubscriptionsExist;
+      mRecognitionIsEnabled = recognitionSubscriptionsExist || segmentationSubscriptionsExist;
+    }
+
+    if (segmentationSubscriptionsExist != mSegmentationIsEnabled) {
+      if (segmentationSubscriptionsExist)
+        wb_camera_recognition_enable_segmentation(mCamera);
+      else
+        wb_camera_recognition_disable_segmentation(mCamera);
+      mSegmentationIsEnabled = segmentationSubscriptionsExist;
     }
 
     // Publish data
     if (mAlwaysOn || imageSubscriptionsExist)
       publishImage();
     if (recognitionSubscriptionsExist)
       publishRecognition();
+    if (segmentationSubscriptionsExist)
+      publishSegmentation();
     if (mCameraInfoPublisher->get_subscription_count() > 0)
       mCameraInfoPublisher->publish(mCameraInfoMessage);
   }
@@ -185,4 +206,13 @@ namespace webots_ros2_driver {
     mWebotsRecognitionPublisher->publish(mWebotsRecognitionMessage);
     mRecognitionPublisher->publish(mRecognitionMessage);
   }
+
+  void Ros2Camera::publishSegmentation() {
+    auto image = wb_camera_recognition_get_segmentation_image(mCamera);
+    if (image) {
+      mSegmentationMessage.header.stamp = mNode->get_clock()->now();
+      memcpy(mSegmentationMessage.data.data(), image, mSegmentationMessage.data.size());
+      mSegmentationPublisher->publish(mSegmentationMessage);
+    }
+  }
 }  // namespace webots_ros2_driver

webots_ros2_tests/test/test_system_driver.py

@@ -129,6 +129,20 @@ def on_objects_received(message):
         self.wait_for_messages(self.__node, CameraRecognitionObjects, '/Pioneer_3_AT/camera/recognitions/webots',
                                condition=on_objects_received)
 
+    def testSegmentation(self):
+        def on_image_received(message):
+            self.assertEqual(message.height, 64)
+            self.assertEqual(message.width, 64)
+
+            pixels = set([tuple(message.data[i:i+4]) for i in range(0,len(message.data),4)])
+            expected = {(0, 0, 0, 255), (0, 0, 255, 255)}
+            assert pixels.issubset( expected )
+
+            return True
+
+        self.wait_for_messages(self.__node, Image, '/Pioneer_3_AT/camera/segmentation', 
+                               condition=on_image_received)
+
     def testPythonPluginService(self):
         client = self.__node.create_client(Trigger, 'move_forward')
         if not client.wait_for_service(timeout_sec=10.0):

webots_ros2_tests/worlds/driver_test.wbt

@@ -86,6 +86,7 @@ Pioneer3at {
       translation -0.25 0 0.25
       rotation 0 0 1 3.14159
       recognition Recognition {
+        segmentation TRUE
       }
     }
   ]