From f9d4383d30df8af991246c5ef5ba75e74bc6eb05 Mon Sep 17 00:00:00 2001 From: Paul-Edouard Sarlin <15985472+sarlinpe@users.noreply.github.com> Date: Tue, 16 Jan 2024 10:17:17 +0100 Subject: [PATCH] Incremental mapping callbacks (#238) * Add incremental mapping callbacks * Add minimal database bindings * Update example * Add missing header * Fix incorrect doc --- example.py | 16 +++++++++++-- pycolmap/pipeline/sfm.h | 20 ++++++++++++++--- pycolmap/scene/bindings.h | 2 ++ pycolmap/scene/database.h | 47 +++++++++++++++++++++++++++++++++++++++ 4 files changed, 80 insertions(+), 5 deletions(-) create mode 100644 pycolmap/scene/database.h diff --git a/example.py b/example.py index 7cc422e7..5c82c98a 100644 --- a/example.py +++ b/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()}") diff --git a/pycolmap/pipeline/sfm.h b/pycolmap/pipeline/sfm.h index e26ad6e3..d76ee319 100644 --- a/pycolmap/pipeline/sfm.h +++ b/pycolmap/pipeline/sfm.h @@ -11,6 +11,7 @@ #include +#include #include #include #include @@ -48,7 +49,9 @@ std::map> 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& initial_image_pair_callback, + const std::function& next_image_callback) { THROW_CHECK_FILE_EXISTS(database_path); THROW_CHECK_DIR_EXISTS(image_path); CreateDirIfNotExists(output_path); @@ -69,7 +72,15 @@ std::map> 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"); } diff --git a/pycolmap/scene/bindings.h b/pycolmap/scene/bindings.h index 11daee06..c586b564 100644 --- a/pycolmap/scene/bindings.h +++ b/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); } diff --git a/pycolmap/scene/database.h b/pycolmap/scene/database.h new file mode 100644 index 00000000..4b429634 --- /dev/null +++ b/pycolmap/scene/database.h @@ -0,0 +1,47 @@ +#include "colmap/scene/database.h" + +#include + +using namespace colmap; +using namespace pybind11::literals; +namespace py = pybind11; + +void BindDatabase(py::module& m) { + py::class_ PyDatabase(m, "Database"); + PyDatabase.def(py::init(), "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_(m, "DatabaseTransaction") + .def(py::init(), "database"_a); +}