Spaces:
Running
Running
| import argparse | |
| import pickle | |
| from collections import defaultdict | |
| from pathlib import Path | |
| from typing import Dict, List, Union | |
| import numpy as np | |
| import pycolmap | |
| from tqdm import tqdm | |
| from . import logger | |
| from .utils.io import get_keypoints, get_matches | |
| from .utils.parsers import parse_image_lists, parse_retrieval | |
| def do_covisibility_clustering( | |
| frame_ids: List[int], reconstruction: pycolmap.Reconstruction | |
| ): | |
| clusters = [] | |
| visited = set() | |
| for frame_id in frame_ids: | |
| # Check if already labeled | |
| if frame_id in visited: | |
| continue | |
| # New component | |
| clusters.append([]) | |
| queue = {frame_id} | |
| while len(queue): | |
| exploration_frame = queue.pop() | |
| # Already part of the component | |
| if exploration_frame in visited: | |
| continue | |
| visited.add(exploration_frame) | |
| clusters[-1].append(exploration_frame) | |
| observed = reconstruction.images[exploration_frame].points2D | |
| connected_frames = { | |
| obs.image_id | |
| for p2D in observed | |
| if p2D.has_point3D() | |
| for obs in reconstruction.points3D[ | |
| p2D.point3D_id | |
| ].track.elements | |
| } | |
| connected_frames &= set(frame_ids) | |
| connected_frames -= visited | |
| queue |= connected_frames | |
| clusters = sorted(clusters, key=len, reverse=True) | |
| return clusters | |
| class QueryLocalizer: | |
| def __init__(self, reconstruction, config=None): | |
| self.reconstruction = reconstruction | |
| self.config = config or {} | |
| def localize(self, points2D_all, points2D_idxs, points3D_id, query_camera): | |
| points2D = points2D_all[points2D_idxs] | |
| points3D = [self.reconstruction.points3D[j].xyz for j in points3D_id] | |
| ret = pycolmap.absolute_pose_estimation( | |
| points2D, | |
| points3D, | |
| query_camera, | |
| estimation_options=self.config.get("estimation", {}), | |
| refinement_options=self.config.get("refinement", {}), | |
| ) | |
| return ret | |
| def pose_from_cluster( | |
| localizer: QueryLocalizer, | |
| qname: str, | |
| query_camera: pycolmap.Camera, | |
| db_ids: List[int], | |
| features_path: Path, | |
| matches_path: Path, | |
| **kwargs, | |
| ): | |
| kpq = get_keypoints(features_path, qname) | |
| kpq += 0.5 # COLMAP coordinates | |
| kp_idx_to_3D = defaultdict(list) | |
| kp_idx_to_3D_to_db = defaultdict(lambda: defaultdict(list)) | |
| num_matches = 0 | |
| for i, db_id in enumerate(db_ids): | |
| image = localizer.reconstruction.images[db_id] | |
| if image.num_points3D == 0: | |
| logger.debug(f"No 3D points found for {image.name}.") | |
| continue | |
| points3D_ids = np.array( | |
| [p.point3D_id if p.has_point3D() else -1 for p in image.points2D] | |
| ) | |
| matches, _ = get_matches(matches_path, qname, image.name) | |
| matches = matches[points3D_ids[matches[:, 1]] != -1] | |
| num_matches += len(matches) | |
| for idx, m in matches: | |
| id_3D = points3D_ids[m] | |
| kp_idx_to_3D_to_db[idx][id_3D].append(i) | |
| # avoid duplicate observations | |
| if id_3D not in kp_idx_to_3D[idx]: | |
| kp_idx_to_3D[idx].append(id_3D) | |
| idxs = list(kp_idx_to_3D.keys()) | |
| mkp_idxs = [i for i in idxs for _ in kp_idx_to_3D[i]] | |
| mp3d_ids = [j for i in idxs for j in kp_idx_to_3D[i]] | |
| ret = localizer.localize(kpq, mkp_idxs, mp3d_ids, query_camera, **kwargs) | |
| if ret is not None: | |
| ret["camera"] = query_camera | |
| # mostly for logging and post-processing | |
| mkp_to_3D_to_db = [ | |
| (j, kp_idx_to_3D_to_db[i][j]) for i in idxs for j in kp_idx_to_3D[i] | |
| ] | |
| log = { | |
| "db": db_ids, | |
| "PnP_ret": ret, | |
| "keypoints_query": kpq[mkp_idxs], | |
| "points3D_ids": mp3d_ids, | |
| "points3D_xyz": None, # we don't log xyz anymore because of file size | |
| "num_matches": num_matches, | |
| "keypoint_index_to_db": (mkp_idxs, mkp_to_3D_to_db), | |
| } | |
| return ret, log | |
| def main( | |
| 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, | |
| ): | |
| assert retrieval.exists(), retrieval | |
| assert features.exists(), features | |
| assert matches.exists(), matches | |
| queries = parse_image_lists(queries, with_intrinsics=True) | |
| retrieval_dict = parse_retrieval(retrieval) | |
| logger.info("Reading the 3D model...") | |
| if not isinstance(reference_sfm, pycolmap.Reconstruction): | |
| reference_sfm = pycolmap.Reconstruction(reference_sfm) | |
| db_name_to_id = {img.name: i for i, img in reference_sfm.images.items()} | |
| config = { | |
| "estimation": {"ransac": {"max_error": ransac_thresh}}, | |
| **(config or {}), | |
| } | |
| localizer = QueryLocalizer(reference_sfm, config) | |
| cam_from_world = {} | |
| logs = { | |
| "features": features, | |
| "matches": matches, | |
| "retrieval": retrieval, | |
| "loc": {}, | |
| } | |
| logger.info("Starting localization...") | |
| for qname, qcam in tqdm(queries): | |
| if qname not in retrieval_dict: | |
| logger.warning( | |
| f"No images retrieved for query image {qname}. Skipping..." | |
| ) | |
| continue | |
| db_names = retrieval_dict[qname] | |
| db_ids = [] | |
| for n in db_names: | |
| if n not in db_name_to_id: | |
| logger.warning(f"Image {n} was retrieved but not in database") | |
| continue | |
| db_ids.append(db_name_to_id[n]) | |
| if covisibility_clustering: | |
| clusters = do_covisibility_clustering(db_ids, reference_sfm) | |
| best_inliers = 0 | |
| best_cluster = None | |
| logs_clusters = [] | |
| for i, cluster_ids in enumerate(clusters): | |
| ret, log = pose_from_cluster( | |
| localizer, qname, qcam, cluster_ids, features, matches | |
| ) | |
| if ret is not None and ret["num_inliers"] > best_inliers: | |
| best_cluster = i | |
| best_inliers = ret["num_inliers"] | |
| logs_clusters.append(log) | |
| if best_cluster is not None: | |
| ret = logs_clusters[best_cluster]["PnP_ret"] | |
| cam_from_world[qname] = ret["cam_from_world"] | |
| logs["loc"][qname] = { | |
| "db": db_ids, | |
| "best_cluster": best_cluster, | |
| "log_clusters": logs_clusters, | |
| "covisibility_clustering": covisibility_clustering, | |
| } | |
| else: | |
| ret, log = pose_from_cluster( | |
| localizer, qname, qcam, db_ids, features, matches | |
| ) | |
| if ret is not None: | |
| cam_from_world[qname] = ret["cam_from_world"] | |
| else: | |
| closest = reference_sfm.images[db_ids[0]] | |
| cam_from_world[qname] = closest.cam_from_world | |
| log["covisibility_clustering"] = covisibility_clustering | |
| logs["loc"][qname] = log | |
| logger.info(f"Localized {len(cam_from_world)} / {len(queries)} images.") | |
| logger.info(f"Writing poses to {results}...") | |
| with open(results, "w") as f: | |
| for query, t in cam_from_world.items(): | |
| qvec = " ".join(map(str, t.rotation.quat[[3, 0, 1, 2]])) | |
| tvec = " ".join(map(str, t.translation)) | |
| name = query.split("/")[-1] | |
| if prepend_camera_name: | |
| name = query.split("/")[-2] + "/" + name | |
| f.write(f"{name} {qvec} {tvec}\n") | |
| logs_path = f"{results}_logs.pkl" | |
| logger.info(f"Writing logs to {logs_path}...") | |
| # TODO: Resolve pickling issue with pycolmap objects. | |
| with open(logs_path, "wb") as f: | |
| pickle.dump(logs, f) | |
| logger.info("Done!") | |
| if __name__ == "__main__": | |
| parser = argparse.ArgumentParser() | |
| parser.add_argument("--reference_sfm", type=Path, required=True) | |
| parser.add_argument("--queries", type=Path, required=True) | |
| parser.add_argument("--features", type=Path, required=True) | |
| parser.add_argument("--matches", type=Path, required=True) | |
| parser.add_argument("--retrieval", type=Path, required=True) | |
| parser.add_argument("--results", type=Path, required=True) | |
| parser.add_argument("--ransac_thresh", type=float, default=12.0) | |
| parser.add_argument("--covisibility_clustering", action="store_true") | |
| parser.add_argument("--prepend_camera_name", action="store_true") | |
| args = parser.parse_args() | |
| main(**args.__dict__) | |