export

cubao · Nov 12, 2023 · d68e73d · d68e73d
1 parent 8beb9bd
commit d68e73d
Show file tree

Hide file tree

Showing 4 changed files with 43 additions and 16 deletions.
diff --git a/src/geobuf/geobuf_plus.hpp b/src/geobuf/geobuf_plus.hpp
@@ -33,11 +33,13 @@ struct GeobufPlus
         return {};
     }
 
-    static std::string encode(const Planet &planet) {
-        return "";
-    }
+    static std::string encode(const Planet &planet) { return ""; }
 
-    static bool encode(const std::string &input_geojson_path, const std::string &output_geobuf_plus_path) {
+    static bool encode(const std::string &input_geojson_path,
+                       const std::string &output_geobuf_plus_path,
+                       uint8_t precision = 8, bool only_xy = false,
+                       std::optional<int> round_z = 3)
+    {
         spdlog::info("Loading {} ...", input_geojson_path);
         auto json = mapbox::geobuf::load_json(input_geojson_path);
         if (json.IsNull()) {
@@ -46,43 +48,64 @@ struct GeobufPlus
         }
         auto geojson = mapbox::geojson::convert(json);
         if (geojson.is<mapbox::geojson::geometry>()) {
-            geojson = mapbox::geojson::feature_collection{
-                mapbox::geojson::feature{geojson.get<mapbox::geojson::geometry>()}
-            };
+            geojson =
+                mapbox::geojson::feature_collection{mapbox::geojson::feature{
+                    geojson.get<mapbox::geojson::geometry>()}};
         } else if (geojson.is<mapbox::geojson::feature>()) {
             geojson = mapbox::geojson::feature_collection{
-                    geojson.get<mapbox::geojson::feature>()
-            };
+                geojson.get<mapbox::geojson::feature>()};
         }
         auto &fc = geojson.get<mapbox::geojson::feature_collection>();
         spdlog::info("encoding {} features...", fc.size());
 
         auto planet = Planet(fc);
-        auto bytes = encode(planet);
         FILE *fp = fopen(output_geobuf_plus_path.c_str(), "wb");
         if (!fp) {
-            spdlog::error("failed to open {} for writing", output_geobuf_plus_path);
+            spdlog::error("failed to open {} for writing",
+                          output_geobuf_plus_path);
             return false;
         }
+        // magic
         fwrite("GeobufPlus", 10, 1, fp);
         int num_features = fc.size();
+        // #features
         fwrite(&num_features, sizeof(num_features), 1, fp);
 
         auto &rtree = planet.packed_rtree();
         auto extent = rtree.getExtent();
+        // spatial
+        //      extent/bbox
         fwrite(&extent, sizeof(extent), 1, fp);
+        //      num_items
         int num_items = rtree.getNumItems();
         fwrite(&num_items, sizeof(num_items), 1, fp);
+        //      num_nodes
         int num_nodes = rtree.getNumNodes();
         fwrite(&num_nodes, sizeof(num_nodes), 1, fp);
+        //      node_size
         int node_size = rtree.getNodeSize();
         fwrite(&node_size, sizeof(node_size), 1, fp);
+        //      tree_size
         int tree_size = rtree.size();
         fwrite(&tree_size, sizeof(tree_size), 1, fp);
+        //      tree_bytes
         rtree.streamWrite([&](const uint8_t *data, size_t size) {
             fwrite(data, 1, size, fp);
         });
 
+        const int padding = 0;
+        fwrite(&padding, sizeof(padding), 1, fp);
+
+        mapbox::geobuf::Encoder encoder(std::pow(10, precision), only_xy,
+                                        round_z);
+        auto bytes = encoder.encode(geojson);
+        std::vector<int> offsets = encoder.__offsets();
+        int num_offsets = offsets.size();
+        fwrite(&num_offsets, sizeof(num_offsets), 1, fp);
+        fwrite(offsets.data(), sizeof(offsets[0]), offsets.size(), fp);
+        fwrite(&padding, sizeof(padding), 1, fp);
+        fwrite(bytes.data(), 1, bytes.size(), fp);
+
         fclose(fp);
         return true;
         // write magic, geobuf_plus

diff --git a/src/geobuf/planet.hpp b/src/geobuf/planet.hpp
@@ -132,9 +132,7 @@ struct Planet
         return fc;
     }
 
-    const FlatGeobuf::PackedRTree &packed_rtree() const {
-        return rtree();
-    }
+    const FlatGeobuf::PackedRTree &packed_rtree() const { return rtree(); }
 
   private:
     FeatureCollection features_;

diff --git a/src/main.cpp b/src/main.cpp
@@ -282,7 +282,12 @@ PYBIND11_MODULE(_pybind11_geobuf, m)
         .def(py::init<>())
         .def("init", &GeobufPlus::init, "header_bytes"_a)
         //
-        .def_static("encode", py::overload_cast<const std::string &, const std::string &>(&GeobufPlus::encode), "input_geojson_path"_a, "output_geobuf_plus_path"_a)
+        .def_static(
+            "encode",
+            py::overload_cast<const std::string &, const std::string &, uint8_t,
+                              bool, std::optional<int>>(&GeobufPlus::encode),
+            "input_geojson_path"_a, "output_geobuf_plus_path"_a, py::kw_only(),
+            "precision"_a = 8, "only_xy"_a = false, "round_z"_a = 3)
         //
         ;
 

diff --git a/tests/test_geobuf.py b/tests/test_geobuf.py
@@ -1891,8 +1891,9 @@ def test_query():
 if __name__ == "__main__":
 
     from pybind11_geobuf import GeobufPlus
+
     ipath = f"{__pwd}/../data/suzhoubeizhan.json"
-    print(GeobufPlus.encode(ipath, f'{__pwd}/../build/export.pbf'))
+    print(GeobufPlus.encode(ipath, f"{__pwd}/../build/export.pbf"))
 
     np.set_printoptions(suppress=True)
     pwd = os.path.abspath(os.path.dirname(__file__))