Segmentation Sensor

include/ignition/sensors/SegmentationCameraSensor.hh

@@ -25,13 +25,15 @@
 #include <ignition/common/PluginMacros.hh>
 #include <ignition/common/SuppressWarning.hh>
 #include <ignition/common/Time.hh>
+#include <ignition/transport/Node.hh>
+#include <ignition/transport/Publisher.hh>
+#include <ignition/msgs.hh>
+#include <sdf/sdf.hh>
 #include "ignition/sensors/CameraSensor.hh"
 #include "ignition/sensors/Export.hh"
-#include "ignition/transport/Node.hh"
-#include "ignition/transport/Publisher.hh"
 #include "ignition/sensors/Sensor.hh"
-#include "ignition/msgs.hh"
-#include <sdf/sdf.hh>
+
+#include "ignition/sensors/segmentation_camera/Export.hh"
 
 namespace ignition
 {
@@ -49,7 +51,8 @@ namespace ignition
     /// It offers both an ignition-transport interface and a direct C++ API
     /// to access the image data. The API works by setting a callback to be
     /// called with image data.
-    class SegmentationCameraSensor : public CameraSensor
+    class IGNITION_SENSORS_SEGMENTATION_CAMERA_VISIBLE
+      SegmentationCameraSensor : public CameraSensor
     {
       /// \brief constructor
       public: SegmentationCameraSensor();

src/SegmentationCameraSensor.cc

@@ -18,32 +18,34 @@
 #include <mutex>
 #include <memory>
 
-#include "ignition/common/Console.hh"
-#include "ignition/common/Profiler.hh"
-#include "ignition/common/Image.hh"
-#include "ignition/msgs.hh"
-#include "ignition/rendering/SegmentationCamera.hh"
+#include <ignition/common/Console.hh>
+#include <ignition/common/Profiler.hh>
+#include <ignition/common/Image.hh>
+#include <ignition/msgs.hh>
+#include <ignition/rendering/SegmentationCamera.hh>
+#include <ignition/transport/Node.hh>
+#include <ignition/transport/Publisher.hh>
+
 #include "ignition/sensors/RenderingEvents.hh"
-#include "ignition/sensors/SensorFactory.hh"
 #include "ignition/sensors/SegmentationCameraSensor.hh"
-#include "ignition/transport/Node.hh"
-#include "ignition/transport/Publisher.hh"
+#include "ignition/sensors/SensorFactory.hh"
 
 using namespace ignition;
 using namespace sensors;
 
 class ignition::sensors::SegmentationCameraSensorPrivate
 {
   /// \brief Save an image
-  /// \param[in] _data the image data to be saved
+  /// \param[in] _coloredBuffer buffer of colored map
+  /// \param[in] _labelsBuffer buffer of labels map
   /// \param[in] _width width of image in pixels
   /// \param[in] _height height of image in pixels
   /// \return True if the image was saved successfully. False can mean
   /// that the path provided to the constructor does exist and creation
   /// of the path was not possible.
   /// \sa ImageSaver
-  public: bool SaveImage(const uint8_t *_data, unsigned int _width,
-    unsigned int _height);
+  public: bool SaveImage(const uint8_t *_coloredBuffer,
+    const uint8_t *_labelsBuffer, unsigned int _width, unsigned int _height);
 
   /// \brief SDF Sensor DOM Object
   public: sdf::Sensor sdfSensor;
@@ -63,17 +65,17 @@ class ignition::sensors::SegmentationCameraSensorPrivate
   /// \brief Publisher to publish segmentation labels image
   public: transport::Node::Publisher labelsMapPublisher;
 
-  /// \brief Segmentation Image Msg
+  /// \brief Segmentation colored image message
   public: msgs::Image coloredMapMsg;
 
-  /// \brief Segmentation Image Msg
+  /// \brief Segmentation labels image message
   public: msgs::Image labelsMapMsg;
 
-  /// \brief Topic to publish the segmentation colored map
-  public: std::string topicColoredMap = "/colored_map";
+  /// \brief Topic suffix to publish the segmentation colored map
+  public: const std::string topicColoredMapSuffix = "/colored_map";
 
- /// \brief Topic to publish the segmentation labels map
-  public: std::string topicLabelsMap = "/labels_map";
+  /// \brief Topic suffix to publish the segmentation labels map
+  public: std::string topicLabelsMapSuffix = "/labels_map";
 
   /// \brief Buffer contains the segmentation colored map data
   public: uint8_t *segmentationColoredBuffer {nullptr};
@@ -184,25 +186,38 @@ bool SegmentationCameraSensor::Load(const sdf::Sensor &_sdf)
       this->dataPtr->type = rendering::SegmentationType::Semantic;
     else if (type == "instance" || type == "panoptic")
       this->dataPtr->type = rendering::SegmentationType::Panoptic;
+    else
+    {
+      igndbg << "Wrong type {" << type <<
+        "}, type should be semantic or instance or panoptic" << std::endl;
+      return false;
+    }
   }
 
   // Create the segmentation colored map image publisher
   this->dataPtr->coloredMapPublisher =
       this->dataPtr->node.Advertise<ignition::msgs::Image>(
-          this->Topic() + this->dataPtr->topicColoredMap);
+          this->Topic() + this->dataPtr->topicColoredMapSuffix);
 
   // Create the segmentation labels map image publisher
   this->dataPtr->labelsMapPublisher =
       this->dataPtr->node.Advertise<ignition::msgs::Image>(
-          this->Topic() + this->dataPtr->topicLabelsMap);
+          this->Topic() + this->dataPtr->topicLabelsMapSuffix);
 
-  if (!this->dataPtr->labelsMapPublisher || !this->dataPtr->coloredMapPublisher)
+  if (!this->dataPtr->labelsMapPublisher)
+  {
+    ignerr << "Unable to create publisher on topic["
+      << this->Topic() + this->dataPtr->topicLabelsMapSuffix << "].\n";
+    return false;
+  }
+  if (!this->dataPtr->coloredMapPublisher)
   {
     ignerr << "Unable to create publisher on topic["
-      << this->Topic() << "].\n";
+      << this->Topic() + this->dataPtr->topicColoredMapSuffix << "].\n";
     return false;
   }
 
+  std::cout << this->InfoTopic() << std::endl;
   if (!this->AdvertiseInfo(this->Topic() + "/camera_info"))
     return false;
 
@@ -355,11 +370,10 @@ bool SegmentationCameraSensor::Update(
     return false;
   }
 
-  // don't render if there is no subscribers
+  // don't render if there are no subscribers
   if (!this->dataPtr->coloredMapPublisher.HasConnections() &&
     !this->dataPtr->labelsMapPublisher.HasConnections() &&
-    this->dataPtr->imageEvent.ConnectionCount() <= 0u &&
-    !this->dataPtr->saveImage)
+    this->dataPtr->imageEvent.ConnectionCount() <= 0u)
   {
     return false;
   }
@@ -426,7 +440,7 @@ bool SegmentationCameraSensor::Update(
   // Save image
   if (this->dataPtr->saveImage)
   {
-    this->dataPtr->SaveImage(
+    this->dataPtr->SaveImage(this->dataPtr->segmentationLabelsBuffer,
       this->dataPtr->segmentationColoredBuffer, width, height);
   }
 
@@ -458,7 +472,8 @@ common::ConnectionPtr SegmentationCameraSensor::ConnectImageCallback(
 
 //////////////////////////////////////////////////
 bool SegmentationCameraSensorPrivate::SaveImage(
-  const uint8_t *_data, unsigned int _width, unsigned int _height)
+  const uint8_t *_coloredBuffer, const uint8_t *_labelsBuffer,
+  unsigned int _width, unsigned int _height)
 {
   // Attempt to create the directory if it doesn't exist
   if (!ignition::common::isDirectory(this->saveImagePath))
@@ -467,16 +482,29 @@ bool SegmentationCameraSensorPrivate::SaveImage(
       return false;
   }
 
-  std::string filename = this->saveImagePrefix +
+  std::string coloredName = this->saveImagePrefix + "labels_" +
                          std::to_string(this->saveImageCounter) + ".png";
+  std::string labelsName = this->saveImagePrefix + "colored_" +
+                         std::to_string(this->saveImageCounter) + ".png";
+
   ++this->saveImageCounter;
 
-  ignition::common::Image localImage;
-  localImage.SetFromData(_data, _width, _height,
+  // save colored map
+  ignition::common::Image localColoredImage;
+  localColoredImage.SetFromData(_coloredBuffer, _width, _height,
     ignition::common::Image::RGB_INT8);
 
-  localImage.SavePNG(
-      ignition::common::joinPaths(this->saveImagePath, filename));
+  localColoredImage.SavePNG(
+      ignition::common::joinPaths(this->saveImagePath, coloredName));
+
+  // save labels map
+  ignition::common::Image localLabelsImage;
+  localLabelsImage.SetFromData(_labelsBuffer, _width, _height,
+    ignition::common::Image::RGB_INT8);
+
+  localLabelsImage.SavePNG(
+      ignition::common::joinPaths(this->saveImagePath, labelsName));
+
   return true;
 }
 

test/integration/segmentation_camera_plugin.cc

@@ -55,10 +55,13 @@ std::mutex g_mutex;
 uint8_t *g_buffer = nullptr;
 
 /// \brief counter of received segmentation msgs
-int g_counter = 0;
+unsigned int g_counter = 0;
 
 /// \brief Label of the hidden box
-uint8_t hiddenLabel = 4;
+const uint8_t hiddenLabel = 4;
+const uint8_t leftBoxLabel = 1;
+const uint8_t rightBoxLabel = 1;
+const uint8_t middleBoxLabel = 2;
 
 /// \brief callback to get the segmentation buffer
 void OnNewSegmentationFrame(const msgs::Image &_msg)
@@ -79,15 +82,15 @@ void OnNewSegmentationFrame(const msgs::Image &_msg)
 /// \brief wait till you read the published frame
 void WaitForNewFrame(ignition::sensors::Manager &mgr)
 {
+  g_counter = 0;
+
   // wait for a few segmentation camera frames
   mgr.RunOnce(std::chrono::steady_clock::duration::zero(), true);
 
-  g_counter = 0;
-
   auto waitTime = std::chrono::duration_cast< std::chrono::milliseconds >(
-  std::chrono::duration< double >(0.001));
+    std::chrono::duration< double >(0.001));
 
-  int counter = 0;
+  unsigned int counter = 0;
   // wait for the counter to increase
   for (int sleep = 0; sleep < 300 && counter == 0; ++sleep)
   {
@@ -96,29 +99,31 @@ void WaitForNewFrame(ignition::sensors::Manager &mgr)
     g_mutex.unlock();
     std::this_thread::sleep_for(waitTime);
   }
-  EXPECT_GT(counter, 0);
+  EXPECT_GT(counter, 0u);
 }
 
 /// \brief Build the scene with 3 boxes besides each other
-/// the 2 outer boxes have the same label & the middle is different
+/// the 2 outer boxes have the same label & the middle is different.
+/// There is also another box with a unique label that is hidden
+/// behind the middle box
 void BuildScene(rendering::ScenePtr scene)
 {
-  math::Vector3d leftPosition(3, -1.5, 0);
-  math::Vector3d rightPosition(3, 1.5, 0);
-  math::Vector3d middlePosition(3, 0, 0);
-  math::Vector3d hiddenPosition(8, 0, 0);
+  const math::Vector3d leftPosition(3, -1.5, 0);
+  const math::Vector3d rightPosition(3, 1.5, 0);
+  const math::Vector3d middlePosition(3, 0, 0);
+  const math::Vector3d hiddenPosition(8, 0, 0);
 
   rendering::VisualPtr root = scene->RootVisual();
 
-  double unitBoxSize = 1.0;
+  const double unitBoxSize = 1.0;
 
   // create box visual
   rendering::VisualPtr box = scene->CreateVisual();
   box->AddGeometry(scene->CreateBox());
   box->SetOrigin(0.0, 0.0, 0.0);
   box->SetLocalPosition(leftPosition);
   box->SetLocalRotation(0, 0, 0);
-  box->SetUserData("label", 1);
+  box->SetUserData("label", leftBoxLabel);
   box->SetLocalScale(unitBoxSize, unitBoxSize, unitBoxSize);
   root->AddChild(box);
 
@@ -128,7 +133,7 @@ void BuildScene(rendering::ScenePtr scene)
   box1->SetOrigin(0.0, 0.0, 0.0);
   box1->SetLocalPosition(rightPosition);
   box1->SetLocalRotation(0, 0, 0);
-  box1->SetUserData("label", 1);
+  box1->SetUserData("label", rightBoxLabel);
   box1->SetLocalScale(unitBoxSize, unitBoxSize, unitBoxSize);
   root->AddChild(box1);
 
@@ -138,7 +143,7 @@ void BuildScene(rendering::ScenePtr scene)
   box2->SetOrigin(0.0, 0.0, 0.0);
   box2->SetLocalPosition(middlePosition);
   box2->SetLocalRotation(0, 0, 0);
-  box2->SetUserData("label", 2);
+  box2->SetUserData("label", middleBoxLabel);
   box2->SetLocalScale(unitBoxSize, unitBoxSize, unitBoxSize);
   root->AddChild(box2);
 
@@ -224,6 +229,7 @@ void SegmentationCameraSensorTest::ImagesWithBuiltinSDF(
   auto camera = sensor->SegmentationCamera();
   ASSERT_NE(camera, nullptr);
 
+  // 23 is a random number that is not used for boxes
   uint32_t backgroundLabel = 23;
   camera->SetSegmentationType(rendering::SegmentationType::Semantic);
   camera->EnableColoredMap(true);
@@ -245,11 +251,10 @@ void SegmentationCameraSensorTest::ImagesWithBuiltinSDF(
   mgr.RunOnce(std::chrono::steady_clock::duration::zero());
 
   EXPECT_TRUE(helper.WaitForMessage()) << helper;
-  EXPECT_TRUE(infoHelper.WaitForMessage()) << infoHelper;
 
   // subscribe to the segmentation camera topic
   ignition::transport::Node node;
-  node.Subscribe(topic, &OnNewSegmentationFrame);
+  EXPECT_TRUE(node.Subscribe(topic, &OnNewSegmentationFrame));
 
   // wait for a new frame
   WaitForNewFrame(mgr);
@@ -265,21 +270,26 @@ void SegmentationCameraSensorTest::ImagesWithBuiltinSDF(
   uint32_t rightIndex = (rightProj.Y() * width + rightProj.X()) * channels;
   uint32_t middleIndex = (middleProj.Y() * width + middleProj.X()) * channels;
 
-  // test
-  g_counter = 0;
-
   uint8_t leftLabel =   g_buffer[leftIndex];
   uint8_t rightLabel =  g_buffer[rightIndex];
   uint8_t middleLabel = g_buffer[middleIndex];
 
   // check the label
-  EXPECT_EQ(leftLabel, 1);
-  EXPECT_EQ(middleLabel, 2);
-  EXPECT_EQ(rightLabel, 1);
+  EXPECT_EQ(leftLabel, leftBoxLabel);
+  EXPECT_EQ(middleLabel, middleBoxLabel);
+  EXPECT_EQ(rightLabel, rightBoxLabel);
 
+  // check a pixel between 2 boxes & a pixel below a box
+  uint32_t betweenBoxesIndex = (120 * width + 230) * channels;
+  uint32_t belowBoxesIndex = (200 * width + 280) * channels;
   // check if the first pixel(background) = the background label
   uint8_t background = g_buffer[0];
+  uint8_t betweenBoxes = g_buffer[betweenBoxesIndex];
+  uint8_t belowBoxes = g_buffer[belowBoxesIndex];
+
   EXPECT_EQ(background, backgroundLabel);
+  EXPECT_EQ(betweenBoxes, backgroundLabel);
+  EXPECT_EQ(belowBoxes, backgroundLabel);
 
   // search for the hidden box label in all pixels
   for (int i = 0; i < height; i++)
@@ -298,23 +308,20 @@ void SegmentationCameraSensorTest::ImagesWithBuiltinSDF(
   // wait for a new frame
   WaitForNewFrame(mgr);
 
-  // the label is in the last channel
-  leftLabel =   g_buffer[leftIndex   + 2];
-  rightLabel =  g_buffer[rightIndex  + 2];
-  middleLabel = g_buffer[middleIndex + 2];
+  const int labelOffset = 2;
+  leftLabel =   g_buffer[leftIndex   + labelOffset];
+  rightLabel =  g_buffer[rightIndex  + labelOffset];
+  middleLabel = g_buffer[middleIndex + labelOffset];
 
   // the instance count is in the first channel
   uint8_t leftCount =   g_buffer[leftIndex];
   uint8_t rightCount =  g_buffer[rightIndex];
   uint8_t middleCount = g_buffer[middleIndex];
 
-  // test
-  g_counter = 0;
-
   // check the label
-  EXPECT_EQ(leftLabel, 1);
-  EXPECT_EQ(middleLabel, 2);
-  EXPECT_EQ(rightLabel, 1);
+  EXPECT_EQ(leftLabel, leftBoxLabel);
+  EXPECT_EQ(middleLabel, middleBoxLabel);
+  EXPECT_EQ(rightLabel, rightBoxLabel);
 
   // instance count
   EXPECT_EQ(middleCount, 1);