Use pycolmap database for rig support and remove COLMAPDatabase.

hloc/__init__.py

@@ -21,7 +21,7 @@
 except ImportError:
     logger.warning("pycolmap is not installed, some features may not work.")
 else:
-    min_version = version.parse("3.12.3")
+    min_version = version.parse("3.12.6")
     found_version = pycolmap.__version__
     if found_version != "dev":
         version = version.parse(found_version)

hloc/reconstruction.py

@@ -15,18 +15,16 @@
     import_matches,
     parse_option_args,
 )
-from .utils.database import COLMAPDatabase
+from .utils.database import open_colmap_database
 
 
 def create_empty_db(database_path: Path):
     if database_path.exists():
         logger.warning("The database already exists, deleting it.")
         database_path.unlink()
     logger.info("Creating an empty database...")
-    db = COLMAPDatabase.connect(database_path)
-    db.create_tables()
-    db.commit()
-    db.close()
+    with open_colmap_database(database_path) as _:
+        pass
 
 
 def import_images(
@@ -53,11 +51,10 @@ def import_images(
 
 
 def get_image_ids(database_path: Path) -> Dict[str, int]:
-    db = COLMAPDatabase.connect(database_path)
     images = {}
-    for name, image_id in db.execute("SELECT name, image_id FROM images;"):
-        images[name] = image_id
-    db.close()
+    with open_colmap_database(database_path) as db:
+        for image_id, image in db.read_all_images().items():
+            images[image.name] = image_id
     return images
 
 

hloc/triangulation.py

@@ -7,7 +7,7 @@
 from tqdm import tqdm
 
 from . import logger
-from .utils.database import COLMAPDatabase
+from .utils.database import open_colmap_database
 from .utils.geometry import compute_epipolar_errors
 from .utils.io import get_keypoints, get_matches
 from .utils.parsers import parse_retrieval
@@ -33,40 +33,27 @@ def create_db_from_model(
         logger.warning("The database already exists, deleting it.")
         database_path.unlink()
 
-    db = COLMAPDatabase.connect(database_path)
-    db.create_tables()
-
-    for i, camera in reconstruction.cameras.items():
-        db.add_camera(
-            camera.model.value,
-            camera.width,
-            camera.height,
-            camera.params,
-            camera_id=i,
-            prior_focal_length=True,
-        )
-
-    for i, image in reconstruction.images.items():
-        db.add_image(image.name, image.camera_id, image_id=i)
-
-    db.commit()
-    db.close()
-    return {image.name: i for i, image in reconstruction.images.items()}
+    with open_colmap_database(database_path) as db:
+        for camera_id, camera in reconstruction.cameras.items():
+            db.write_camera(camera, use_camera_id=True)
+        for rig_id, rig in reconstruction.rigs.items():
+            db.write_rig(rig, use_rig_id=True)
+        for frame_id, frame in reconstruction.frames.items():
+            db.write_frame(frame, use_frame_id=True)
+        for image_id, image in reconstruction.images.items():
+            db.write_image(image, use_image_id=True)
+    return {image.name: image_id for image_id, image in reconstruction.images.items()}
 
 
 def import_features(
     image_ids: Dict[str, int], database_path: Path, features_path: Path
 ):
     logger.info("Importing features into the database...")
-    db = COLMAPDatabase.connect(database_path)
-
-    for image_name, image_id in tqdm(image_ids.items()):
-        keypoints = get_keypoints(features_path, image_name)
-        keypoints += 0.5  # COLMAP origin
-        db.add_keypoints(image_id, keypoints)
-
-    db.commit()
-    db.close()
+    with open_colmap_database(database_path) as db:
+        for image_name, image_id in tqdm(image_ids.items()):
+            keypoints = get_keypoints(features_path, image_name)
+            keypoints += 0.5  # COLMAP origin
+            db.write_keypoints(image_id, keypoints)
 
 
 def import_matches(
@@ -82,24 +69,22 @@ def import_matches(
     with open(str(pairs_path), "r") as f:
         pairs = [p.split() for p in f.readlines()]
 
-    db = COLMAPDatabase.connect(database_path)
-
-    matched = set()
-    for name0, name1 in tqdm(pairs):
-        id0, id1 = image_ids[name0], image_ids[name1]
-        if len({(id0, id1), (id1, id0)} & matched) > 0:
-            continue
-        matches, scores = get_matches(matches_path, name0, name1)
-        if min_match_score:
-            matches = matches[scores > min_match_score]
-        db.add_matches(id0, id1, matches)
-        matched |= {(id0, id1), (id1, id0)}
-
-        if skip_geometric_verification:
-            db.add_two_view_geometry(id0, id1, matches)
+    with open_colmap_database(database_path) as db:
+        matched = set()
+        for name0, name1 in tqdm(pairs):
+            id0, id1 = image_ids[name0], image_ids[name1]
+            if len({(id0, id1), (id1, id0)} & matched) > 0:
+                continue
+            matches, scores = get_matches(matches_path, name0, name1)
+            if min_match_score:
+                matches = matches[scores > min_match_score]
+            db.write_matches(id0, id1, matches)
+            matched |= {(id0, id1), (id1, id0)}
 
-    db.commit()
-    db.close()
+            if skip_geometric_verification:
+                db.write_two_view_geometry(
+                    id0, id1, pycolmap.TwoViewGeometry(inlier_matches=matches)
+                )
 
 
 def estimation_and_geometric_verification(
@@ -126,54 +111,62 @@ def geometric_verification(
     logger.info("Performing geometric verification of the matches...")
 
     pairs = parse_retrieval(pairs_path)
-    db = COLMAPDatabase.connect(database_path)
-
-    inlier_ratios = []
-    matched = set()
-    for name0 in tqdm(pairs):
-        id0 = image_ids[name0]
-        image0 = reference.images[id0]
-        cam0 = reference.cameras[image0.camera_id]
-        kps0, noise0 = get_keypoints(features_path, name0, return_uncertainty=True)
-        noise0 = 1.0 if noise0 is None else noise0
-        if len(kps0) > 0:
-            kps0 = np.stack(cam0.cam_from_img(kps0))
-        else:
-            kps0 = np.zeros((0, 2))
-
-        for name1 in pairs[name0]:
-            id1 = image_ids[name1]
-            image1 = reference.images[id1]
-            cam1 = reference.cameras[image1.camera_id]
-            kps1, noise1 = get_keypoints(features_path, name1, return_uncertainty=True)
-            noise1 = 1.0 if noise1 is None else noise1
-            if len(kps1) > 0:
-                kps1 = np.stack(cam1.cam_from_img(kps1))
-            else:
-                kps1 = np.zeros((0, 2))
 
-            matches = get_matches(matches_path, name0, name1)[0]
-
-            if len({(id0, id1), (id1, id0)} & matched) > 0:
-                continue
-            matched |= {(id0, id1), (id1, id0)}
-
-            if matches.shape[0] == 0:
-                db.add_two_view_geometry(id0, id1, matches)
-                continue
-
-            cam1_from_cam0 = image1.cam_from_world() * image0.cam_from_world().inverse()
-            errors0, errors1 = compute_epipolar_errors(
-                cam1_from_cam0, kps0[matches[:, 0]], kps1[matches[:, 1]]
-            )
-            valid_matches = np.logical_and(
-                errors0 <= cam0.cam_from_img_threshold(noise0 * max_error),
-                errors1 <= cam1.cam_from_img_threshold(noise1 * max_error),
-            )
-            # TODO: We could also add E to the database, but we need
-            # to reverse the transformations if id0 > id1 in utils/database.py.
-            db.add_two_view_geometry(id0, id1, matches[valid_matches, :])
-            inlier_ratios.append(np.mean(valid_matches))
+    with open_colmap_database(database_path) as db:
+        inlier_ratios = []
+        matched = set()
+        for name0 in tqdm(pairs):
+            id0 = image_ids[name0]
+            image0 = reference.images[id0]
+            cam0 = reference.cameras[image0.camera_id]
+            kps0, noise0 = get_keypoints(features_path, name0, return_uncertainty=True)
+            noise0 = 1.0 if noise0 is None else noise0
+            if len(kps0) > 0:
+                kps0 = np.stack(cam0.cam_from_img(kps0))
+            else:
+                kps0 = np.zeros((0, 2))
+
+            for name1 in pairs[name0]:
+                id1 = image_ids[name1]
+                image1 = reference.images[id1]
+                cam1 = reference.cameras[image1.camera_id]
+                kps1, noise1 = get_keypoints(
+                    features_path, name1, return_uncertainty=True
+                )
+                noise1 = 1.0 if noise1 is None else noise1
+                if len(kps1) > 0:
+                    kps1 = np.stack(cam1.cam_from_img(kps1))
+                else:
+                    kps1 = np.zeros((0, 2))
+
+                matches = get_matches(matches_path, name0, name1)[0]
+
+                if len({(id0, id1), (id1, id0)} & matched) > 0:
+                    continue
+                matched |= {(id0, id1), (id1, id0)}
+
+                if matches.shape[0] == 0:
+                    db.write_two_view_geometry(id0, id1, pycolmap.TwoViewGeometry())
+                    continue
+
+                cam1_from_cam0 = (
+                    image1.cam_from_world() * image0.cam_from_world().inverse()
+                )
+                errors0, errors1 = compute_epipolar_errors(
+                    cam1_from_cam0, kps0[matches[:, 0]], kps1[matches[:, 1]]
+                )
+                valid_matches = np.logical_and(
+                    errors0 <= cam0.cam_from_img_threshold(noise0 * max_error),
+                    errors1 <= cam1.cam_from_img_threshold(noise1 * max_error),
+                )
+                # TODO: We could also add E to the database, but we need
+                # to reverse the transformations if id0 > id1 in utils/database.py.
+                db.write_two_view_geometry(
+                    id0,
+                    id1,
+                    pycolmap.TwoViewGeometry(inlier_matches=matches[valid_matches, :]),
+                )
+                inlier_ratios.append(np.mean(valid_matches))
     logger.info(
         "mean/med/min/max valid matches %.2f/%.2f/%.2f/%.2f%%.",
         np.mean(inlier_ratios) * 100,
@@ -182,9 +175,6 @@ def geometric_verification(
         np.max(inlier_ratios) * 100,
     )
 
-    db.commit()
-    db.close()
-
 
 def run_triangulation(
     model_path: Path,