diff --git a/webots_ros2_driver/include/webots_ros2_driver/plugins/static/Ros2Camera.hpp b/webots_ros2_driver/include/webots_ros2_driver/plugins/static/Ros2Camera.hpp index 73b82229d..65c0c3939 100644 --- a/webots_ros2_driver/include/webots_ros2_driver/plugins/static/Ros2Camera.hpp +++ b/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,6 +60,9 @@ namespace webots_ros2_driver { rclcpp::Publisher::SharedPtr mCameraInfoPublisher; sensor_msgs::msg::CameraInfo mCameraInfoMessage; + rclcpp::Publisher::SharedPtr mSegmentationPublisher; + sensor_msgs::msg::Image mSegmentationMessage; + rclcpp::Publisher::SharedPtr mRecognitionPublisher; rclcpp::Publisher::SharedPtr mWebotsRecognitionPublisher; vision_msgs::msg::Detection2DArray mRecognitionMessage; @@ -66,6 +70,7 @@ namespace webots_ros2_driver { bool mIsEnabled; bool mRecognitionIsEnabled; + bool mSegmentationIsEnabled; }; } // namespace webots_ros2_driver diff --git a/webots_ros2_driver/src/plugins/static/Ros2Camera.cpp b/webots_ros2_driver/src/plugins/static/Ros2Camera.cpp index 9c2067229..f5a565f01 100644 --- a/webots_ros2_driver/src/plugins/static/Ros2Camera.cpp +++ b/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(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,12 +110,20 @@ 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 @@ -112,6 +131,8 @@ namespace webots_ros2_driver { 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 diff --git a/webots_ros2_tests/test/test_system_driver.py b/webots_ros2_tests/test/test_system_driver.py index 80c1463d2..a15037b61 100644 --- a/webots_ros2_tests/test/test_system_driver.py +++ b/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): diff --git a/webots_ros2_tests/worlds/driver_test.wbt b/webots_ros2_tests/worlds/driver_test.wbt index 40ef7a4d2..3485c03c6 100644 --- a/webots_ros2_tests/worlds/driver_test.wbt +++ b/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 } } ]