Expose SfM from a given image list

kaiyi98 · Jan 4, 2022 · d3c4272 · d3c4272
1 parent 14438dd
commit d3c4272
Showing 1 changed file with 7 additions and 4 deletions.
diff --git a/hloc/reconstruction.py b/hloc/reconstruction.py
@@ -1,5 +1,6 @@
 import argparse
 import shutil
+from typing import Optional, List
 import multiprocessing
 from pathlib import Path
 import pycolmap
@@ -21,13 +22,14 @@ def create_empty_db(database_path):
     db.close()
 
 
-def import_images(image_dir, database_path, camera_mode):
+def import_images(image_dir, database_path, camera_mode, image_list=None):
     logger.info('Importing images into the database...')
     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)
+        pycolmap.import_images(database_path, image_dir, camera_mode,
+                               image_list=image_list or [])
 
 
 def get_image_ids(database_path):
@@ -75,7 +77,8 @@ def run_reconstruction(sfm_dir, database_path, image_dir, verbose=False):
 
 def main(sfm_dir, image_dir, pairs, features, matches,
          camera_mode=pycolmap.CameraMode.AUTO, verbose=False,
-         skip_geometric_verification=False, min_match_score=None):
+         skip_geometric_verification=False, min_match_score=None,
+         image_list: Optional[List[str]] = None):
 
     assert features.exists(), features
     assert pairs.exists(), pairs
@@ -85,7 +88,7 @@ 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)
+    import_images(image_dir, database, camera_mode, image_list)
     image_ids = get_image_ids(database)
     import_features(image_ids, database, features)
     import_matches(image_ids, database, pairs, matches,