Incremental mapping callbacks (#238)

* Add incremental mapping callbacks

* Add minimal database bindings

* Update example

* Add missing header

* Fix incorrect doc

example.py

@@ -3,6 +3,8 @@
 import zipfile
 from pathlib import Path
 
+import enlighten
+
 import pycolmap
 from pycolmap import logging
 
@@ -28,13 +30,23 @@ def run():
     if database_path.exists():
         database_path.unlink()
     pycolmap.extract_features(database_path, image_path)
-
     pycolmap.match_exhaustive(database_path)
+    num_images = pycolmap.Database(database_path).num_images
 
     if sfm_path.exists():
         shutil.rmtree(sfm_path)
     sfm_path.mkdir(exist_ok=True)
-    recs = pycolmap.incremental_mapping(database_path, image_path, sfm_path)
+
+    with enlighten.Manager() as manager:
+        with manager.counter(total=num_images, desc="Images registered:") as pbar:
+            pbar.update(0, force=True)
+            recs = pycolmap.incremental_mapping(
+                database_path,
+                image_path,
+                sfm_path,
+                initial_image_pair_callback=lambda: pbar.update(2),
+                next_image_callback=lambda: pbar.update(1),
+            )
     for idx, rec in recs.items():
         logging.info(f"#{idx} {rec.summary()}")
 

pycolmap/pipeline/sfm.h

@@ -11,6 +11,7 @@
 
 #include <memory>
 
+#include <pybind11/functional.h>
 #include <pybind11/iostream.h>
 #include <pybind11/pybind11.h>
 #include <pybind11/stl.h>
@@ -48,7 +49,9 @@ std::map<size_t, std::shared_ptr<Reconstruction>> IncrementalMapping(
     const std::string& image_path,
     const std::string& output_path,
     const IncrementalMapperOptions& options,
-    const std::string& input_path) {
+    const std::string& input_path,
+    const std::function<void()>& initial_image_pair_callback,
+    const std::function<void()>& next_image_callback) {
   THROW_CHECK_FILE_EXISTS(database_path);
   THROW_CHECK_DIR_EXISTS(image_path);
   CreateDirIfNotExists(output_path);
@@ -69,7 +72,15 @@ std::map<size_t, std::shared_ptr<Reconstruction>> IncrementalMapping(
                        if (py_interrupt.Raised()) {
                          throw py::error_already_set();
                        }
+                       if (next_image_callback) {
+                         next_image_callback();
+                       }
                      });
+  if (initial_image_pair_callback) {
+    mapper.AddCallback(
+        IncrementalMapperController::INITIAL_IMAGE_PAIR_REG_CALLBACK,
+        initial_image_pair_callback);
+  }
 
   mapper.Start();
   mapper.Wait();
@@ -341,10 +352,13 @@ void BindSfM(py::module& m) {
         "output_path"_a,
         "options"_a = mapper_options,
         "input_path"_a = py::str(""),
-        "Triangulate 3D points from known poses");
+        "initial_image_pair_callback"_a = py::none(),
+        "next_image_callback"_a = py::none(),
+        "Recover 3D points and unknown camera poses");
 
   m.def("bundle_adjustment",
         &BundleAdjustment,
         "reconstruction"_a,
-        "options"_a = ba_options);
+        "options"_a = ba_options,
+        "Jointly refine 3D points and camera poses");
 }

pycolmap/scene/bindings.h

@@ -1,5 +1,6 @@
 #include "pycolmap/scene/camera.h"
 #include "pycolmap/scene/correspondence_graph.h"
+#include "pycolmap/scene/database.h"
 #include "pycolmap/scene/image.h"
 #include "pycolmap/scene/point2D.h"
 #include "pycolmap/scene/point3D.h"
@@ -18,4 +19,5 @@ void BindScene(py::module& m) {
   BindPoint3D(m);
   BindCorrespondenceGraph(m);
   BindReconstruction(m);
+  BindDatabase(m);
 }

pycolmap/scene/database.h

@@ -0,0 +1,47 @@
+#include "colmap/scene/database.h"
+
+#include <pybind11/pybind11.h>
+
+using namespace colmap;
+using namespace pybind11::literals;
+namespace py = pybind11;
+
+void BindDatabase(py::module& m) {
+  py::class_<Database> PyDatabase(m, "Database");
+  PyDatabase.def(py::init<const std::string&>(), "path"_a)
+      .def("open", &Database::Open, "path"_a)
+      .def("close", &Database::Close)
+      .def_property_readonly("num_cameras", &Database::NumCameras)
+      .def_property_readonly("num_images", &Database::NumImages)
+      .def_property_readonly("num_keypoints", &Database::NumKeypoints)
+      .def_property_readonly("num_keypoints_for_image",
+                             &Database::NumKeypointsForImage)
+      .def_property_readonly("num_descriptors", &Database::NumDescriptors)
+      .def_property_readonly("num_descriptors_for_image",
+                             &Database::NumDescriptorsForImage)
+      .def_property_readonly("num_matches", &Database::NumMatches)
+      .def_property_readonly("num_inlier_matches", &Database::NumInlierMatches)
+      .def_property_readonly("num_matched_image_pairs",
+                             &Database::NumMatchedImagePairs)
+      .def_property_readonly("num_verified_image_pairs",
+                             &Database::NumVerifiedImagePairs)
+      .def("image_pair_to_pair_id", &Database::ImagePairToPairId)
+      .def("pair_id_to_image_pair", &Database::PairIdToImagePair)
+      .def("read_camera", &Database::ReadCamera)
+      .def("read_all_cameras", &Database::ReadAllCameras)
+      .def("read_image", &Database::ReadImage)
+      .def("read_image_with_name", &Database::ReadImageWithName)
+      .def("read_all_images", &Database::ReadAllImages)
+      // ReadKeypoints
+      // ReadDescriptors
+      // ReadMatches
+      // ReadAllMatches
+      .def("read_two_view_geometry", &Database::ReadTwoViewGeometry)
+      // ReadTwoViewGeometries
+      // ReadTwoViewGeometryNumInliers
+      .def("write_camera", &Database::WriteCamera)
+      .def("write_image", &Database::WriteImage);
+
+  py::class_<DatabaseTransaction>(m, "DatabaseTransaction")
+      .def(py::init<Database*>(), "database"_a);
+}