Writing probability grid as proto with probability grid points processor (#1069)

The probability grid points processor can now be configured to write the probability grid as pbstream. Before it was only possible to write the probability grid as png.

cartographer/io/probability_grid_points_processor.cc

@@ -24,6 +24,7 @@
 #include "cartographer/io/image.h"
 #include "cartographer/io/points_batch.h"
 #include "cartographer/mapping/probability_values.h"
+#include "glog/logging.h"
 
 namespace cartographer {
 namespace io {
@@ -53,22 +54,49 @@ uint8 ProbabilityToColor(float probability_from_grid) {
              (mapping::kMaxProbability - mapping::kMinProbability)));
 }
 
+std::string FileExtensionFromOutputType(
+    const ProbabilityGridPointsProcessor::OutputType& output_type) {
+  if (output_type == ProbabilityGridPointsProcessor::OutputType::kPng) {
+    return ".png";
+  } else if (output_type == ProbabilityGridPointsProcessor::OutputType::kPb) {
+    return ".pb";
+  }
+  LOG(FATAL) << "OutputType does not exist!";
+}
+
+ProbabilityGridPointsProcessor::OutputType OutputTypeFromString(
+    const std::string& output_type) {
+  if (output_type == "png") {
+    return ProbabilityGridPointsProcessor::OutputType::kPng;
+  } else if (output_type == "pb") {
+    return ProbabilityGridPointsProcessor::OutputType::kPb;
+  } else {
+    LOG(FATAL) << "OutputType " << output_type << " does not exist!";
+  }
+}
+
 }  // namespace
 
 ProbabilityGridPointsProcessor::ProbabilityGridPointsProcessor(
     const double resolution,
     const mapping::proto::ProbabilityGridRangeDataInserterOptions2D&
         probability_grid_range_data_inserter_options,
-    const DrawTrajectories& draw_trajectories,
+    const DrawTrajectories& draw_trajectories, const OutputType& output_type,
     std::unique_ptr<FileWriter> file_writer,
     const std::vector<mapping::proto::Trajectory>& trajectories,
     PointsProcessor* const next)
     : draw_trajectories_(draw_trajectories),
+      output_type_(output_type),
       trajectories_(trajectories),
       file_writer_(std::move(file_writer)),
       next_(next),
       range_data_inserter_(probability_grid_range_data_inserter_options),
-      probability_grid_(CreateProbabilityGrid(resolution)) {}
+      probability_grid_(CreateProbabilityGrid(resolution)) {
+  LOG_IF(WARNING, output_type == OutputType::kPb &&
+                      draw_trajectories_ == DrawTrajectories::kYes)
+      << "Drawing the trajectories is not supported when writing the "
+         "probability grid as protobuf.";
+}
 
 std::unique_ptr<ProbabilityGridPointsProcessor>
 ProbabilityGridPointsProcessor::FromDictionary(
@@ -80,12 +108,17 @@ ProbabilityGridPointsProcessor::FromDictionary(
                                   dictionary->GetBool("draw_trajectories"))
                                      ? DrawTrajectories::kYes
                                      : DrawTrajectories::kNo;
+  const auto output_type =
+      dictionary->HasKey("output_type")
+          ? OutputTypeFromString(dictionary->GetString("output_type"))
+          : OutputType::kPng;
   return common::make_unique<ProbabilityGridPointsProcessor>(
       dictionary->GetDouble("resolution"),
       mapping::CreateProbabilityGridRangeDataInserterOptions2D(
           dictionary->GetDictionary("range_data_inserter").get()),
-      draw_trajectories,
-      file_writer_factory(dictionary->GetString("filename") + ".png"),
+      draw_trajectories, output_type,
+      file_writer_factory(dictionary->GetString("filename") +
+                          FileExtensionFromOutputType(output_type)),
       trajectories, next);
 }
 
@@ -97,17 +130,29 @@ void ProbabilityGridPointsProcessor::Process(
 }
 
 PointsProcessor::FlushResult ProbabilityGridPointsProcessor::Flush() {
-  Eigen::Array2i offset;
-  std::unique_ptr<Image> image =
-      DrawProbabilityGrid(probability_grid_, &offset);
-  if (image != nullptr) {
-    if (draw_trajectories_ ==
-        ProbabilityGridPointsProcessor::DrawTrajectories::kYes) {
-      DrawTrajectoriesIntoImage(probability_grid_, offset, trajectories_,
-                                image->GetCairoSurface().get());
+  if (output_type_ == OutputType::kPng) {
+    Eigen::Array2i offset;
+    std::unique_ptr<Image> image =
+        DrawProbabilityGrid(probability_grid_, &offset);
+    if (image != nullptr) {
+      if (draw_trajectories_ ==
+          ProbabilityGridPointsProcessor::DrawTrajectories::kYes) {
+        DrawTrajectoriesIntoImage(probability_grid_, offset, trajectories_,
+                                  image->GetCairoSurface().get());
+      }
+      image->WritePng(file_writer_.get());
+      CHECK(file_writer_->Close());
     }
-    image->WritePng(file_writer_.get());
+  } else if (output_type_ == OutputType::kPb) {
+    const auto probability_grid_proto = probability_grid_.ToProto();
+    std::string probability_grid_serialized;
+    probability_grid_proto.SerializeToString(&probability_grid_serialized);
+    file_writer_->Write(probability_grid_serialized.data(),
+                        probability_grid_serialized.size());
     CHECK(file_writer_->Close());
+  } else {
+    LOG(FATAL) << "Output Type " << FileExtensionFromOutputType(output_type_)
+               << " is not supported.";
   }
 
   switch (next_->Flush()) {

cartographer/io/probability_grid_points_processor.h

@@ -40,13 +40,14 @@ class ProbabilityGridPointsProcessor : public PointsProcessor {
   constexpr static const char* kConfigurationFileActionName =
       "write_probability_grid";
   enum class DrawTrajectories { kNo, kYes };
+  enum class OutputType { kPng, kPb };
   ProbabilityGridPointsProcessor(
       double resolution,
       const mapping::proto::ProbabilityGridRangeDataInserterOptions2D&
           probability_grid_range_data_inserter_options,
-      const DrawTrajectories& draw_trajectories,
+      const DrawTrajectories& draw_trajectories, const OutputType& output_type,
       std::unique_ptr<FileWriter> file_writer,
-      const std::vector<mapping::proto::Trajectory>& trajectorios,
+      const std::vector<mapping::proto::Trajectory>& trajectories,
       PointsProcessor* next);
   ProbabilityGridPointsProcessor(const ProbabilityGridPointsProcessor&) =
       delete;
@@ -65,6 +66,7 @@ class ProbabilityGridPointsProcessor : public PointsProcessor {
 
  private:
   const DrawTrajectories draw_trajectories_;
+  const OutputType output_type_;
   const std::vector<mapping::proto::Trajectory> trajectories_;
   std::unique_ptr<FileWriter> file_writer_;
   PointsProcessor* const next_;

cartographer/io/probability_grid_points_processor_test.cc

@@ -0,0 +1,156 @@
+/*
+ * Copyright 2018 The Cartographer Authors
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#include "cartographer/io/probability_grid_points_processor.h"
+
+#include <string>
+#include "cartographer/common/lua_parameter_dictionary.h"
+#include "cartographer/common/lua_parameter_dictionary_test_helpers.h"
+#include "cartographer/common/port.h"
+#include "cartographer/io/fake_file_writer.h"
+#include "cartographer/io/points_processor_pipeline_builder.h"
+#include "cartographer/mapping/2d/probability_grid_range_data_inserter_2d.h"
+#include "gmock/gmock.h"
+#include "gtest/gtest.h"
+
+namespace cartographer {
+namespace io {
+namespace {
+
+std::unique_ptr<PointsBatch> CreatePointsBatch() {
+  auto points_batch = cartographer::common::make_unique<PointsBatch>();
+  points_batch->origin = Eigen::Vector3f(0, 0, 0);
+  points_batch->points.emplace_back(0.0f, 0.0f, 0.0f);
+  points_batch->points.emplace_back(0.0f, 1.0f, 2.0f);
+  points_batch->points.emplace_back(1.0f, 2.0f, 4.0f);
+  points_batch->points.emplace_back(0.0f, 3.0f, 5.0f);
+  points_batch->points.emplace_back(3.0f, 0.0f, 6.0f);
+  return points_batch;
+}
+
+::cartographer::io::FileWriterFactory CreateFakeFileWriterFactory(
+    const std::string& expected_filename,
+    std::shared_ptr<std::vector<char>> fake_file_writer_output) {
+  return [&fake_file_writer_output,
+          &expected_filename](const std::string& full_filename) {
+    EXPECT_EQ(expected_filename, full_filename);
+    return ::cartographer::common::make_unique<
+        ::cartographer::io::FakeFileWriter>(full_filename,
+                                            fake_file_writer_output);
+  };
+}
+
+std::vector<std::unique_ptr<::cartographer::io::PointsProcessor>>
+CreatePipelineFromDictionary(
+    common::LuaParameterDictionary* const pipeline_dictionary,
+    const std::vector<mapping::proto::Trajectory>& trajectories,
+    ::cartographer::io::FileWriterFactory file_writer_factory) {
+  auto builder = ::cartographer::common::make_unique<
+      ::cartographer::io::PointsProcessorPipelineBuilder>();
+  builder->Register(
+      ProbabilityGridPointsProcessor::kConfigurationFileActionName,
+      [&trajectories, file_writer_factory](
+          common::LuaParameterDictionary* const dictionary,
+          PointsProcessor* const next) -> std::unique_ptr<PointsProcessor> {
+        return ProbabilityGridPointsProcessor::FromDictionary(
+            trajectories, file_writer_factory, dictionary, next);
+      });
+
+  return builder->CreatePipeline(pipeline_dictionary);
+}
+
+std::vector<char> CreateExpectedProbabilityGrid(
+    std::unique_ptr<PointsBatch> points_batch,
+    common::LuaParameterDictionary* const probability_grid_options) {
+  ::cartographer::mapping::ProbabilityGridRangeDataInserter2D
+      range_data_inserter(
+          cartographer::mapping::
+              CreateProbabilityGridRangeDataInserterOptions2D(
+                  probability_grid_options->GetDictionary("range_data_inserter")
+                      .get()));
+  auto probability_grid =
+      CreateProbabilityGrid(probability_grid_options->GetDouble("resolution"));
+  range_data_inserter.Insert({points_batch->origin, points_batch->points, {}},
+                             &probability_grid);
+
+  std::vector<char> probability_grid_proto(
+      probability_grid.ToProto().ByteSize());
+  probability_grid.ToProto().SerializePartialToArray(
+      probability_grid_proto.data(), probability_grid_proto.size());
+  return probability_grid_proto;
+}
+
+std::unique_ptr<common::LuaParameterDictionary> CreateParameterDictionary() {
+  auto parameter_dictionary =
+      cartographer::common::LuaParameterDictionary::NonReferenceCounted(
+          R"text(
+          pipeline = { 
+            { 
+              action = "write_probability_grid", 
+              resolution = 0.05, 
+              range_data_inserter = { 
+                insert_free_space = true, 
+                hit_probability = 0.55, 
+                miss_probability = 0.49, 
+              }, 
+              draw_trajectories = false, 
+              output_type = "pb", 
+              filename = "map" 
+            }
+          } 
+          return pipeline
+    )text",
+          common::make_unique<cartographer::common::DummyFileResolver>());
+  return parameter_dictionary;
+}
+
+class ProbabilityGridPointsProcessorTest : public ::testing::Test {
+ protected:
+  ProbabilityGridPointsProcessorTest()
+      : pipeline_dictionary_(CreateParameterDictionary()) {}
+
+  void Run(const std::string& expected_filename) {
+    const auto pipeline = CreatePipelineFromDictionary(
+        pipeline_dictionary_.get(), dummy_trajectories_,
+        CreateFakeFileWriterFactory(expected_filename,
+                                    fake_file_writer_output_));
+    EXPECT_TRUE(pipeline.size() > 0);
+
+    do {
+      pipeline.back()->Process(CreatePointsBatch());
+    } while (pipeline.back()->Flush() ==
+             cartographer::io::PointsProcessor::FlushResult::kRestartStream);
+  }
+
+  std::shared_ptr<std::vector<char>> fake_file_writer_output_ =
+      std::make_shared<std::vector<char>>();
+  std::unique_ptr<cartographer::common::LuaParameterDictionary>
+      pipeline_dictionary_;
+  const std::vector<mapping::proto::Trajectory> dummy_trajectories_;
+};
+
+TEST_F(ProbabilityGridPointsProcessorTest, WriteProto) {
+  const auto expected_prob_grid_proto = CreateExpectedProbabilityGrid(
+      CreatePointsBatch(),
+      pipeline_dictionary_->GetArrayValuesAsDictionaries().front().get());
+  Run("map.pb");
+  EXPECT_THAT(*fake_file_writer_output_,
+              ::testing::ContainerEq(expected_prob_grid_proto));
+}
+
+}  // namespace
+}  // namespace io
+}  // namespace cartographer