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__)