colmap · sarlinpe · May 1, 2024 · Apr 8, 2024 · Apr 8, 2024 · Apr 8, 2024
diff --git a/pycolmap/custom_bundle_adjustment.py b/pycolmap/custom_bundle_adjustment.py
diff --git a/pycolmap/custom_incremental_mapping.py b/pycolmap/custom_incremental_mapping.py
@@ -12,6 +12,7 @@
 
 import pycolmap
 from pycolmap import logging
+import custom_bundle_adjustment
 
 
 def extract_colors(image_path, image_id, reconstruction):
@@ -30,7 +31,9 @@ def write_snapshot(reconstruction, snapshot_path):
 
 def iterative_global_refinement(options, mapper_options, mapper):
     logging.info("Retriangulation and Global bundle adjustment")
-    mapper.iterative_global_refinement(
+    # The following is equivalent to mapper.iterative_global_refinement(...)
+    custom_bundle_adjustment.iterative_global_refinement(
+        mapper,
         options.ba_global_max_refinements,
         options.ba_global_max_refinement_change,
         mapper_options,
@@ -70,8 +73,9 @@ def initialize_reconstruction(
         mapper_options, two_view_geometry, *init_pair
     )
     logging.info("Global bundle adjustment")
-    mapper.adjust_global_bundle(
-        mapper_options, options.get_global_bundle_adjustment()
+    # The following is equivalent to: mapper.adjust_global_bundle(...)
+    custom_bundle_adjustment.adjust_global_bundle(
+        mapper, mapper_options, options.get_global_bundle_adjustment()
     )
     reconstruction.normalize()
     mapper.filter_points(mapper_options)
@@ -144,7 +148,9 @@ def reconstruct_sub_model(controller, mapper, mapper_options, reconstruction):
                 break
         if reg_next_success:
             mapper.triangulate_image(options.get_triangulation(), next_image_id)
-            mapper.iterative_local_refinement(
+            # The following is equivalent to mapper.iterative_local_refinement(...)
+            custom_bundle_adjustment.iterative_local_refinement(
+                mapper,
                 options.ba_local_max_refinements,
                 options.ba_local_max_refinement_change,
                 mapper_options,

diff --git a/src/colmap/estimators/bundle_adjustment.cc b/src/colmap/estimators/bundle_adjustment.cc
@@ -194,6 +194,11 @@ bool BundleAdjustmentConfig::HasConstantCamPositions(
          constant_cam_positions_.end();
 }
 
+const std::unordered_set<camera_t> BundleAdjustmentConfig::ConstantIntrinsics()
+    const {
+  return constant_intrinsics_;
+}
+
 const std::unordered_set<image_t>& BundleAdjustmentConfig::Images() const {
   return image_ids_;
 }
@@ -208,6 +213,11 @@ const std::unordered_set<point3D_t>& BundleAdjustmentConfig::ConstantPoints()
   return constant_point3D_ids_;
 }
 
+const std::unordered_set<image_t>& BundleAdjustmentConfig::ConstantCamPoses()
+    const {
+  return constant_cam_poses_;
+}
+
 const std::vector<int>& BundleAdjustmentConfig::ConstantCamPositions(
     const image_t image_id) const {
   return constant_cam_positions_.at(image_id);
@@ -256,22 +266,68 @@ BundleAdjuster::BundleAdjuster(const BundleAdjustmentOptions& options,
 }
 
 bool BundleAdjuster::Solve(Reconstruction* reconstruction) {
-  THROW_CHECK_NOTNULL(reconstruction);
-  THROW_CHECK(!problem_) << "Cannot use the same BundleAdjuster multiple times";
-
-  ceres::Problem::Options problem_options;
-  problem_options.loss_function_ownership = ceres::DO_NOT_TAKE_OWNERSHIP;
-  problem_ = std::make_unique<ceres::Problem>(problem_options);
-
   const auto loss_function =
       std::unique_ptr<ceres::LossFunction>(options_.CreateLossFunction());
-  SetUp(reconstruction, loss_function.get());
+  SetUpProblem(reconstruction, loss_function.get());
 
   if (problem_->NumResiduals() == 0) {
     return false;
   }
 
-  ceres::Solver::Options solver_options = options_.solver_options;
+  ceres::Solver::Options solver_options =
+      SetUpSolverOptions(*problem_, options_.solver_options);
+
+  ceres::Solve(solver_options, problem_.get(), &summary_);
+
+  if (options_.print_summary || VLOG_IS_ON(1)) {
+    PrintSolverSummary(summary_, "Bundle adjustment report");
+  }
+
+  return true;
+}
+
+const BundleAdjustmentOptions& BundleAdjuster::Options() const {
+  return options_;
+}
+
+const BundleAdjustmentConfig& BundleAdjuster::Config() const { return config_; }
+
+std::shared_ptr<ceres::Problem> BundleAdjuster::Problem() { return problem_; }
+
+const ceres::Solver::Summary& BundleAdjuster::Summary() const {
+  return summary_;
+}
+
+void BundleAdjuster::SetUpProblem(Reconstruction* reconstruction,
+                                  ceres::LossFunction* loss_function) {
+  THROW_CHECK_NOTNULL(reconstruction);
+
+  // Initialize an empty problem
+  ceres::Problem::Options problem_options;
+  problem_options.loss_function_ownership = ceres::DO_NOT_TAKE_OWNERSHIP;
+  problem_ = std::make_shared<ceres::Problem>(problem_options);
+
+  // Set up problem
+  // Warning: AddPointsToProblem assumes that AddImageToProblem is called first.
+  // Do not change order of instructions!
+  for (const image_t image_id : config_.Images()) {
+    AddImageToProblem(image_id, reconstruction, loss_function);
+  }
+  for (const auto point3D_id : config_.VariablePoints()) {
+    AddPointToProblem(point3D_id, reconstruction, loss_function);
+  }
+  for (const auto point3D_id : config_.ConstantPoints()) {
+    AddPointToProblem(point3D_id, reconstruction, loss_function);
+  }
+
+  ParameterizeCameras(reconstruction);
+  ParameterizePoints(reconstruction);
+}
+
+ceres::Solver::Options BundleAdjuster::SetUpSolverOptions(
+    const ceres::Problem& problem,
+    const ceres::Solver::Options& input_solver_options) const {
+  ceres::Solver::Options solver_options = input_solver_options;
   const bool has_sparse =
       solver_options.sparse_linear_algebra_library_type != ceres::NO_SPARSE;
 
@@ -288,8 +344,7 @@ bool BundleAdjuster::Solve(Reconstruction* reconstruction) {
     solver_options.preconditioner_type = ceres::SCHUR_JACOBI;
   }
 
-  if (problem_->NumResiduals() <
-      options_.min_num_residuals_for_multi_threading) {
+  if (problem.NumResiduals() < options_.min_num_residuals_for_multi_threading) {
     solver_options.num_threads = 1;
 #if CERES_VERSION_MAJOR < 2
     solver_options.num_linear_solver_threads = 1;
@@ -305,42 +360,7 @@ bool BundleAdjuster::Solve(Reconstruction* reconstruction) {
 
   std::string solver_error;
  THROW_CHECK(solver_options.IsValid(&solver_error)) << solver_error;
-
-  ceres::Solve(solver_options, problem_.get(), &summary_);
-
-  if (options_.print_summary || VLOG_IS_ON(1)) {
-    PrintSolverSummary(summary_, "Bundle adjustment report");
-  }
-
-  TearDown(reconstruction);
-
-  return true;
-}
-
-const ceres::Solver::Summary& BundleAdjuster::Summary() const {
-  return summary_;
-}
-
-void BundleAdjuster::SetUp(Reconstruction* reconstruction,
-                           ceres::LossFunction* loss_function) {
-  // Warning: AddPointsToProblem assumes that AddImageToProblem is called first.
-  // Do not change order of instructions!
-  for (const image_t image_id : config_.Images()) {
-    AddImageToProblem(image_id, reconstruction, loss_function);
-  }
-  for (const auto point3D_id : config_.VariablePoints()) {
-    AddPointToProblem(point3D_id, reconstruction, loss_function);
-  }
-  for (const auto point3D_id : config_.ConstantPoints()) {
-    AddPointToProblem(point3D_id, reconstruction, loss_function);
-  }
-
-  ParameterizeCameras(reconstruction);
-  ParameterizePoints(reconstruction);
-}
-
-void BundleAdjuster::TearDown(Reconstruction*) {
-  // Nothing to do
+  return solver_options;
 }
 
 void BundleAdjuster::AddImageToProblem(const image_t image_id,
@@ -516,9 +536,35 @@ RigBundleAdjuster::RigBundleAdjuster(const BundleAdjustmentOptions& options,
 
 bool RigBundleAdjuster::Solve(Reconstruction* reconstruction,
                               std::vector<CameraRig>* camera_rigs) {
+  const auto loss_function =
+      std::unique_ptr<ceres::LossFunction>(options_.CreateLossFunction());
+  SetUpProblem(reconstruction, camera_rigs, loss_function.get());
+
+  if (problem_->NumResiduals() == 0) {
+    return false;
+  }
+
+  ceres::Solver::Options solver_options =
+      SetUpSolverOptions(*problem_, options_.solver_options);
+
+  ceres::Solve(solver_options, problem_.get(), &summary_);
+
+  if (options_.print_summary || VLOG_IS_ON(1)) {
+    PrintSolverSummary(summary_, "Rig Bundle adjustment report");
+  }
+
+  TearDown(reconstruction, *camera_rigs);
+
+  return true;
+}
+
+void RigBundleAdjuster::SetUpProblem(Reconstruction* reconstruction,
+                                     std::vector<CameraRig>* camera_rigs,
+                                     ceres::LossFunction* loss_function) {
   THROW_CHECK_NOTNULL(reconstruction);
   THROW_CHECK_NOTNULL(camera_rigs);
-  THROW_CHECK(!problem_) << "Cannot use the same BundleAdjuster multiple times";
+  THROW_CHECK(!problem_)
+      << "Cannot set up problem from the same BundleAdjuster multiple times";
 
   // Check the validity of the provided camera rigs.
   std::unordered_set<camera_t> rig_camera_ids;
@@ -539,61 +585,12 @@ bool RigBundleAdjuster::Solve(Reconstruction* reconstruction,
     }
   }
 
-  problem_ = std::make_unique<ceres::Problem>();
-
+  // Initialize an empty problem
   ceres::Problem::Options problem_options;
   problem_options.loss_function_ownership = ceres::DO_NOT_TAKE_OWNERSHIP;
-  problem_ = std::make_unique<ceres::Problem>(problem_options);
-
-  const auto loss_function =
-      std::unique_ptr<ceres::LossFunction>(options_.CreateLossFunction());
-  SetUp(reconstruction, camera_rigs, loss_function.get());
-
-  if (problem_->NumResiduals() == 0) {
-    return false;
-  }
-
-  ceres::Solver::Options solver_options = options_.solver_options;
-  const bool has_sparse =
-      solver_options.sparse_linear_algebra_library_type != ceres::NO_SPARSE;
-
-  // Empirical choice.
-  const size_t kMaxNumImagesDirectDenseSolver = 50;
-  const size_t kMaxNumImagesDirectSparseSolver = 1000;
-  const size_t num_images = config_.NumImages();
-  if (num_images <= kMaxNumImagesDirectDenseSolver) {
-    solver_options.linear_solver_type = ceres::DENSE_SCHUR;
-  } else if (num_images <= kMaxNumImagesDirectSparseSolver && has_sparse) {
-    solver_options.linear_solver_type = ceres::SPARSE_SCHUR;
-  } else {  // Indirect sparse (preconditioned CG) solver.
-    solver_options.linear_solver_type = ceres::ITERATIVE_SCHUR;
-    solver_options.preconditioner_type = ceres::SCHUR_JACOBI;
-  }
-
-  solver_options.num_threads =
-      GetEffectiveNumThreads(solver_options.num_threads);
-#if CERES_VERSION_MAJOR < 2
-  solver_options.num_linear_solver_threads =
-      GetEffectiveNumThreads(solver_options.num_linear_solver_threads);
-#endif  // CERES_VERSION_MAJOR
-
-  std::string solver_error;
-  THROW_CHECK(solver_options.IsValid(&solver_error)) << solver_error;
-
-  ceres::Solve(solver_options, problem_.get(), &summary_);
-
-  if (options_.print_summary || VLOG_IS_ON(1)) {
-    PrintSolverSummary(summary_, "Rig Bundle adjustment report");
-  }
-
-  TearDown(reconstruction, *camera_rigs);
-
-  return true;
-}
+  problem_ = std::make_shared<ceres::Problem>(problem_options);
 
-void RigBundleAdjuster::SetUp(Reconstruction* reconstruction,
-                              std::vector<CameraRig>* camera_rigs,
-                              ceres::LossFunction* loss_function) {
+  // Set up problem
   ComputeCameraRigPoses(*reconstruction, *camera_rigs);
 
   for (const image_t image_id : config_.Images()) {

diff --git a/src/colmap/estimators/bundle_adjustment.h b/src/colmap/estimators/bundle_adjustment.h
@@ -147,9 +147,11 @@ class BundleAdjustmentConfig {
   void RemoveConstantPoint(point3D_t point3D_id);
 
   // Access configuration data.
+  const std::unordered_set<camera_t> ConstantIntrinsics() const;
  const std::unordered_set<image_t>& Images() const;
   const std::unordered_set<point3D_t>& VariablePoints() const;
   const std::unordered_set<point3D_t>& ConstantPoints() const;
+  const std::unordered_set<image_t>& ConstantCamPoses() const;
   const std::vector<int>& ConstantCamPositions(image_t image_id) const;
 
  private:
@@ -170,14 +172,22 @@ class BundleAdjuster {
 
   bool Solve(Reconstruction* reconstruction);
 
-  // Get the Ceres solver summary for the last call to `Solve`.
+  // Set up the problem
+  void SetUpProblem(Reconstruction* reconstruction,
+                    ceres::LossFunction* loss_function);
+  ceres::Solver::Options SetUpSolverOptions(
+      const ceres::Problem& problem,
+      const ceres::Solver::Options& input_solver_options) const;
+
+  // Getter functions below
+  const BundleAdjustmentOptions& Options() const;
+  const BundleAdjustmentConfig& Config() const;
+  // Get the Ceres problem after the last call to "set_up"
+  std::shared_ptr<ceres::Problem> Problem();
+  // Get the Ceres solver summary after the last call to `Solve`.
   const ceres::Solver::Summary& Summary() const;
 
  private:
-  void SetUp(Reconstruction* reconstruction,
-             ceres::LossFunction* loss_function);
-  void TearDown(Reconstruction* reconstruction);
-
   void AddImageToProblem(image_t image_id,
                          Reconstruction* reconstruction,
                          ceres::LossFunction* loss_function);
@@ -192,7 +202,7 @@ class BundleAdjuster {
 
   const BundleAdjustmentOptions options_;
   BundleAdjustmentConfig config_;
-  std::unique_ptr<ceres::Problem> problem_;
+  std::shared_ptr<ceres::Problem> problem_;
   ceres::Solver::Summary summary_;
   std::unordered_set<camera_t> camera_ids_;
   std::unordered_map<point3D_t, size_t> point3D_num_observations_;
@@ -219,13 +229,14 @@ class RigBundleAdjuster : public BundleAdjuster {
   bool Solve(Reconstruction* reconstruction,
              std::vector<CameraRig>* camera_rigs);
 
- private:
-  void SetUp(Reconstruction* reconstruction,
-             std::vector<CameraRig>* camera_rigs,
-             ceres::LossFunction* loss_function);
+  void SetUpProblem(Reconstruction* reconstruction,
+                    std::vector<CameraRig>* camera_rigs,
+                    ceres::LossFunction* loss_function);
+
   void TearDown(Reconstruction* reconstruction,
                 const std::vector<CameraRig>& camera_rigs);
 
+ private:
   void AddImageToProblem(image_t image_id,
                          Reconstruction* reconstruction,
                          std::vector<CameraRig>* camera_rigs,