(reference_sfm: Union[Path, pycolmap.Reconstruction],
queries: Path,
retrieval: Path,
features: Path,
matches: Path,
results: Path,
ransac_thresh: int = 12,
covisibility_clustering: bool = False,
prepend_camera_name: bool = False,
config: Dict = None)
| 124 | |
| 125 | |
| 126 | def main(reference_sfm: Union[Path, pycolmap.Reconstruction], |
| 127 | queries: Path, |
| 128 | retrieval: Path, |
| 129 | features: Path, |
| 130 | matches: Path, |
| 131 | results: Path, |
| 132 | ransac_thresh: int = 12, |
| 133 | covisibility_clustering: bool = False, |
| 134 | prepend_camera_name: bool = False, |
| 135 | config: Dict = None): |
| 136 | |
| 137 | assert retrieval.exists(), retrieval |
| 138 | assert features.exists(), features |
| 139 | assert matches.exists(), matches |
| 140 | |
| 141 | queries = parse_image_lists(queries, with_intrinsics=True) |
| 142 | retrieval_dict = parse_retrieval(retrieval) |
| 143 | |
| 144 | logger.info('Reading the 3D model...') |
| 145 | if not isinstance(reference_sfm, pycolmap.Reconstruction): |
| 146 | reference_sfm = pycolmap.Reconstruction(reference_sfm) |
| 147 | db_name_to_id = {img.name: i for i, img in reference_sfm.images.items()} |
| 148 | |
| 149 | config = {"estimation": {"ransac": {"max_error": ransac_thresh}}, |
| 150 | **(config or {})} |
| 151 | localizer = QueryLocalizer(reference_sfm, config) |
| 152 | |
| 153 | poses = {} |
| 154 | logs = { |
| 155 | 'features': features, |
| 156 | 'matches': matches, |
| 157 | 'retrieval': retrieval, |
| 158 | 'loc': {}, |
| 159 | } |
| 160 | logger.info('Starting localization...') |
| 161 | for qname, qcam in tqdm(queries): |
| 162 | if qname not in retrieval_dict: |
| 163 | logger.warning( |
| 164 | f'No images retrieved for query image {qname}. Skipping...') |
| 165 | continue |
| 166 | db_names = retrieval_dict[qname] |
| 167 | db_ids = [] |
| 168 | for n in db_names: |
| 169 | if n not in db_name_to_id: |
| 170 | logger.warning(f'Image {n} was retrieved but not in database') |
| 171 | continue |
| 172 | db_ids.append(db_name_to_id[n]) |
| 173 | |
| 174 | if covisibility_clustering: |
| 175 | clusters = do_covisibility_clustering(db_ids, reference_sfm) |
| 176 | best_inliers = 0 |
| 177 | best_cluster = None |
| 178 | logs_clusters = [] |
| 179 | for i, cluster_ids in enumerate(clusters): |
| 180 | ret, log = pose_from_cluster( |
| 181 | localizer, qname, qcam, cluster_ids, features, matches) |
| 182 | if ret['success'] and ret['num_inliers'] > best_inliers: |
| 183 | best_cluster = i |
no test coverage detected