Add serialization and deserialization methods to Octree

cpp/open3d/geometry/Octree.cpp

@@ -49,6 +49,35 @@ std::shared_ptr<OctreeNode> OctreeNode::ConstructFromJsonValue(
     return node;
 }
 
+std::shared_ptr<OctreeNode> OctreeNode::ConstructFromBinaryStream(
+        const std::string& in, size_t& offset) {
+    if (in.size() < offset + 18) return nullptr;
+
+    std::shared_ptr<OctreeNode> node = nullptr;
+    if (std::string(&in[offset], 18) == "OctreeInternalNode") {
+        node = std::make_shared<OctreeInternalNode>();
+    } else if (std::string(&in[offset], 23) == "OctreeInternalPointNode") {
+        node = std::make_shared<OctreeInternalPointNode>();
+    } else if (std::string(&in[offset], 19) == "OctreeColorLeafNode") {
+        node = std::make_shared<OctreeColorLeafNode>();
+    } else if (std::string(&in[offset], 24) == "OctreePointColorLeafNode") {
+        node = std::make_shared<OctreePointColorLeafNode>();
+    } else {
+        utility::LogError("Unhandled class id {}",
+                          std::string(&in[offset], 24));
+    }
+
+    // Convert from binary
+    if (node != nullptr) {
+        bool deserialization_success =
+                node->DeserializeFromBinaryStream(in, offset);
+        if (!deserialization_success) {
+            node = nullptr;
+        }
+    }
+    return node;
+}
+
 std::shared_ptr<OctreeNodeInfo> OctreeInternalNode::GetInsertionNodeInfo(
         const std::shared_ptr<OctreeNodeInfo>& node_info,
         const Eigen::Vector3d& point) {
@@ -120,6 +149,46 @@ bool OctreeInternalNode::ConvertFromJsonValue(const Json::Value& value) {
     return rc;
 }
 
+bool OctreeInternalNode::SerializeToBinaryStream(std::string& out) const {
+    bool rc = true;
+    // Write class identifier
+    const char class_id[] = "OctreeInternalNode";
+    out.append(class_id, sizeof(class_id));
+    // Write children
+    for (int cid = 0; cid < 8; ++cid) {
+        uint8_t has_child = (children_[cid] != nullptr);
+        out.append(reinterpret_cast<const char*>(&has_child),
+                   sizeof(has_child));
+        if (has_child) {
+            rc = rc && children_[cid]->SerializeToBinaryStream(out);
+        }
+    }
+    return rc;
+}
+
+bool OctreeInternalNode::DeserializeFromBinaryStream(const std::string& in,
+                                                     size_t& offset) {
+    // Read and check class identifier
+    if (in.size() < offset + 19) return false;
+    if (std::string(&in[offset], 18) != "OctreeInternalNode") return false;
+    offset += 19;
+    // Read children
+    for (int cid = 0; cid < 8; ++cid) {
+        if (in.size() < offset + sizeof(uint8_t)) return false;
+        uint8_t has_child = *reinterpret_cast<const uint8_t*>(&in[offset]);
+        offset += sizeof(uint8_t);
+        if (has_child) {
+            children_[cid] = OctreeNode::ConstructFromBinaryStream(in, offset);
+            if (!children_[cid]) {
+                return false;
+            }
+        } else {
+            children_[cid] = nullptr;
+        }
+    }
+    return true;
+}
+
 std::function<std::shared_ptr<OctreeInternalNode>()>
 OctreeInternalPointNode::GetInitFunction() {
     return []() -> std::shared_ptr<geometry::OctreeInternalNode> {
@@ -190,6 +259,68 @@ bool OctreeInternalPointNode::ConvertFromJsonValue(const Json::Value& value) {
     return rc;
 }
 
+bool OctreeInternalPointNode::SerializeToBinaryStream(std::string& out) const {
+    bool rc = true;
+    // Write class identifier
+    const char class_id[] = "OctreeInternalPointNode";
+    out.append(class_id, sizeof(class_id));
+    // Write children
+    for (int cid = 0; cid < 8; ++cid) {
+        uint8_t has_child = (children_[cid] != nullptr);
+        out.append(reinterpret_cast<const char*>(&has_child),
+                   sizeof(has_child));
+        if (has_child) {
+            rc = rc && children_[cid]->SerializeToBinaryStream(out);
+        }
+    }
+    // Write indices
+    uint64_t num_indices = indices_.size();
+    out.append(reinterpret_cast<const char*>(&num_indices),
+               sizeof(num_indices));
+    for (uint64_t idx : indices_) {
+        out.append(reinterpret_cast<const char*>(&idx), sizeof(idx));
+    }
+    return rc;
+}
+
+bool OctreeInternalPointNode::DeserializeFromBinaryStream(const std::string& in,
+                                                          size_t& offset) {
+    // Read and check class identifier
+    if (in.size() < offset + 24) return false;
+    if (std::string(&in[offset], 23) != "OctreeInternalPointNode") return false;
+    offset += 24;
+
+    // Read children
+    for (int cid = 0; cid < 8; ++cid) {
+        if (in.size() < offset + sizeof(uint8_t)) return false;
+        uint8_t has_child = *reinterpret_cast<const uint8_t*>(&in[offset]);
+        offset += sizeof(uint8_t);
+        if (has_child) {
+            children_[cid] = OctreeNode::ConstructFromBinaryStream(in, offset);
+            if (!children_[cid]) {
+                return false;
+            }
+        } else {
+            children_[cid] = nullptr;
+        }
+    }
+
+    // Read indices size
+    if (in.size() < offset + sizeof(uint64_t)) return false;
+    uint64_t indices_size = *reinterpret_cast<const uint64_t*>(&in[offset]);
+    offset += sizeof(uint64_t);
+    // Read indices
+    indices_.clear();
+    for (uint64_t i = 0; i < indices_size; ++i) {
+        if (in.size() < offset + sizeof(uint64_t)) return false;
+        uint64_t idx = *reinterpret_cast<const uint64_t*>(&in[offset]);
+        offset += sizeof(uint64_t);
+        indices_.push_back(idx);
+    }
+
+    return true;
+}
+
 std::function<std::shared_ptr<OctreeLeafNode>()>
 OctreeColorLeafNode::GetInitFunction() {
     return []() -> std::shared_ptr<geometry::OctreeLeafNode> {
@@ -247,6 +378,29 @@ bool OctreeColorLeafNode::ConvertFromJsonValue(const Json::Value& value) {
     return EigenVector3dFromJsonArray(color_, value["color"]);
 }
 
+bool OctreeColorLeafNode::SerializeToBinaryStream(std::string& out) const {
+    // Write class identifier
+    const char class_id[] = "OctreeColorLeafNode";
+    out.append(class_id, sizeof(class_id));
+    // Write color_ (Eigen::Vector3d)
+    out.append(reinterpret_cast<const char*>(color_.data()),
+               sizeof(double) * 3);
+    return true;
+}
+
+bool OctreeColorLeafNode::DeserializeFromBinaryStream(const std::string& in,
+                                                      size_t& offset) {
+    // Read and check class identifier
+    if (in.size() < offset + 20) return false;
+    if (std::string(&in[offset], 19) != "OctreeColorLeafNode") return false;
+    offset += 20;
+    // Read color_ (Eigen::Vector3d)
+    if (in.size() < offset + sizeof(double) * 3) return false;
+    std::memcpy(color_.data(), &in[offset], sizeof(double) * 3);
+    offset += sizeof(double) * 3;
+    return true;
+}
+
 std::function<std::shared_ptr<OctreeLeafNode>()>
 OctreePointColorLeafNode::GetInitFunction() {
     return []() -> std::shared_ptr<geometry::OctreeLeafNode> {
@@ -322,6 +476,52 @@ bool OctreePointColorLeafNode::ConvertFromJsonValue(const Json::Value& value) {
     return true;
 }
 
+bool OctreePointColorLeafNode::SerializeToBinaryStream(std::string& out) const {
+    bool rc = true;
+    // Write class identifier
+    const char class_id[] = "OctreePointColorLeafNode";
+    out.append(class_id, sizeof(class_id));
+    // Write color_ (Eigen::Vector3d)
+    out.append(reinterpret_cast<const char*>(color_.data()),
+               sizeof(double) * 3);
+    // Write indices
+    uint64_t num_indices = indices_.size();
+    out.append(reinterpret_cast<const char*>(&num_indices),
+               sizeof(num_indices));
+    for (uint64_t idx : indices_) {
+        out.append(reinterpret_cast<const char*>(&idx), sizeof(idx));
+    }
+    return rc;
+}
+
+bool OctreePointColorLeafNode::DeserializeFromBinaryStream(
+        const std::string& in, size_t& offset) {
+    // Read and check class identifier
+    if (in.size() < offset + 25) return false;
+    if (std::string(&in[offset], 24) != "OctreePointColorLeafNode")
+        return false;
+    offset += 25;
+    // Read color_ (Eigen::Vector3d)
+    if (in.size() < offset + sizeof(double) * 3) return false;
+    std::memcpy(color_.data(), &in[offset], sizeof(double) * 3);
+    offset += sizeof(double) * 3;
+
+    // Read indices size
+    if (in.size() < offset + sizeof(uint64_t)) return false;
+    uint64_t indices_size = *reinterpret_cast<const uint64_t*>(&in[offset]);
+    offset += sizeof(uint64_t);
+    // Read indices
+    indices_.clear();
+    for (uint64_t i = 0; i < indices_size; ++i) {
+        if (in.size() < offset + sizeof(uint64_t)) return false;
+        uint64_t idx = *reinterpret_cast<const uint64_t*>(&in[offset]);
+        offset += sizeof(uint64_t);
+        indices_.push_back(idx);
+    }
+
+    return true;
+}
+
 Octree::Octree(const Octree& src_octree)
     : Geometry3D(Geometry::GeometryType::Octree),
       origin_(src_octree.origin_),
@@ -786,6 +986,41 @@ bool Octree::ConvertFromJsonValue(const Json::Value& value) {
     return rc;
 }
 
+bool Octree::SerializeToBinaryStream(std::string& out) const {
+    bool rc = true;
+    // Serialize the basic attributes
+    out.clear();
+    out.append(reinterpret_cast<const char*>(origin_.data()),
+               sizeof(double) * 3);
+    out.append(reinterpret_cast<const char*>(&size_), sizeof(double));
+    out.append(reinterpret_cast<const char*>(&max_depth_), sizeof(uint64_t));
+    // Serialize the root node
+    if (root_node_) {
+        rc = rc && root_node_->SerializeToBinaryStream(out);
+    }
+    return rc;
+}
+
+bool Octree::DeserializeFromBinaryStream(const std::string& in) {
+    size_t offset = 0;
+    // Deserialize the basic attributes
+    if (in.size() < offset + sizeof(double) * 3) return false;
+    std::memcpy(origin_.data(), &in[offset], sizeof(double) * 3);
+    offset += sizeof(double) * 3;
+
+    if (in.size() < offset + sizeof(size_)) return false;
+    std::memcpy(&size_, &in[offset], sizeof(size_));
+    offset += sizeof(double);
+
+    if (in.size() < offset + sizeof(uint64_t)) return false;
+    std::memcpy(&max_depth_, &in[offset], sizeof(uint64_t));
+    offset += sizeof(uint64_t);
+
+    // Deserialize the root node
+    root_node_ = OctreeNode::ConstructFromBinaryStream(in, offset);
+    return true;
+}
+
 void Octree::CreateFromVoxelGrid(const geometry::VoxelGrid& voxel_grid) {
     origin_ = voxel_grid.origin_;
     size_ = (voxel_grid.GetMaxBound() - origin_).maxCoeff();

cpp/open3d/geometry/Octree.h

@@ -9,6 +9,7 @@
 
 #include <cstddef>
 #include <memory>
+#include <string>
 #include <vector>
 
 #include "open3d/geometry/Geometry3D.h"
@@ -77,6 +78,14 @@ class OctreeNode : public utility::IJsonConvertible {
     /// Factory function to construct an OctreeNode by parsing the json value.
     static std::shared_ptr<OctreeNode> ConstructFromJsonValue(
             const Json::Value& value);
+    /// Factory function to construct an OctreeNode by reading from a binary
+    /// stream.
+    static std::shared_ptr<OctreeNode> ConstructFromBinaryStream(
+            const std::string& in, size_t& offset);
+
+    virtual bool SerializeToBinaryStream(std::string& out) const = 0;
+    virtual bool DeserializeFromBinaryStream(const std::string& in,
+                                             size_t& offset) = 0;
 };
 
 /// \class OctreeInternalNode
@@ -121,6 +130,8 @@ class OctreeInternalNode : public OctreeNode {
 
     bool ConvertToJsonValue(Json::Value& value) const override;
     bool ConvertFromJsonValue(const Json::Value& value) override;
+    bool SerializeToBinaryStream(std::string& out) const;
+    bool DeserializeFromBinaryStream(const std::string& in, size_t& offset);
 
 public:
     /// Use vector instead of C-array for Pybind11, otherwise, need to define
@@ -156,6 +167,8 @@ class OctreeInternalPointNode : public OctreeInternalNode {
 
     bool ConvertToJsonValue(Json::Value& value) const override;
     bool ConvertFromJsonValue(const Json::Value& value) override;
+    bool SerializeToBinaryStream(std::string& out) const;
+    bool DeserializeFromBinaryStream(const std::string& in, size_t& offset);
 
 public:
     /// Indices of points associated with any children of this node
@@ -199,6 +212,8 @@ class OctreeColorLeafNode : public OctreeLeafNode {
 
     bool ConvertToJsonValue(Json::Value& value) const override;
     bool ConvertFromJsonValue(const Json::Value& value) override;
+    bool SerializeToBinaryStream(std::string& out) const;
+    bool DeserializeFromBinaryStream(const std::string& in, size_t& offset);
     /// TODO: flexible data, with lambda function for handling node
     /// Color of the node.
     Eigen::Vector3d color_ = Eigen::Vector3d(0, 0, 0);
@@ -233,6 +248,8 @@ class OctreePointColorLeafNode : public OctreeColorLeafNode {
 
     bool ConvertToJsonValue(Json::Value& value) const override;
     bool ConvertFromJsonValue(const Json::Value& value) override;
+    bool SerializeToBinaryStream(std::string& out) const;
+    bool DeserializeFromBinaryStream(const std::string& in, size_t& offset);
 
     /// Associated point indices with this node.
     std::vector<size_t> indices_;
@@ -301,6 +318,8 @@ class Octree : public Geometry3D, public utility::IJsonConvertible {
                    const Eigen::Vector3d& center) override;
     bool ConvertToJsonValue(Json::Value& value) const override;
     bool ConvertFromJsonValue(const Json::Value& value) override;
+    bool SerializeToBinaryStream(std::string& out) const;
+    bool DeserializeFromBinaryStream(const std::string& in);
 
 public:
     /// \brief Convert octree from point cloud.

cpp/open3d/io/OctreeIO.cpp

@@ -21,13 +21,15 @@ static const std::unordered_map<
         std::function<bool(const std::string &, geometry::Octree &)>>
         file_extension_to_octree_read_function{
                 {"json", ReadOctreeFromJson},
+                {"bin", ReadOctreeFromBIN},
         };
 
 static const std::unordered_map<
         std::string,
         std::function<bool(const std::string &, const geometry::Octree &)>>
         file_extension_to_octree_write_function{
                 {"json", WriteOctreeToJson},
+                {"bin", WriteOctreeToBIN},
         };
 
 std::shared_ptr<geometry::Octree> CreateOctreeFromFile(
@@ -90,5 +92,22 @@ bool WriteOctreeToJson(const std::string &filename,
                        const geometry::Octree &octree) {
     return WriteIJsonConvertibleToJSON(filename, octree);
 }
+
+bool ReadOctreeFromBIN(const std::string &filename, geometry::Octree &octree) {
+    std::string bin_data;
+    if (!ReadOctreeBinaryStreamFromBIN(filename, bin_data)) {
+        return false;
+    }
+    // Deserialize bin_data into octree
+    octree.DeserializeFromBinaryStream(bin_data);
+    return true;
+}
+
+bool WriteOctreeToBIN(const std::string &filename,
+                      const geometry::Octree &octree) {
+    std::string bin_data;
+    octree.SerializeToBinaryStream(bin_data);
+    return WriteOctreeBinaryStreamToBIN(filename, bin_data);
+}
 }  // namespace io
 }  // namespace open3d