|
| 1 | +# colmap_pipeline_options.py |
| 2 | +# Helpers for COLMAP pipeline presets and options |
| 3 | + |
| 4 | +from __future__ import annotations |
| 5 | + |
| 6 | +import os |
| 7 | + |
| 8 | +VOCAB_TREE_PRIORITY_NAMES = ( |
| 9 | + "vocab_tree_faiss_flickr100K_words256K.bin", |
| 10 | + "vocab_tree_faiss_flickr100K_words1M.bin", |
| 11 | + "vocab_tree_flickr100K_words256K.bin", |
| 12 | + "vocab_tree_flickr100K_words1M.bin", |
| 13 | +) |
| 14 | + |
| 15 | + |
| 16 | +def _copy_options(options): |
| 17 | + return {section: dict(values) for section, values in options.items()} |
| 18 | + |
| 19 | + |
| 20 | +BALANCED_OPTIONS = { |
| 21 | + "feature": { |
| 22 | + "SiftExtraction.max_num_features": 16384, |
| 23 | + "SiftExtraction.estimate_affine_shape": 1, |
| 24 | + "SiftExtraction.domain_size_pooling": 1, |
| 25 | + }, |
| 26 | + "matcher": { |
| 27 | + "FeatureMatching.guided_matching": 1, |
| 28 | + }, |
| 29 | + "mapper": { |
| 30 | + "Mapper.ba_refine_sensor_from_rig": 1, |
| 31 | + }, |
| 32 | +} |
| 33 | + |
| 34 | +ULTRA_OPTIONS = _copy_options(BALANCED_OPTIONS) |
| 35 | +ULTRA_OPTIONS["feature"]["SiftExtraction.max_num_features"] = 32768 |
| 36 | + |
| 37 | + |
| 38 | +COLMAP_PRESETS = { |
| 39 | + "standard": { |
| 40 | + "matcher": "sequential", |
| 41 | + "options": {}, |
| 42 | + }, |
| 43 | + "balanced": { |
| 44 | + "matcher": "sequential", |
| 45 | + "options": BALANCED_OPTIONS, |
| 46 | + }, |
| 47 | + "ultra": { |
| 48 | + "matcher": "sequential", |
| 49 | + "options": ULTRA_OPTIONS, |
| 50 | + }, |
| 51 | + "multi_path": { |
| 52 | + "matcher": "vocab_tree", |
| 53 | + "options": BALANCED_OPTIONS, |
| 54 | + }, |
| 55 | +} |
| 56 | + |
| 57 | + |
| 58 | +def merge_options(base, preset, overrides): |
| 59 | + merged = {} |
| 60 | + for source in (base or {}, preset or {}, overrides or {}): |
| 61 | + for section, options in source.items(): |
| 62 | + if not options: |
| 63 | + continue |
| 64 | + merged_section = merged.setdefault(section, {}) |
| 65 | + merged_section.update(options) |
| 66 | + return merged |
| 67 | + |
| 68 | + |
| 69 | +def _stringify_option_value(value): |
| 70 | + if isinstance(value, bool): |
| 71 | + return "1" if value else "0" |
| 72 | + return str(value) |
| 73 | + |
| 74 | + |
| 75 | +def build_colmap_command(base_command, options, supported_options=None, alias_map=None, skipped=None): |
| 76 | + command = list(base_command) |
| 77 | + if not options: |
| 78 | + return command |
| 79 | + for key, value in options.items(): |
| 80 | + if value is None: |
| 81 | + continue |
| 82 | + if isinstance(value, str) and not value.strip(): |
| 83 | + continue |
| 84 | + option_key = key |
| 85 | + if supported_options is not None and key not in supported_options: |
| 86 | + option_key = None |
| 87 | + if alias_map and key in alias_map: |
| 88 | + candidates = alias_map[key] |
| 89 | + if isinstance(candidates, str): |
| 90 | + candidates = [candidates] |
| 91 | + for candidate in candidates: |
| 92 | + if candidate in supported_options: |
| 93 | + option_key = candidate |
| 94 | + break |
| 95 | + if option_key is None: |
| 96 | + if skipped is not None: |
| 97 | + skipped.append(key) |
| 98 | + continue |
| 99 | + command.extend([f"--{option_key}", _stringify_option_value(value)]) |
| 100 | + return command |
| 101 | + |
| 102 | + |
| 103 | +def find_vocab_tree_path(colmap_exec_path): |
| 104 | + if not colmap_exec_path or not os.path.isfile(colmap_exec_path): |
| 105 | + return None |
| 106 | + exe_dir = os.path.dirname(os.path.abspath(colmap_exec_path)) |
| 107 | + candidate_dirs = [ |
| 108 | + exe_dir, |
| 109 | + os.path.abspath(os.path.join(exe_dir, os.pardir)), |
| 110 | + os.path.abspath(os.path.join(exe_dir, os.pardir, "share", "colmap")), |
| 111 | + os.path.abspath(os.path.join(exe_dir, os.pardir, os.pardir, "share", "colmap")), |
| 112 | + ] |
| 113 | + seen = set() |
| 114 | + search_dirs = [] |
| 115 | + for directory in candidate_dirs: |
| 116 | + if directory in seen: |
| 117 | + continue |
| 118 | + seen.add(directory) |
| 119 | + search_dirs.append(directory) |
| 120 | + |
| 121 | + for name in VOCAB_TREE_PRIORITY_NAMES: |
| 122 | + for directory in search_dirs: |
| 123 | + candidate = os.path.join(directory, name) |
| 124 | + if os.path.isfile(candidate): |
| 125 | + return candidate |
| 126 | + |
| 127 | + for directory in search_dirs: |
| 128 | + if not os.path.isdir(directory): |
| 129 | + continue |
| 130 | + try: |
| 131 | + entries = sorted(os.listdir(directory)) |
| 132 | + except OSError: |
| 133 | + continue |
| 134 | + for entry in entries: |
| 135 | + if entry.startswith("vocab_tree_") and entry.endswith(".bin"): |
| 136 | + candidate = os.path.join(directory, entry) |
| 137 | + if os.path.isfile(candidate): |
| 138 | + return candidate |
| 139 | + return None |
0 commit comments