Control over COLMAP import and reconstruction options

hloc/__init__.py

@@ -20,7 +20,7 @@
 except ImportError:
     logger.warning('pycolmap is not installed, some features may not work.')
 else:
-    minimal_version = version.parse('0.2.0')
+    minimal_version = version.parse('0.3.0')
     found_version = version.parse(getattr(pycolmap, '__version__'))
     if found_version < minimal_version:
         logger.warning(

hloc/reconstruction.py

@@ -1,6 +1,6 @@
 import argparse
 import shutil
-from typing import Optional, List
+from typing import Optional, List, Dict, Any
 import multiprocessing
 from pathlib import Path
 import pycolmap
@@ -9,10 +9,10 @@
 from .utils.database import COLMAPDatabase
 from .triangulation import (
     import_features, import_matches, estimation_and_geometric_verification,
-    OutputCapture)
+    OutputCapture, parse_option_args)
 
 
-def create_empty_db(database_path):
+def create_empty_db(database_path: Path):
     if database_path.exists():
         logger.warning('The database already exists, deleting it.')
         database_path.unlink()
@@ -23,17 +23,24 @@ def create_empty_db(database_path):
     db.close()
 
 
-def import_images(image_dir, database_path, camera_mode, image_list=None):
+def import_images(image_dir: Path,
+                  database_path: Path,
+                  camera_mode: pycolmap.CameraMode,
+                  image_list: Optional[List[str]] = None,
+                  options: Optional[Dict[str, Any]] = None):
     logger.info('Importing images into the database...')
+    if options is None:
+        options = {}
     images = list(image_dir.iterdir())
     if len(images) == 0:
         raise IOError(f'No images found in {image_dir}.')
     with pycolmap.ostream():
         pycolmap.import_images(database_path, image_dir, camera_mode,
-                               image_list=image_list or [])
+                               image_list=image_list or [],
+                               options=options)
 
 
-def get_image_ids(database_path):
+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;"):
@@ -42,15 +49,22 @@ def get_image_ids(database_path):
     return images
 
 
-def run_reconstruction(sfm_dir, database_path, image_dir, verbose=False):
+def run_reconstruction(sfm_dir: Path,
+                       database_path: Path,
+                       image_dir: Path,
+                       verbose: bool = False,
+                       options: Optional[Dict[str, Any]] = None,
+                       ) -> pycolmap.Reconstruction:
     models_path = sfm_dir / 'models'
     models_path.mkdir(exist_ok=True, parents=True)
     logger.info('Running 3D reconstruction...')
+    if options is None:
+        options = {}
+    options = {'num_threads': min(multiprocessing.cpu_count(), 16), **options}
     with OutputCapture(verbose):
         with pycolmap.ostream():
             reconstructions = pycolmap.incremental_mapping(
-                database_path, image_dir, models_path,
-                num_threads=min(multiprocessing.cpu_count(), 16))
+                database_path, image_dir, models_path, options=options)
 
     if len(reconstructions) == 0:
         logger.error('Could not reconstruct any model!')
@@ -76,10 +90,19 @@ def run_reconstruction(sfm_dir, database_path, image_dir, verbose=False):
     return reconstructions[largest_index]
 
 
-def main(sfm_dir, image_dir, pairs, features, matches,
-         camera_mode=pycolmap.CameraMode.AUTO, verbose=False,
-         skip_geometric_verification=False, min_match_score=None,
-         image_list: Optional[List[str]] = None):
+def main(sfm_dir: Path,
+         image_dir: Path,
+         pairs: Path,
+         features: Path,
+         matches: Path,
+         camera_mode: pycolmap.CameraMode = pycolmap.CameraMode.AUTO,
+         verbose: bool = False,
+         skip_geometric_verification: bool = False,
+         min_match_score: Optional[float] = None,
+         image_list: Optional[List[str]] = None,
+         image_options: Optional[Dict[str, Any]] = None,
+         mapper_options: Optional[Dict[str, Any]] = None,
+         ) -> pycolmap.Reconstruction:
 
     assert features.exists(), features
     assert pairs.exists(), pairs
@@ -89,14 +112,15 @@ def main(sfm_dir, image_dir, pairs, features, matches,
     database = sfm_dir / 'database.db'
 
     create_empty_db(database)
-    import_images(image_dir, database, camera_mode, image_list)
+    import_images(image_dir, database, camera_mode, image_list, image_options)
     image_ids = get_image_ids(database)
     import_features(image_ids, database, features)
     import_matches(image_ids, database, pairs, matches,
                    min_match_score, skip_geometric_verification)
     if not skip_geometric_verification:
         estimation_and_geometric_verification(database, pairs, verbose)
-    reconstruction = run_reconstruction(sfm_dir, database, image_dir, verbose)
+    reconstruction = run_reconstruction(
+        sfm_dir, database, image_dir, verbose, mapper_options)
     if reconstruction is not None:
         logger.info(f'Reconstruction statistics:\n{reconstruction.summary()}'
                     + f'\n\tnum_input_images = {len(image_ids)}')
@@ -117,6 +141,18 @@ def main(sfm_dir, image_dir, pairs, features, matches,
     parser.add_argument('--skip_geometric_verification', action='store_true')
     parser.add_argument('--min_match_score', type=float)
     parser.add_argument('--verbose', action='store_true')
-    args = parser.parse_args()
 
-    main(**args.__dict__)
+    parser.add_argument('--image_options', nargs='+', default=[],
+                        help='List of key=value from {}'.format(
+                            pycolmap.ImageReaderOptions().todict()))
+    parser.add_argument('--mapper_options', nargs='+', default=[],
+                        help='List of key=value from {}'.format(
+                            pycolmap.IncrementalMapperOptions().todict()))
+    args = parser.parse_args().__dict__
+
+    image_options = parse_option_args(
+        args.pop("image_options"), pycolmap.ImageReaderOptions())
+    mapper_options = parse_option_args(
+        args.pop("mapper_options"), pycolmap.IncrementalMapperOptions())
+
+    main(**args, image_options=image_options, mapper_options=mapper_options)

hloc/triangulation.py

@@ -1,5 +1,6 @@
 import argparse
 import contextlib
+from typing import Optional, List, Dict, Any
 import io
 import sys
 from pathlib import Path
@@ -15,7 +16,7 @@
 
 
 class OutputCapture:
-    def __init__(self, verbose):
+    def __init__(self, verbose: bool):
         self.verbose = verbose
 
     def __enter__(self):
@@ -31,7 +32,8 @@ def __exit__(self, exc_type, *args):
         sys.stdout.flush()
 
 
-def create_db_from_model(reconstruction, database_path):
+def create_db_from_model(reconstruction: pycolmap.Reconstruction,
+                         database_path: Path) -> Dict[str, int]:
     if database_path.exists():
         logger.warning('The database already exists, deleting it.')
         database_path.unlink()
@@ -52,7 +54,9 @@ def create_db_from_model(reconstruction, database_path):
     return {image.name: i for i, image in reconstruction.images.items()}
 
 
-def import_features(image_ids, database_path, features_path):
+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)
 
@@ -65,8 +69,12 @@ def import_features(image_ids, database_path, features_path):
     db.close()
 
 
-def import_matches(image_ids, database_path, pairs_path, matches_path,
-                   min_match_score=None, skip_geometric_verification=False):
+def import_matches(image_ids: Dict[str, int],
+                   database_path: Path,
+                   pairs_path: Path,
+                   matches_path: Path,
+                   min_match_score: Optional[float] = None,
+                   skip_geometric_verification: bool = False):
     logger.info('Importing matches into the database...')
 
     with open(str(pairs_path), 'r') as f:
@@ -92,8 +100,9 @@ def import_matches(image_ids, database_path, pairs_path, matches_path,
     db.close()
 
 
-def estimation_and_geometric_verification(database_path, pairs_path,
-                                          verbose=False):
+def estimation_and_geometric_verification(database_path: Path,
+                                          pairs_path: Path,
+                                          verbose: bool = False):
     logger.info('Performing geometric verification of the matches...')
     with OutputCapture(verbose):
         with pycolmap.ostream():
@@ -102,8 +111,13 @@ def estimation_and_geometric_verification(database_path, pairs_path,
                 max_num_trials=20000, min_inlier_ratio=0.1)
 
 
-def geometric_verification(image_ids, reference, database_path, features_path,
-                           pairs_path, matches_path, max_error=4.0):
+def geometric_verification(image_ids: Dict[str, int],
+                           reference: pycolmap.Reconstruction,
+                           database_path: Path,
+                           features_path: Path,
+                           pairs_path: Path,
+                           matches_path: Path,
+                           max_error: float = 4.0):
     logger.info('Performing geometric verification of the matches...')
 
     pairs = parse_retrieval(pairs_path)
@@ -156,20 +170,37 @@ def geometric_verification(image_ids, reference, database_path, features_path,
     db.close()
 
 
-def run_triangulation(model_path, database_path, image_dir, reference_model,
-                      verbose=False):
+def run_triangulation(model_path: Path,
+                      database_path: Path,
+                      image_dir: Path,
+                      reference_model: pycolmap.Reconstruction,
+                      verbose: bool = False,
+                      options: Optional[Dict[str, Any]] = None,
+                      ) -> pycolmap.Reconstruction:
     model_path.mkdir(parents=True, exist_ok=True)
     logger.info('Running 3D triangulation...')
+    if options is None:
+        options = {}
     with OutputCapture(verbose):
         with pycolmap.ostream():
             reconstruction = pycolmap.triangulate_points(
-                reference_model, database_path, image_dir, model_path)
+                reference_model, database_path, image_dir, model_path,
+                options=options)
     return reconstruction
 
 
-def main(sfm_dir, reference_model, image_dir, pairs, features, matches,
-         skip_geometric_verification=False, estimate_two_view_geometries=False,
-         min_match_score=None, verbose=False):
+def main(sfm_dir: Path,
+         reference_model: Path,
+         image_dir: Path,
+         pairs: Path,
+         features: Path,
+         matches: Path,
+         skip_geometric_verification: bool = False,
+         estimate_two_view_geometries: bool = False,
+         min_match_score: Optional[float] = None,
+         verbose: bool = False,
+         mapper_options: Optional[Dict[str, Any]] = None,
+         ) -> pycolmap.Reconstruction:
 
     assert reference_model.exists(), reference_model
     assert features.exists(), features
@@ -191,12 +222,32 @@ def main(sfm_dir, reference_model, image_dir, pairs, features, matches,
             geometric_verification(
                 image_ids, reference, database, features, pairs, matches)
     reconstruction = run_triangulation(sfm_dir, database, image_dir, reference,
-                                       verbose)
+                                       verbose, mapper_options)
     logger.info('Finished the triangulation with statistics:\n%s',
                 reconstruction.summary())
     return reconstruction
 
 
+def parse_option_args(args: List[str], default_options) -> Dict[str, Any]:
+    options = {}
+    for arg in args:
+        idx = arg.find('=')
+        if idx == -1:
+            raise ValueError('Options format: key1=value1 key2=value2 etc.')
+        key, value = arg[:idx], arg[idx+1:]
+        if not hasattr(default_options, key):
+            raise ValueError(
+                f'Unknown option "{key}", allowed options and default values'
+                f' for {default_options.summary()}')
+        value = eval(value)
+        target_type = type(getattr(default_options, key))
+        if not isinstance(value, target_type):
+            raise ValueError(f'Incorrect type for option "{key}":'
+                             f' {type(value)} vs {target_type}')
+        options[key] = value
+    return options
+
+
 if __name__ == '__main__':
     parser = argparse.ArgumentParser()
     parser.add_argument('--sfm_dir', type=Path, required=True)
@@ -210,6 +261,9 @@ def main(sfm_dir, reference_model, image_dir, pairs, features, matches,
     parser.add_argument('--skip_geometric_verification', action='store_true')
     parser.add_argument('--min_match_score', type=float)
     parser.add_argument('--verbose', action='store_true')
-    args = parser.parse_args()
+    args = parser.parse_args().__dict__
+
+    mapper_options = parse_option_args(
+        args.pop("mapper_options"), pycolmap.IncrementalMapperOptions())
 
-    main(**args.__dict__)
+    main(**args, mapper_options=mapper_options)

requirements.txt

@@ -7,6 +7,6 @@ matplotlib
 plotly
 scipy
 h5py
-pycolmap>=0.2.0
+pycolmap>=0.3.0
 kornia>=0.6.4
 gdown