Cuda bindings (MVS, SIFT GPU)

CMakeLists.txt

@@ -13,6 +13,13 @@ endif()
 
 add_subdirectory(pybind11)
 
+
+if (${COLMAP_CUDA_ENABLED})
+    message("Compiling bindings with CUDA support.")
+    add_definitions(-DCUDA_ENABLED)
+else()
+    message("Compiling bindings without CUDA support.")
+endif()
 pybind11_add_module(pycolmap main.cc)
 target_link_libraries(pycolmap PRIVATE ${COLMAP_LIBRARIES})
 if(MSVC)

README.md

@@ -258,15 +258,37 @@ camera_dict = {
 
 ## Reconstruction pipeline
 
-We provide bindings for multiple step of the standard reconstruction pipeline. They are defined in `pipeline.cc` and include:
+We provide bindings for multiple step of the standard reconstruction pipeline. They are defined in `pipeline/` and include:
 
+- extracting and matching SIFT features
 - importing an image folder into a COLMAP database
 - inferring the camera parameters from the EXIF metadata of an image file
 - running two-view geometric verification of matches on a COLMAP database
 - triangulating points into an existing COLMAP model
 - running incremental reconstruction from a COLMAP database
+- dense reconstruction with multi-view stereo
 
-For an example of usage, see [`hloc/reconstruction.py`](https://github.com/cvg/Hierarchical-Localization/blob/master/hloc/reconstruction.py).
+Dense reconstruction from a folder of images can be performed with:
+```python
+output_path: pathlib.Path
+image_dir: pathlib.Path
+
+output_path.mkdir()
+mvs_path = output_path / "mvs"
+database_path = output_path / "database.db"
+
+pycolmap.extract_features(database_path, image_dir)
+pycolmap.match_exhaustive(database_path)
+maps = pycolmap.incremental_mapping(database_path, image_dir, output_path)
+maps[0].write(output_path)
+pycolmap.undistort_images(mvs_path, output_path, image_dir)
+pycolmap.patch_match_stereo(mvs_path)
+pycolmap.stereo_fusion(mvs_path / "dense.ply", mvs_path)
+```
+
+PyCOLMAP can leverage the GPU for feature extraction, matching, and multi-view stereo if COLMAP was compiled with CUDA support. This requires to build the package from source and is not available with the PyPI wheels.
+
+For another example of usage, see [`hloc/reconstruction.py`](https://github.com/cvg/Hierarchical-Localization/blob/master/hloc/reconstruction.py).
 
 ## SIFT feature extraction
 
@@ -276,22 +298,18 @@ import pycolmap
 from PIL import Image, ImageOps
 
 # Input should be grayscale image with range [0, 1].
-with open('image.jpg', 'rb') as f:
-    img = Image.open(f)
-    img = img.convert('RGB')
-    img = ImageOps.grayscale(img)
-    img = np.array(img).astype(np.float) / 255.
+img = Image.open('image.jpg').convert('RGB')
+img = ImageOps.grayscale(img)
+img = np.array(img).astype(np.float) / 255.
+
+# Optional parameters:
+# - options: dict or pycolmap.SiftExtractionOptions
+# - device: default pycolmap.Device.auto uses the GPU if available
+sift = pycolmap.Sift()
 
 # Parameters:
 # - image: HxW float array
-# Named parameters:
-# - num_octaves: int (4)
-# - octave_resolution: int (3)
-# - first_octave: int (0)
-# - edge_thresh: float (10)
-# - peak_thresh: float (0.01)
-# - upright: bool (False)
-keypoints, scores, descriptors = pycolmap.extract_sift(img)
+keypoints, scores, descriptors = sift.extract(img)
 # Returns:
 # - keypoints: Nx4 array; format: x (j), y (i), sigma, angle
 # - scores: N array; DoG scores

helpers.h

@@ -2,6 +2,8 @@
 
 #pragma once
 
+#include <colmap/util/threading.h>
+
 #include <pybind11/embed.h>
 #include <pybind11/eval.h>
 #include <pybind11/numpy.h>
@@ -102,6 +104,8 @@ inline void make_dataclass(py::class_<T> cls) {
                     <<py::str(ex.value()).cast<std::string>()<<std::endl;
                     throw;
                 } else {
+                    std::cerr<<"Internal Error: "
+                    <<py::str(ex.value()).cast<std::string>()<<std::endl;
                     throw;
                 }
             }
@@ -176,3 +180,69 @@ inline void make_dataclass(py::class_<T> cls) {
         return dict;
     });
 }
+
+
+// Catch python keyboard interrupts
+
+/*
+// single
+if (PyInterrupt().Raised()) {
+    // stop the execution and raise an exception
+    throw py::error_already_set();
+}
+
+// loop
+PyInterrupt py_interrupt = PyInterrupt(2.0)
+for (...) {
+    if (py_interrupt.Raised()) {
+        // stop the execution and raise an exception
+        throw py::error_already_set();
+    }
+    // Do your workload here
+}
+
+
+*/
+struct PyInterrupt {
+  using clock = std::chrono::steady_clock;
+  using sec = std::chrono::duration<double>;
+  PyInterrupt(double gap = -1.0);
+
+  inline bool Raised();
+
+ private:
+  std::mutex mutex_;
+  bool found = false;
+  colmap::Timer timer_;
+  clock::time_point start;
+  double gap_;
+};
+
+PyInterrupt::PyInterrupt(double gap) : gap_(gap), start(clock::now()) {}
+
+bool PyInterrupt::Raised() {
+  const sec duration = clock::now() - start;
+  if (!found && duration.count() > gap_) {
+    std::lock_guard<std::mutex> lock(mutex_);
+    py::gil_scoped_acquire acq;
+    found = (PyErr_CheckSignals() != 0);
+    start = clock::now();
+  }
+  return found;
+}
+
+
+// Instead of thread.Wait() call this to allow interrupts through python
+void PyWait(Thread* thread, double gap=2.0) {
+    PyInterrupt py_interrupt(gap);
+    while(thread->IsRunning()) {
+        if (py_interrupt.Raised()) {
+            std::cerr<<"Stopping thread..."<<std::endl;
+            thread->Stop();
+            thread->Wait();
+            throw py::error_already_set();
+        }
+    }
+    // after finishing join the thread to avoid abort
+    thread->Wait();
+}

main.cc

@@ -16,8 +16,11 @@ namespace py = pybind11;
 #include "homography_decomposition.cc"
 #include "transformations.cc"
 #include "sift.cc"
-#include "pipeline.cc"
 #include "helpers.h"
+#include "utils.h"
+
+#include "pipeline/sfm.cc"
+#include "pipeline/mvs.cc"
 
 #include "reconstruction/reconstruction.cc"
 #include "reconstruction/correspondence_graph.cc"
@@ -34,6 +37,14 @@ PYBIND11_MODULE(pycolmap, m) {
     m.attr("__version__") = py::str("dev");
 #endif
 
+    auto PyDevice = py::enum_<Device>(m, "Device")
+        .value("auto", Device::AUTO)
+        .value("cpu", Device::CPU)
+        .value("cuda", Device::CUDA);
+    AddStringToEnumConstructor(PyDevice);
+
+    m.attr("has_cuda") = IsGPU(Device::AUTO);
+
     // Estimators
     auto PyRANSACOptions = py::class_<RANSACOptions>(m, "RANSACOptions")
         .def(py::init<>([]() {
@@ -69,13 +80,6 @@ PYBIND11_MODULE(pycolmap, m) {
           py::arg("points2"),
           "Analytical Homography Decomposition.");
 
-    // SIFT.
-    m.def("extract_sift", &extract_sift,
-          py::arg("image"),
-          py::arg("num_octaves") = 4, py::arg("octave_resolution") = 3, py::arg("first_octave") = 0,
-          py::arg("edge_thresh") = 10.0, py::arg("peak_thresh") = 0.01, py::arg("upright") = false,
-          "Extract SIFT features.");
-
     // Reconstruction bindings
     init_reconstruction(m);
 
@@ -92,7 +96,11 @@ PYBIND11_MODULE(pycolmap, m) {
     init_transforms(m);
 
     // Main reconstruction steps
-    init_pipeline(m);
+    init_sfm(m);
+    init_mvs(m);
+
+    // SIFT feature detector and descriptor
+    init_sift(m);
 
     py::add_ostream_redirect(m, "ostream");
 }