Spaces:
Running
Running
File size: 6,177 Bytes
499e141 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 |
import argparse
import shutil
from typing import Optional, List, Dict, Any
import multiprocessing
from pathlib import Path
import pycolmap
from . import logger
from .utils.database import COLMAPDatabase
from .triangulation import (
import_features, import_matches, estimation_and_geometric_verification,
OutputCapture, parse_option_args)
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()
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 [],
options=options)
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()
return images
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, options=options)
if len(reconstructions) == 0:
logger.error('Could not reconstruct any model!')
return None
logger.info(f'Reconstructed {len(reconstructions)} model(s).')
largest_index = None
largest_num_images = 0
for index, rec in reconstructions.items():
num_images = rec.num_reg_images()
if num_images > largest_num_images:
largest_index = index
largest_num_images = num_images
assert largest_index is not None
logger.info(f'Largest model is #{largest_index} '
f'with {largest_num_images} images.')
for filename in ['images.bin', 'cameras.bin', 'points3D.bin']:
if (sfm_dir / filename).exists():
(sfm_dir / filename).unlink()
shutil.move(
str(models_path / str(largest_index) / filename), str(sfm_dir))
return reconstructions[largest_index]
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
assert matches.exists(), matches
sfm_dir.mkdir(parents=True, exist_ok=True)
database = sfm_dir / 'database.db'
create_empty_db(database)
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, mapper_options)
if reconstruction is not None:
logger.info(f'Reconstruction statistics:\n{reconstruction.summary()}'
+ f'\n\tnum_input_images = {len(image_ids)}')
return reconstruction
if __name__ == '__main__':
parser = argparse.ArgumentParser()
parser.add_argument('--sfm_dir', type=Path, required=True)
parser.add_argument('--image_dir', type=Path, required=True)
parser.add_argument('--pairs', type=Path, required=True)
parser.add_argument('--features', type=Path, required=True)
parser.add_argument('--matches', type=Path, required=True)
parser.add_argument('--camera_mode', type=str, default="AUTO",
choices=list(pycolmap.CameraMode.__members__.keys()))
parser.add_argument('--skip_geometric_verification', action='store_true')
parser.add_argument('--min_match_score', type=float)
parser.add_argument('--verbose', action='store_true')
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)
|