From 2f3a9f4dccf3ce0e46346017621a32c809e04fb6 Mon Sep 17 00:00:00 2001 From: Wolfgang Hess Date: Fri, 28 Apr 2017 15:05:42 +0200 Subject: [PATCH] Add support for files containing multiple protos. (#241) This adds code for reading and writing files that contain sequences of compressed protocol buffer messages. --- cartographer/CMakeLists.txt | 1 + cartographer/io/CMakeLists.txt | 18 ++++++ cartographer/io/proto_stream.cc | 88 ++++++++++++++++++++++++++++ cartographer/io/proto_stream.h | 83 ++++++++++++++++++++++++++ cartographer/io/proto_stream_test.cc | 72 +++++++++++++++++++++++ 5 files changed, 262 insertions(+) create mode 100644 cartographer/io/CMakeLists.txt create mode 100644 cartographer/io/proto_stream.cc create mode 100644 cartographer/io/proto_stream.h create mode 100644 cartographer/io/proto_stream_test.cc diff --git a/cartographer/CMakeLists.txt b/cartographer/CMakeLists.txt index d9e35ec..ee520ed 100644 --- a/cartographer/CMakeLists.txt +++ b/cartographer/CMakeLists.txt @@ -14,6 +14,7 @@ add_subdirectory("common") add_subdirectory("ground_truth") +add_subdirectory("io") add_subdirectory("kalman_filter") add_subdirectory("mapping") add_subdirectory("mapping_2d") diff --git a/cartographer/io/CMakeLists.txt b/cartographer/io/CMakeLists.txt new file mode 100644 index 0000000..4d6723a --- /dev/null +++ b/cartographer/io/CMakeLists.txt @@ -0,0 +1,18 @@ +# Copyright 2016 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. + +google_test(io_proto_stream_test + SRCS + proto_stream_test.cc +) diff --git a/cartographer/io/proto_stream.cc b/cartographer/io/proto_stream.cc new file mode 100644 index 0000000..7f47553 --- /dev/null +++ b/cartographer/io/proto_stream.cc @@ -0,0 +1,88 @@ +/* + * Copyright 2016 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/proto_stream.h" + +namespace cartographer { +namespace io { + +namespace { + +// First eight bytes to identify our proto stream format. +const size_t kMagic = 0x7b1d1f7b5bf501db; + +void WriteSizeAsLittleEndian(size_t size, std::ostream* out) { + for (int i = 0; i != 8; ++i) { + out->put(size & 0xff); + size >>= 8; + } +} + +bool ReadSizeAsLittleEndian(std::istream* in, size_t* size) { + *size = 0; + for (int i = 0; i != 8; ++i) { + *size >>= 8; + *size += static_cast(in->get()) << 56; + } + return !in->fail(); +} + +} // namespace + +ProtoStreamWriter::ProtoStreamWriter(const string& filename) + : out_(filename, std::ios::out | std::ios::binary) { + WriteSizeAsLittleEndian(kMagic, &out_); +} + +ProtoStreamWriter::~ProtoStreamWriter() {} + +void ProtoStreamWriter::Write(const string& uncompressed_data) { + string compressed_data; + common::FastGzipString(uncompressed_data, &compressed_data); + WriteSizeAsLittleEndian(compressed_data.size(), &out_); + out_.write(compressed_data.data(), compressed_data.size()); +} + +bool ProtoStreamWriter::Close() { + out_.close(); + return !out_.fail(); +} + +ProtoStreamReader::ProtoStreamReader(const string& filename) + : in_(filename, std::ios::in | std::ios::binary) { + size_t magic; + if (!ReadSizeAsLittleEndian(&in_, &magic) || magic != kMagic) { + in_.setstate(std::ios::failbit); + } +} + +ProtoStreamReader::~ProtoStreamReader() {} + +bool ProtoStreamReader::Read(string* decompressed_data) { + size_t compressed_size; + if (!ReadSizeAsLittleEndian(&in_, &compressed_size)) { + return false; + } + string compressed_data(compressed_size, '\0'); + if (!in_.read(&compressed_data.front(), compressed_size)) { + return false; + } + common::FastGunzipString(compressed_data, decompressed_data); + return true; +} + +} // namespace io +} // namespace cartographer diff --git a/cartographer/io/proto_stream.h b/cartographer/io/proto_stream.h new file mode 100644 index 0000000..5525041 --- /dev/null +++ b/cartographer/io/proto_stream.h @@ -0,0 +1,83 @@ +/* + * Copyright 2016 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. + */ + +#ifndef CARTOGRAPHER_IO_PROTO_STREAM_H_ +#define CARTOGRAPHER_IO_PROTO_STREAM_H_ + +#include + +#include "cartographer/common/port.h" + +namespace cartographer { +namespace io { + +// A simple writer of a compressed sequence of protocol buffer messages to a +// file. The format is not intended to be compatible with any other format used +// outside of Cartographer. +// +// TODO(whess): Compress the file instead of individual messages for better +// compression performance? Should we use LZ4? +class ProtoStreamWriter { + public: + ProtoStreamWriter(const string& filename); + ~ProtoStreamWriter(); + + ProtoStreamWriter(const ProtoStreamWriter&) = delete; + ProtoStreamWriter& operator=(const ProtoStreamWriter&) = delete; + + // Serializes, compressed and writes the 'proto' to the file. + template + void WriteProto(const MessageType& proto) { + string uncompressed_data; + proto.SerializeToString(&uncompressed_data); + Write(uncompressed_data); + } + + // This should be called to check whether writing was successful. + bool Close(); + + private: + void Write(const string& uncompressed_data); + + std::ofstream out_; +}; + +// A reader of the format produced by ProtoStreamWriter. +class ProtoStreamReader { + public: + ProtoStreamReader(const string& filename); + ~ProtoStreamReader(); + + ProtoStreamReader(const ProtoStreamReader&) = delete; + ProtoStreamReader& operator=(const ProtoStreamReader&) = delete; + + template + bool ReadProto(MessageType* proto) { + string decompressed_data; + return Read(&decompressed_data) && + proto->ParseFromString(decompressed_data); + } + + private: + bool Read(string* decompressed_data); + + std::ifstream in_; +}; + +} // namespace io +} // namespace cartographer + +#endif // CARTOGRAPHER_IO_PROTO_STREAM_H_ diff --git a/cartographer/io/proto_stream_test.cc b/cartographer/io/proto_stream_test.cc new file mode 100644 index 0000000..a849ba6 --- /dev/null +++ b/cartographer/io/proto_stream_test.cc @@ -0,0 +1,72 @@ +/* + * Copyright 2016 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/proto_stream.h" + +#include +#include +#include +#include + +#include "cartographer/common/port.h" +#include "cartographer/mapping/proto/trajectory.pb.h" +#include "gtest/gtest.h" + +namespace cartographer { +namespace io { +namespace { + +class ProtoStreamTest : public ::testing::Test { + protected: + void SetUp() override { + const string tmpdir = P_tmpdir; + test_directory_ = tmpdir + "/proto_stream_test_XXXXXX"; + ASSERT_NE(mkdtemp(&test_directory_[0]), nullptr) << strerror(errno); + } + + void TearDown() override { remove(test_directory_.c_str()); } + + string test_directory_; +}; + +TEST_F(ProtoStreamTest, WriteAndReadBack) { + const string test_file = test_directory_ + "/test_trajectory.pbstream"; + { + ProtoStreamWriter writer(test_file); + for (int i = 0; i != 10; ++i) { + mapping::proto::Trajectory trajectory; + trajectory.add_node()->set_timestamp(i); + writer.WriteProto(trajectory); + } + ASSERT_TRUE(writer.Close()); + } + { + ProtoStreamReader reader(test_file); + for (int i = 0; i != 10; ++i) { + mapping::proto::Trajectory trajectory; + ASSERT_TRUE(reader.ReadProto(&trajectory)); + ASSERT_EQ(1, trajectory.node_size()); + EXPECT_EQ(i, trajectory.node(0).timestamp()); + } + mapping::proto::Trajectory trajectory; + EXPECT_FALSE(reader.ReadProto(&trajectory)); + } + remove(test_file.c_str()); +} + +} // namespace +} // namespace io +} // namespace cartographer