Allow passing custom estimation_options in localize_inloc.py

hloc/localize_inloc.py

@@ -65,7 +65,15 @@ def get_scan_pose(dataset_dir, rpath):
     return P_after_GICP
 
 
-def pose_from_cluster(dataset_dir, q, retrieved, feature_file, match_file, skip=None):
+def pose_from_cluster(
+    dataset_dir,
+    q,
+    retrieved,
+    feature_file,
+    match_file,
+    skip=None,
+    estimation_options=None,
+):
     height, width = cv2.imread(str(dataset_dir / q)).shape[:2]
     cx = 0.5 * width
     cy = 0.5 * height
@@ -111,16 +119,22 @@ def pose_from_cluster(dataset_dir, q, retrieved, feature_file, match_file, skip=
         "height": height,
         "params": [focal_length, cx, cy],
     }
-    estimation_options = pycolmap.AbsolutePoseEstimationOptions()
-    estimation_options.ransac.max_error = 48
     ret = pycolmap.estimate_and_refine_absolute_pose(
         all_mkpq, all_mkp3d, cfg, estimation_options
     )
     ret["cfg"] = cfg
     return ret, all_mkpq, all_mkpr, all_mkp3d, all_indices, num_matches
 
 
-def main(dataset_dir, retrieval, features, matches, results, skip_matches=None):
+def main(
+    dataset_dir,
+    retrieval,
+    features,
+    matches,
+    results,
+    skip_matches=None,
+    estimation_options=None,
+):
     assert retrieval.exists(), retrieval
     assert features.exists(), features
     assert matches.exists(), matches
@@ -138,11 +152,22 @@ def main(dataset_dir, retrieval, features, matches, results, skip_matches=None):
         "retrieval": retrieval,
         "loc": {},
     }
+
+    if estimation_options is None:
+        estimation_options = pycolmap.AbsolutePoseEstimationOptions()
+        estimation_options.ransac.max_error = 48
+
     logger.info("Starting localization...")
     for q in tqdm(queries):
         db = retrieval_dict[q]
         ret, mkpq, mkpr, mkp3d, indices, num_matches = pose_from_cluster(
-            dataset_dir, q, db, feature_file, match_file, skip_matches
+            dataset_dir,
+            q,
+            db,
+            feature_file,
+            match_file,
+            skip_matches,
+            estimation_options,
         )
 
         poses[q] = ret["cam_from_world"]