diff --git a/hloc/__init__.py b/hloc/__init__.py index 0aa271bd..52eff7d7 100644 --- a/hloc/__init__.py +++ b/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( diff --git a/hloc/reconstruction.py b/hloc/reconstruction.py index a6705af9..94392054 100644 --- a/hloc/reconstruction.py +++ b/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) diff --git a/hloc/triangulation.py b/hloc/triangulation.py index b903e6be..79c94f34 100644 --- a/hloc/triangulation.py +++ b/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) diff --git a/requirements.txt b/requirements.txt index 2ad6630e..77bec656 100644 --- a/requirements.txt +++ b/requirements.txt @@ -7,6 +7,6 @@ matplotlib plotly scipy h5py -pycolmap>=0.2.0 +pycolmap>=0.3.0 kornia>=0.6.4 gdown