Merge pull request #46 from cyberbotics/port-camera-segmentation-test…

…-to-master Port camera segmentation tests to master branch
cyberbotics · Dec 16, 2020 · 8b17bf4 · 8b17bf4
2 parents e8cf5a3 + 0cabc91
commit 8b17bf4
Show file tree

Hide file tree

Showing 2 changed files with 106 additions and 1 deletion.
diff --git a/src/complete_test.cpp b/src/complete_test.cpp
@@ -173,6 +173,17 @@ void cameraRecognitionCallback(const webots_ros::RecognitionObject::ConstPtr &ob
   callbackCalled = true;
 }
 
+void segmentationCallback(const sensor_msgs::Image::ConstPtr &values) {
+  ROS_INFO("Segmentation callback called.");
+  int i = 0;
+  imageColor.resize(values->step * values->height);
+  for (std::vector<unsigned char>::const_iterator it = values->data.begin(); it != values->data.end(); ++it) {
+    imageColor[i] = *it;
+    i++;
+  }
+  callbackCalled = true;
+}
+
 void joystickCallback(const webots_ros::Int8Stamped::ConstPtr &value) {
   ROS_INFO("Joystick button pressed: %d (time: %d:%d).", value->data, value->header.stamp.sec, value->header.stamp.nsec);
   callbackCalled = true;
@@ -800,6 +811,93 @@ int main(int argc, char **argv) {
   enable_camera_recognition_client.shutdown();
   time_step_client.call(time_step_srv);
 
+  // camera recognition segmentation is enabled
+  ros::ServiceClient has_segmentation_client =
+    n.serviceClient<webots_ros::get_bool>(model_name + "/camera/recognition_has_segmentation");
+  webots_ros::get_bool has_segmentation_srv;
+  if (has_segmentation_client.call(has_segmentation_srv)) {
+    if (has_segmentation_srv.response.value)
+      ROS_INFO("Camera recognition segmentation field is TRUE.");
+    else
+      ROS_INFO("Camera recognition segmentation field is FALSE.");
+  } else
+    ROS_ERROR("Failed to get segmentation field value.");
+  has_segmentation_client.shutdown();
+
+  // camera recognition enable segmentation
+  ros::ServiceClient enable_segmentation_client =
+    n.serviceClient<webots_ros::get_bool>(model_name + "/camera/recognition_enable_segmentation");
+  webots_ros::get_bool enable_segmentation_srv;
+  if (enable_segmentation_client.call(enable_segmentation_srv) && enable_segmentation_srv.response.value)
+    ROS_INFO("Segmentation correctly available.");
+  else {
+    if (!enable_segmentation_srv.response.value)
+      ROS_ERROR("Segmentation value could not be retrieved correctly.");
+    ROS_ERROR("Failed to retrieve segmentation value.");
+    return 1;
+  }
+
+  enable_segmentation_client.shutdown();
+  time_step_client.call(time_step_srv);
+
+  ros::ServiceClient is_segmentation_enabled_client =
+    n.serviceClient<webots_ros::get_bool>(model_name + "/camera/recognition_is_segmentation_enabled");
+  webots_ros::get_bool is_segmentation_enabled_srv;
+  if (is_segmentation_enabled_client.call(is_segmentation_enabled_srv) && is_segmentation_enabled_srv.response.value)
+    ROS_INFO("Segmentation correctly enabled.");
+  else {
+    if (!enable_segmentation_srv.response.value)
+      ROS_ERROR("Failed to enable segmentation.");
+    ROS_ERROR("Failed to query segmentation enabled status.");
+    return 1;
+  }
+
+  is_segmentation_enabled_client.shutdown();
+  time_step_client.call(time_step_srv);
+
+  ros::Subscriber sub_segmentation =
+    n.subscribe(model_name + "/camera/recognition_segmentation_image", 1, segmentationCallback);
+  ROS_INFO("Topic for camera recognition segmentation initialized.");
+  callbackCalled = false;
+  while (sub_segmentation.getNumPublishers() == 0 && !callbackCalled) {
+    ros::spinOnce();
+    time_step_client.call(time_step_srv);
+  }
+  ROS_INFO("Topic for camera recognition segmentation connected.");
+  sub_segmentation.shutdown();
+  time_step_client.call(time_step_srv);
+
+  sub_camera_color.shutdown();
+  enable_camera_client.shutdown();
+  time_step_client.call(time_step_srv);
+
+  // camera recognition save segmentation image
+  ros::ServiceClient save_segmentation_image_client =
+    n.serviceClient<webots_ros::save_image>(model_name + "/camera/recognition_save_segmentation_image");
+  webots_ros::save_image save_segmentation_image_srv;
+  save_segmentation_image_srv.request.filename = std::string(getenv("HOME")) + std::string("/test_image_segmentation.png");
+  save_segmentation_image_srv.request.quality = 100;
+  if (save_segmentation_image_client.call(save_segmentation_image_srv) && save_segmentation_image_srv.response.success == 1)
+    ROS_INFO("Segmentation image saved.");
+  else
+    ROS_ERROR("Failed to call save segmentation image.");
+
+  // camera recognition disable segmentation
+  ros::ServiceClient disable_segmentation_client =
+    n.serviceClient<webots_ros::get_bool>(model_name + "/camera/recognition_disable_segmentation");
+  webots_ros::get_bool disable_segmentation_srv;
+  if (disable_segmentation_client.call(disable_segmentation_srv) && disable_segmentation_srv.response.value)
+    ROS_INFO("Segmentation correctly disabled.");
+  else {
+    if (!disable_segmentation_srv.response.value)
+      ROS_ERROR("Segmentation value could not be disabled.");
+    ROS_ERROR("Failed to set segmentation.");
+    return 1;
+  }
+
+  enable_segmentation_client.shutdown();
+  time_step_client.call(time_step_srv);
+
   // camera_save_image
   ros::ServiceClient save_image_client = n.serviceClient<webots_ros::save_image>(model_name + "/camera/save_image");
   webots_ros::save_image save_image_srv;

diff --git a/worlds/complete_test.wbt b/worlds/complete_test.wbt
@@ -50,6 +50,9 @@ DEF CONE Solid {
     }
   ]
   name "solid(1)"
+  recognitionColors [
+    0.988235 0.913725 0.309804
+  ]
 }
 DEF RADAR_TARGET Solid {
   translation 0 0 -1.86
@@ -176,6 +179,7 @@ Robot {
         minFieldOfView 0.49
       }
       recognition Recognition {
+        segmentation TRUE
       }
     }
     RangeFinder {
@@ -268,6 +272,9 @@ Robot {
   boundingObject Box {
     size 0.2 0.2 0.2
   }
+  recognitionColors [
+    0.678431 0.498039 0.658824
+  ]
   controller "ros"
   controllerArgs [
     "--name=my_robot"
@@ -277,7 +284,7 @@ Robot {
   customData "INITIAL"
   supervisor TRUE
   battery [
-    10, 10, 2
+    50, 50, 2
   ]
   window "info"
 }