From 900af51069cfa3e232f420ee95fbcf4ba196ed77 Mon Sep 17 00:00:00 2001 From: Paul-Edouard Sarlin Date: Thu, 11 Jan 2024 16:55:49 +0100 Subject: [PATCH 1/9] Add config files --- .flake8 | 4 ++++ .isort.cfg | 2 ++ 2 files changed, 6 insertions(+) create mode 100644 .flake8 create mode 100644 .isort.cfg diff --git a/.flake8 b/.flake8 new file mode 100644 index 00000000..899119f2 --- /dev/null +++ b/.flake8 @@ -0,0 +1,4 @@ +[flake8] +max-line-length = 88 +extend-ignore = E203 +exclude = .git,__pycache__,build,.venv/ diff --git a/.isort.cfg b/.isort.cfg new file mode 100644 index 00000000..b9fb3f3e --- /dev/null +++ b/.isort.cfg @@ -0,0 +1,2 @@ +[settings] +profile=black From aea5d7765a2b65422c4e7f182734a2d81401ddbb Mon Sep 17 00:00:00 2001 From: Paul-Edouard Sarlin Date: Thu, 11 Jan 2024 16:56:48 +0100 Subject: [PATCH 2/9] Format --- hloc/__init__.py | 24 +- hloc/colmap_from_nvm.py | 108 +++-- hloc/extract_features.py | 327 +++++++------- hloc/extractors/d2net.py | 52 ++- hloc/extractors/dir.py | 66 +-- hloc/extractors/disk.py | 30 +- hloc/extractors/dog.py | 76 ++-- hloc/extractors/eigenplaces.py | 27 +- hloc/extractors/netvlad.py | 64 ++- hloc/extractors/openibl.py | 11 +- hloc/extractors/r2d2.py | 60 +-- hloc/extractors/superpoint.py | 26 +- hloc/localize_inloc.py | 149 ++++--- hloc/localize_sfm.py | 199 +++++---- hloc/match_dense.py | 439 ++++++++++--------- hloc/match_features.py | 240 +++++----- hloc/matchers/__init__.py | 4 +- hloc/matchers/adalam.py | 78 ++-- hloc/matchers/lightglue.py | 34 +- hloc/matchers/loftr.py | 46 +- hloc/matchers/nearest_neighbor.py | 42 +- hloc/matchers/superglue.py | 20 +- hloc/pairs_from_covisibility.py | 21 +- hloc/pairs_from_exhaustive.py | 39 +- hloc/pairs_from_poses.py | 31 +- hloc/pairs_from_retrieval.py | 83 ++-- hloc/pipelines/4Seasons/localize.py | 87 ++-- hloc/pipelines/4Seasons/prepare_reference.py | 44 +- hloc/pipelines/4Seasons/utils.py | 124 +++--- hloc/pipelines/7Scenes/create_gt_sfm.py | 66 +-- hloc/pipelines/7Scenes/pipeline.py | 137 +++--- hloc/pipelines/7Scenes/utils.py | 15 +- hloc/pipelines/Aachen/pipeline.py | 112 +++-- hloc/pipelines/Aachen_v1_1/pipeline.py | 101 +++-- hloc/pipelines/Aachen_v1_1/pipeline_loftr.py | 107 +++-- hloc/pipelines/CMU/pipeline.py | 146 +++--- hloc/pipelines/Cambridge/pipeline.py | 143 +++--- hloc/pipelines/Cambridge/utils.py | 85 ++-- hloc/pipelines/RobotCar/colmap_from_nvm.py | 103 +++-- hloc/pipelines/RobotCar/pipeline.py | 138 +++--- hloc/reconstruction.py | 188 ++++---- hloc/triangulation.py | 215 +++++---- hloc/utils/base_model.py | 11 +- hloc/utils/database.py | 282 +++++------- hloc/utils/geometry.py | 24 +- hloc/utils/io.py | 35 +- hloc/utils/parsers.py | 23 +- hloc/utils/read_write_model.py | 272 ++++++++---- hloc/utils/viz.py | 68 ++- hloc/utils/viz_3d.py | 161 ++++--- hloc/visualization.py | 130 +++--- 51 files changed, 2795 insertions(+), 2318 deletions(-) diff --git a/hloc/__init__.py b/hloc/__init__.py index 61c41829..36db24b8 100644 --- a/hloc/__init__.py +++ b/hloc/__init__.py @@ -1,11 +1,12 @@ import logging + from packaging import version -__version__ = '1.5' +__version__ = "1.5" formatter = logging.Formatter( - fmt='[%(asctime)s %(name)s %(levelname)s] %(message)s', - datefmt='%Y/%m/%d %H:%M:%S') + fmt="[%(asctime)s %(name)s %(levelname)s] %(message)s", datefmt="%Y/%m/%d %H:%M:%S" +) handler = logging.StreamHandler() handler.setFormatter(formatter) handler.setLevel(logging.INFO) @@ -18,16 +19,19 @@ try: import pycolmap except ImportError: - logger.warning('pycolmap is not installed, some features may not work.') + logger.warning("pycolmap is not installed, some features may not work.") else: - min_version = version.parse('0.3.0') - max_version = version.parse('0.4.0') + min_version = version.parse("0.3.0") + max_version = version.parse("0.4.0") found_version = pycolmap.__version__ - if found_version != 'dev': + if found_version != "dev": version = version.parse(found_version) if version < min_version or version > max_version: - s = f'pycolmap>={min_version},<={max_version}' + s = f"pycolmap>={min_version},<={max_version}" logger.warning( - 'hloc now requires %s but found pycolmap==%s, ' + "hloc now requires %s but found pycolmap==%s, " 'please upgrade with `pip install --upgrade "%s"`', - s, found_version, s) + s, + found_version, + s, + ) diff --git a/hloc/colmap_from_nvm.py b/hloc/colmap_from_nvm.py index 9afee5ed..50b7401b 100644 --- a/hloc/colmap_from_nvm.py +++ b/hloc/colmap_from_nvm.py @@ -1,13 +1,19 @@ import argparse import sqlite3 -from tqdm import tqdm from collections import defaultdict -import numpy as np from pathlib import Path +import numpy as np +from tqdm import tqdm + from . import logger -from .utils.read_write_model import Camera, Image, Point3D, CAMERA_MODEL_NAMES -from .utils.read_write_model import write_model +from .utils.read_write_model import ( + CAMERA_MODEL_NAMES, + Camera, + Image, + Point3D, + write_model, +) def recover_database_images_and_ids(database_path): @@ -19,18 +25,20 @@ def recover_database_images_and_ids(database_path): images[name] = image_id cameras[name] = camera_id db.close() - logger.info( - f'Found {len(images)} images and {len(cameras)} cameras in database.') + logger.info(f"Found {len(images)} images and {len(cameras)} cameras in database.") return images, cameras def quaternion_to_rotation_matrix(qvec): qvec = qvec / np.linalg.norm(qvec) w, x, y, z = qvec - R = np.array([ - [1 - 2 * y * y - 2 * z * z, 2 * x * y - 2 * z * w, 2 * x * z + 2 * y * w], - [2 * x * y + 2 * z * w, 1 - 2 * x * x - 2 * z * z, 2 * y * z - 2 * x * w], - [2 * x * z - 2 * y * w, 2 * y * z + 2 * x * w, 1 - 2 * x * x - 2 * y * y]]) + R = np.array( + [ + [1 - 2 * y * y - 2 * z * z, 2 * x * y - 2 * z * w, 2 * x * z + 2 * y * w], + [2 * x * y + 2 * z * w, 1 - 2 * x * x - 2 * z * z, 2 * y * z - 2 * x * w], + [2 * x * z - 2 * y * w, 2 * y * z + 2 * x * w, 1 - 2 * x * x - 2 * y * y], + ] + ) return R @@ -39,73 +47,76 @@ def camera_center_to_translation(c, qvec): return (-1) * np.matmul(R, c) -def read_nvm_model( - nvm_path, intrinsics_path, image_ids, camera_ids, skip_points=False): - - with open(intrinsics_path, 'r') as f: +def read_nvm_model(nvm_path, intrinsics_path, image_ids, camera_ids, skip_points=False): + with open(intrinsics_path, "r") as f: raw_intrinsics = f.readlines() - logger.info(f'Reading {len(raw_intrinsics)} cameras...') + logger.info(f"Reading {len(raw_intrinsics)} cameras...") cameras = {} for intrinsics in raw_intrinsics: - intrinsics = intrinsics.strip('\n').split(' ') + intrinsics = intrinsics.strip("\n").split(" ") name, camera_model, width, height = intrinsics[:4] params = [float(p) for p in intrinsics[4:]] camera_model = CAMERA_MODEL_NAMES[camera_model] assert len(params) == camera_model.num_params camera_id = camera_ids[name] camera = Camera( - id=camera_id, model=camera_model.model_name, - width=int(width), height=int(height), params=params) + id=camera_id, + model=camera_model.model_name, + width=int(width), + height=int(height), + params=params, + ) cameras[camera_id] = camera - nvm_f = open(nvm_path, 'r') + nvm_f = open(nvm_path, "r") line = nvm_f.readline() - while line == '\n' or line.startswith('NVM_V3'): + while line == "\n" or line.startswith("NVM_V3"): line = nvm_f.readline() num_images = int(line) assert num_images == len(cameras) - logger.info(f'Reading {num_images} images...') + logger.info(f"Reading {num_images} images...") image_idx_to_db_image_id = [] image_data = [] i = 0 while i < num_images: line = nvm_f.readline() - if line == '\n': + if line == "\n": continue - data = line.strip('\n').split(' ') + data = line.strip("\n").split(" ") image_data.append(data) image_idx_to_db_image_id.append(image_ids[data[0]]) i += 1 line = nvm_f.readline() - while line == '\n': + while line == "\n": line = nvm_f.readline() num_points = int(line) if skip_points: - logger.info(f'Skipping {num_points} points.') + logger.info(f"Skipping {num_points} points.") num_points = 0 else: - logger.info(f'Reading {num_points} points...') + logger.info(f"Reading {num_points} points...") points3D = {} image_idx_to_keypoints = defaultdict(list) i = 0 - pbar = tqdm(total=num_points, unit='pts') + pbar = tqdm(total=num_points, unit="pts") while i < num_points: line = nvm_f.readline() - if line == '\n': + if line == "\n": continue - data = line.strip('\n').split(' ') + data = line.strip("\n").split(" ") x, y, z, r, g, b, num_observations = data[:7] obs_image_ids, point2D_idxs = [], [] for j in range(int(num_observations)): - s = 7 + 4*j - img_index, kp_index, kx, ky = data[s:s+4] + s = 7 + 4 * j + img_index, kp_index, kx, ky = data[s : s + 4] image_idx_to_keypoints[int(img_index)].append( - (int(kp_index), float(kx), float(ky), i)) + (int(kp_index), float(kx), float(ky), i) + ) db_image_id = image_idx_to_db_image_id[int(img_index)] obs_image_ids.append(db_image_id) point2D_idxs.append(kp_index) @@ -114,16 +125,17 @@ def read_nvm_model( id=i, xyz=np.array([x, y, z], float), rgb=np.array([r, g, b], int), - error=1., # fake + error=1.0, # fake image_ids=np.array(obs_image_ids, int), - point2D_idxs=np.array(point2D_idxs, int)) + point2D_idxs=np.array(point2D_idxs, int), + ) points3D[i] = point i += 1 pbar.update(1) pbar.close() - logger.info('Parsing image data...') + logger.info("Parsing image data...") images = {} for i, data in enumerate(image_data): # Skip the focal length. Skip the distortion and terminal 0. @@ -156,7 +168,8 @@ def read_nvm_model( camera_id=camera_ids[name], name=name, xys=xys, - point3D_ids=point3D_ids) + point3D_ids=point3D_ids, + ) images[image_id] = image return cameras, images, points3D @@ -169,22 +182,23 @@ def main(nvm, intrinsics, database, output, skip_points=False): image_ids, camera_ids = recover_database_images_and_ids(database) - logger.info('Reading the NVM model...') + logger.info("Reading the NVM model...") model = read_nvm_model( - nvm, intrinsics, image_ids, camera_ids, skip_points=skip_points) + nvm, intrinsics, image_ids, camera_ids, skip_points=skip_points + ) - logger.info('Writing the COLMAP model...') + logger.info("Writing the COLMAP model...") output.mkdir(exist_ok=True, parents=True) - write_model(*model, path=str(output), ext='.bin') - logger.info('Done.') + write_model(*model, path=str(output), ext=".bin") + logger.info("Done.") -if __name__ == '__main__': +if __name__ == "__main__": parser = argparse.ArgumentParser() - parser.add_argument('--nvm', required=True, type=Path) - parser.add_argument('--intrinsics', required=True, type=Path) - parser.add_argument('--database', required=True, type=Path) - parser.add_argument('--output', required=True, type=Path) - parser.add_argument('--skip_points', action='store_true') + parser.add_argument("--nvm", required=True, type=Path) + parser.add_argument("--intrinsics", required=True, type=Path) + parser.add_argument("--database", required=True, type=Path) + parser.add_argument("--output", required=True, type=Path) + parser.add_argument("--skip_points", action="store_true") args = parser.parse_args() main(**args.__dict__) diff --git a/hloc/extract_features.py b/hloc/extract_features.py index 053fddb9..53130ab5 100644 --- a/hloc/extract_features.py +++ b/hloc/extract_features.py @@ -1,173 +1,167 @@ import argparse -import torch +import collections.abc as collections +import glob +import pprint from pathlib import Path -from typing import Dict, List, Union, Optional -import h5py from types import SimpleNamespace +from typing import Dict, List, Optional, Union + import cv2 +import h5py import numpy as np -from tqdm import tqdm -import pprint -import collections.abc as collections import PIL.Image -import glob +import torch +from tqdm import tqdm from . import extractors, logger from .utils.base_model import dynamic_load +from .utils.io import list_h5_names, read_image from .utils.parsers import parse_image_lists -from .utils.io import read_image, list_h5_names - -''' +""" A set of standard configurations that can be directly selected from the command line using their name. Each is a dictionary with the following entries: - output: the name of the feature file that will be generated. - model: the model configuration, as passed to a feature extractor. - preprocessing: how to preprocess the images read from disk. -''' +""" confs = { - 'superpoint_aachen': { - 'output': 'feats-superpoint-n4096-r1024', - 'model': { - 'name': 'superpoint', - 'nms_radius': 3, - 'max_keypoints': 4096, + "superpoint_aachen": { + "output": "feats-superpoint-n4096-r1024", + "model": { + "name": "superpoint", + "nms_radius": 3, + "max_keypoints": 4096, }, - 'preprocessing': { - 'grayscale': True, - 'resize_max': 1024, + "preprocessing": { + "grayscale": True, + "resize_max": 1024, }, }, # Resize images to 1600px even if they are originally smaller. # Improves the keypoint localization if the images are of good quality. - 'superpoint_max': { - 'output': 'feats-superpoint-n4096-rmax1600', - 'model': { - 'name': 'superpoint', - 'nms_radius': 3, - 'max_keypoints': 4096, + "superpoint_max": { + "output": "feats-superpoint-n4096-rmax1600", + "model": { + "name": "superpoint", + "nms_radius": 3, + "max_keypoints": 4096, }, - 'preprocessing': { - 'grayscale': True, - 'resize_max': 1600, - 'resize_force': True, + "preprocessing": { + "grayscale": True, + "resize_max": 1600, + "resize_force": True, }, }, - 'superpoint_inloc': { - 'output': 'feats-superpoint-n4096-r1600', - 'model': { - 'name': 'superpoint', - 'nms_radius': 4, - 'max_keypoints': 4096, + "superpoint_inloc": { + "output": "feats-superpoint-n4096-r1600", + "model": { + "name": "superpoint", + "nms_radius": 4, + "max_keypoints": 4096, }, - 'preprocessing': { - 'grayscale': True, - 'resize_max': 1600, + "preprocessing": { + "grayscale": True, + "resize_max": 1600, }, }, - 'r2d2': { - 'output': 'feats-r2d2-n5000-r1024', - 'model': { - 'name': 'r2d2', - 'max_keypoints': 5000, + "r2d2": { + "output": "feats-r2d2-n5000-r1024", + "model": { + "name": "r2d2", + "max_keypoints": 5000, }, - 'preprocessing': { - 'grayscale': False, - 'resize_max': 1024, + "preprocessing": { + "grayscale": False, + "resize_max": 1024, }, }, - 'd2net-ss': { - 'output': 'feats-d2net-ss', - 'model': { - 'name': 'd2net', - 'multiscale': False, + "d2net-ss": { + "output": "feats-d2net-ss", + "model": { + "name": "d2net", + "multiscale": False, }, - 'preprocessing': { - 'grayscale': False, - 'resize_max': 1600, + "preprocessing": { + "grayscale": False, + "resize_max": 1600, }, }, - 'sift': { - 'output': 'feats-sift', - 'model': { - 'name': 'dog' - }, - 'preprocessing': { - 'grayscale': True, - 'resize_max': 1600, + "sift": { + "output": "feats-sift", + "model": {"name": "dog"}, + "preprocessing": { + "grayscale": True, + "resize_max": 1600, }, }, - 'sosnet': { - 'output': 'feats-sosnet', - 'model': { - 'name': 'dog', - 'descriptor': 'sosnet' - }, - 'preprocessing': { - 'grayscale': True, - 'resize_max': 1600, + "sosnet": { + "output": "feats-sosnet", + "model": {"name": "dog", "descriptor": "sosnet"}, + "preprocessing": { + "grayscale": True, + "resize_max": 1600, }, }, - 'disk': { - 'output': 'feats-disk', - 'model': { - 'name': 'disk', - 'max_keypoints': 5000, + "disk": { + "output": "feats-disk", + "model": { + "name": "disk", + "max_keypoints": 5000, }, - 'preprocessing': { - 'grayscale': False, - 'resize_max': 1600, + "preprocessing": { + "grayscale": False, + "resize_max": 1600, }, }, # Global descriptors - 'dir': { - 'output': 'global-feats-dir', - 'model': {'name': 'dir'}, - 'preprocessing': {'resize_max': 1024}, + "dir": { + "output": "global-feats-dir", + "model": {"name": "dir"}, + "preprocessing": {"resize_max": 1024}, }, - 'netvlad': { - 'output': 'global-feats-netvlad', - 'model': {'name': 'netvlad'}, - 'preprocessing': {'resize_max': 1024}, + "netvlad": { + "output": "global-feats-netvlad", + "model": {"name": "netvlad"}, + "preprocessing": {"resize_max": 1024}, }, - 'openibl': { - 'output': 'global-feats-openibl', - 'model': {'name': 'openibl'}, - 'preprocessing': {'resize_max': 1024}, + "openibl": { + "output": "global-feats-openibl", + "model": {"name": "openibl"}, + "preprocessing": {"resize_max": 1024}, + }, + "eigenplaces": { + "output": "global-feats-eigenplaces", + "model": {"name": "eigenplaces"}, + "preprocessing": {"resize_max": 1024}, }, - 'eigenplaces': { - 'output': 'global-feats-eigenplaces', - 'model': {'name': 'eigenplaces'}, - 'preprocessing': {'resize_max': 1024}, - } } def resize_image(image, size, interp): - if interp.startswith('cv2_'): - interp = getattr(cv2, 'INTER_'+interp[len('cv2_'):].upper()) + if interp.startswith("cv2_"): + interp = getattr(cv2, "INTER_" + interp[len("cv2_") :].upper()) h, w = image.shape[:2] if interp == cv2.INTER_AREA and (w < size[0] or h < size[1]): interp = cv2.INTER_LINEAR resized = cv2.resize(image, size, interpolation=interp) - elif interp.startswith('pil_'): - interp = getattr(PIL.Image, interp[len('pil_'):].upper()) + elif interp.startswith("pil_"): + interp = getattr(PIL.Image, interp[len("pil_") :].upper()) resized = PIL.Image.fromarray(image.astype(np.uint8)) resized = resized.resize(size, resample=interp) resized = np.asarray(resized, dtype=image.dtype) else: - raise ValueError( - f'Unknown interpolation {interp}.') + raise ValueError(f"Unknown interpolation {interp}.") return resized class ImageDataset(torch.utils.data.Dataset): default_conf = { - 'globs': ['*.jpg', '*.png', '*.jpeg', '*.JPG', '*.PNG'], - 'grayscale': False, - 'resize_max': None, - 'resize_force': False, - 'interpolation': 'cv2_area', # pil_linear is more accurate but slower + "globs": ["*.jpg", "*.png", "*.jpeg", "*.JPG", "*.PNG"], + "grayscale": False, + "resize_max": None, + "resize_force": False, + "interpolation": "cv2_area", # pil_linear is more accurate but slower } def __init__(self, root, conf, paths=None): @@ -177,26 +171,23 @@ def __init__(self, root, conf, paths=None): if paths is None: paths = [] for g in conf.globs: - paths += glob.glob( - (Path(root) / '**' / g).as_posix(), recursive=True) + paths += glob.glob((Path(root) / "**" / g).as_posix(), recursive=True) if len(paths) == 0: - raise ValueError(f'Could not find any image in root: {root}.') + raise ValueError(f"Could not find any image in root: {root}.") paths = sorted(set(paths)) self.names = [Path(p).relative_to(root).as_posix() for p in paths] - logger.info(f'Found {len(self.names)} images in root {root}.') + logger.info(f"Found {len(self.names)} images in root {root}.") else: if isinstance(paths, (Path, str)): self.names = parse_image_lists(paths) elif isinstance(paths, collections.Iterable): - self.names = [p.as_posix() if isinstance(p, Path) else p - for p in paths] + self.names = [p.as_posix() if isinstance(p, Path) else p for p in paths] else: - raise ValueError(f'Unknown format for path argument {paths}.') + raise ValueError(f"Unknown format for path argument {paths}.") for name in self.names: if not (root / name).exists(): - raise ValueError( - f'Image {name} does not exists in root: {root}.') + raise ValueError(f"Image {name} does not exists in root: {root}.") def __getitem__(self, idx): name = self.names[idx] @@ -204,21 +195,22 @@ def __getitem__(self, idx): image = image.astype(np.float32) size = image.shape[:2][::-1] - if self.conf.resize_max and (self.conf.resize_force - or max(size) > self.conf.resize_max): + if self.conf.resize_max and ( + self.conf.resize_force or max(size) > self.conf.resize_max + ): scale = self.conf.resize_max / max(size) - size_new = tuple(int(round(x*scale)) for x in size) + size_new = tuple(int(round(x * scale)) for x in size) image = resize_image(image, size_new, self.conf.interpolation) if self.conf.grayscale: image = image[None] else: image = image.transpose((2, 0, 1)) # HxWxC to CxHxW - image = image / 255. + image = image / 255.0 data = { - 'image': image, - 'original_size': np.array(size), + "image": image, + "original_size": np.array(size), } return data @@ -227,47 +219,52 @@ def __len__(self): @torch.no_grad() -def main(conf: Dict, - image_dir: Path, - export_dir: Optional[Path] = None, - as_half: bool = True, - image_list: Optional[Union[Path, List[str]]] = None, - feature_path: Optional[Path] = None, - overwrite: bool = False) -> Path: - logger.info('Extracting local features with configuration:' - f'\n{pprint.pformat(conf)}') +def main( + conf: Dict, + image_dir: Path, + export_dir: Optional[Path] = None, + as_half: bool = True, + image_list: Optional[Union[Path, List[str]]] = None, + feature_path: Optional[Path] = None, + overwrite: bool = False, +) -> Path: + logger.info( + "Extracting local features with configuration:" f"\n{pprint.pformat(conf)}" + ) - dataset = ImageDataset(image_dir, conf['preprocessing'], image_list) + dataset = ImageDataset(image_dir, conf["preprocessing"], image_list) if feature_path is None: - feature_path = Path(export_dir, conf['output']+'.h5') + feature_path = Path(export_dir, conf["output"] + ".h5") feature_path.parent.mkdir(exist_ok=True, parents=True) - skip_names = set(list_h5_names(feature_path) - if feature_path.exists() and not overwrite else ()) + skip_names = set( + list_h5_names(feature_path) if feature_path.exists() and not overwrite else () + ) dataset.names = [n for n in dataset.names if n not in skip_names] if len(dataset.names) == 0: - logger.info('Skipping the extraction.') + logger.info("Skipping the extraction.") return feature_path - device = 'cuda' if torch.cuda.is_available() else 'cpu' - Model = dynamic_load(extractors, conf['model']['name']) - model = Model(conf['model']).eval().to(device) + device = "cuda" if torch.cuda.is_available() else "cpu" + Model = dynamic_load(extractors, conf["model"]["name"]) + model = Model(conf["model"]).eval().to(device) loader = torch.utils.data.DataLoader( - dataset, num_workers=1, shuffle=False, pin_memory=True) + dataset, num_workers=1, shuffle=False, pin_memory=True + ) for idx, data in enumerate(tqdm(loader)): name = dataset.names[idx] - pred = model({'image': data['image'].to(device, non_blocking=True)}) + pred = model({"image": data["image"].to(device, non_blocking=True)}) pred = {k: v[0].cpu().numpy() for k, v in pred.items()} - pred['image_size'] = original_size = data['original_size'][0].numpy() - if 'keypoints' in pred: - size = np.array(data['image'].shape[-2:][::-1]) + pred["image_size"] = original_size = data["original_size"][0].numpy() + if "keypoints" in pred: + size = np.array(data["image"].shape[-2:][::-1]) scales = (original_size / size).astype(np.float32) - pred['keypoints'] = (pred['keypoints'] + .5) * scales[None] - .5 - if 'scales' in pred: - pred['scales'] *= scales.mean() + pred["keypoints"] = (pred["keypoints"] + 0.5) * scales[None] - 0.5 + if "scales" in pred: + pred["scales"] *= scales.mean() # add keypoint uncertainties scaled to the original resolution - uncertainty = getattr(model, 'detection_noise', 1) * scales.mean() + uncertainty = getattr(model, "detection_noise", 1) * scales.mean() if as_half: for k in pred: @@ -275,37 +272,39 @@ def main(conf: Dict, if (dt == np.float32) and (dt != np.float16): pred[k] = pred[k].astype(np.float16) - with h5py.File(str(feature_path), 'a', libver='latest') as fd: + with h5py.File(str(feature_path), "a", libver="latest") as fd: try: if name in fd: del fd[name] grp = fd.create_group(name) for k, v in pred.items(): grp.create_dataset(k, data=v) - if 'keypoints' in pred: - grp['keypoints'].attrs['uncertainty'] = uncertainty + if "keypoints" in pred: + grp["keypoints"].attrs["uncertainty"] = uncertainty except OSError as error: - if 'No space left on device' in error.args[0]: + if "No space left on device" in error.args[0]: logger.error( - 'Out of disk space: storing features on disk can take ' - 'significant space, did you enable the as_half flag?') + "Out of disk space: storing features on disk can take " + "significant space, did you enable the as_half flag?" + ) del grp, fd[name] raise error del pred - logger.info('Finished exporting features.') + logger.info("Finished exporting features.") return feature_path -if __name__ == '__main__': +if __name__ == "__main__": parser = argparse.ArgumentParser() - parser.add_argument('--image_dir', type=Path, required=True) - parser.add_argument('--export_dir', type=Path, required=True) - parser.add_argument('--conf', type=str, default='superpoint_aachen', - choices=list(confs.keys())) - parser.add_argument('--as_half', action='store_true') - parser.add_argument('--image_list', type=Path) - parser.add_argument('--feature_path', type=Path) + parser.add_argument("--image_dir", type=Path, required=True) + parser.add_argument("--export_dir", type=Path, required=True) + parser.add_argument( + "--conf", type=str, default="superpoint_aachen", choices=list(confs.keys()) + ) + parser.add_argument("--as_half", action="store_true") + parser.add_argument("--image_list", type=Path) + parser.add_argument("--feature_path", type=Path) args = parser.parse_args() main(confs[args.conf], args.image_dir, args.export_dir, args.as_half) diff --git a/hloc/extractors/d2net.py b/hloc/extractors/d2net.py index 64af5f9e..040b0dd5 100644 --- a/hloc/extractors/d2net.py +++ b/hloc/extractors/d2net.py @@ -1,54 +1,58 @@ +import subprocess import sys from pathlib import Path -import subprocess + import torch from ..utils.base_model import BaseModel -d2net_path = Path(__file__).parent / '../../third_party/d2net' +d2net_path = Path(__file__).parent / "../../third_party/d2net" sys.path.append(str(d2net_path)) -from lib.model_test import D2Net as _D2Net -from lib.pyramid import process_multiscale +from lib.model_test import D2Net as _D2Net # noqa: E402 +from lib.pyramid import process_multiscale # noqa: E402 class D2Net(BaseModel): default_conf = { - 'model_name': 'd2_tf.pth', - 'checkpoint_dir': d2net_path / 'models', - 'use_relu': True, - 'multiscale': False, + "model_name": "d2_tf.pth", + "checkpoint_dir": d2net_path / "models", + "use_relu": True, + "multiscale": False, } - required_inputs = ['image'] + required_inputs = ["image"] def _init(self, conf): - model_file = conf['checkpoint_dir'] / conf['model_name'] + model_file = conf["checkpoint_dir"] / conf["model_name"] if not model_file.exists(): model_file.parent.mkdir(exist_ok=True) - cmd = ['wget', 'https://dsmn.ml/files/d2-net/'+conf['model_name'], - '-O', str(model_file)] + cmd = [ + "wget", + "https://dsmn.ml/files/d2-net/" + conf["model_name"], + "-O", + str(model_file), + ] subprocess.run(cmd, check=True) self.net = _D2Net( - model_file=model_file, - use_relu=conf['use_relu'], - use_cuda=False) + model_file=model_file, use_relu=conf["use_relu"], use_cuda=False + ) def _forward(self, data): - image = data['image'] + image = data["image"] image = image.flip(1) # RGB -> BGR norm = image.new_tensor([103.939, 116.779, 123.68]) - image = (image * 255 - norm.view(1, 3, 1, 1)) # caffe normalization + image = image * 255 - norm.view(1, 3, 1, 1) # caffe normalization - if self.conf['multiscale']: - keypoints, scores, descriptors = process_multiscale( - image, self.net) + if self.conf["multiscale"]: + keypoints, scores, descriptors = process_multiscale(image, self.net) else: keypoints, scores, descriptors = process_multiscale( - image, self.net, scales=[1]) + image, self.net, scales=[1] + ) keypoints = keypoints[:, [1, 0]] # (x, y) and remove the scale return { - 'keypoints': torch.from_numpy(keypoints)[None], - 'scores': torch.from_numpy(scores)[None], - 'descriptors': torch.from_numpy(descriptors.T)[None], + "keypoints": torch.from_numpy(keypoints)[None], + "scores": torch.from_numpy(scores)[None], + "descriptors": torch.from_numpy(descriptors.T)[None], } diff --git a/hloc/extractors/dir.py b/hloc/extractors/dir.py index 4290135d..5b8f42b4 100644 --- a/hloc/extractors/dir.py +++ b/hloc/extractors/dir.py @@ -1,77 +1,77 @@ +import os import sys from pathlib import Path -import torch from zipfile import ZipFile -import os -import sklearn + import gdown +import sklearn +import torch from ..utils.base_model import BaseModel -sys.path.append(str( - Path(__file__).parent / '../../third_party/deep-image-retrieval')) -os.environ['DB_ROOT'] = '' # required by dirtorch +sys.path.append(str(Path(__file__).parent / "../../third_party/deep-image-retrieval")) +os.environ["DB_ROOT"] = "" # required by dirtorch -from dirtorch.utils import common # noqa: E402 from dirtorch.extract_features import load_model # noqa: E402 +from dirtorch.utils import common # noqa: E402 # The DIR model checkpoints (pickle files) include sklearn.decomposition.pca, # which has been deprecated in sklearn v0.24 # and must be explicitly imported with `from sklearn.decomposition import PCA`. # This is a hacky workaround to maintain forward compatibility. -sys.modules['sklearn.decomposition.pca'] = sklearn.decomposition._pca +sys.modules["sklearn.decomposition.pca"] = sklearn.decomposition._pca class DIR(BaseModel): default_conf = { - 'model_name': 'Resnet-101-AP-GeM', - 'whiten_name': 'Landmarks_clean', - 'whiten_params': { - 'whitenp': 0.25, - 'whitenv': None, - 'whitenm': 1.0, + "model_name": "Resnet-101-AP-GeM", + "whiten_name": "Landmarks_clean", + "whiten_params": { + "whitenp": 0.25, + "whitenv": None, + "whitenm": 1.0, }, - 'pooling': 'gem', - 'gemp': 3, + "pooling": "gem", + "gemp": 3, } - required_inputs = ['image'] + required_inputs = ["image"] dir_models = { - 'Resnet-101-AP-GeM': 'https://docs.google.com/uc?export=download&id=1UWJGDuHtzaQdFhSMojoYVQjmCXhIwVvy', + "Resnet-101-AP-GeM": "https://docs.google.com/uc?export=download&id=1UWJGDuHtzaQdFhSMojoYVQjmCXhIwVvy", # noqa: E501 } def _init(self, conf): - checkpoint = Path( - torch.hub.get_dir(), 'dirtorch', conf['model_name'] + '.pt') + checkpoint = Path(torch.hub.get_dir(), "dirtorch", conf["model_name"] + ".pt") if not checkpoint.exists(): checkpoint.parent.mkdir(exist_ok=True, parents=True) - link = self.dir_models[conf['model_name']] - gdown.download(str(link), str(checkpoint)+'.zip', quiet=False) - zf = ZipFile(str(checkpoint)+'.zip', 'r') + link = self.dir_models[conf["model_name"]] + gdown.download(str(link), str(checkpoint) + ".zip", quiet=False) + zf = ZipFile(str(checkpoint) + ".zip", "r") zf.extractall(checkpoint.parent) zf.close() - os.remove(str(checkpoint)+'.zip') + os.remove(str(checkpoint) + ".zip") self.net = load_model(checkpoint, False) # first load on CPU - if conf['whiten_name']: - assert conf['whiten_name'] in self.net.pca + if conf["whiten_name"]: + assert conf["whiten_name"] in self.net.pca def _forward(self, data): - image = data['image'] + image = data["image"] assert image.shape[1] == 3 - mean = self.net.preprocess['mean'] - std = self.net.preprocess['std'] + mean = self.net.preprocess["mean"] + std = self.net.preprocess["std"] image = image - image.new_tensor(mean)[:, None, None] image = image / image.new_tensor(std)[:, None, None] desc = self.net(image) desc = desc.unsqueeze(0) # batch dimension - if self.conf['whiten_name']: - pca = self.net.pca[self.conf['whiten_name']] + if self.conf["whiten_name"]: + pca = self.net.pca[self.conf["whiten_name"]] desc = common.whiten_features( - desc.cpu().numpy(), pca, **self.conf['whiten_params']) + desc.cpu().numpy(), pca, **self.conf["whiten_params"] + ) desc = torch.from_numpy(desc) return { - 'global_descriptor': desc, + "global_descriptor": desc, } diff --git a/hloc/extractors/disk.py b/hloc/extractors/disk.py index dc04280c..bd21d559 100644 --- a/hloc/extractors/disk.py +++ b/hloc/extractors/disk.py @@ -5,28 +5,28 @@ class DISK(BaseModel): default_conf = { - 'weights': 'depth', - 'max_keypoints': None, - 'nms_window_size': 5, - 'detection_threshold': 0.0, - 'pad_if_not_divisible': True, + "weights": "depth", + "max_keypoints": None, + "nms_window_size": 5, + "detection_threshold": 0.0, + "pad_if_not_divisible": True, } - required_inputs = ['image'] + required_inputs = ["image"] def _init(self, conf): - self.model = kornia.feature.DISK.from_pretrained(conf['weights']) + self.model = kornia.feature.DISK.from_pretrained(conf["weights"]) def _forward(self, data): - image = data['image'] + image = data["image"] features = self.model( image, - n=self.conf['max_keypoints'], - window_size=self.conf['nms_window_size'], - score_threshold=self.conf['detection_threshold'], - pad_if_not_divisible=self.conf['pad_if_not_divisible'], + n=self.conf["max_keypoints"], + window_size=self.conf["nms_window_size"], + score_threshold=self.conf["detection_threshold"], + pad_if_not_divisible=self.conf["pad_if_not_divisible"], ) return { - 'keypoints': [f.keypoints for f in features], - 'keypoint_scores': [f.detection_scores for f in features], - 'descriptors': [f.descriptors.t() for f in features], + "keypoints": [f.keypoints for f in features], + "keypoint_scores": [f.detection_scores for f in features], + "descriptors": [f.descriptors.t() for f in features], } diff --git a/hloc/extractors/dog.py b/hloc/extractors/dog.py index 33320a81..a3085384 100644 --- a/hloc/extractors/dog.py +++ b/hloc/extractors/dog.py @@ -1,13 +1,11 @@ import kornia -from kornia.feature.laf import ( - laf_from_center_scale_ori, extract_patches_from_pyramid) import numpy as np -import torch import pycolmap +import torch +from kornia.feature.laf import extract_patches_from_pyramid, laf_from_center_scale_ori from ..utils.base_model import BaseModel - EPS = 1e-6 @@ -20,74 +18,78 @@ def sift_to_rootsift(x): class DoG(BaseModel): default_conf = { - 'options': { - 'first_octave': 0, - 'peak_threshold': 0.01, + "options": { + "first_octave": 0, + "peak_threshold": 0.01, }, - 'descriptor': 'rootsift', - 'max_keypoints': -1, - 'patch_size': 32, - 'mr_size': 12, + "descriptor": "rootsift", + "max_keypoints": -1, + "patch_size": 32, + "mr_size": 12, } - required_inputs = ['image'] + required_inputs = ["image"] detection_noise = 1.0 max_batch_size = 1024 def _init(self, conf): - if conf['descriptor'] == 'sosnet': + if conf["descriptor"] == "sosnet": self.describe = kornia.feature.SOSNet(pretrained=True) - elif conf['descriptor'] == 'hardnet': + elif conf["descriptor"] == "hardnet": self.describe = kornia.feature.HardNet(pretrained=True) - elif conf['descriptor'] not in ['sift', 'rootsift']: + elif conf["descriptor"] not in ["sift", "rootsift"]: raise ValueError(f'Unknown descriptor: {conf["descriptor"]}') self.sift = None # lazily instantiated on the first image self.dummy_param = torch.nn.Parameter(torch.empty(0)) def _forward(self, data): - image = data['image'] + image = data["image"] image_np = image.cpu().numpy()[0, 0] assert image.shape[1] == 1 assert image_np.min() >= -EPS and image_np.max() <= 1 + EPS if self.sift is None: device = self.dummy_param.device - use_gpu = pycolmap.has_cuda and device.type == 'cuda' - options = {**self.conf['options']} - if self.conf['descriptor'] == 'rootsift': - options['normalization'] = pycolmap.Normalization.L1_ROOT + use_gpu = pycolmap.has_cuda and device.type == "cuda" + options = {**self.conf["options"]} + if self.conf["descriptor"] == "rootsift": + options["normalization"] = pycolmap.Normalization.L1_ROOT else: - options['normalization'] = pycolmap.Normalization.L2 + options["normalization"] = pycolmap.Normalization.L2 self.sift = pycolmap.Sift( options=pycolmap.SiftExtractionOptions(options), - device=getattr(pycolmap.Device, 'cuda' if use_gpu else 'cpu')) + device=getattr(pycolmap.Device, "cuda" if use_gpu else "cpu"), + ) keypoints, scores, descriptors = self.sift.extract(image_np) scales = keypoints[:, 2] oris = np.rad2deg(keypoints[:, 3]) - if self.conf['descriptor'] in ['sift', 'rootsift']: + if self.conf["descriptor"] in ["sift", "rootsift"]: # We still renormalize because COLMAP does not normalize well, # maybe due to numerical errors - if self.conf['descriptor'] == 'rootsift': + if self.conf["descriptor"] == "rootsift": descriptors = sift_to_rootsift(descriptors) descriptors = torch.from_numpy(descriptors) - elif self.conf['descriptor'] in ('sosnet', 'hardnet'): + elif self.conf["descriptor"] in ("sosnet", "hardnet"): center = keypoints[:, :2] + 0.5 - laf_scale = scales * self.conf['mr_size'] / 2 + laf_scale = scales * self.conf["mr_size"] / 2 laf_ori = -oris lafs = laf_from_center_scale_ori( torch.from_numpy(center)[None], torch.from_numpy(laf_scale)[None, :, None, None], - torch.from_numpy(laf_ori)[None, :, None]).to(image.device) + torch.from_numpy(laf_ori)[None, :, None], + ).to(image.device) patches = extract_patches_from_pyramid( - image, lafs, PS=self.conf['patch_size'])[0] + image, lafs, PS=self.conf["patch_size"] + )[0] descriptors = patches.new_zeros((len(patches), 128)) if len(patches) > 0: for start_idx in range(0, len(patches), self.max_batch_size): - end_idx = min(len(patches), start_idx+self.max_batch_size) + end_idx = min(len(patches), start_idx + self.max_batch_size) descriptors[start_idx:end_idx] = self.describe( - patches[start_idx:end_idx]) + patches[start_idx:end_idx] + ) else: raise ValueError(f'Unknown descriptor: {self.conf["descriptor"]}') @@ -96,10 +98,10 @@ def _forward(self, data): oris = torch.from_numpy(oris) scores = torch.from_numpy(scores) - if self.conf['max_keypoints'] != -1: + if self.conf["max_keypoints"] != -1: # TODO: check that the scores from PyCOLMAP are 100% correct, # follow https://github.com/mihaidusmanu/pycolmap/issues/8 - indices = torch.topk(scores, self.conf['max_keypoints']) + indices = torch.topk(scores, self.conf["max_keypoints"]) keypoints = keypoints[indices] scales = scales[indices] oris = oris[indices] @@ -107,9 +109,9 @@ def _forward(self, data): descriptors = descriptors[indices] return { - 'keypoints': keypoints[None], - 'scales': scales[None], - 'oris': oris[None], - 'scores': scores[None], - 'descriptors': descriptors.T[None], + "keypoints": keypoints[None], + "scales": scales[None], + "oris": oris[None], + "scores": scores[None], + "descriptors": descriptors.T[None], } diff --git a/hloc/extractors/eigenplaces.py b/hloc/extractors/eigenplaces.py index b37149bf..fd9953b2 100644 --- a/hloc/extractors/eigenplaces.py +++ b/hloc/extractors/eigenplaces.py @@ -1,4 +1,4 @@ -''' +""" Code for loading models trained with EigenPlaces (or CosPlace) as a global features extractor for geolocalization through image retrieval. Multiple models are available with different backbones. Below is a summary of @@ -21,7 +21,7 @@ EigenPlaces paper (ICCV 2023): https://arxiv.org/abs/2308.10832 CosPlace paper (CVPR 2022): https://arxiv.org/abs/2204.02287 -''' +""" import torch import torchvision.transforms as tvf @@ -31,26 +31,27 @@ class EigenPlaces(BaseModel): default_conf = { - 'variant': 'EigenPlaces', - 'backbone': 'ResNet101', - 'fc_output_dim' : 2048 + "variant": "EigenPlaces", + "backbone": "ResNet101", + "fc_output_dim": 2048, } - required_inputs = ['image'] + required_inputs = ["image"] + def _init(self, conf): self.net = torch.hub.load( - 'gmberton/' + conf['variant'], - 'get_trained_model', - backbone=conf['backbone'], - fc_output_dim=conf['fc_output_dim'] + "gmberton/" + conf["variant"], + "get_trained_model", + backbone=conf["backbone"], + fc_output_dim=conf["fc_output_dim"], ).eval() - + mean = [0.485, 0.456, 0.406] std = [0.229, 0.224, 0.225] self.norm_rgb = tvf.Normalize(mean=mean, std=std) def _forward(self, data): - image = self.norm_rgb(data['image']) + image = self.norm_rgb(data["image"]) desc = self.net(image) return { - 'global_descriptor': desc, + "global_descriptor": desc, } diff --git a/hloc/extractors/netvlad.py b/hloc/extractors/netvlad.py index a78b57be..c1051e0d 100644 --- a/hloc/extractors/netvlad.py +++ b/hloc/extractors/netvlad.py @@ -1,5 +1,6 @@ -from pathlib import Path import logging +from pathlib import Path + import numpy as np import torch import torch.nn as nn @@ -17,11 +18,10 @@ class NetVLADLayer(nn.Module): def __init__(self, input_dim=512, K=64, score_bias=False, intranorm=True): super().__init__() - self.score_proj = nn.Conv1d( - input_dim, K, kernel_size=1, bias=score_bias) + self.score_proj = nn.Conv1d(input_dim, K, kernel_size=1, bias=score_bias) centers = nn.parameter.Parameter(torch.empty([input_dim, K])) nn.init.xavier_uniform_(centers) - self.register_parameter('centers', centers) + self.register_parameter("centers", centers) self.intranorm = intranorm self.output_dim = input_dim * K @@ -29,7 +29,7 @@ def forward(self, x): b = x.size(0) scores = self.score_proj(x) scores = F.softmax(scores, dim=1) - diff = (x.unsqueeze(2) - self.centers.unsqueeze(0).unsqueeze(-1)) + diff = x.unsqueeze(2) - self.centers.unsqueeze(0).unsqueeze(-1) desc = (scores.unsqueeze(1) * diff).sum(dim=-1) if self.intranorm: # From the official MATLAB implementation. @@ -40,49 +40,47 @@ def forward(self, x): class NetVLAD(BaseModel): - default_conf = { - 'model_name': 'VGG16-NetVLAD-Pitts30K', - 'whiten': True - } - required_inputs = ['image'] + default_conf = {"model_name": "VGG16-NetVLAD-Pitts30K", "whiten": True} + required_inputs = ["image"] # Models exported using # https://github.com/uzh-rpg/netvlad_tf_open/blob/master/matlab/net_class2struct.m. checkpoint_urls = { - 'VGG16-NetVLAD-Pitts30K': 'https://cvg-data.inf.ethz.ch/hloc/netvlad/Pitts30K_struct.mat', - 'VGG16-NetVLAD-TokyoTM': 'https://cvg-data.inf.ethz.ch/hloc/netvlad/TokyoTM_struct.mat' + "VGG16-NetVLAD-Pitts30K": "https://cvg-data.inf.ethz.ch/hloc/netvlad/Pitts30K_struct.mat", # noqa: E501 + "VGG16-NetVLAD-TokyoTM": "https://cvg-data.inf.ethz.ch/hloc/netvlad/TokyoTM_struct.mat", # noqa: E501 } def _init(self, conf): - if conf['model_name'] not in self.checkpoint_urls: + if conf["model_name"] not in self.checkpoint_urls: raise ValueError( - f'{conf["model_name"]} not in {self.checkpoint_urls.keys()}.') + f'{conf["model_name"]} not in {self.checkpoint_urls.keys()}.' + ) # Download the checkpoint. checkpoint_path = Path( - torch.hub.get_dir(), 'netvlad', conf['model_name'] + '.mat') + torch.hub.get_dir(), "netvlad", conf["model_name"] + ".mat" + ) if not checkpoint_path.exists(): checkpoint_path.parent.mkdir(exist_ok=True, parents=True) - url = self.checkpoint_urls[conf['model_name']] + url = self.checkpoint_urls[conf["model_name"]] torch.hub.download_url_to_file(url, checkpoint_path) # Create the network. # Remove classification head. backbone = list(models.vgg16().children())[0] # Remove last ReLU + MaxPool2d. - self.backbone = nn.Sequential(*list(backbone.children())[: -2]) + self.backbone = nn.Sequential(*list(backbone.children())[:-2]) self.netvlad = NetVLADLayer() - if conf['whiten']: + if conf["whiten"]: self.whiten = nn.Linear(self.netvlad.output_dim, 4096) # Parse MATLAB weights using https://github.com/uzh-rpg/netvlad_tf_open mat = loadmat(checkpoint_path, struct_as_record=False, squeeze_me=True) # CNN weights. - for layer, mat_layer in zip(self.backbone.children(), - mat['net'].layers): + for layer, mat_layer in zip(self.backbone.children(), mat["net"].layers): if isinstance(layer, nn.Conv2d): w = mat_layer.weights[0] # Shape: S x S x IN x OUT b = mat_layer.weights[1] # Shape: OUT @@ -96,9 +94,9 @@ def _init(self, conf): layer.bias = nn.Parameter(b) # NetVLAD weights. - score_w = mat['net'].layers[30].weights[0] # D x K + score_w = mat["net"].layers[30].weights[0] # D x K # centers are stored as opposite in official MATLAB code - center_w = -mat['net'].layers[30].weights[1] # D x K + center_w = -mat["net"].layers[30].weights[1] # D x K # Prepare for PyTorch - make sure it is float32 and has right shape. # score_w should have shape K x D x 1 # center_w should have shape D x K @@ -109,9 +107,9 @@ def _init(self, conf): self.netvlad.centers = nn.Parameter(center_w) # Whitening weights. - if conf['whiten']: - w = mat['net'].layers[33].weights[0] # Shape: 1 x 1 x IN x OUT - b = mat['net'].layers[33].weights[1] # Shape: OUT + if conf["whiten"]: + w = mat["net"].layers[33].weights[0] # Shape: 1 x 1 x IN x OUT + b = mat["net"].layers[33].weights[1] # Shape: OUT # Prepare for PyTorch - make sure it is float32 and has right shape w = torch.tensor(w).float().squeeze().permute([1, 0]) # OUT x IN b = torch.tensor(b.squeeze()).float() # Shape: OUT @@ -121,17 +119,17 @@ def _init(self, conf): # Preprocessing parameters. self.preprocess = { - 'mean': mat['net'].meta.normalization.averageImage[0, 0], - 'std': np.array([1, 1, 1], dtype=np.float32) + "mean": mat["net"].meta.normalization.averageImage[0, 0], + "std": np.array([1, 1, 1], dtype=np.float32), } def _forward(self, data): - image = data['image'] + image = data["image"] assert image.shape[1] == 3 assert image.min() >= -EPS and image.max() <= 1 + EPS image = torch.clamp(image * 255, 0.0, 255.0) # Input should be 0-255. - mean = self.preprocess['mean'] - std = self.preprocess['std'] + mean = self.preprocess["mean"] + std = self.preprocess["std"] image = image - image.new_tensor(mean).view(1, -1, 1, 1) image = image / image.new_tensor(std).view(1, -1, 1, 1) @@ -145,10 +143,8 @@ def _forward(self, data): desc = self.netvlad(descriptors) # Whiten if needed. - if hasattr(self, 'whiten'): + if hasattr(self, "whiten"): desc = self.whiten(desc) desc = F.normalize(desc, dim=1) # Final L2 normalization. - return { - 'global_descriptor': desc - } + return {"global_descriptor": desc} diff --git a/hloc/extractors/openibl.py b/hloc/extractors/openibl.py index a27ce523..9e332a4e 100644 --- a/hloc/extractors/openibl.py +++ b/hloc/extractors/openibl.py @@ -6,20 +6,21 @@ class OpenIBL(BaseModel): default_conf = { - 'model_name': 'vgg16_netvlad', + "model_name": "vgg16_netvlad", } - required_inputs = ['image'] + required_inputs = ["image"] def _init(self, conf): self.net = torch.hub.load( - 'yxgeee/OpenIBL', conf['model_name'], pretrained=True).eval() + "yxgeee/OpenIBL", conf["model_name"], pretrained=True + ).eval() mean = [0.48501960784313836, 0.4579568627450961, 0.4076039215686255] std = [0.00392156862745098, 0.00392156862745098, 0.00392156862745098] self.norm_rgb = tvf.Normalize(mean=mean, std=std) def _forward(self, data): - image = self.norm_rgb(data['image']) + image = self.norm_rgb(data["image"]) desc = self.net(image) return { - 'global_descriptor': desc, + "global_descriptor": desc, } diff --git a/hloc/extractors/r2d2.py b/hloc/extractors/r2d2.py index e1465ee8..e5f49c98 100644 --- a/hloc/extractors/r2d2.py +++ b/hloc/extractors/r2d2.py @@ -1,56 +1,62 @@ import sys from pathlib import Path + import torchvision.transforms as tvf from ..utils.base_model import BaseModel r2d2_path = Path(__file__).parent / "../../third_party/r2d2" sys.path.append(str(r2d2_path)) -from extract import load_network, NonMaxSuppression, extract_multiscale +from extract import NonMaxSuppression, extract_multiscale, load_network # noqa: E402 class R2D2(BaseModel): default_conf = { - 'model_name': 'r2d2_WASF_N16.pt', - 'max_keypoints': 5000, - 'scale_factor': 2**0.25, - 'min_size': 256, - 'max_size': 1024, - 'min_scale': 0, - 'max_scale': 1, - 'reliability_threshold': 0.7, - 'repetability_threshold': 0.7, + "model_name": "r2d2_WASF_N16.pt", + "max_keypoints": 5000, + "scale_factor": 2**0.25, + "min_size": 256, + "max_size": 1024, + "min_scale": 0, + "max_scale": 1, + "reliability_threshold": 0.7, + "repetability_threshold": 0.7, } - required_inputs = ['image'] + required_inputs = ["image"] def _init(self, conf): - model_fn = r2d2_path / "models" / conf['model_name'] - self.norm_rgb = tvf.Normalize(mean=[0.485, 0.456, 0.406], - std=[0.229, 0.224, 0.225]) + model_fn = r2d2_path / "models" / conf["model_name"] + self.norm_rgb = tvf.Normalize( + mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225] + ) self.net = load_network(model_fn) self.detector = NonMaxSuppression( - rel_thr=conf['reliability_threshold'], - rep_thr=conf['repetability_threshold'] + rel_thr=conf["reliability_threshold"], + rep_thr=conf["repetability_threshold"], ) def _forward(self, data): - img = data['image'] + img = data["image"] img = self.norm_rgb(img) xys, desc, scores = extract_multiscale( - self.net, img, self.detector, - scale_f=self.conf['scale_factor'], - min_size=self.conf['min_size'], - max_size=self.conf['max_size'], - min_scale=self.conf['min_scale'], - max_scale=self.conf['max_scale'], + self.net, + img, + self.detector, + scale_f=self.conf["scale_factor"], + min_size=self.conf["min_size"], + max_size=self.conf["max_size"], + min_scale=self.conf["min_scale"], + max_scale=self.conf["max_scale"], ) - idxs = scores.argsort()[-self.conf['max_keypoints'] or None:] + idxs = scores.argsort()[-self.conf["max_keypoints"] or None :] xy = xys[idxs, :2] desc = desc[idxs].t() scores = scores[idxs] - pred = {'keypoints': xy[None], - 'descriptors': desc[None], - 'scores': scores[None]} + pred = { + "keypoints": xy[None], + "descriptors": desc[None], + "scores": scores[None], + } return pred diff --git a/hloc/extractors/superpoint.py b/hloc/extractors/superpoint.py index 739246a1..a3c2a534 100644 --- a/hloc/extractors/superpoint.py +++ b/hloc/extractors/superpoint.py @@ -1,41 +1,43 @@ import sys from pathlib import Path + import torch from ..utils.base_model import BaseModel -sys.path.append(str(Path(__file__).parent / '../../third_party')) +sys.path.append(str(Path(__file__).parent / "../../third_party")) from SuperGluePretrainedNetwork.models import superpoint # noqa E402 # The original keypoint sampling is incorrect. We patch it here but # we don't fix it upstream to not impact exisiting evaluations. def sample_descriptors_fix_sampling(keypoints, descriptors, s: int = 8): - """ Interpolate descriptors at keypoint locations """ + """Interpolate descriptors at keypoint locations""" b, c, h, w = descriptors.shape keypoints = (keypoints + 0.5) / (keypoints.new_tensor([w, h]) * s) keypoints = keypoints * 2 - 1 # normalize to (-1, 1) descriptors = torch.nn.functional.grid_sample( - descriptors, keypoints.view(b, 1, -1, 2), - mode='bilinear', align_corners=False) + descriptors, keypoints.view(b, 1, -1, 2), mode="bilinear", align_corners=False + ) descriptors = torch.nn.functional.normalize( - descriptors.reshape(b, c, -1), p=2, dim=1) + descriptors.reshape(b, c, -1), p=2, dim=1 + ) return descriptors class SuperPoint(BaseModel): default_conf = { - 'nms_radius': 4, - 'keypoint_threshold': 0.005, - 'max_keypoints': -1, - 'remove_borders': 4, - 'fix_sampling': False, + "nms_radius": 4, + "keypoint_threshold": 0.005, + "max_keypoints": -1, + "remove_borders": 4, + "fix_sampling": False, } - required_inputs = ['image'] + required_inputs = ["image"] detection_noise = 2.0 def _init(self, conf): - if conf['fix_sampling']: + if conf["fix_sampling"]: superpoint.sample_descriptors = sample_descriptors_fix_sampling self.net = superpoint.SuperPoint(conf) diff --git a/hloc/localize_inloc.py b/hloc/localize_inloc.py index c3cfc81c..acda7520 100644 --- a/hloc/localize_inloc.py +++ b/hloc/localize_inloc.py @@ -1,21 +1,22 @@ import argparse +import pickle from pathlib import Path -import numpy as np + +import cv2 import h5py -from scipy.io import loadmat +import numpy as np +import pycolmap import torch +from scipy.io import loadmat from tqdm import tqdm -import pickle -import cv2 -import pycolmap from . import logger -from .utils.parsers import parse_retrieval, names_to_pair +from .utils.parsers import names_to_pair, parse_retrieval def interpolate_scan(scan, kp): h, w, c = scan.shape - kp = kp / np.array([[w-1, h-1]]) * 2 - 1 + kp = kp / np.array([[w - 1, h - 1]]) * 2 - 1 assert np.all(kp > -1) and np.all(kp < 1) scan = torch.from_numpy(scan).permute(2, 0, 1)[None] kp = torch.from_numpy(kp)[None, None] @@ -23,10 +24,10 @@ def interpolate_scan(scan, kp): # To maximize the number of points that have depth: # do bilinear interpolation first and then nearest for the remaining points - interp_lin = grid_sample( - scan, kp, align_corners=True, mode='bilinear')[0, :, 0] + interp_lin = grid_sample(scan, kp, align_corners=True, mode="bilinear")[0, :, 0] interp_nn = torch.nn.functional.grid_sample( - scan, kp, align_corners=True, mode='nearest')[0, :, 0] + scan, kp, align_corners=True, mode="nearest" + )[0, :, 0] interp = torch.where(torch.isnan(interp_lin), interp_nn, interp_lin) valid = ~torch.any(torch.isnan(interp), 0) @@ -36,47 +37,51 @@ def interpolate_scan(scan, kp): def get_scan_pose(dataset_dir, rpath): - split_image_rpath = rpath.split('/') + split_image_rpath = rpath.split("/") floor_name = split_image_rpath[-3] scan_id = split_image_rpath[-2] image_name = split_image_rpath[-1] building_name = image_name[:3] path = Path( - dataset_dir, 'database/alignments', floor_name, - f'transformations/{building_name}_trans_{scan_id}.txt') + dataset_dir, + "database/alignments", + floor_name, + f"transformations/{building_name}_trans_{scan_id}.txt", + ) with open(path) as f: raw_lines = f.readlines() - P_after_GICP = np.array([ - np.fromstring(raw_lines[7], sep=' '), - np.fromstring(raw_lines[8], sep=' '), - np.fromstring(raw_lines[9], sep=' '), - np.fromstring(raw_lines[10], sep=' ') - ]) + P_after_GICP = np.array( + [ + np.fromstring(raw_lines[7], sep=" "), + np.fromstring(raw_lines[8], sep=" "), + np.fromstring(raw_lines[9], sep=" "), + np.fromstring(raw_lines[10], sep=" "), + ] + ) return P_after_GICP -def pose_from_cluster(dataset_dir, q, retrieved, feature_file, match_file, - skip=None): +def pose_from_cluster(dataset_dir, q, retrieved, feature_file, match_file, skip=None): height, width = cv2.imread(str(dataset_dir / q)).shape[:2] - cx = .5 * width - cy = .5 * height - focal_length = 4032. * 28. / 36. + cx = 0.5 * width + cy = 0.5 * height + focal_length = 4032.0 * 28.0 / 36.0 all_mkpq = [] all_mkpr = [] all_mkp3d = [] all_indices = [] - kpq = feature_file[q]['keypoints'].__array__() + kpq = feature_file[q]["keypoints"].__array__() num_matches = 0 for i, r in enumerate(retrieved): - kpr = feature_file[r]['keypoints'].__array__() + kpr = feature_file[r]["keypoints"].__array__() pair = names_to_pair(q, r) - m = match_file[pair]['matches0'].__array__() - v = (m > -1) + m = match_file[pair]["matches0"].__array__() + v = m > -1 if skip and (np.count_nonzero(v) < skip): continue @@ -84,7 +89,7 @@ def pose_from_cluster(dataset_dir, q, retrieved, feature_file, match_file, mkpq, mkpr = kpq[v], kpr[m[v]] num_matches += len(mkpq) - scan_r = loadmat(Path(dataset_dir, r + '.mat'))["XYZcut"] + scan_r = loadmat(Path(dataset_dir, r + ".mat"))["XYZcut"] mkp3d, valid = interpolate_scan(scan_r, mkpr) Tr = get_scan_pose(dataset_dir, r) mkp3d = (Tr[:3, :3] @ mkp3d.T + Tr[:3, -1:]).T @@ -100,20 +105,17 @@ def pose_from_cluster(dataset_dir, q, retrieved, feature_file, match_file, all_indices = np.concatenate(all_indices, 0) cfg = { - 'model': 'SIMPLE_PINHOLE', - 'width': width, - 'height': height, - 'params': [focal_length, cx, cy] + "model": "SIMPLE_PINHOLE", + "width": width, + "height": height, + "params": [focal_length, cx, cy], } - ret = pycolmap.absolute_pose_estimation( - all_mkpq, all_mkp3d, cfg, 48.00) - ret['cfg'] = cfg + ret = pycolmap.absolute_pose_estimation(all_mkpq, all_mkp3d, cfg, 48.00) + ret["cfg"] = cfg return ret, all_mkpq, all_mkpr, all_mkp3d, all_indices, num_matches -def main(dataset_dir, retrieval, features, matches, results, - skip_matches=None): - +def main(dataset_dir, retrieval, features, matches, results, skip_matches=None): assert retrieval.exists(), retrieval assert features.exists(), features assert matches.exists(), matches @@ -121,56 +123,57 @@ def main(dataset_dir, retrieval, features, matches, results, retrieval_dict = parse_retrieval(retrieval) queries = list(retrieval_dict.keys()) - feature_file = h5py.File(features, 'r', libver='latest') - match_file = h5py.File(matches, 'r', libver='latest') + feature_file = h5py.File(features, "r", libver="latest") + match_file = h5py.File(matches, "r", libver="latest") poses = {} logs = { - 'features': features, - 'matches': matches, - 'retrieval': retrieval, - 'loc': {}, + "features": features, + "matches": matches, + "retrieval": retrieval, + "loc": {}, } - logger.info('Starting localization...') + logger.info("Starting localization...") for q in tqdm(queries): db = retrieval_dict[q] ret, mkpq, mkpr, mkp3d, indices, num_matches = pose_from_cluster( - dataset_dir, q, db, feature_file, match_file, skip_matches) - - poses[q] = (ret['qvec'], ret['tvec']) - logs['loc'][q] = { - 'db': db, - 'PnP_ret': ret, - 'keypoints_query': mkpq, - 'keypoints_db': mkpr, - '3d_points': mkp3d, - 'indices_db': indices, - 'num_matches': num_matches, + dataset_dir, q, db, feature_file, match_file, skip_matches + ) + + poses[q] = (ret["qvec"], ret["tvec"]) + logs["loc"][q] = { + "db": db, + "PnP_ret": ret, + "keypoints_query": mkpq, + "keypoints_db": mkpr, + "3d_points": mkp3d, + "indices_db": indices, + "num_matches": num_matches, } - logger.info(f'Writing poses to {results}...') - with open(results, 'w') as f: + logger.info(f"Writing poses to {results}...") + with open(results, "w") as f: for q in queries: qvec, tvec = poses[q] - qvec = ' '.join(map(str, qvec)) - tvec = ' '.join(map(str, tvec)) + qvec = " ".join(map(str, qvec)) + tvec = " ".join(map(str, tvec)) name = q.split("/")[-1] - f.write(f'{name} {qvec} {tvec}\n') + f.write(f"{name} {qvec} {tvec}\n") - logs_path = f'{results}_logs.pkl' - logger.info(f'Writing logs to {logs_path}...') - with open(logs_path, 'wb') as f: + logs_path = f"{results}_logs.pkl" + logger.info(f"Writing logs to {logs_path}...") + with open(logs_path, "wb") as f: pickle.dump(logs, f) - logger.info('Done!') + logger.info("Done!") -if __name__ == '__main__': +if __name__ == "__main__": parser = argparse.ArgumentParser() - parser.add_argument('--dataset_dir', type=Path, required=True) - parser.add_argument('--retrieval', type=Path, required=True) - parser.add_argument('--features', type=Path, required=True) - parser.add_argument('--matches', type=Path, required=True) - parser.add_argument('--results', type=Path, required=True) - parser.add_argument('--skip_matches', type=int) + parser.add_argument("--dataset_dir", type=Path, required=True) + parser.add_argument("--retrieval", type=Path, required=True) + parser.add_argument("--features", type=Path, required=True) + parser.add_argument("--matches", type=Path, required=True) + parser.add_argument("--results", type=Path, required=True) + parser.add_argument("--skip_matches", type=int) args = parser.parse_args() main(**args.__dict__) diff --git a/hloc/localize_sfm.py b/hloc/localize_sfm.py index ba68d0e5..76e6037e 100644 --- a/hloc/localize_sfm.py +++ b/hloc/localize_sfm.py @@ -1,19 +1,21 @@ import argparse -import numpy as np -from pathlib import Path +import pickle from collections import defaultdict +from pathlib import Path from typing import Dict, List, Union -from tqdm import tqdm -import pickle + +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): +def do_covisibility_clustering( + frame_ids: List[int], reconstruction: pycolmap.Reconstruction +): clusters = [] visited = set() for frame_id in frame_ids: @@ -36,9 +38,9 @@ def do_covisibility_clustering(frame_ids: List[int], 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 + 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 @@ -57,22 +59,24 @@ 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', {}), + 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): - + 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 @@ -82,10 +86,11 @@ def pose_from_cluster( 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}.') + 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]) + 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] @@ -101,39 +106,41 @@ def pose_from_cluster( 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) - ret['camera'] = { - 'model': query_camera.model_name, - 'width': query_camera.width, - 'height': query_camera.height, - 'params': query_camera.params, + ret["camera"] = { + "model": query_camera.model_name, + "width": query_camera.width, + "height": query_camera.height, + "params": query_camera.params, } # 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]] + 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), + "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): - +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 @@ -141,33 +148,31 @@ def main(reference_sfm: Union[Path, pycolmap.Reconstruction], queries = parse_image_lists(queries, with_intrinsics=True) retrieval_dict = parse_retrieval(retrieval) - logger.info('Reading the 3D model...') + 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 {})} + config = {"estimation": {"ransac": {"max_error": ransac_thresh}}, **(config or {})} localizer = QueryLocalizer(reference_sfm, config) poses = {} logs = { - 'features': features, - 'matches': matches, - 'retrieval': retrieval, - 'loc': {}, + "features": features, + "matches": matches, + "retrieval": retrieval, + "loc": {}, } - logger.info('Starting localization...') + 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...') + 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') + logger.warning(f"Image {n} was retrieved but not in database") continue db_ids.append(db_name_to_id[n]) @@ -178,60 +183,62 @@ def main(reference_sfm: Union[Path, pycolmap.Reconstruction], logs_clusters = [] for i, cluster_ids in enumerate(clusters): ret, log = pose_from_cluster( - localizer, qname, qcam, cluster_ids, features, matches) - if ret['success'] and ret['num_inliers'] > best_inliers: + localizer, qname, qcam, cluster_ids, features, matches + ) + if ret["success"] and ret["num_inliers"] > best_inliers: best_cluster = i - best_inliers = ret['num_inliers'] + best_inliers = ret["num_inliers"] logs_clusters.append(log) if best_cluster is not None: - ret = logs_clusters[best_cluster]['PnP_ret'] - poses[qname] = (ret['qvec'], ret['tvec']) - logs['loc'][qname] = { - 'db': db_ids, - 'best_cluster': best_cluster, - 'log_clusters': logs_clusters, - 'covisibility_clustering': covisibility_clustering, + ret = logs_clusters[best_cluster]["PnP_ret"] + poses[qname] = (ret["qvec"], ret["tvec"]) + 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['success']: - poses[qname] = (ret['qvec'], ret['tvec']) + localizer, qname, qcam, db_ids, features, matches + ) + if ret["success"]: + poses[qname] = (ret["qvec"], ret["tvec"]) else: closest = reference_sfm.images[db_ids[0]] poses[qname] = (closest.qvec, closest.tvec) - log['covisibility_clustering'] = covisibility_clustering - logs['loc'][qname] = log + log["covisibility_clustering"] = covisibility_clustering + logs["loc"][qname] = log - logger.info(f'Localized {len(poses)} / {len(queries)} images.') - logger.info(f'Writing poses to {results}...') - with open(results, 'w') as f: + logger.info(f"Localized {len(poses)} / {len(queries)} images.") + logger.info(f"Writing poses to {results}...") + with open(results, "w") as f: for q in poses: qvec, tvec = poses[q] - qvec = ' '.join(map(str, qvec)) - tvec = ' '.join(map(str, tvec)) - name = q.split('/')[-1] + qvec = " ".join(map(str, qvec)) + tvec = " ".join(map(str, tvec)) + name = q.split("/")[-1] if prepend_camera_name: - name = q.split('/')[-2] + '/' + name - f.write(f'{name} {qvec} {tvec}\n') + name = q.split("/")[-2] + "/" + name + f.write(f"{name} {qvec} {tvec}\n") - logs_path = f'{results}_logs.pkl' - logger.info(f'Writing logs to {logs_path}...') - with open(logs_path, 'wb') as f: + logs_path = f"{results}_logs.pkl" + logger.info(f"Writing logs to {logs_path}...") + with open(logs_path, "wb") as f: pickle.dump(logs, f) - logger.info('Done!') + logger.info("Done!") -if __name__ == '__main__': +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') + 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__) diff --git a/hloc/match_dense.py b/hloc/match_dense.py index 18fb930b..96062880 100644 --- a/hloc/match_dense.py +++ b/hloc/match_dense.py @@ -1,25 +1,24 @@ -from tqdm import tqdm -import numpy as np +import argparse +import pprint +from collections import Counter, defaultdict +from itertools import chain +from pathlib import Path +from types import SimpleNamespace +from typing import Dict, Iterable, List, Optional, Set, Tuple, Union + import h5py +import numpy as np import torch -from pathlib import Path -from typing import Dict, Iterable, Optional, List, Tuple, Union, Set -import pprint -import argparse import torchvision.transforms.functional as F -from types import SimpleNamespace -from collections import defaultdict from scipy.spatial import KDTree -from collections import Counter -from itertools import chain +from tqdm import tqdm -from . import matchers, logger -from .utils.base_model import dynamic_load -from .utils.parsers import parse_retrieval, names_to_pair -from .match_features import find_unique_new_pairs +from . import logger, matchers from .extract_features import read_image, resize_image +from .match_features import find_unique_new_pairs +from .utils.base_model import dynamic_load from .utils.io import list_h5_names - +from .utils.parsers import names_to_pair, parse_retrieval # Default usage: # dense_conf = confs['loftr'] @@ -38,49 +37,28 @@ confs = { # Best quality but loads of points. Only use for small scenes - 'loftr': { - 'output': 'matches-loftr', - 'model': { - 'name': 'loftr', - 'weights': 'outdoor' - }, - 'preprocessing': { - 'grayscale': True, - 'resize_max': 1024, - 'dfactor': 8 - }, - 'max_error': 1, # max error for assigned keypoints (in px) - 'cell_size': 1, # size of quantization patch (max 1 kp/patch) + "loftr": { + "output": "matches-loftr", + "model": {"name": "loftr", "weights": "outdoor"}, + "preprocessing": {"grayscale": True, "resize_max": 1024, "dfactor": 8}, + "max_error": 1, # max error for assigned keypoints (in px) + "cell_size": 1, # size of quantization patch (max 1 kp/patch) }, # Semi-scalable loftr which limits detected keypoints - 'loftr_aachen': { - 'output': 'matches-loftr_aachen', - 'model': { - 'name': 'loftr', - 'weights': 'outdoor' - }, - 'preprocessing': { - 'grayscale': True, - 'resize_max': 1024, - 'dfactor': 8 - }, - 'max_error': 2, # max error for assigned keypoints (in px) - 'cell_size': 8, # size of quantization patch (max 1 kp/patch) + "loftr_aachen": { + "output": "matches-loftr_aachen", + "model": {"name": "loftr", "weights": "outdoor"}, + "preprocessing": {"grayscale": True, "resize_max": 1024, "dfactor": 8}, + "max_error": 2, # max error for assigned keypoints (in px) + "cell_size": 8, # size of quantization patch (max 1 kp/patch) }, # Use for matching superpoint feats with loftr - 'loftr_superpoint': { - 'output': 'matches-loftr_aachen', - 'model': { - 'name': 'loftr', - 'weights': 'outdoor' - }, - 'preprocessing': { - 'grayscale': True, - 'resize_max': 1024, - 'dfactor': 8 - }, - 'max_error': 4, # max error for assigned keypoints (in px) - 'cell_size': 4, # size of quantization patch (max 1 kp/patch) + "loftr_superpoint": { + "output": "matches-loftr_aachen", + "model": {"name": "loftr", "weights": "outdoor"}, + "preprocessing": {"grayscale": True, "resize_max": 1024, "dfactor": 8}, + "max_error": 4, # max error for assigned keypoints (in px) + "cell_size": 4, # size of quantization patch (max 1 kp/patch) }, } @@ -91,19 +69,21 @@ def to_cpts(kpts, ps): return [tuple(cpt) for cpt in kpts] -def assign_keypoints(kpts: np.ndarray, - other_cpts: Union[List[Tuple], np.ndarray], - max_error: float, - update: bool = False, - ref_bins: Optional[List[Counter]] = None, - scores: Optional[np.ndarray] = None, - cell_size: Optional[int] = None): +def assign_keypoints( + kpts: np.ndarray, + other_cpts: Union[List[Tuple], np.ndarray], + max_error: float, + update: bool = False, + ref_bins: Optional[List[Counter]] = None, + scores: Optional[np.ndarray] = None, + cell_size: Optional[int] = None, +): if not update: # Without update this is just a NN search if len(other_cpts) == 0 or len(kpts) == 0: return np.full(len(kpts), -1) dist, kpt_ids = KDTree(np.array(other_cpts)).query(kpts) - valid = (dist <= max_error) + valid = dist <= max_error kpt_ids[~valid] = -1 return kpt_ids else: @@ -136,8 +116,7 @@ def get_grouped_ids(array): # all duplicates are grouped as a set idx_sort = np.argsort(array) sorted_array = array[idx_sort] - _, ids, _ = np.unique(sorted_array, return_counts=True, - return_index=True) + _, ids, _ = np.unique(sorted_array, return_counts=True, return_index=True) res = np.split(idx_sort, ids[1:]) return res @@ -184,10 +163,10 @@ def scale_keypoints(kpts, scale): class ImagePairDataset(torch.utils.data.Dataset): default_conf = { - 'grayscale': True, - 'resize_max': 1024, - 'dfactor': 8, - 'cache_images': False, + "grayscale": True, + "resize_max": 1024, + "dfactor": 8, + "cache_images": False, } def __init__(self, image_dir, conf, pairs): @@ -196,8 +175,7 @@ def __init__(self, image_dir, conf, pairs): self.pairs = pairs if self.conf.cache_images: image_names = set(sum(pairs, ())) # unique image names in pairs - logger.info( - f'Loading and caching {len(image_names)} unique images.') + logger.info(f"Loading and caching {len(image_names)} unique images.") self.images = {} self.scales = {} for name in tqdm(image_names): @@ -212,8 +190,8 @@ def preprocess(self, image: np.ndarray): if self.conf.resize_max: scale = self.conf.resize_max / max(size) if scale < 1.0: - size_new = tuple(int(round(x*scale)) for x in size) - image = resize_image(image, size_new, 'cv2_area') + size_new = tuple(int(round(x * scale)) for x in size) + image = resize_image(image, size_new, "cv2_area") scale = np.array(size) / np.array(size_new) if self.conf.grayscale: @@ -224,9 +202,12 @@ def preprocess(self, image: np.ndarray): image = torch.from_numpy(image / 255.0).float() # assure that the size is divisible by dfactor - size_new = tuple(map( + size_new = tuple( + map( lambda x: int(x // self.conf.dfactor * self.conf.dfactor), - image.shape[-2:])) + image.shape[-2:], + ) + ) image = F.resize(image, size=size_new) scale = np.array(size) / np.array(size_new)[::-1] return image, scale @@ -248,23 +229,25 @@ def __getitem__(self, idx): @torch.no_grad() -def match_dense(conf: Dict, - pairs: List[Tuple[str, str]], - image_dir: Path, - match_path: Path, # out - existing_refs: Optional[List] = []): - - device = 'cuda' if torch.cuda.is_available() else 'cpu' - Model = dynamic_load(matchers, conf['model']['name']) - model = Model(conf['model']).eval().to(device) +def match_dense( + conf: Dict, + pairs: List[Tuple[str, str]], + image_dir: Path, + match_path: Path, # out + existing_refs: Optional[List] = [], +): + device = "cuda" if torch.cuda.is_available() else "cpu" + Model = dynamic_load(matchers, conf["model"]["name"]) + model = Model(conf["model"]).eval().to(device) dataset = ImagePairDataset(image_dir, conf["preprocessing"], pairs) loader = torch.utils.data.DataLoader( - dataset, num_workers=16, batch_size=1, shuffle=False) + dataset, num_workers=16, batch_size=1, shuffle=False + ) logger.info("Performing dense matching...") - with h5py.File(str(match_path), 'a') as fd: - for data in tqdm(loader, smoothing=.1): + with h5py.File(str(match_path), "a") as fd: + for data in tqdm(loader, smoothing=0.1): # load image-pair data image0, image1, scale0, scale1, (name0,), (name1,) = data scale0, scale1 = scale0[0].numpy(), scale1[0].numpy() @@ -274,21 +257,23 @@ def match_dense(conf: Dict, # for consistency with pairs_from_*: refine kpts of image0 if name0 in existing_refs: # special case: flip to enable refinement in query image - pred = model({'image0': image1, 'image1': image0}) - pred = {**pred, - 'keypoints0': pred['keypoints1'], - 'keypoints1': pred['keypoints0']} + pred = model({"image0": image1, "image1": image0}) + pred = { + **pred, + "keypoints0": pred["keypoints1"], + "keypoints1": pred["keypoints0"], + } else: # usual case - pred = model({'image0': image0, 'image1': image1}) + pred = model({"image0": image0, "image1": image1}) # Rescale keypoints and move to cpu - kpts0, kpts1 = pred['keypoints0'], pred['keypoints1'] + kpts0, kpts1 = pred["keypoints0"], pred["keypoints1"] kpts0 = scale_keypoints(kpts0 + 0.5, scale0) - 0.5 kpts1 = scale_keypoints(kpts1 + 0.5, scale1) - 0.5 kpts0 = kpts0.cpu().numpy() kpts1 = kpts1.cpu().numpy() - scores = pred['scores'].cpu().numpy() + scores = pred["scores"].cpu().numpy() # Write matches and matching scores in hloc format pair = names_to_pair(name0, name1) @@ -297,66 +282,72 @@ def match_dense(conf: Dict, grp = fd.create_group(pair) # Write dense matching output - grp.create_dataset('keypoints0', data=kpts0) - grp.create_dataset('keypoints1', data=kpts1) - grp.create_dataset('scores', data=scores) + grp.create_dataset("keypoints0", data=kpts0) + grp.create_dataset("keypoints1", data=kpts1) + grp.create_dataset("scores", data=scores) del model, loader # default: quantize all! -def load_keypoints(conf: Dict, - feature_paths_refs: List[Path], - quantize: Optional[set] = None): - name2ref = {n: i for i, p in enumerate(feature_paths_refs) - for n in list_h5_names(p)} +def load_keypoints( + conf: Dict, feature_paths_refs: List[Path], quantize: Optional[set] = None +): + name2ref = { + n: i for i, p in enumerate(feature_paths_refs) for n in list_h5_names(p) + } existing_refs = set(name2ref.keys()) if quantize is None: quantize = existing_refs # quantize all if len(existing_refs) > 0: - logger.info(f'Loading keypoints from {len(existing_refs)} images.') + logger.info(f"Loading keypoints from {len(existing_refs)} images.") # Load query keypoints cpdict = defaultdict(list) bindict = defaultdict(list) for name in existing_refs: - with h5py.File(str(feature_paths_refs[name2ref[name]]), 'r') as fd: - kps = fd[name]['keypoints'].__array__() + with h5py.File(str(feature_paths_refs[name2ref[name]]), "r") as fd: + kps = fd[name]["keypoints"].__array__() if name not in quantize: cpdict[name] = kps else: - if 'scores' in fd[name].keys(): - kp_scores = fd[name]['scores'].__array__() + if "scores" in fd[name].keys(): + kp_scores = fd[name]["scores"].__array__() else: # we set the score to 1.0 if not provided # increase for more weight on reference keypoints for # stronger anchoring - kp_scores = \ - [1.0 for _ in range(kps.shape[0])] + kp_scores = [1.0 for _ in range(kps.shape[0])] # bin existing keypoints of reference images for association assign_keypoints( - kps, cpdict[name], conf['max_error'], True, bindict[name], - kp_scores, conf['cell_size']) + kps, + cpdict[name], + conf["max_error"], + True, + bindict[name], + kp_scores, + conf["cell_size"], + ) return cpdict, bindict def aggregate_matches( - conf: Dict, - pairs: List[Tuple[str, str]], - match_path: Path, - feature_path: Path, - required_queries: Optional[Set[str]] = None, - max_kps: Optional[int] = None, - cpdict: Dict[str, Iterable] = defaultdict(list), - bindict: Dict[str, List[Counter]] = defaultdict(list)): + conf: Dict, + pairs: List[Tuple[str, str]], + match_path: Path, + feature_path: Path, + required_queries: Optional[Set[str]] = None, + max_kps: Optional[int] = None, + cpdict: Dict[str, Iterable] = defaultdict(list), + bindict: Dict[str, List[Counter]] = defaultdict(list), +): if required_queries is None: required_queries = set(sum(pairs, ())) # default: do not overwrite existing features in feature_path! required_queries -= set(list_h5_names(feature_path)) # if an entry in cpdict is provided as np.ndarray we assume it is fixed - required_queries -= set( - [k for k, v in cpdict.items() if isinstance(v, np.ndarray)]) + required_queries -= set([k for k, v in cpdict.items() if isinstance(v, np.ndarray)]) # sort pairs for reduced RAM pairs_per_q = Counter(list(chain(*pairs))) @@ -364,15 +355,15 @@ def aggregate_matches( pairs = [p for _, p in sorted(zip(pairs_score, pairs))] if len(required_queries) > 0: - logger.info(f'Aggregating keypoints for {len(required_queries)} images.') + logger.info(f"Aggregating keypoints for {len(required_queries)} images.") n_kps = 0 - with h5py.File(str(match_path), 'a') as fd: - for name0, name1 in tqdm(pairs, smoothing=.1): + with h5py.File(str(match_path), "a") as fd: + for name0, name1 in tqdm(pairs, smoothing=0.1): pair = names_to_pair(name0, name1) grp = fd[pair] - kpts0 = grp['keypoints0'].__array__() - kpts1 = grp['keypoints1'].__array__() - scores = grp['scores'].__array__() + kpts0 = grp["keypoints0"].__array__() + kpts1 = grp["keypoints1"].__array__() + scores = grp["scores"].__array__() # Aggregate local features update0 = name0 in required_queries @@ -383,23 +374,35 @@ def aggregate_matches( if update0 and not update1 and max_kps is None: max_error0 = cell_size0 = 0.0 else: - max_error0 = conf['max_error'] - cell_size0 = conf['cell_size'] + max_error0 = conf["max_error"] + cell_size0 = conf["cell_size"] # Get match ids and extend query keypoints (cpdict) - mkp_ids0 = assign_keypoints(kpts0, cpdict[name0], max_error0, - update0, bindict[name0], scores, - cell_size0) - mkp_ids1 = assign_keypoints(kpts1, cpdict[name1], conf['max_error'], - update1, bindict[name1], scores, - conf['cell_size']) + mkp_ids0 = assign_keypoints( + kpts0, + cpdict[name0], + max_error0, + update0, + bindict[name0], + scores, + cell_size0, + ) + mkp_ids1 = assign_keypoints( + kpts1, + cpdict[name1], + conf["max_error"], + update1, + bindict[name1], + scores, + conf["cell_size"], + ) # Build matches from assignments matches0, scores0 = kpids_to_matches0(mkp_ids0, mkp_ids1, scores) assert kpts0.shape[0] == scores.shape[0] - grp.create_dataset('matches0', data=matches0) - grp.create_dataset('matching_scores0', data=scores0) + grp.create_dataset("matches0", data=matches0) + grp.create_dataset("matching_scores0", data=scores0) # Convert bins to kps if finished, and store them for name in (name0, name1): @@ -418,70 +421,75 @@ def aggregate_matches( kp_score = np.array(kp_score)[top_k] # Write query keypoints - with h5py.File(feature_path, 'a') as kfd: + with h5py.File(feature_path, "a") as kfd: if name in kfd: del kfd[name] kgrp = kfd.create_group(name) - kgrp.create_dataset('keypoints', data=cpdict[name]) - kgrp.create_dataset('score', data=kp_score) + kgrp.create_dataset("keypoints", data=cpdict[name]) + kgrp.create_dataset("score", data=kp_score) n_kps += cpdict[name].shape[0] del bindict[name] if len(required_queries) > 0: avg_kp_per_image = round(n_kps / len(required_queries), 1) - logger.info(f'Finished assignment, found {avg_kp_per_image} ' - f'keypoints/image (avg.), total {n_kps}.') + logger.info( + f"Finished assignment, found {avg_kp_per_image} " + f"keypoints/image (avg.), total {n_kps}." + ) return cpdict def assign_matches( - pairs: List[Tuple[str, str]], - match_path: Path, - keypoints: Union[List[Path], Dict[str, np.array]], - max_error: float): + pairs: List[Tuple[str, str]], + match_path: Path, + keypoints: Union[List[Path], Dict[str, np.array]], + max_error: float, +): if isinstance(keypoints, list): keypoints = load_keypoints({}, keypoints, kpts_as_bin=set([])) assert len(set(sum(pairs, ())) - set(keypoints.keys())) == 0 - with h5py.File(str(match_path), 'a') as fd: + with h5py.File(str(match_path), "a") as fd: for name0, name1 in tqdm(pairs): pair = names_to_pair(name0, name1) grp = fd[pair] - kpts0 = grp['keypoints0'].__array__() - kpts1 = grp['keypoints1'].__array__() - scores = grp['scores'].__array__() + kpts0 = grp["keypoints0"].__array__() + kpts1 = grp["keypoints1"].__array__() + scores = grp["scores"].__array__() # NN search across cell boundaries mkp_ids0 = assign_keypoints(kpts0, keypoints[name0], max_error) mkp_ids1 = assign_keypoints(kpts1, keypoints[name1], max_error) - matches0, scores0 = kpids_to_matches0(mkp_ids0, mkp_ids1, - scores) + matches0, scores0 = kpids_to_matches0(mkp_ids0, mkp_ids1, scores) # overwrite matches0 and matching_scores0 - del grp['matches0'], grp['matching_scores0'] - grp.create_dataset('matches0', data=matches0) - grp.create_dataset('matching_scores0', data=scores0) + del grp["matches0"], grp["matching_scores0"] + grp.create_dataset("matches0", data=matches0) + grp.create_dataset("matching_scores0", data=scores0) @torch.no_grad() -def match_and_assign(conf: Dict, - pairs_path: Path, - image_dir: Path, - match_path: Path, # out - feature_path_q: Path, # out - feature_paths_refs: Optional[List[Path]] = [], - max_kps: Optional[int] = 8192, - overwrite: bool = False) -> Path: +def match_and_assign( + conf: Dict, + pairs_path: Path, + image_dir: Path, + match_path: Path, # out + feature_path_q: Path, # out + feature_paths_refs: Optional[List[Path]] = [], + max_kps: Optional[int] = 8192, + overwrite: bool = False, +) -> Path: for path in feature_paths_refs: if not path.exists(): - raise FileNotFoundError(f'Reference feature file {path}.') + raise FileNotFoundError(f"Reference feature file {path}.") pairs = parse_retrieval(pairs_path) pairs = [(q, r) for q, rs in pairs.items() for r in rs] pairs = find_unique_new_pairs(pairs, None if overwrite else match_path) required_queries = set(sum(pairs, ())) - name2ref = {n: i for i, p in enumerate(feature_paths_refs) - for n in list_h5_names(p)} + name2ref = { + n: i for i, p in enumerate(feature_paths_refs) for n in list_h5_names(p) + } existing_refs = required_queries.intersection(set(name2ref.keys())) # images which require feature extraction @@ -499,59 +507,67 @@ def match_and_assign(conf: Dict, return # extract semi-dense matches - match_dense(conf, pairs, image_dir, match_path, - existing_refs=existing_refs) + match_dense(conf, pairs, image_dir, match_path, existing_refs=existing_refs) logger.info("Assigning matches...") # Pre-load existing keypoints cpdict, bindict = load_keypoints( - conf, feature_paths_refs, - quantize=required_queries) + conf, feature_paths_refs, quantize=required_queries + ) # Reassign matches by aggregation cpdict = aggregate_matches( - conf, pairs, match_path, feature_path=feature_path_q, - required_queries=required_queries, max_kps=max_kps, cpdict=cpdict, - bindict=bindict) + conf, + pairs, + match_path, + feature_path=feature_path_q, + required_queries=required_queries, + max_kps=max_kps, + cpdict=cpdict, + bindict=bindict, + ) # Invalidate matches that are far from selected bin by reassignment if max_kps is not None: logger.info(f'Reassign matches with max_error={conf["max_error"]}.') - assign_matches(pairs, match_path, cpdict, - max_error=conf['max_error']) + assign_matches(pairs, match_path, cpdict, max_error=conf["max_error"]) @torch.no_grad() -def main(conf: Dict, - pairs: Path, - image_dir: Path, - export_dir: Optional[Path] = None, - matches: Optional[Path] = None, # out - features: Optional[Path] = None, # out - features_ref: Optional[Path] = None, - max_kps: Optional[int] = 8192, - overwrite: bool = False) -> Path: - logger.info('Extracting semi-dense features with configuration:' - f'\n{pprint.pformat(conf)}') +def main( + conf: Dict, + pairs: Path, + image_dir: Path, + export_dir: Optional[Path] = None, + matches: Optional[Path] = None, # out + features: Optional[Path] = None, # out + features_ref: Optional[Path] = None, + max_kps: Optional[int] = 8192, + overwrite: bool = False, +) -> Path: + logger.info( + "Extracting semi-dense features with configuration:" f"\n{pprint.pformat(conf)}" + ) if features is None: - features = 'feats_' + features = "feats_" if isinstance(features, Path): features_q = features if matches is None: - raise ValueError('Either provide both features and matches as Path' - ' or both as names.') + raise ValueError( + "Either provide both features and matches as Path" " or both as names." + ) else: if export_dir is None: - raise ValueError('Provide an export_dir if features and matches' - f' are not file paths: {features}, {matches}.') - features_q = Path(export_dir, - f'{features}{conf["output"]}.h5') + raise ValueError( + "Provide an export_dir if features and matches" + f" are not file paths: {features}, {matches}." + ) + features_q = Path(export_dir, f'{features}{conf["output"]}.h5') if matches is None: - matches = Path( - export_dir, f'{conf["output"]}_{pairs.stem}.h5') + matches = Path(export_dir, f'{conf["output"]}_{pairs.stem}.h5') if features_ref is None: features_ref = [] @@ -562,24 +578,29 @@ def main(conf: Dict, else: raise TypeError(str(features_ref)) - match_and_assign(conf, pairs, image_dir, matches, - features_q, features_ref, - max_kps, overwrite) + match_and_assign( + conf, pairs, image_dir, matches, features_q, features_ref, max_kps, overwrite + ) return features_q, matches -if __name__ == '__main__': +if __name__ == "__main__": parser = argparse.ArgumentParser() - parser.add_argument('--pairs', type=Path, required=True) - parser.add_argument('--image_dir', type=Path, required=True) - parser.add_argument('--export_dir', type=Path, required=True) - parser.add_argument('--matches', type=Path, - default=confs['loftr']['output']) - parser.add_argument('--features', type=str, - default='feats_' + confs['loftr']['output']) - parser.add_argument('--conf', type=str, default='loftr', - choices=list(confs.keys())) + parser.add_argument("--pairs", type=Path, required=True) + parser.add_argument("--image_dir", type=Path, required=True) + parser.add_argument("--export_dir", type=Path, required=True) + parser.add_argument("--matches", type=Path, default=confs["loftr"]["output"]) + parser.add_argument( + "--features", type=str, default="feats_" + confs["loftr"]["output"] + ) + parser.add_argument("--conf", type=str, default="loftr", choices=list(confs.keys())) args = parser.parse_args() - main(confs[args.conf], args.pairs, args.image_dir, args.export_dir, - args.matches, args.features) + main( + confs[args.conf], + args.pairs, + args.image_dir, + args.export_dir, + args.matches, + args.features, + ) diff --git a/hloc/match_features.py b/hloc/match_features.py index f9ef3e5f..be8cc682 100644 --- a/hloc/match_features.py +++ b/hloc/match_features.py @@ -1,94 +1,91 @@ import argparse -from typing import Union, Optional, Dict, List, Tuple -from pathlib import Path import pprint +from functools import partial +from pathlib import Path from queue import Queue from threading import Thread -from functools import partial -from tqdm import tqdm +from typing import Dict, List, Optional, Tuple, Union + import h5py import torch +from tqdm import tqdm -from . import matchers, logger +from . import logger, matchers from .utils.base_model import dynamic_load from .utils.parsers import names_to_pair, names_to_pair_old, parse_retrieval - -''' +""" A set of standard configurations that can be directly selected from the command line using their name. Each is a dictionary with the following entries: - output: the name of the match file that will be generated. - model: the model configuration, as passed to a feature matcher. -''' +""" confs = { - 'superpoint+lightglue': { - 'output': 'matches-superpoint-lightglue', - 'model': { - 'name': 'lightglue', - 'features': 'superpoint', + "superpoint+lightglue": { + "output": "matches-superpoint-lightglue", + "model": { + "name": "lightglue", + "features": "superpoint", }, }, - 'disk+lightglue': { - 'output': 'matches-disk-lightglue', - 'model': { - 'name': 'lightglue', - 'features': 'disk', + "disk+lightglue": { + "output": "matches-disk-lightglue", + "model": { + "name": "lightglue", + "features": "disk", }, }, - 'superglue': { - 'output': 'matches-superglue', - 'model': { - 'name': 'superglue', - 'weights': 'outdoor', - 'sinkhorn_iterations': 50, + "superglue": { + "output": "matches-superglue", + "model": { + "name": "superglue", + "weights": "outdoor", + "sinkhorn_iterations": 50, }, }, - 'superglue-fast': { - 'output': 'matches-superglue-it5', - 'model': { - 'name': 'superglue', - 'weights': 'outdoor', - 'sinkhorn_iterations': 5, + "superglue-fast": { + "output": "matches-superglue-it5", + "model": { + "name": "superglue", + "weights": "outdoor", + "sinkhorn_iterations": 5, }, }, - 'NN-superpoint': { - 'output': 'matches-NN-mutual-dist.7', - 'model': { - 'name': 'nearest_neighbor', - 'do_mutual_check': True, - 'distance_threshold': 0.7, + "NN-superpoint": { + "output": "matches-NN-mutual-dist.7", + "model": { + "name": "nearest_neighbor", + "do_mutual_check": True, + "distance_threshold": 0.7, }, }, - 'NN-ratio': { - 'output': 'matches-NN-mutual-ratio.8', - 'model': { - 'name': 'nearest_neighbor', - 'do_mutual_check': True, - 'ratio_threshold': 0.8, - } - }, - 'NN-mutual': { - 'output': 'matches-NN-mutual', - 'model': { - 'name': 'nearest_neighbor', - 'do_mutual_check': True, + "NN-ratio": { + "output": "matches-NN-mutual-ratio.8", + "model": { + "name": "nearest_neighbor", + "do_mutual_check": True, + "ratio_threshold": 0.8, }, }, - 'adalam': { - 'output': 'matches-adalam', - 'model': { - 'name': 'adalam' + "NN-mutual": { + "output": "matches-NN-mutual", + "model": { + "name": "nearest_neighbor", + "do_mutual_check": True, }, - } + }, + "adalam": { + "output": "matches-adalam", + "model": {"name": "adalam"}, + }, } -class WorkQueue(): +class WorkQueue: def __init__(self, work_fn, num_threads=1): self.queue = Queue(num_threads) self.threads = [ - Thread(target=self.thread_fn, args=(work_fn,)) - for _ in range(num_threads) + Thread(target=self.thread_fn, args=(work_fn,)) for _ in range(num_threads) ] for thread in self.threads: thread.start() @@ -118,17 +115,17 @@ def __init__(self, pairs, feature_path_q, feature_path_r): def __getitem__(self, idx): name0, name1 = self.pairs[idx] data = {} - with h5py.File(self.feature_path_q, 'r') as fd: + with h5py.File(self.feature_path_q, "r") as fd: grp = fd[name0] for k, v in grp.items(): - data[k+'0'] = torch.from_numpy(v.__array__()).float() + data[k + "0"] = torch.from_numpy(v.__array__()).float() # some matchers might expect an image but only use its size - data['image0'] = torch.empty((1,)+tuple(grp['image_size'])[::-1]) - with h5py.File(self.feature_path_r, 'r') as fd: + data["image0"] = torch.empty((1,) + tuple(grp["image_size"])[::-1]) + with h5py.File(self.feature_path_r, "r") as fd: grp = fd[name1] for k, v in grp.items(): - data[k+'1'] = torch.from_numpy(v.__array__()).float() - data['image1'] = torch.empty((1,)+tuple(grp['image_size'])[::-1]) + data[k + "1"] = torch.from_numpy(v.__array__()).float() + data["image1"] = torch.empty((1,) + tuple(grp["image_size"])[::-1]) return data def __len__(self): @@ -137,37 +134,40 @@ def __len__(self): def writer_fn(inp, match_path): pair, pred = inp - with h5py.File(str(match_path), 'a', libver='latest') as fd: + with h5py.File(str(match_path), "a", libver="latest") as fd: if pair in fd: del fd[pair] grp = fd.create_group(pair) - matches = pred['matches0'][0].cpu().short().numpy() - grp.create_dataset('matches0', data=matches) - if 'matching_scores0' in pred: - scores = pred['matching_scores0'][0].cpu().half().numpy() - grp.create_dataset('matching_scores0', data=scores) - + matches = pred["matches0"][0].cpu().short().numpy() + grp.create_dataset("matches0", data=matches) + if "matching_scores0" in pred: + scores = pred["matching_scores0"][0].cpu().half().numpy() + grp.create_dataset("matching_scores0", data=scores) -def main(conf: Dict, - pairs: Path, features: Union[Path, str], - export_dir: Optional[Path] = None, - matches: Optional[Path] = None, - features_ref: Optional[Path] = None, - overwrite: bool = False) -> Path: +def main( + conf: Dict, + pairs: Path, + features: Union[Path, str], + export_dir: Optional[Path] = None, + matches: Optional[Path] = None, + features_ref: Optional[Path] = None, + overwrite: bool = False, +) -> Path: if isinstance(features, Path) or Path(features).exists(): features_q = features if matches is None: - raise ValueError('Either provide both features and matches as Path' - ' or both as names.') + raise ValueError( + "Either provide both features and matches as Path" " or both as names." + ) else: if export_dir is None: - raise ValueError('Provide an export_dir if features is not' - f' a file path: {features}.') - features_q = Path(export_dir, features+'.h5') + raise ValueError( + "Provide an export_dir if features is not" f" a file path: {features}." + ) + features_q = Path(export_dir, features + ".h5") if matches is None: - matches = Path( - export_dir, f'{features}_{conf["output"]}_{pairs.stem}.h5') + matches = Path(export_dir, f'{features}_{conf["output"]}_{pairs.stem}.h5') if features_ref is None: features_ref = features_q @@ -177,20 +177,22 @@ def main(conf: Dict, def find_unique_new_pairs(pairs_all: List[Tuple[str]], match_path: Path = None): - '''Avoid to recompute duplicates to save time.''' + """Avoid to recompute duplicates to save time.""" pairs = set() for i, j in pairs_all: if (j, i) not in pairs: pairs.add((i, j)) pairs = list(pairs) if match_path is not None and match_path.exists(): - with h5py.File(str(match_path), 'r', libver='latest') as fd: + with h5py.File(str(match_path), "r", libver="latest") as fd: pairs_filtered = [] for i, j in pairs: - if (names_to_pair(i, j) in fd or - names_to_pair(j, i) in fd or - names_to_pair_old(i, j) in fd or - names_to_pair_old(j, i) in fd): + if ( + names_to_pair(i, j) in fd + or names_to_pair(j, i) in fd + or names_to_pair_old(i, j) in fd + or names_to_pair_old(j, i) in fd + ): continue pairs_filtered.append((i, j)) return pairs_filtered @@ -198,19 +200,22 @@ def find_unique_new_pairs(pairs_all: List[Tuple[str]], match_path: Path = None): @torch.no_grad() -def match_from_paths(conf: Dict, - pairs_path: Path, - match_path: Path, - feature_path_q: Path, - feature_path_ref: Path, - overwrite: bool = False) -> Path: - logger.info('Matching local features with configuration:' - f'\n{pprint.pformat(conf)}') +def match_from_paths( + conf: Dict, + pairs_path: Path, + match_path: Path, + feature_path_q: Path, + feature_path_ref: Path, + overwrite: bool = False, +) -> Path: + logger.info( + "Matching local features with configuration:" f"\n{pprint.pformat(conf)}" + ) if not feature_path_q.exists(): - raise FileNotFoundError(f'Query feature file {feature_path_q}.') + raise FileNotFoundError(f"Query feature file {feature_path_q}.") if not feature_path_ref.exists(): - raise FileNotFoundError(f'Reference feature file {feature_path_ref}.') + raise FileNotFoundError(f"Reference feature file {feature_path_ref}.") match_path.parent.mkdir(exist_ok=True, parents=True) assert pairs_path.exists(), pairs_path @@ -218,36 +223,39 @@ def match_from_paths(conf: Dict, pairs = [(q, r) for q, rs in pairs.items() for r in rs] pairs = find_unique_new_pairs(pairs, None if overwrite else match_path) if len(pairs) == 0: - logger.info('Skipping the matching.') + logger.info("Skipping the matching.") return - device = 'cuda' if torch.cuda.is_available() else 'cpu' - Model = dynamic_load(matchers, conf['model']['name']) - model = Model(conf['model']).eval().to(device) + device = "cuda" if torch.cuda.is_available() else "cpu" + Model = dynamic_load(matchers, conf["model"]["name"]) + model = Model(conf["model"]).eval().to(device) dataset = FeaturePairsDataset(pairs, feature_path_q, feature_path_ref) loader = torch.utils.data.DataLoader( - dataset, num_workers=5, batch_size=1, shuffle=False, pin_memory=True) + dataset, num_workers=5, batch_size=1, shuffle=False, pin_memory=True + ) writer_queue = WorkQueue(partial(writer_fn, match_path=match_path), 5) - for idx, data in enumerate(tqdm(loader, smoothing=.1)): - data = {k: v if k.startswith('image') - else v.to(device, non_blocking=True) for k, v in data.items()} + for idx, data in enumerate(tqdm(loader, smoothing=0.1)): + data = { + k: v if k.startswith("image") else v.to(device, non_blocking=True) + for k, v in data.items() + } pred = model(data) pair = names_to_pair(*pairs[idx]) writer_queue.put((pair, pred)) writer_queue.join() - logger.info('Finished exporting matches.') + logger.info("Finished exporting matches.") -if __name__ == '__main__': +if __name__ == "__main__": parser = argparse.ArgumentParser() - parser.add_argument('--pairs', type=Path, required=True) - parser.add_argument('--export_dir', type=Path) - parser.add_argument('--features', type=str, - default='feats-superpoint-n4096-r1024') - parser.add_argument('--matches', type=Path) - parser.add_argument('--conf', type=str, default='superglue', - choices=list(confs.keys())) + parser.add_argument("--pairs", type=Path, required=True) + parser.add_argument("--export_dir", type=Path) + parser.add_argument("--features", type=str, default="feats-superpoint-n4096-r1024") + parser.add_argument("--matches", type=Path) + parser.add_argument( + "--conf", type=str, default="superglue", choices=list(confs.keys()) + ) args = parser.parse_args() main(confs[args.conf], args.pairs, args.features, args.export_dir) diff --git a/hloc/matchers/__init__.py b/hloc/matchers/__init__.py index 7edac76f..a9fd381e 100644 --- a/hloc/matchers/__init__.py +++ b/hloc/matchers/__init__.py @@ -1,3 +1,3 @@ def get_matcher(matcher): - mod = __import__(f'{__name__}.{matcher}', fromlist=['']) - return getattr(mod, 'Model') + mod = __import__(f"{__name__}.{matcher}", fromlist=[""]) + return getattr(mod, "Model") diff --git a/hloc/matchers/adalam.py b/hloc/matchers/adalam.py index 3ce1cb52..d2d4e61d 100644 --- a/hloc/matchers/adalam.py +++ b/hloc/matchers/adalam.py @@ -1,55 +1,67 @@ import torch - -from ..utils.base_model import BaseModel - from kornia.feature.adalam import AdalamFilter from kornia.utils.helpers import get_cuda_device_if_available +from ..utils.base_model import BaseModel + class AdaLAM(BaseModel): - # See https://kornia.readthedocs.io/en/latest/_modules/kornia/feature/adalam/adalam.html. default_conf = { - 'area_ratio': 100, - 'search_expansion': 4, - 'ransac_iters': 128, - 'min_inliers': 6, - 'min_confidence': 200, - 'orientation_difference_threshold': 30, - 'scale_rate_threshold': 1.5, - 'detected_scale_rate_threshold': 5, - 'refit': True, - 'force_seed_mnn': True, - 'device': get_cuda_device_if_available() + "area_ratio": 100, + "search_expansion": 4, + "ransac_iters": 128, + "min_inliers": 6, + "min_confidence": 200, + "orientation_difference_threshold": 30, + "scale_rate_threshold": 1.5, + "detected_scale_rate_threshold": 5, + "refit": True, + "force_seed_mnn": True, + "device": get_cuda_device_if_available(), } required_inputs = [ - 'image0', 'image1', - 'descriptors0', 'descriptors1', - 'keypoints0', 'keypoints1', - 'scales0', 'scales1', - 'oris0', 'oris1'] + "image0", + "image1", + "descriptors0", + "descriptors1", + "keypoints0", + "keypoints1", + "scales0", + "scales1", + "oris0", + "oris1", + ] def _init(self, conf): self.adalam = AdalamFilter(conf) def _forward(self, data): - assert data['keypoints0'].size(0) == 1 - if data['keypoints0'].size(1) < 2 or data['keypoints1'].size(1) < 2: + assert data["keypoints0"].size(0) == 1 + if data["keypoints0"].size(1) < 2 or data["keypoints1"].size(1) < 2: matches = torch.zeros( - (0, 2), dtype=torch.int64, - device=data['keypoints0'].device) + (0, 2), dtype=torch.int64, device=data["keypoints0"].device + ) else: matches = self.adalam.match_and_filter( - data['keypoints0'][0], data['keypoints1'][0], - data['descriptors0'][0].T, data['descriptors1'][0].T, - data['image0'].shape[2:], data['image1'].shape[2:], - data['oris0'][0], data['oris1'][0], - data['scales0'][0], data['scales1'][0] + data["keypoints0"][0], + data["keypoints1"][0], + data["descriptors0"][0].T, + data["descriptors1"][0].T, + data["image0"].shape[2:], + data["image1"].shape[2:], + data["oris0"][0], + data["oris1"][0], + data["scales0"][0], + data["scales1"][0], ) matches_new = torch.full( - (data['keypoints0'].size(1),), -1, - dtype=torch.int64, device=data['keypoints0'].device) + (data["keypoints0"].size(1),), + -1, + dtype=torch.int64, + device=data["keypoints0"].device, + ) matches_new[matches[:, 0]] = matches[:, 1] return { - 'matches0': matches_new.unsqueeze(0), - 'matching_scores0': torch.zeros(matches_new.size(0)).unsqueeze(0) + "matches0": matches_new.unsqueeze(0), + "matching_scores0": torch.zeros(matches_new.size(0)).unsqueeze(0), } diff --git a/hloc/matchers/lightglue.py b/hloc/matchers/lightglue.py index 4663115e..59264036 100644 --- a/hloc/matchers/lightglue.py +++ b/hloc/matchers/lightglue.py @@ -1,25 +1,33 @@ -from ..utils.base_model import BaseModel from lightglue import LightGlue as LightGlue_ +from ..utils.base_model import BaseModel + + class LightGlue(BaseModel): default_conf = { - 'features': 'superpoint', - 'depth_confidence': 0.95, - 'width_confidence': 0.99, + "features": "superpoint", + "depth_confidence": 0.95, + "width_confidence": 0.99, } required_inputs = [ - 'image0', 'keypoints0', 'descriptors0', - 'image1', 'keypoints1', 'descriptors1', + "image0", + "keypoints0", + "descriptors0", + "image1", + "keypoints1", + "descriptors1", ] def _init(self, conf): - self.net = LightGlue_(conf.pop('features'), **conf) + self.net = LightGlue_(conf.pop("features"), **conf) def _forward(self, data): - data['descriptors0'] = data['descriptors0'].transpose(-1, -2) - data['descriptors1'] = data['descriptors1'].transpose(-1, -2) + data["descriptors0"] = data["descriptors0"].transpose(-1, -2) + data["descriptors1"] = data["descriptors1"].transpose(-1, -2) - return self.net({ - 'image0': {k[:-1]: v for k, v in data.items() if k[-1] == '0'}, - 'image1': {k[:-1]: v for k, v in data.items() if k[-1] == '1'} - }) + return self.net( + { + "image0": {k[:-1]: v for k, v in data.items() if k[-1] == "0"}, + "image1": {k[:-1]: v for k, v in data.items() if k[-1] == "1"}, + } + ) diff --git a/hloc/matchers/loftr.py b/hloc/matchers/loftr.py index c833bb6d..7e12bac4 100644 --- a/hloc/matchers/loftr.py +++ b/hloc/matchers/loftr.py @@ -1,53 +1,53 @@ -import torch import warnings -from kornia.feature.loftr.loftr import default_cfg + +import torch from kornia.feature import LoFTR as LoFTR_ +from kornia.feature.loftr.loftr import default_cfg from ..utils.base_model import BaseModel class LoFTR(BaseModel): default_conf = { - 'weights': 'outdoor', - 'match_threshold': 0.2, - 'max_num_matches': None, + "weights": "outdoor", + "match_threshold": 0.2, + "max_num_matches": None, } - required_inputs = [ - 'image0', - 'image1' - ] + required_inputs = ["image0", "image1"] def _init(self, conf): cfg = default_cfg - cfg['match_coarse']['thr'] = conf['match_threshold'] - self.net = LoFTR_(pretrained=conf['weights'], config=cfg) + cfg["match_coarse"]["thr"] = conf["match_threshold"] + self.net = LoFTR_(pretrained=conf["weights"], config=cfg) def _forward(self, data): # For consistency with hloc pairs, we refine kpts in image0! rename = { - 'keypoints0': 'keypoints1', - 'keypoints1': 'keypoints0', - 'image0': 'image1', - 'image1': 'image0', - 'mask0': 'mask1', - 'mask1': 'mask0', + "keypoints0": "keypoints1", + "keypoints1": "keypoints0", + "image0": "image1", + "image1": "image0", + "mask0": "mask1", + "mask1": "mask0", } data_ = {rename[k]: v for k, v in data.items()} with warnings.catch_warnings(): warnings.simplefilter("ignore") pred = self.net(data_) - scores = pred['confidence'] + scores = pred["confidence"] - top_k = self.conf['max_num_matches'] + top_k = self.conf["max_num_matches"] if top_k is not None and len(scores) > top_k: keep = torch.argsort(scores, descending=True)[:top_k] - pred['keypoints0'], pred['keypoints1'] =\ - pred['keypoints0'][keep], pred['keypoints1'][keep] + pred["keypoints0"], pred["keypoints1"] = ( + pred["keypoints0"][keep], + pred["keypoints1"][keep], + ) scores = scores[keep] # Switch back indices pred = {(rename[k] if k in rename else k): v for k, v in pred.items()} - pred['scores'] = scores - del pred['confidence'] + pred["scores"] = scores + del pred["confidence"] return pred diff --git a/hloc/matchers/nearest_neighbor.py b/hloc/matchers/nearest_neighbor.py index 996bfe22..c036f8e7 100644 --- a/hloc/matchers/nearest_neighbor.py +++ b/hloc/matchers/nearest_neighbor.py @@ -8,11 +8,11 @@ def find_nn(sim, ratio_thresh, distance_thresh): dist_nn = 2 * (1 - sim_nn) mask = torch.ones(ind_nn.shape[:-1], dtype=torch.bool, device=sim.device) if ratio_thresh: - mask = mask & (dist_nn[..., 0] <= (ratio_thresh**2)*dist_nn[..., 1]) + mask = mask & (dist_nn[..., 0] <= (ratio_thresh**2) * dist_nn[..., 1]) if distance_thresh: mask = mask & (dist_nn[..., 0] <= distance_thresh**2) matches = torch.where(mask, ind_nn[..., 0], ind_nn.new_tensor(-1)) - scores = torch.where(mask, (sim_nn[..., 0]+1)/2, sim_nn.new_tensor(0)) + scores = torch.where(mask, (sim_nn[..., 0] + 1) / 2, sim_nn.new_tensor(0)) return matches, scores @@ -26,37 +26,37 @@ def mutual_check(m0, m1): class NearestNeighbor(BaseModel): default_conf = { - 'ratio_threshold': None, - 'distance_threshold': None, - 'do_mutual_check': True, + "ratio_threshold": None, + "distance_threshold": None, + "do_mutual_check": True, } - required_inputs = ['descriptors0', 'descriptors1'] + required_inputs = ["descriptors0", "descriptors1"] def _init(self, conf): pass def _forward(self, data): - if data['descriptors0'].size(-1) == 0 or data['descriptors1'].size(-1) == 0: + if data["descriptors0"].size(-1) == 0 or data["descriptors1"].size(-1) == 0: matches0 = torch.full( - data['descriptors0'].shape[:2], -1, - device=data['descriptors0'].device) + data["descriptors0"].shape[:2], -1, device=data["descriptors0"].device + ) return { - 'matches0': matches0, - 'matching_scores0': torch.zeros_like(matches0) + "matches0": matches0, + "matching_scores0": torch.zeros_like(matches0), } - ratio_threshold = self.conf['ratio_threshold'] - if data['descriptors0'].size(-1) == 1 or data['descriptors1'].size(-1) == 1: + ratio_threshold = self.conf["ratio_threshold"] + if data["descriptors0"].size(-1) == 1 or data["descriptors1"].size(-1) == 1: ratio_threshold = None - sim = torch.einsum( - 'bdn,bdm->bnm', data['descriptors0'], data['descriptors1']) + sim = torch.einsum("bdn,bdm->bnm", data["descriptors0"], data["descriptors1"]) matches0, scores0 = find_nn( - sim, ratio_threshold, self.conf['distance_threshold']) - if self.conf['do_mutual_check']: + sim, ratio_threshold, self.conf["distance_threshold"] + ) + if self.conf["do_mutual_check"]: matches1, scores1 = find_nn( - sim.transpose(1, 2), ratio_threshold, - self.conf['distance_threshold']) + sim.transpose(1, 2), ratio_threshold, self.conf["distance_threshold"] + ) matches0 = mutual_check(matches0, matches1) return { - 'matches0': matches0, - 'matching_scores0': scores0, + "matches0": matches0, + "matching_scores0": scores0, } diff --git a/hloc/matchers/superglue.py b/hloc/matchers/superglue.py index 916f9785..7e157236 100644 --- a/hloc/matchers/superglue.py +++ b/hloc/matchers/superglue.py @@ -3,19 +3,25 @@ from ..utils.base_model import BaseModel -sys.path.append(str(Path(__file__).parent / '../../third_party')) -from SuperGluePretrainedNetwork.models.superglue import SuperGlue as SG +sys.path.append(str(Path(__file__).parent / "../../third_party")) +from SuperGluePretrainedNetwork.models.superglue import SuperGlue as SG # noqa: E402 class SuperGlue(BaseModel): default_conf = { - 'weights': 'outdoor', - 'sinkhorn_iterations': 100, - 'match_threshold': 0.2, + "weights": "outdoor", + "sinkhorn_iterations": 100, + "match_threshold": 0.2, } required_inputs = [ - 'image0', 'keypoints0', 'scores0', 'descriptors0', - 'image1', 'keypoints1', 'scores1', 'descriptors1', + "image0", + "keypoints0", + "scores0", + "descriptors0", + "image1", + "keypoints1", + "scores1", + "descriptors1", ] def _init(self, conf): diff --git a/hloc/pairs_from_covisibility.py b/hloc/pairs_from_covisibility.py index 60fff2a2..49f3e57f 100644 --- a/hloc/pairs_from_covisibility.py +++ b/hloc/pairs_from_covisibility.py @@ -1,18 +1,19 @@ import argparse +from collections import defaultdict from pathlib import Path + import numpy as np from tqdm import tqdm -from collections import defaultdict from . import logger from .utils.read_write_model import read_model def main(model, output, num_matched): - logger.info('Reading the COLMAP model...') + logger.info("Reading the COLMAP model...") cameras, images, points3D = read_model(model) - logger.info('Extracting image pairs from covisibility info...') + logger.info("Extracting image pairs from covisibility info...") pairs = [] for image_id, image in tqdm(images.items()): matched = image.point3D_ids != -1 @@ -25,7 +26,7 @@ def main(model, output, num_matched): covis[image_covis_id] += 1 if len(covis) == 0: - logger.info(f'Image {image_id} does not have any covisibility.') + logger.info(f"Image {image_id} does not have any covisibility.") continue covis_ids = np.array(list(covis.keys())) @@ -45,15 +46,15 @@ def main(model, output, num_matched): pair = (image.name, images[i].name) pairs.append(pair) - logger.info(f'Found {len(pairs)} pairs.') - with open(output, 'w') as f: - f.write('\n'.join(' '.join([i, j]) for i, j in pairs)) + logger.info(f"Found {len(pairs)} pairs.") + with open(output, "w") as f: + f.write("\n".join(" ".join([i, j]) for i, j in pairs)) if __name__ == "__main__": parser = argparse.ArgumentParser() - parser.add_argument('--model', required=True, type=Path) - parser.add_argument('--output', required=True, type=Path) - parser.add_argument('--num_matched', required=True, type=int) + parser.add_argument("--model", required=True, type=Path) + parser.add_argument("--output", required=True, type=Path) + parser.add_argument("--num_matched", required=True, type=int) args = parser.parse_args() main(**args.__dict__) diff --git a/hloc/pairs_from_exhaustive.py b/hloc/pairs_from_exhaustive.py index 5858b0fd..0d54ed1d 100644 --- a/hloc/pairs_from_exhaustive.py +++ b/hloc/pairs_from_exhaustive.py @@ -1,31 +1,31 @@ import argparse import collections.abc as collections from pathlib import Path -from typing import Optional, Union, List +from typing import List, Optional, Union from . import logger -from .utils.parsers import parse_image_lists from .utils.io import list_h5_names +from .utils.parsers import parse_image_lists def main( - output: Path, - image_list: Optional[Union[Path, List[str]]] = None, - features: Optional[Path] = None, - ref_list: Optional[Union[Path, List[str]]] = None, - ref_features: Optional[Path] = None): - + output: Path, + image_list: Optional[Union[Path, List[str]]] = None, + features: Optional[Path] = None, + ref_list: Optional[Union[Path, List[str]]] = None, + ref_features: Optional[Path] = None, +): if image_list is not None: if isinstance(image_list, (str, Path)): names_q = parse_image_lists(image_list) elif isinstance(image_list, collections.Iterable): names_q = list(image_list) else: - raise ValueError(f'Unknown type for image list: {image_list}') + raise ValueError(f"Unknown type for image list: {image_list}") elif features is not None: names_q = list_h5_names(features) else: - raise ValueError('Provide either a list of images or a feature file.') + raise ValueError("Provide either a list of images or a feature file.") self_matching = False if ref_list is not None: @@ -34,8 +34,7 @@ def main( elif isinstance(image_list, collections.Iterable): names_ref = list(ref_list) else: - raise ValueError( - f'Unknown type for reference image list: {ref_list}') + raise ValueError(f"Unknown type for reference image list: {ref_list}") elif ref_features is not None: names_ref = list_h5_names(ref_features) else: @@ -49,17 +48,17 @@ def main( continue pairs.append((n1, n2)) - logger.info(f'Found {len(pairs)} pairs.') - with open(output, 'w') as f: - f.write('\n'.join(' '.join([i, j]) for i, j in pairs)) + logger.info(f"Found {len(pairs)} pairs.") + with open(output, "w") as f: + f.write("\n".join(" ".join([i, j]) for i, j in pairs)) if __name__ == "__main__": parser = argparse.ArgumentParser() - parser.add_argument('--output', required=True, type=Path) - parser.add_argument('--image_list', type=Path) - parser.add_argument('--features', type=Path) - parser.add_argument('--ref_list', type=Path) - parser.add_argument('--ref_features', type=Path) + parser.add_argument("--output", required=True, type=Path) + parser.add_argument("--image_list", type=Path) + parser.add_argument("--features", type=Path) + parser.add_argument("--ref_list", type=Path) + parser.add_argument("--ref_features", type=Path) args = parser.parse_args() main(**args.__dict__) diff --git a/hloc/pairs_from_poses.py b/hloc/pairs_from_poses.py index 9643d463..83ee1b8c 100644 --- a/hloc/pairs_from_poses.py +++ b/hloc/pairs_from_poses.py @@ -1,11 +1,12 @@ import argparse from pathlib import Path + import numpy as np import scipy.spatial from . import logger -from .utils.read_write_model import read_images_binary from .pairs_from_retrieval import pairs_from_score_matrix +from .utils.read_write_model import read_images_binary DEFAULT_ROT_THRESH = 30 # in degrees @@ -33,37 +34,35 @@ def get_pairwise_distances(images): # we compute the angle between the principal axes, as two images rotated # around their principal axis still observe the same scene. axes = Rs[:, :, -1] - dots = np.einsum('mi,ni->mn', axes, axes, optimize=True) - dR = np.rad2deg(np.arccos(np.clip(dots, -1., 1.))) + dots = np.einsum("mi,ni->mn", axes, axes, optimize=True) + dR = np.rad2deg(np.arccos(np.clip(dots, -1.0, 1.0))) return ids, dist, dR def main(model, output, num_matched, rotation_threshold=DEFAULT_ROT_THRESH): - logger.info('Reading the COLMAP model...') - images = read_images_binary(model / 'images.bin') + logger.info("Reading the COLMAP model...") + images = read_images_binary(model / "images.bin") - logger.info( - f'Obtaining pairwise distances between {len(images)} images...') + logger.info(f"Obtaining pairwise distances between {len(images)} images...") ids, dist, dR = get_pairwise_distances(images) scores = -dist - invalid = (dR >= rotation_threshold) + invalid = dR >= rotation_threshold np.fill_diagonal(invalid, True) pairs = pairs_from_score_matrix(scores, invalid, num_matched) pairs = [(images[ids[i]].name, images[ids[j]].name) for i, j in pairs] - logger.info(f'Found {len(pairs)} pairs.') - with open(output, 'w') as f: - f.write('\n'.join(' '.join(p) for p in pairs)) + logger.info(f"Found {len(pairs)} pairs.") + with open(output, "w") as f: + f.write("\n".join(" ".join(p) for p in pairs)) if __name__ == "__main__": parser = argparse.ArgumentParser() - parser.add_argument('--model', required=True, type=Path) - parser.add_argument('--output', required=True, type=Path) - parser.add_argument('--num_matched', required=True, type=int) - parser.add_argument('--rotation_threshold', - default=DEFAULT_ROT_THRESH, type=float) + parser.add_argument("--model", required=True, type=Path) + parser.add_argument("--output", required=True, type=Path) + parser.add_argument("--num_matched", required=True, type=int) + parser.add_argument("--rotation_threshold", default=DEFAULT_ROT_THRESH, type=float) args = parser.parse_args() main(**args.__dict__) diff --git a/hloc/pairs_from_retrieval.py b/hloc/pairs_from_retrieval.py index 4329547a..32336801 100644 --- a/hloc/pairs_from_retrieval.py +++ b/hloc/pairs_from_retrieval.py @@ -1,15 +1,16 @@ import argparse +import collections.abc as collections from pathlib import Path from typing import Optional + import h5py import numpy as np import torch -import collections.abc as collections from . import logger +from .utils.io import list_h5_names from .utils.parsers import parse_image_lists from .utils.read_write_model import read_images_binary -from .utils.io import list_h5_names def parse_names(prefix, names, names_all): @@ -18,44 +19,47 @@ def parse_names(prefix, names, names_all): prefix = tuple(prefix) names = [n for n in names_all if n.startswith(prefix)] if len(names) == 0: - raise ValueError( - f'Could not find any image with the prefix `{prefix}`.') + raise ValueError(f"Could not find any image with the prefix `{prefix}`.") elif names is not None: if isinstance(names, (str, Path)): names = parse_image_lists(names) elif isinstance(names, collections.Iterable): names = list(names) else: - raise ValueError(f'Unknown type of image list: {names}.' - 'Provide either a list or a path to a list file.') + raise ValueError( + f"Unknown type of image list: {names}." + "Provide either a list or a path to a list file." + ) else: names = names_all return names -def get_descriptors(names, path, name2idx=None, key='global_descriptor'): +def get_descriptors(names, path, name2idx=None, key="global_descriptor"): if name2idx is None: - with h5py.File(str(path), 'r', libver='latest') as fd: + with h5py.File(str(path), "r", libver="latest") as fd: desc = [fd[n][key].__array__() for n in names] else: desc = [] for n in names: - with h5py.File(str(path[name2idx[n]]), 'r', libver='latest') as fd: + with h5py.File(str(path[name2idx[n]]), "r", libver="latest") as fd: desc.append(fd[n][key].__array__()) return torch.from_numpy(np.stack(desc, 0)).float() -def pairs_from_score_matrix(scores: torch.Tensor, - invalid: np.array, - num_select: int, - min_score: Optional[float] = None): +def pairs_from_score_matrix( + scores: torch.Tensor, + invalid: np.array, + num_select: int, + min_score: Optional[float] = None, +): assert scores.shape == invalid.shape if isinstance(scores, np.ndarray): scores = torch.from_numpy(scores) invalid = torch.from_numpy(invalid).to(scores.device) if min_score is not None: invalid |= scores < min_score - scores.masked_fill_(invalid, float('-inf')) + scores.masked_fill_(invalid, float("-inf")) topk = torch.topk(scores, num_select, dim=1) indices = topk.indices.cpu().numpy() @@ -67,10 +71,18 @@ def pairs_from_score_matrix(scores: torch.Tensor, return pairs -def main(descriptors, output, num_matched, - query_prefix=None, query_list=None, - db_prefix=None, db_list=None, db_model=None, db_descriptors=None): - logger.info('Extracting image pairs from a retrieval database.') +def main( + descriptors, + output, + num_matched, + query_prefix=None, + query_list=None, + db_prefix=None, + db_list=None, + db_model=None, + db_descriptors=None, +): + logger.info("Extracting image pairs from a retrieval database.") # We handle multiple reference feature files. # We only assume that names are unique among them and map names to files. @@ -78,45 +90,44 @@ def main(descriptors, output, num_matched, db_descriptors = descriptors if isinstance(db_descriptors, (Path, str)): db_descriptors = [db_descriptors] - name2db = {n: i for i, p in enumerate(db_descriptors) - for n in list_h5_names(p)} + name2db = {n: i for i, p in enumerate(db_descriptors) for n in list_h5_names(p)} db_names_h5 = list(name2db.keys()) query_names_h5 = list_h5_names(descriptors) if db_model: - images = read_images_binary(db_model / 'images.bin') + images = read_images_binary(db_model / "images.bin") db_names = [i.name for i in images.values()] else: db_names = parse_names(db_prefix, db_list, db_names_h5) if len(db_names) == 0: - raise ValueError('Could not find any database image.') + raise ValueError("Could not find any database image.") query_names = parse_names(query_prefix, query_list, query_names_h5) - device = 'cuda' if torch.cuda.is_available() else 'cpu' + device = "cuda" if torch.cuda.is_available() else "cpu" db_desc = get_descriptors(db_names, db_descriptors, name2db) query_desc = get_descriptors(query_names, descriptors) - sim = torch.einsum('id,jd->ij', query_desc.to(device), db_desc.to(device)) + sim = torch.einsum("id,jd->ij", query_desc.to(device), db_desc.to(device)) # Avoid self-matching self = np.array(query_names)[:, None] == np.array(db_names)[None] pairs = pairs_from_score_matrix(sim, self, num_matched, min_score=0) pairs = [(query_names[i], db_names[j]) for i, j in pairs] - logger.info(f'Found {len(pairs)} pairs.') - with open(output, 'w') as f: - f.write('\n'.join(' '.join([i, j]) for i, j in pairs)) + logger.info(f"Found {len(pairs)} pairs.") + with open(output, "w") as f: + f.write("\n".join(" ".join([i, j]) for i, j in pairs)) if __name__ == "__main__": parser = argparse.ArgumentParser() - parser.add_argument('--descriptors', type=Path, required=True) - parser.add_argument('--output', type=Path, required=True) - parser.add_argument('--num_matched', type=int, required=True) - parser.add_argument('--query_prefix', type=str, nargs='+') - parser.add_argument('--query_list', type=Path) - parser.add_argument('--db_prefix', type=str, nargs='+') - parser.add_argument('--db_list', type=Path) - parser.add_argument('--db_model', type=Path) - parser.add_argument('--db_descriptors', type=Path) + parser.add_argument("--descriptors", type=Path, required=True) + parser.add_argument("--output", type=Path, required=True) + parser.add_argument("--num_matched", type=int, required=True) + parser.add_argument("--query_prefix", type=str, nargs="+") + parser.add_argument("--query_list", type=Path) + parser.add_argument("--db_prefix", type=str, nargs="+") + parser.add_argument("--db_list", type=Path) + parser.add_argument("--db_model", type=Path) + parser.add_argument("--db_descriptors", type=Path) args = parser.parse_args() main(**args.__dict__) diff --git a/hloc/pipelines/4Seasons/localize.py b/hloc/pipelines/4Seasons/localize.py index aa268a50..50ed957b 100644 --- a/hloc/pipelines/4Seasons/localize.py +++ b/hloc/pipelines/4Seasons/localize.py @@ -1,50 +1,67 @@ -from pathlib import Path import argparse +from pathlib import Path -from ... import extract_features, match_features, localize_sfm, logger -from .utils import get_timestamps, delete_unused_images -from .utils import generate_query_lists, generate_localization_pairs -from .utils import prepare_submission, evaluate_submission +from ... import extract_features, localize_sfm, logger, match_features +from .utils import ( + delete_unused_images, + evaluate_submission, + generate_localization_pairs, + generate_query_lists, + get_timestamps, + prepare_submission, +) relocalization_files = { - 'training': 'RelocalizationFilesTrain//relocalizationFile_recording_2020-03-24_17-36-22.txt', - 'validation': 'RelocalizationFilesVal/relocalizationFile_recording_2020-03-03_12-03-23.txt', - 'test0': 'RelocalizationFilesTest/relocalizationFile_recording_2020-03-24_17-45-31_*.txt', - 'test1': 'RelocalizationFilesTest/relocalizationFile_recording_2020-04-23_19-37-00_*.txt', + "training": "RelocalizationFilesTrain//relocalizationFile_recording_2020-03-24_17-36-22.txt", # noqa: E501 + "validation": "RelocalizationFilesVal/relocalizationFile_recording_2020-03-03_12-03-23.txt", # noqa: E501 + "test0": "RelocalizationFilesTest/relocalizationFile_recording_2020-03-24_17-45-31_*.txt", # noqa: E501 + "test1": "RelocalizationFilesTest/relocalizationFile_recording_2020-04-23_19-37-00_*.txt", # noqa: E501 } parser = argparse.ArgumentParser() -parser.add_argument('--sequence', type=str, required=True, - choices=['training', 'validation', 'test0', 'test1'], - help='Sequence to be relocalized.') -parser.add_argument('--dataset', type=Path, default='datasets/4Seasons', - help='Path to the dataset, default: %(default)s') -parser.add_argument('--outputs', type=Path, default='outputs/4Seasons', - help='Path to the output directory, default: %(default)s') +parser.add_argument( + "--sequence", + type=str, + required=True, + choices=["training", "validation", "test0", "test1"], + help="Sequence to be relocalized.", +) +parser.add_argument( + "--dataset", + type=Path, + default="datasets/4Seasons", + help="Path to the dataset, default: %(default)s", +) +parser.add_argument( + "--outputs", + type=Path, + default="outputs/4Seasons", + help="Path to the output directory, default: %(default)s", +) args = parser.parse_args() sequence = args.sequence data_dir = args.dataset -ref_dir = data_dir / 'reference' -assert ref_dir.exists(), f'{ref_dir} does not exist' +ref_dir = data_dir / "reference" +assert ref_dir.exists(), f"{ref_dir} does not exist" seq_dir = data_dir / sequence -assert seq_dir.exists(), f'{seq_dir} does not exist' -seq_images = seq_dir / 'undistorted_images' +assert seq_dir.exists(), f"{seq_dir} does not exist" +seq_images = seq_dir / "undistorted_images" reloc = ref_dir / relocalization_files[sequence] output_dir = args.outputs output_dir.mkdir(exist_ok=True, parents=True) -query_list = output_dir / f'{sequence}_queries_with_intrinsics.txt' -ref_pairs = output_dir / 'pairs-db-dist20.txt' -ref_sfm = output_dir / 'sfm_superpoint+superglue' -results_path = output_dir / f'localization_{sequence}_hloc+superglue.txt' -submission_dir = output_dir / 'submission_hloc+superglue' +query_list = output_dir / f"{sequence}_queries_with_intrinsics.txt" +ref_pairs = output_dir / "pairs-db-dist20.txt" +ref_sfm = output_dir / "sfm_superpoint+superglue" +results_path = output_dir / f"localization_{sequence}_hloc+superglue.txt" +submission_dir = output_dir / "submission_hloc+superglue" num_loc_pairs = 10 -loc_pairs = output_dir / f'pairs-query-{sequence}-dist{num_loc_pairs}.txt' +loc_pairs = output_dir / f"pairs-query-{sequence}-dist{num_loc_pairs}.txt" -fconf = extract_features.confs['superpoint_max'] -mconf = match_features.confs['superglue'] +fconf = extract_features.confs["superpoint_max"] +mconf = match_features.confs["superglue"] # Not all query images that are used for the evaluation # To save time in feature extraction, we delete unsused images. @@ -55,20 +72,18 @@ generate_query_lists(timestamps, seq_dir, query_list) # Generate the localization pairs from the given reference frames. -generate_localization_pairs( - sequence, reloc, num_loc_pairs, ref_pairs, loc_pairs) +generate_localization_pairs(sequence, reloc, num_loc_pairs, ref_pairs, loc_pairs) # Extract, match, amd localize. ffile = extract_features.main(fconf, seq_images, output_dir) -mfile = match_features.main(mconf, loc_pairs, fconf['output'], output_dir) -localize_sfm.main( - ref_sfm, query_list, loc_pairs, ffile, mfile, results_path) +mfile = match_features.main(mconf, loc_pairs, fconf["output"], output_dir) +localize_sfm.main(ref_sfm, query_list, loc_pairs, ffile, mfile, results_path) # Convert the absolute poses to relative poses with the reference frames. submission_dir.mkdir(exist_ok=True) -prepare_submission(results_path, reloc, ref_dir / 'poses.txt', submission_dir) +prepare_submission(results_path, reloc, ref_dir / "poses.txt", submission_dir) # If not a test sequence: evaluation the localization accuracy -if 'test' not in sequence: - logger.info('Evaluating the relocalization submission...') +if "test" not in sequence: + logger.info("Evaluating the relocalization submission...") evaluate_submission(submission_dir, reloc) diff --git a/hloc/pipelines/4Seasons/prepare_reference.py b/hloc/pipelines/4Seasons/prepare_reference.py index 1a1de52e..f47aee77 100644 --- a/hloc/pipelines/4Seasons/prepare_reference.py +++ b/hloc/pipelines/4Seasons/prepare_reference.py @@ -1,36 +1,42 @@ -from pathlib import Path import argparse +from pathlib import Path -from ... import extract_features, match_features -from ... import pairs_from_poses, triangulation -from .utils import get_timestamps, delete_unused_images -from .utils import build_empty_colmap_model +from ... import extract_features, match_features, pairs_from_poses, triangulation +from .utils import build_empty_colmap_model, delete_unused_images, get_timestamps parser = argparse.ArgumentParser() -parser.add_argument('--dataset', type=Path, default='datasets/4Seasons', - help='Path to the dataset, default: %(default)s') -parser.add_argument('--outputs', type=Path, default='outputs/4Seasons', - help='Path to the output directory, default: %(default)s') +parser.add_argument( + "--dataset", + type=Path, + default="datasets/4Seasons", + help="Path to the dataset, default: %(default)s", +) +parser.add_argument( + "--outputs", + type=Path, + default="outputs/4Seasons", + help="Path to the output directory, default: %(default)s", +) args = parser.parse_args() -ref_dir = args.dataset / 'reference' -assert ref_dir.exists(), f'{ref_dir} does not exist' -ref_images = ref_dir / 'undistorted_images' +ref_dir = args.dataset / "reference" +assert ref_dir.exists(), f"{ref_dir} does not exist" +ref_images = ref_dir / "undistorted_images" output_dir = args.outputs output_dir.mkdir(exist_ok=True, parents=True) -ref_sfm_empty = output_dir / 'sfm_reference_empty' -ref_sfm = output_dir / 'sfm_superpoint+superglue' +ref_sfm_empty = output_dir / "sfm_reference_empty" +ref_sfm = output_dir / "sfm_superpoint+superglue" num_ref_pairs = 20 -ref_pairs = output_dir / f'pairs-db-dist{num_ref_pairs}.txt' +ref_pairs = output_dir / f"pairs-db-dist{num_ref_pairs}.txt" -fconf = extract_features.confs['superpoint_max'] -mconf = match_features.confs['superglue'] +fconf = extract_features.confs["superpoint_max"] +mconf = match_features.confs["superglue"] # Only reference images that have a pose are used in the pipeline. # To save time in feature extraction, we delete unsused images. -delete_unused_images(ref_images, get_timestamps(ref_dir / 'poses.txt', 0)) +delete_unused_images(ref_images, get_timestamps(ref_dir / "poses.txt", 0)) # Build an empty COLMAP model containing only camera and images # from the provided poses and intrinsics. @@ -41,5 +47,5 @@ # Extract, match, and triangulate the reference SfM model. ffile = extract_features.main(fconf, ref_images, output_dir) -mfile = match_features.main(mconf, ref_pairs, fconf['output'], output_dir) +mfile = match_features.main(mconf, ref_pairs, fconf["output"], output_dir) triangulation.main(ref_sfm, ref_sfm_empty, ref_images, ref_pairs, ffile, mfile) diff --git a/hloc/pipelines/4Seasons/utils.py b/hloc/pipelines/4Seasons/utils.py index 95e7a3d9..e5aace9d 100644 --- a/hloc/pipelines/4Seasons/utils.py +++ b/hloc/pipelines/4Seasons/utils.py @@ -1,12 +1,18 @@ -import os -import numpy as np +import glob import logging +import os from pathlib import Path -import glob -from ...utils.read_write_model import qvec2rotmat, rotmat2qvec -from ...utils.read_write_model import Image, write_model, Camera +import numpy as np + from ...utils.parsers import parse_retrieval +from ...utils.read_write_model import ( + Camera, + Image, + qvec2rotmat, + rotmat2qvec, + write_model, +) logger = logging.getLogger(__name__) @@ -19,38 +25,38 @@ def get_timestamps(files, idx): lines += f.readlines() timestamps = set() for line in lines: - line = line.rstrip('\n') - if line[0] == '#' or line == '': + line = line.rstrip("\n") + if line[0] == "#" or line == "": continue - ts = line.replace(',', ' ').split()[idx] + ts = line.replace(",", " ").split()[idx] timestamps.add(ts) return timestamps def delete_unused_images(root, timestamps): """Delete all images in root if they are not contained in timestamps.""" - images = glob.glob((root / '**/*.png').as_posix(), recursive=True) + images = glob.glob((root / "**/*.png").as_posix(), recursive=True) deleted = 0 for image in images: ts = Path(image).stem if ts not in timestamps: os.remove(image) deleted += 1 - logger.info(f'Deleted {deleted} images in {root}.') + logger.info(f"Deleted {deleted} images in {root}.") def camera_from_calibration_file(id_, path): """Create a COLMAP camera from an MLAD calibration file.""" - with open(path, 'r') as f: + with open(path, "r") as f: data = f.readlines() model, fx, fy, cx, cy = data[0].split()[:5] width, height = data[1].split() - assert model == 'Pinhole' - model_name = 'PINHOLE' + assert model == "Pinhole" + model_name = "PINHOLE" params = [float(i) for i in [fx, fy, cx, cy]] camera = Camera( - id=id_, model=model_name, - width=int(width), height=int(height), params=params) + id=id_, model=model_name, width=int(width), height=int(height), params=params + ) return camera @@ -59,10 +65,10 @@ def parse_poses(path, colmap=False): poses = [] with open(path) as f: for line in f.readlines(): - line = line.rstrip('\n') - if line[0] == '#' or line == '': + line = line.rstrip("\n") + if line[0] == "#" or line == "": continue - data = line.replace(',', ' ').split() + data = line.replace(",", " ").split() ts, p = data[0], np.array(data[1:], float) if colmap: q, t = np.split(p, [4]) @@ -79,10 +85,10 @@ def parse_relocalization(path, has_poses=False): reloc = [] with open(path) as f: for line in f.readlines(): - line = line.rstrip('\n') - if line[0] == '#' or line == '': + line = line.rstrip("\n") + if line[0] == "#" or line == "": continue - data = line.replace(',', ' ').split() + data = line.replace(",", " ").split() out = data[:2] # ref_ts, q_ts if has_poses: assert len(data) == 9 @@ -96,13 +102,13 @@ def parse_relocalization(path, has_poses=False): def build_empty_colmap_model(root, sfm_dir): """Build a COLMAP model with images and cameras only.""" - calibration = 'Calibration/undistorted_calib_{}.txt' + calibration = "Calibration/undistorted_calib_{}.txt" cam0 = camera_from_calibration_file(0, root / calibration.format(0)) cam1 = camera_from_calibration_file(1, root / calibration.format(1)) cameras = {0: cam0, 1: cam1} - T_0to1 = np.loadtxt(root / 'Calibration/undistorted_calib_stereo.txt') - poses = parse_poses(root / 'poses.txt') + T_0to1 = np.loadtxt(root / "Calibration/undistorted_calib_stereo.txt") + poses = parse_poses(root / "poses.txt") images = {} id_ = 0 for ts, R_cam0_to_w, t_cam0_to_w in poses: @@ -113,63 +119,65 @@ def build_empty_colmap_model(root, sfm_dir): t_w_to_cam1 = T_0to1[:3, :3] @ t_w_to_cam0 + T_0to1[:3, 3] for idx, (R_w_to_cam, t_w_to_cam) in enumerate( - zip([R_w_to_cam0, R_w_to_cam1], [t_w_to_cam0, t_w_to_cam1])): + zip([R_w_to_cam0, R_w_to_cam1], [t_w_to_cam0, t_w_to_cam1]) + ): image = Image( id=id_, qvec=rotmat2qvec(R_w_to_cam), tvec=t_w_to_cam, camera_id=idx, - name=f'cam{idx}/{ts}.png', + name=f"cam{idx}/{ts}.png", xys=np.zeros((0, 2), float), - point3D_ids=np.full(0, -1, int)) + point3D_ids=np.full(0, -1, int), + ) images[id_] = image id_ += 1 sfm_dir.mkdir(exist_ok=True, parents=True) - write_model(cameras, images, {}, path=str(sfm_dir), ext='.bin') + write_model(cameras, images, {}, path=str(sfm_dir), ext=".bin") def generate_query_lists(timestamps, seq_dir, out_path): """Create a list of query images with intrinsics from timestamps.""" cam0 = camera_from_calibration_file( - 0, seq_dir / 'Calibration/undistorted_calib_0.txt') + 0, seq_dir / "Calibration/undistorted_calib_0.txt" + ) intrinsics = [cam0.model, cam0.width, cam0.height] + cam0.params intrinsics = [str(p) for p in intrinsics] - data = map(lambda ts: ' '.join([f'cam0/{ts}.png']+intrinsics), timestamps) - with open(out_path, 'w') as f: - f.write('\n'.join(data)) + data = map(lambda ts: " ".join([f"cam0/{ts}.png"] + intrinsics), timestamps) + with open(out_path, "w") as f: + f.write("\n".join(data)) def generate_localization_pairs(sequence, reloc, num, ref_pairs, out_path): """Create the matching pairs for the localization. - We simply lookup the corresponding reference frame - and extract its `num` closest frames from the existing pair list. + We simply lookup the corresponding reference frame + and extract its `num` closest frames from the existing pair list. """ - if 'test' in sequence: + if "test" in sequence: # hard pairs will be overwritten by easy ones if available - relocs = [ - str(reloc).replace('*', d) for d in ['hard', 'moderate', 'easy']] + relocs = [str(reloc).replace("*", d) for d in ["hard", "moderate", "easy"]] else: relocs = [reloc] query_to_ref_ts = {} for reloc in relocs: - with open(reloc, 'r') as f: + with open(reloc, "r") as f: for line in f.readlines(): - line = line.rstrip('\n') - if line[0] == '#' or line == '': + line = line.rstrip("\n") + if line[0] == "#" or line == "": continue ref_ts, q_ts = line.split()[:2] query_to_ref_ts[q_ts] = ref_ts - ts_to_name = 'cam0/{}.png'.format + ts_to_name = "cam0/{}.png".format ref_pairs = parse_retrieval(ref_pairs) loc_pairs = [] for q_ts, ref_ts in query_to_ref_ts.items(): ref_name = ts_to_name(ref_ts) - selected = [ref_name] + ref_pairs[ref_name][:num-1] - loc_pairs.extend([' '.join((ts_to_name(q_ts), s)) for s in selected]) - with open(out_path, 'w') as f: - f.write('\n'.join(loc_pairs)) + selected = [ref_name] + ref_pairs[ref_name][: num - 1] + loc_pairs.extend([" ".join((ts_to_name(q_ts), s)) for s in selected]) + with open(out_path, "w") as f: + f.write("\n".join(loc_pairs)) def prepare_submission(results, relocs, poses_path, out_dir): @@ -194,24 +202,20 @@ def prepare_submission(results, relocs, poses_path, out_dir): qvec = rotmat2qvec(R_ref0_to_q0)[[1, 2, 3, 0]] # wxyz to xyzw out = [ref_ts, q_ts] + list(map(str, tvec)) + list(map(str, qvec)) - relative_poses.append(' '.join(out)) + relative_poses.append(" ".join(out)) out_path = out_dir / reloc.name - with open(out_path, 'w') as f: - f.write('\n'.join(relative_poses)) - logger.info(f'Submission file written to {out_path}.') + with open(out_path, "w") as f: + f.write("\n".join(relative_poses)) + logger.info(f"Submission file written to {out_path}.") def evaluate_submission(submission_dir, relocs, ths=[0.1, 0.2, 0.5]): - """Compute the relocalization recall from predicted and ground truth poses. - """ + """Compute the relocalization recall from predicted and ground truth poses.""" for reloc in relocs.parent.glob(relocs.name): - poses_gt = parse_relocalization( - reloc, has_poses=True) - poses_pred = parse_relocalization( - submission_dir / reloc.name, has_poses=True) - poses_pred = { - (ref_ts, q_ts): (R, t) for ref_ts, q_ts, R, t in poses_pred} + poses_gt = parse_relocalization(reloc, has_poses=True) + poses_pred = parse_relocalization(submission_dir / reloc.name, has_poses=True) + poses_pred = {(ref_ts, q_ts): (R, t) for ref_ts, q_ts, R, t in poses_pred} error = [] for ref_ts, q_ts, R_gt, t_gt in poses_gt: @@ -221,7 +225,7 @@ def evaluate_submission(submission_dir, relocs, ths=[0.1, 0.2, 0.5]): error = np.array(error) recall = [np.mean(error <= th) for th in ths] - s = f'Relocalization evaluation {submission_dir.name}/{reloc.name}\n' - s += ' / '.join([f'{th:>7}m' for th in ths]) + '\n' - s += ' / '.join([f'{100*r:>7.3f}%' for r in recall]) + s = f"Relocalization evaluation {submission_dir.name}/{reloc.name}\n" + s += " / ".join([f"{th:>7}m" for th in ths]) + "\n" + s += " / ".join([f"{100*r:>7.3f}%" for r in recall]) logger.info(s) diff --git a/hloc/pipelines/7Scenes/create_gt_sfm.py b/hloc/pipelines/7Scenes/create_gt_sfm.py index 3eccb5c5..95dfa461 100644 --- a/hloc/pipelines/7Scenes/create_gt_sfm.py +++ b/hloc/pipelines/7Scenes/create_gt_sfm.py @@ -1,11 +1,12 @@ from pathlib import Path + import numpy as np -import torch import PIL.Image -from tqdm import tqdm import pycolmap +import torch +from tqdm import tqdm -from ...utils.read_write_model import write_model, read_model +from ...utils.read_write_model import read_model, write_model def scene_coordinates(p2D, R_w2c, t_w2c, depth, camera): @@ -19,7 +20,7 @@ def scene_coordinates(p2D, R_w2c, t_w2c, depth, camera): def interpolate_depth(depth, kp): h, w = depth.shape - kp = kp / np.array([[w-1, h-1]]) * 2 - 1 + kp = kp / np.array([[w - 1, h - 1]]) * 2 - 1 assert np.all(kp > -1) and np.all(kp < 1) depth = torch.from_numpy(depth)[None, None] kp = torch.from_numpy(kp)[None, None] @@ -27,10 +28,10 @@ def interpolate_depth(depth, kp): # To maximize the number of points that have depth: # do bilinear interpolation first and then nearest for the remaining points - interp_lin = grid_sample( - depth, kp, align_corners=True, mode='bilinear')[0, :, 0] + interp_lin = grid_sample(depth, kp, align_corners=True, mode="bilinear")[0, :, 0] interp_nn = torch.nn.functional.grid_sample( - depth, kp, align_corners=True, mode='nearest')[0, :, 0] + depth, kp, align_corners=True, mode="nearest" + )[0, :, 0] interp = torch.where(torch.isnan(interp_lin), interp_nn, interp_lin) valid = ~torch.any(torch.isnan(interp), 0) @@ -40,10 +41,10 @@ def interpolate_depth(depth, kp): def image_path_to_rendered_depth_path(image_name): - parts = image_name.split('/') - name = '_'.join([''.join(parts[0].split('-')), parts[1]]) - name = name.replace('color', 'pose') - name = name.replace('png', 'depth.tiff') + parts = image_name.split("/") + name = "_".join(["".join(parts[0].split("-")), parts[1]]) + name = name.replace("color", "pose") + name = name.replace("png", "depth.tiff") return name @@ -65,8 +66,8 @@ def correct_sfm_with_gt_depth(sfm_path, depth_folder_path, output_path): depth_name = image_path_to_rendered_depth_path(image_name) depth = PIL.Image.open(Path(depth_folder_path) / depth_name) - depth = np.array(depth).astype('float64') - depth = depth/1000. # mm to meter + depth = np.array(depth).astype("float64") + depth = depth / 1000.0 # mm to meter depth[(depth == 0.0) | (depth > 1000.0)] = np.nan R_w2c, t_w2c = img.qvec2rotmat(), img.tvec @@ -77,12 +78,17 @@ def correct_sfm_with_gt_depth(sfm_path, depth_folder_path, output_path): p2Ds, valids_projected = project_to_image(p3Ds, R_w2c, t_w2c, camera) invalid_p3D_ids = p3D_ids[p3D_ids != -1][~valids_projected] interp_depth, valids_backprojected = interpolate_depth(depth, p2Ds) - scs = scene_coordinates(p2Ds[valids_backprojected], R_w2c, t_w2c, - interp_depth[valids_backprojected], - camera) + scs = scene_coordinates( + p2Ds[valids_backprojected], + R_w2c, + t_w2c, + interp_depth[valids_backprojected], + camera, + ) invalid_p3D_ids = np.append( invalid_p3D_ids, - p3D_ids[p3D_ids != -1][valids_projected][~valids_backprojected]) + p3D_ids[p3D_ids != -1][valids_projected][~valids_backprojected], + ) for p3did in invalid_p3D_ids: if p3did == -1: continue @@ -91,8 +97,10 @@ def correct_sfm_with_gt_depth(sfm_path, depth_folder_path, output_path): invalid_imgids = list(np.where(obs_imgids == img.id)[0]) points3D[p3did] = points3D[p3did]._replace( image_ids=np.delete(obs_imgids, invalid_imgids), - point2D_idxs=np.delete(points3D[p3did].point2D_idxs, - invalid_imgids)) + point2D_idxs=np.delete( + points3D[p3did].point2D_idxs, invalid_imgids + ), + ) new_p3D_ids = p3D_ids.copy() sub_p3D_ids = new_p3D_ids[new_p3D_ids != -1] @@ -103,8 +111,9 @@ def correct_sfm_with_gt_depth(sfm_path, depth_folder_path, output_path): new_p3D_ids[new_p3D_ids != -1] = sub_p3D_ids img = img._replace(point3D_ids=new_p3D_ids) - assert len(img.point3D_ids[img.point3D_ids != -1]) == len(scs), ( - f"{len(scs)}, {len(img.point3D_ids[img.point3D_ids != -1])}") + assert len(img.point3D_ids[img.point3D_ids != -1]) == len( + scs + ), f"{len(scs)}, {len(img.point3D_ids[img.point3D_ids != -1])}" for i, p3did in enumerate(img.point3D_ids[img.point3D_ids != -1]): points3D[p3did] = points3D[p3did]._replace(xyz=scs[i]) images[imgid] = img @@ -113,14 +122,13 @@ def correct_sfm_with_gt_depth(sfm_path, depth_folder_path, output_path): write_model(cameras, images, points3D, output_path) -if __name__ == '__main__': - dataset = Path('datasets/7scenes') - outputs = Path('outputs/7Scenes') +if __name__ == "__main__": + dataset = Path("datasets/7scenes") + outputs = Path("outputs/7Scenes") - SCENES = ['chess', 'fire', 'heads', 'office', 'pumpkin', - 'redkitchen', 'stairs'] + SCENES = ["chess", "fire", "heads", "office", "pumpkin", "redkitchen", "stairs"] for scene in SCENES: - sfm_path = outputs / scene / 'sfm_superpoint+superglue' - depth_path = dataset / f'depth/7scenes_{scene}/train/depth' - output_path = outputs / scene / 'sfm_superpoint+superglue+depth' + sfm_path = outputs / scene / "sfm_superpoint+superglue" + depth_path = dataset / f"depth/7scenes_{scene}/train/depth" + output_path = outputs / scene / "sfm_superpoint+superglue+depth" correct_sfm_with_gt_depth(sfm_path, depth_path, output_path) diff --git a/hloc/pipelines/7Scenes/pipeline.py b/hloc/pipelines/7Scenes/pipeline.py index 754e06d8..6fc28c6d 100644 --- a/hloc/pipelines/7Scenes/pipeline.py +++ b/hloc/pipelines/7Scenes/pipeline.py @@ -1,66 +1,76 @@ -from pathlib import Path import argparse +from pathlib import Path -from .utils import create_reference_sfm -from .create_gt_sfm import correct_sfm_with_gt_depth +from ... import ( + extract_features, + localize_sfm, + logger, + match_features, + pairs_from_covisibility, + triangulation, +) from ..Cambridge.utils import create_query_list_with_intrinsics, evaluate -from ... import extract_features, match_features, pairs_from_covisibility -from ... import triangulation, localize_sfm, logger +from .create_gt_sfm import correct_sfm_with_gt_depth +from .utils import create_reference_sfm -SCENES = ['chess', 'fire', 'heads', 'office', 'pumpkin', - 'redkitchen', 'stairs'] +SCENES = ["chess", "fire", "heads", "office", "pumpkin", "redkitchen", "stairs"] -def run_scene(images, gt_dir, retrieval, outputs, results, num_covis, - use_dense_depth, depth_dir=None): +def run_scene( + images, + gt_dir, + retrieval, + outputs, + results, + num_covis, + use_dense_depth, + depth_dir=None, +): outputs.mkdir(exist_ok=True, parents=True) - ref_sfm_sift = outputs / 'sfm_sift' - ref_sfm = outputs / 'sfm_superpoint+superglue' - query_list = outputs / 'query_list_with_intrinsics.txt' + ref_sfm_sift = outputs / "sfm_sift" + ref_sfm = outputs / "sfm_superpoint+superglue" + query_list = outputs / "query_list_with_intrinsics.txt" feature_conf = { - 'output': 'feats-superpoint-n4096-r1024', - 'model': { - 'name': 'superpoint', - 'nms_radius': 3, - 'max_keypoints': 4096, + "output": "feats-superpoint-n4096-r1024", + "model": { + "name": "superpoint", + "nms_radius": 3, + "max_keypoints": 4096, }, - 'preprocessing': { - 'globs': ['*.color.png'], - 'grayscale': True, - 'resize_max': 1024, + "preprocessing": { + "globs": ["*.color.png"], + "grayscale": True, + "resize_max": 1024, }, } - matcher_conf = match_features.confs['superglue'] - matcher_conf['model']['sinkhorn_iterations'] = 5 + matcher_conf = match_features.confs["superglue"] + matcher_conf["model"]["sinkhorn_iterations"] = 5 - test_list = gt_dir / 'list_test.txt' + test_list = gt_dir / "list_test.txt" create_reference_sfm(gt_dir, ref_sfm_sift, test_list) create_query_list_with_intrinsics(gt_dir, query_list, test_list) - features = extract_features.main( - feature_conf, images, outputs, as_half=True) + features = extract_features.main(feature_conf, images, outputs, as_half=True) - sfm_pairs = outputs / f'pairs-db-covis{num_covis}.txt' - pairs_from_covisibility.main( - ref_sfm_sift, sfm_pairs, num_matched=num_covis) + sfm_pairs = outputs / f"pairs-db-covis{num_covis}.txt" + pairs_from_covisibility.main(ref_sfm_sift, sfm_pairs, num_matched=num_covis) sfm_matches = match_features.main( - matcher_conf, sfm_pairs, feature_conf['output'], outputs) + matcher_conf, sfm_pairs, feature_conf["output"], outputs + ) if not (use_dense_depth and ref_sfm.exists()): triangulation.main( - ref_sfm, ref_sfm_sift, - images, - sfm_pairs, - features, - sfm_matches) + ref_sfm, ref_sfm_sift, images, sfm_pairs, features, sfm_matches + ) if use_dense_depth: assert depth_dir is not None - ref_sfm_fix = outputs / 'sfm_superpoint+superglue+depth' + ref_sfm_fix = outputs / "sfm_superpoint+superglue+depth" correct_sfm_with_gt_depth(ref_sfm, depth_dir, ref_sfm_fix) ref_sfm = ref_sfm_fix loc_matches = match_features.main( - matcher_conf, retrieval, feature_conf['output'], outputs) + matcher_conf, retrieval, feature_conf["output"], outputs + ) localize_sfm.main( ref_sfm, @@ -70,43 +80,60 @@ def run_scene(images, gt_dir, retrieval, outputs, results, num_covis, loc_matches, results, covisibility_clustering=False, - prepend_camera_name=True) + prepend_camera_name=True, + ) -if __name__ == '__main__': +if __name__ == "__main__": parser = argparse.ArgumentParser() - parser.add_argument('--scenes', default=SCENES, choices=SCENES, nargs='+') - parser.add_argument('--overwrite', action='store_true') - parser.add_argument('--dataset', type=Path, default='datasets/7scenes', - help='Path to the dataset, default: %(default)s') - parser.add_argument('--outputs', type=Path, default='outputs/7scenes', - help='Path to the output directory, default: %(default)s') - parser.add_argument('--use_dense_depth', action='store_true') - parser.add_argument('--num_covis', type=int, default=30, - help='Number of image pairs for SfM, default: %(default)s') + parser.add_argument("--scenes", default=SCENES, choices=SCENES, nargs="+") + parser.add_argument("--overwrite", action="store_true") + parser.add_argument( + "--dataset", + type=Path, + default="datasets/7scenes", + help="Path to the dataset, default: %(default)s", + ) + parser.add_argument( + "--outputs", + type=Path, + default="outputs/7scenes", + help="Path to the output directory, default: %(default)s", + ) + parser.add_argument("--use_dense_depth", action="store_true") + parser.add_argument( + "--num_covis", + type=int, + default=30, + help="Number of image pairs for SfM, default: %(default)s", + ) args = parser.parse_args() - gt_dirs = args.dataset / '7scenes_sfm_triangulated/{scene}/triangulated' - retrieval_dirs = args.dataset / '7scenes_densevlad_retrieval_top_10' + gt_dirs = args.dataset / "7scenes_sfm_triangulated/{scene}/triangulated" + retrieval_dirs = args.dataset / "7scenes_densevlad_retrieval_top_10" all_results = {} for scene in args.scenes: logger.info(f'Working on scene "{scene}".') - results = args.outputs / scene / 'results_{}.txt'.format( - "dense" if args.use_dense_depth else "sparse") + results = ( + args.outputs + / scene + / "results_{}.txt".format("dense" if args.use_dense_depth else "sparse") + ) if args.overwrite or not results.exists(): run_scene( args.dataset / scene, Path(str(gt_dirs).format(scene=scene)), - retrieval_dirs / f'{scene}_top10.txt', + retrieval_dirs / f"{scene}_top10.txt", args.outputs / scene, results, args.num_covis, args.use_dense_depth, - depth_dir=args.dataset / f'depth/7scenes_{scene}/train/depth') + depth_dir=args.dataset / f"depth/7scenes_{scene}/train/depth", + ) all_results[scene] = results for scene in args.scenes: logger.info(f'Evaluate scene "{scene}".') gt_dir = Path(str(gt_dirs).format(scene=scene)) - evaluate(gt_dir, all_results[scene], gt_dir / 'list_test.txt') + evaluate(gt_dir, all_results[scene], gt_dir / "list_test.txt") diff --git a/hloc/pipelines/7Scenes/utils.py b/hloc/pipelines/7Scenes/utils.py index 871343fb..1cb02128 100644 --- a/hloc/pipelines/7Scenes/utils.py +++ b/hloc/pipelines/7Scenes/utils.py @@ -1,4 +1,5 @@ import logging + import numpy as np from hloc.utils.read_write_model import read_model, write_model @@ -6,15 +7,15 @@ logger = logging.getLogger(__name__) -def create_reference_sfm(full_model, ref_model, blacklist=None, ext='.bin'): - '''Create a new COLMAP model with only training images.''' - logger.info('Creating the reference model.') +def create_reference_sfm(full_model, ref_model, blacklist=None, ext=".bin"): + """Create a new COLMAP model with only training images.""" + logger.info("Creating the reference model.") ref_model.mkdir(exist_ok=True) cameras, images, points3D = read_model(full_model, ext) if blacklist is not None: - with open(blacklist, 'r') as f: - blacklist = f.read().rstrip().split('\n') + with open(blacklist, "r") as f: + blacklist = f.read().rstrip().split("\n") images_ref = dict() for id_, image in images.items(): @@ -29,5 +30,5 @@ def create_reference_sfm(full_model, ref_model, blacklist=None, ext='.bin'): continue points3D_ref[id_] = point3D._replace(image_ids=np.array(ref_ids)) - write_model(cameras, images_ref, points3D_ref, ref_model, '.bin') - logger.info(f'Kept {len(images_ref)} images out of {len(images)}.') + write_model(cameras, images_ref, points3D_ref, ref_model, ".bin") + logger.info(f"Kept {len(images_ref)} images out of {len(images)}.") diff --git a/hloc/pipelines/Aachen/pipeline.py b/hloc/pipelines/Aachen/pipeline.py index 7f582cf6..340c1d17 100644 --- a/hloc/pipelines/Aachen/pipeline.py +++ b/hloc/pipelines/Aachen/pipeline.py @@ -1,75 +1,101 @@ +import argparse from pathlib import Path from pprint import pformat -import argparse - -from ... import extract_features, match_features -from ... import pairs_from_covisibility, pairs_from_retrieval -from ... import colmap_from_nvm, triangulation, localize_sfm +from ... import ( + colmap_from_nvm, + extract_features, + localize_sfm, + match_features, + pairs_from_covisibility, + pairs_from_retrieval, + triangulation, +) parser = argparse.ArgumentParser() -parser.add_argument('--dataset', type=Path, default='datasets/aachen', - help='Path to the dataset, default: %(default)s') -parser.add_argument('--outputs', type=Path, default='outputs/aachen', - help='Path to the output directory, default: %(default)s') -parser.add_argument('--num_covis', type=int, default=20, - help='Number of image pairs for SfM, default: %(default)s') -parser.add_argument('--num_loc', type=int, default=50, - help='Number of image pairs for loc, default: %(default)s') +parser.add_argument( + "--dataset", + type=Path, + default="datasets/aachen", + help="Path to the dataset, default: %(default)s", +) +parser.add_argument( + "--outputs", + type=Path, + default="outputs/aachen", + help="Path to the output directory, default: %(default)s", +) +parser.add_argument( + "--num_covis", + type=int, + default=20, + help="Number of image pairs for SfM, default: %(default)s", +) +parser.add_argument( + "--num_loc", + type=int, + default=50, + help="Number of image pairs for loc, default: %(default)s", +) args = parser.parse_args() # Setup the paths dataset = args.dataset -images = dataset / 'images_upright/' +images = dataset / "images_upright/" outputs = args.outputs # where everything will be saved -sift_sfm = outputs / 'sfm_sift' # from which we extract the reference poses -reference_sfm = outputs / 'sfm_superpoint+superglue' # the SfM model we will build -sfm_pairs = outputs / f'pairs-db-covis{args.num_covis}.txt' # top-k most covisible in SIFT model -loc_pairs = outputs / f'pairs-query-netvlad{args.num_loc}.txt' # top-k retrieved by NetVLAD -results = outputs / f'Aachen_hloc_superpoint+superglue_netvlad{args.num_loc}.txt' +sift_sfm = outputs / "sfm_sift" # from which we extract the reference poses +reference_sfm = outputs / "sfm_superpoint+superglue" # the SfM model we will build +sfm_pairs = ( + outputs / f"pairs-db-covis{args.num_covis}.txt" +) # top-k most covisible in SIFT model +loc_pairs = ( + outputs / f"pairs-query-netvlad{args.num_loc}.txt" +) # top-k retrieved by NetVLAD +results = outputs / f"Aachen_hloc_superpoint+superglue_netvlad{args.num_loc}.txt" # list the standard configurations available -print(f'Configs for feature extractors:\n{pformat(extract_features.confs)}') -print(f'Configs for feature matchers:\n{pformat(match_features.confs)}') +print(f"Configs for feature extractors:\n{pformat(extract_features.confs)}") +print(f"Configs for feature matchers:\n{pformat(match_features.confs)}") # pick one of the configurations for extraction and matching -retrieval_conf = extract_features.confs['netvlad'] -feature_conf = extract_features.confs['superpoint_aachen'] -matcher_conf = match_features.confs['superglue'] +retrieval_conf = extract_features.confs["netvlad"] +feature_conf = extract_features.confs["superpoint_aachen"] +matcher_conf = match_features.confs["superglue"] features = extract_features.main(feature_conf, images, outputs) colmap_from_nvm.main( - dataset / '3D-models/aachen_cvpr2018_db.nvm', - dataset / '3D-models/database_intrinsics.txt', - dataset / 'aachen.db', - sift_sfm) -pairs_from_covisibility.main( - sift_sfm, sfm_pairs, num_matched=args.num_covis) + dataset / "3D-models/aachen_cvpr2018_db.nvm", + dataset / "3D-models/database_intrinsics.txt", + dataset / "aachen.db", + sift_sfm, +) +pairs_from_covisibility.main(sift_sfm, sfm_pairs, num_matched=args.num_covis) sfm_matches = match_features.main( - matcher_conf, sfm_pairs, feature_conf['output'], outputs) + matcher_conf, sfm_pairs, feature_conf["output"], outputs +) -triangulation.main( - reference_sfm, - sift_sfm, - images, - sfm_pairs, - features, - sfm_matches) +triangulation.main(reference_sfm, sift_sfm, images, sfm_pairs, features, sfm_matches) global_descriptors = extract_features.main(retrieval_conf, images, outputs) pairs_from_retrieval.main( - global_descriptors, loc_pairs, args.num_loc, - query_prefix='query', db_model=reference_sfm) + global_descriptors, + loc_pairs, + args.num_loc, + query_prefix="query", + db_model=reference_sfm, +) loc_matches = match_features.main( - matcher_conf, loc_pairs, feature_conf['output'], outputs) + matcher_conf, loc_pairs, feature_conf["output"], outputs +) localize_sfm.main( reference_sfm, - dataset / 'queries/*_time_queries_with_intrinsics.txt', + dataset / "queries/*_time_queries_with_intrinsics.txt", loc_pairs, features, loc_matches, results, - covisibility_clustering=False) # not required with SuperPoint+SuperGlue + covisibility_clustering=False, +) # not required with SuperPoint+SuperGlue diff --git a/hloc/pipelines/Aachen_v1_1/pipeline.py b/hloc/pipelines/Aachen_v1_1/pipeline.py index 12e29500..bbce73b7 100644 --- a/hloc/pipelines/Aachen_v1_1/pipeline.py +++ b/hloc/pipelines/Aachen_v1_1/pipeline.py @@ -1,69 +1,94 @@ +import argparse from pathlib import Path from pprint import pformat -import argparse - -from ... import extract_features, match_features, triangulation -from ... import pairs_from_covisibility, pairs_from_retrieval, localize_sfm +from ... import ( + extract_features, + localize_sfm, + match_features, + pairs_from_covisibility, + pairs_from_retrieval, + triangulation, +) parser = argparse.ArgumentParser() -parser.add_argument('--dataset', type=Path, default='datasets/aachen_v1.1', - help='Path to the dataset, default: %(default)s') -parser.add_argument('--outputs', type=Path, default='outputs/aachen_v1.1', - help='Path to the output directory, default: %(default)s') -parser.add_argument('--num_covis', type=int, default=20, - help='Number of image pairs for SfM, default: %(default)s') -parser.add_argument('--num_loc', type=int, default=50, - help='Number of image pairs for loc, default: %(default)s') +parser.add_argument( + "--dataset", + type=Path, + default="datasets/aachen_v1.1", + help="Path to the dataset, default: %(default)s", +) +parser.add_argument( + "--outputs", + type=Path, + default="outputs/aachen_v1.1", + help="Path to the output directory, default: %(default)s", +) +parser.add_argument( + "--num_covis", + type=int, + default=20, + help="Number of image pairs for SfM, default: %(default)s", +) +parser.add_argument( + "--num_loc", + type=int, + default=50, + help="Number of image pairs for loc, default: %(default)s", +) args = parser.parse_args() # Setup the paths dataset = args.dataset -images = dataset / 'images_upright/' -sift_sfm = dataset / '3D-models/aachen_v_1_1' +images = dataset / "images_upright/" +sift_sfm = dataset / "3D-models/aachen_v_1_1" outputs = args.outputs # where everything will be saved -reference_sfm = outputs / 'sfm_superpoint+superglue' # the SfM model we will build -sfm_pairs = outputs / f'pairs-db-covis{args.num_covis}.txt' # top-k most covisible in SIFT model -loc_pairs = outputs / f'pairs-query-netvlad{args.num_loc}.txt' # top-k retrieved by NetVLAD -results = outputs / f'Aachen-v1.1_hloc_superpoint+superglue_netvlad{args.num_loc}.txt' +reference_sfm = outputs / "sfm_superpoint+superglue" # the SfM model we will build +sfm_pairs = ( + outputs / f"pairs-db-covis{args.num_covis}.txt" +) # top-k most covisible in SIFT model +loc_pairs = ( + outputs / f"pairs-query-netvlad{args.num_loc}.txt" +) # top-k retrieved by NetVLAD +results = outputs / f"Aachen-v1.1_hloc_superpoint+superglue_netvlad{args.num_loc}.txt" # list the standard configurations available -print(f'Configs for feature extractors:\n{pformat(extract_features.confs)}') -print(f'Configs for feature matchers:\n{pformat(match_features.confs)}') +print(f"Configs for feature extractors:\n{pformat(extract_features.confs)}") +print(f"Configs for feature matchers:\n{pformat(match_features.confs)}") # pick one of the configurations for extraction and matching -retrieval_conf = extract_features.confs['netvlad'] -feature_conf = extract_features.confs['superpoint_max'] -matcher_conf = match_features.confs['superglue'] +retrieval_conf = extract_features.confs["netvlad"] +feature_conf = extract_features.confs["superpoint_max"] +matcher_conf = match_features.confs["superglue"] features = extract_features.main(feature_conf, images, outputs) -pairs_from_covisibility.main( - sift_sfm, sfm_pairs, num_matched=args.num_covis) +pairs_from_covisibility.main(sift_sfm, sfm_pairs, num_matched=args.num_covis) sfm_matches = match_features.main( - matcher_conf, sfm_pairs, feature_conf['output'], outputs) + matcher_conf, sfm_pairs, feature_conf["output"], outputs +) -triangulation.main( - reference_sfm, - sift_sfm, - images, - sfm_pairs, - features, - sfm_matches) +triangulation.main(reference_sfm, sift_sfm, images, sfm_pairs, features, sfm_matches) global_descriptors = extract_features.main(retrieval_conf, images, outputs) pairs_from_retrieval.main( - global_descriptors, loc_pairs, args.num_loc, - query_prefix='query', db_model=reference_sfm) + global_descriptors, + loc_pairs, + args.num_loc, + query_prefix="query", + db_model=reference_sfm, +) loc_matches = match_features.main( - matcher_conf, loc_pairs, feature_conf['output'], outputs) + matcher_conf, loc_pairs, feature_conf["output"], outputs +) localize_sfm.main( reference_sfm, - dataset / 'queries/*_time_queries_with_intrinsics.txt', + dataset / "queries/*_time_queries_with_intrinsics.txt", loc_pairs, features, loc_matches, results, - covisibility_clustering=False) # not required with SuperPoint+SuperGlue + covisibility_clustering=False, +) # not required with SuperPoint+SuperGlue diff --git a/hloc/pipelines/Aachen_v1_1/pipeline_loftr.py b/hloc/pipelines/Aachen_v1_1/pipeline_loftr.py index 299b5660..08b7a9b9 100644 --- a/hloc/pipelines/Aachen_v1_1/pipeline_loftr.py +++ b/hloc/pipelines/Aachen_v1_1/pipeline_loftr.py @@ -1,68 +1,97 @@ +import argparse from pathlib import Path from pprint import pformat -import argparse - -from ... import extract_features, match_dense, triangulation -from ... import pairs_from_covisibility, pairs_from_retrieval, localize_sfm +from ... import ( + extract_features, + localize_sfm, + match_dense, + pairs_from_covisibility, + pairs_from_retrieval, + triangulation, +) parser = argparse.ArgumentParser() -parser.add_argument('--dataset', type=Path, default='datasets/aachen_v1.1', - help='Path to the dataset, default: %(default)s') -parser.add_argument('--outputs', type=Path, default='outputs/aachen_v1.1', - help='Path to the output directory, default: %(default)s') -parser.add_argument('--num_covis', type=int, default=20, - help='Number of image pairs for SfM, default: %(default)s') -parser.add_argument('--num_loc', type=int, default=50, - help='Number of image pairs for loc, default: %(default)s') +parser.add_argument( + "--dataset", + type=Path, + default="datasets/aachen_v1.1", + help="Path to the dataset, default: %(default)s", +) +parser.add_argument( + "--outputs", + type=Path, + default="outputs/aachen_v1.1", + help="Path to the output directory, default: %(default)s", +) +parser.add_argument( + "--num_covis", + type=int, + default=20, + help="Number of image pairs for SfM, default: %(default)s", +) +parser.add_argument( + "--num_loc", + type=int, + default=50, + help="Number of image pairs for loc, default: %(default)s", +) args = parser.parse_args() # Setup the paths dataset = args.dataset -images = dataset / 'images_upright/' -sift_sfm = dataset / '3D-models/aachen_v_1_1' +images = dataset / "images_upright/" +sift_sfm = dataset / "3D-models/aachen_v_1_1" outputs = args.outputs # where everything will be saved outputs.mkdir() -reference_sfm = outputs / 'sfm_loftr' # the SfM model we will build -sfm_pairs = outputs / f'pairs-db-covis{args.num_covis}.txt' # top-k most covisible in SIFT model -loc_pairs = outputs / f'pairs-query-netvlad{args.num_loc}.txt' # top-k retrieved by NetVLAD -results = outputs / f'Aachen-v1.1_hloc_loftr_netvlad{args.num_loc}.txt' +reference_sfm = outputs / "sfm_loftr" # the SfM model we will build +sfm_pairs = ( + outputs / f"pairs-db-covis{args.num_covis}.txt" +) # top-k most covisible in SIFT model +loc_pairs = ( + outputs / f"pairs-query-netvlad{args.num_loc}.txt" +) # top-k retrieved by NetVLAD +results = outputs / f"Aachen-v1.1_hloc_loftr_netvlad{args.num_loc}.txt" # list the standard configurations available -print(f'Configs for dense feature matchers:\n{pformat(match_dense.confs)}') +print(f"Configs for dense feature matchers:\n{pformat(match_dense.confs)}") # pick one of the configurations for extraction and matching -retrieval_conf = extract_features.confs['netvlad'] -matcher_conf = match_dense.confs['loftr_aachen'] +retrieval_conf = extract_features.confs["netvlad"] +matcher_conf = match_dense.confs["loftr_aachen"] -pairs_from_covisibility.main( - sift_sfm, sfm_pairs, num_matched=args.num_covis) -features, sfm_matches = match_dense.main(matcher_conf, sfm_pairs, images, - outputs, max_kps=8192, - overwrite=False) +pairs_from_covisibility.main(sift_sfm, sfm_pairs, num_matched=args.num_covis) +features, sfm_matches = match_dense.main( + matcher_conf, sfm_pairs, images, outputs, max_kps=8192, overwrite=False +) -triangulation.main( - reference_sfm, - sift_sfm, - images, - sfm_pairs, - features, - sfm_matches) +triangulation.main(reference_sfm, sift_sfm, images, sfm_pairs, features, sfm_matches) global_descriptors = extract_features.main(retrieval_conf, images, outputs) pairs_from_retrieval.main( - global_descriptors, loc_pairs, args.num_loc, - query_prefix='query', db_model=reference_sfm) + global_descriptors, + loc_pairs, + args.num_loc, + query_prefix="query", + db_model=reference_sfm, +) features, loc_matches = match_dense.main( - matcher_conf, loc_pairs, images, outputs, features=features, max_kps=None, - matches=sfm_matches) + matcher_conf, + loc_pairs, + images, + outputs, + features=features, + max_kps=None, + matches=sfm_matches, +) localize_sfm.main( reference_sfm, - dataset / 'queries/*_time_queries_with_intrinsics.txt', + dataset / "queries/*_time_queries_with_intrinsics.txt", loc_pairs, features, loc_matches, results, - covisibility_clustering=False) # not required with loftr + covisibility_clustering=False, +) # not required with loftr diff --git a/hloc/pipelines/CMU/pipeline.py b/hloc/pipelines/CMU/pipeline.py index 8e9e52a0..4706a05c 100644 --- a/hloc/pipelines/CMU/pipeline.py +++ b/hloc/pipelines/CMU/pipeline.py @@ -1,115 +1,133 @@ -from pathlib import Path import argparse +from pathlib import Path -from ... import extract_features, match_features, triangulation, logger -from ... import pairs_from_covisibility, pairs_from_retrieval, localize_sfm +from ... import ( + extract_features, + localize_sfm, + logger, + match_features, + pairs_from_covisibility, + pairs_from_retrieval, + triangulation, +) TEST_SLICES = [2, 3, 4, 5, 6, 13, 14, 15, 16, 17, 18, 19, 20, 21] def generate_query_list(dataset, path, slice_): cameras = {} - with open(dataset / 'intrinsics.txt', 'r') as f: + with open(dataset / "intrinsics.txt", "r") as f: for line in f.readlines(): - if line[0] == '#' or line == '\n': + if line[0] == "#" or line == "\n": continue data = line.split() cameras[data[0]] = data[1:] assert len(cameras) == 2 - queries = dataset / f'{slice_}/test-images-{slice_}.txt' - with open(queries, 'r') as f: - queries = [q.rstrip('\n') for q in f.readlines()] + queries = dataset / f"{slice_}/test-images-{slice_}.txt" + with open(queries, "r") as f: + queries = [q.rstrip("\n") for q in f.readlines()] - out = [[q] + cameras[q.split('_')[2]] for q in queries] - with open(path, 'w') as f: - f.write('\n'.join(map(' '.join, out))) + out = [[q] + cameras[q.split("_")[2]] for q in queries] + with open(path, "w") as f: + f.write("\n".join(map(" ".join, out))) def run_slice(slice_, root, outputs, num_covis, num_loc): dataset = root / slice_ - ref_images = dataset / 'database' - query_images = dataset / 'query' - sift_sfm = dataset / 'sparse' + ref_images = dataset / "database" + query_images = dataset / "query" + sift_sfm = dataset / "sparse" outputs = outputs / slice_ outputs.mkdir(exist_ok=True, parents=True) - query_list = dataset / 'queries_with_intrinsics.txt' - sfm_pairs = outputs / f'pairs-db-covis{num_covis}.txt' - loc_pairs = outputs / f'pairs-query-netvlad{num_loc}.txt' - ref_sfm = outputs / 'sfm_superpoint+superglue' - results = outputs / f'CMU_hloc_superpoint+superglue_netvlad{num_loc}.txt' + query_list = dataset / "queries_with_intrinsics.txt" + sfm_pairs = outputs / f"pairs-db-covis{num_covis}.txt" + loc_pairs = outputs / f"pairs-query-netvlad{num_loc}.txt" + ref_sfm = outputs / "sfm_superpoint+superglue" + results = outputs / f"CMU_hloc_superpoint+superglue_netvlad{num_loc}.txt" # pick one of the configurations for extraction and matching - retrieval_conf = extract_features.confs['netvlad'] - feature_conf = extract_features.confs['superpoint_aachen'] - matcher_conf = match_features.confs['superglue'] + retrieval_conf = extract_features.confs["netvlad"] + feature_conf = extract_features.confs["superpoint_aachen"] + matcher_conf = match_features.confs["superglue"] pairs_from_covisibility.main(sift_sfm, sfm_pairs, num_matched=num_covis) - features = extract_features.main( - feature_conf, ref_images, outputs, as_half=True) + features = extract_features.main(feature_conf, ref_images, outputs, as_half=True) sfm_matches = match_features.main( - matcher_conf, sfm_pairs, feature_conf['output'], outputs) - triangulation.main( - ref_sfm, - sift_sfm, - ref_images, - sfm_pairs, - features, - sfm_matches) + matcher_conf, sfm_pairs, feature_conf["output"], outputs + ) + triangulation.main(ref_sfm, sift_sfm, ref_images, sfm_pairs, features, sfm_matches) generate_query_list(root, query_list, slice_) - global_descriptors = extract_features.main( - retrieval_conf, ref_images, outputs) - global_descriptors = extract_features.main( - retrieval_conf, query_images, outputs) + global_descriptors = extract_features.main(retrieval_conf, ref_images, outputs) + global_descriptors = extract_features.main(retrieval_conf, query_images, outputs) pairs_from_retrieval.main( - global_descriptors, loc_pairs, num_loc, - query_list=query_list, db_model=ref_sfm) + global_descriptors, loc_pairs, num_loc, query_list=query_list, db_model=ref_sfm + ) - features = extract_features.main( - feature_conf, query_images, outputs, as_half=True) + features = extract_features.main(feature_conf, query_images, outputs, as_half=True) loc_matches = match_features.main( - matcher_conf, loc_pairs, feature_conf['output'], outputs) + matcher_conf, loc_pairs, feature_conf["output"], outputs + ) localize_sfm.main( ref_sfm, - dataset / 'queries/*_time_queries_with_intrinsics.txt', + dataset / "queries/*_time_queries_with_intrinsics.txt", loc_pairs, features, loc_matches, - results) + results, + ) -if __name__ == '__main__': +if __name__ == "__main__": parser = argparse.ArgumentParser() - parser.add_argument('--slices', type=str, default='*', - help='a single number, an interval (e.g. 2-6), ' - 'or a Python-style list or int (e.g. [2, 3, 4]') - parser.add_argument('--dataset', type=Path, - default='datasets/cmu_extended', - help='Path to the dataset, default: %(default)s') - parser.add_argument('--outputs', type=Path, - default='outputs/aachen_extended', - help='Path to the output directory, default: %(default)s') - parser.add_argument('--num_covis', type=int, default=20, - help='Number of image pairs for SfM, default: %(default)s') - parser.add_argument('--num_loc', type=int, default=10, - help='Number of image pairs for loc, default: %(default)s') + parser.add_argument( + "--slices", + type=str, + default="*", + help="a single number, an interval (e.g. 2-6), " + "or a Python-style list or int (e.g. [2, 3, 4]", + ) + parser.add_argument( + "--dataset", + type=Path, + default="datasets/cmu_extended", + help="Path to the dataset, default: %(default)s", + ) + parser.add_argument( + "--outputs", + type=Path, + default="outputs/aachen_extended", + help="Path to the output directory, default: %(default)s", + ) + parser.add_argument( + "--num_covis", + type=int, + default=20, + help="Number of image pairs for SfM, default: %(default)s", + ) + parser.add_argument( + "--num_loc", + type=int, + default=10, + help="Number of image pairs for loc, default: %(default)s", + ) args = parser.parse_args() - if args.slice == '*': + if args.slice == "*": slices = TEST_SLICES - if '-' in args.slices: - min_, max_ = args.slices.split('-') - slices = list(range(int(min_), int(max_)+1)) + if "-" in args.slices: + min_, max_ = args.slices.split("-") + slices = list(range(int(min_), int(max_) + 1)) else: slices = eval(args.slices) if isinstance(slices, int): slices = [slices] for slice_ in slices: - logger.info('Working on slice %s.', slice_) + logger.info("Working on slice %s.", slice_) run_slice( - f'slice{slice_}', args.dataset, args.outputs, - args.num_covis, args.num_loc) + f"slice{slice_}", args.dataset, args.outputs, args.num_covis, args.num_loc + ) diff --git a/hloc/pipelines/Cambridge/pipeline.py b/hloc/pipelines/Cambridge/pipeline.py index 7971ccd5..3a676e5a 100644 --- a/hloc/pipelines/Cambridge/pipeline.py +++ b/hloc/pipelines/Cambridge/pipeline.py @@ -1,69 +1,75 @@ -from pathlib import Path import argparse +from pathlib import Path -from .utils import ( - create_query_list_with_intrinsics, scale_sfm_images, evaluate) -from ... import extract_features, match_features, pairs_from_covisibility -from ... import triangulation, localize_sfm, pairs_from_retrieval, logger +from ... import ( + extract_features, + localize_sfm, + logger, + match_features, + pairs_from_covisibility, + pairs_from_retrieval, + triangulation, +) +from .utils import create_query_list_with_intrinsics, evaluate, scale_sfm_images -SCENES = ['KingsCollege', 'OldHospital', 'ShopFacade', 'StMarysChurch', - 'GreatCourt'] +SCENES = ["KingsCollege", "OldHospital", "ShopFacade", "StMarysChurch", "GreatCourt"] def run_scene(images, gt_dir, outputs, results, num_covis, num_loc): - ref_sfm_sift = gt_dir / 'model_train' - test_list = gt_dir / 'list_query.txt' + ref_sfm_sift = gt_dir / "model_train" + test_list = gt_dir / "list_query.txt" outputs.mkdir(exist_ok=True, parents=True) - ref_sfm = outputs / 'sfm_superpoint+superglue' - ref_sfm_scaled = outputs / 'sfm_sift_scaled' - query_list = outputs / 'query_list_with_intrinsics.txt' - sfm_pairs = outputs / f'pairs-db-covis{num_covis}.txt' - loc_pairs = outputs / f'pairs-query-netvlad{num_loc}.txt' + ref_sfm = outputs / "sfm_superpoint+superglue" + ref_sfm_scaled = outputs / "sfm_sift_scaled" + query_list = outputs / "query_list_with_intrinsics.txt" + sfm_pairs = outputs / f"pairs-db-covis{num_covis}.txt" + loc_pairs = outputs / f"pairs-query-netvlad{num_loc}.txt" feature_conf = { - 'output': 'feats-superpoint-n4096-r1024', - 'model': { - 'name': 'superpoint', - 'nms_radius': 3, - 'max_keypoints': 4096, + "output": "feats-superpoint-n4096-r1024", + "model": { + "name": "superpoint", + "nms_radius": 3, + "max_keypoints": 4096, }, - 'preprocessing': { - 'grayscale': True, - 'resize_max': 1024, + "preprocessing": { + "grayscale": True, + "resize_max": 1024, }, } - matcher_conf = match_features.confs['superglue'] - retrieval_conf = extract_features.confs['netvlad'] + matcher_conf = match_features.confs["superglue"] + retrieval_conf = extract_features.confs["netvlad"] create_query_list_with_intrinsics( - gt_dir / 'empty_all', query_list, test_list, - ext='.txt', image_dir=images) - with open(test_list, 'r') as f: - query_seqs = {q.split('/')[0] for q in f.read().rstrip().split('\n')} + gt_dir / "empty_all", query_list, test_list, ext=".txt", image_dir=images + ) + with open(test_list, "r") as f: + query_seqs = {q.split("/")[0] for q in f.read().rstrip().split("\n")} global_descriptors = extract_features.main(retrieval_conf, images, outputs) pairs_from_retrieval.main( - global_descriptors, loc_pairs, num_loc, - db_model=ref_sfm_sift, query_prefix=query_seqs) + global_descriptors, + loc_pairs, + num_loc, + db_model=ref_sfm_sift, + query_prefix=query_seqs, + ) - features = extract_features.main( - feature_conf, images, outputs, as_half=True) - pairs_from_covisibility.main( - ref_sfm_sift, sfm_pairs, num_matched=num_covis) + features = extract_features.main(feature_conf, images, outputs, as_half=True) + pairs_from_covisibility.main(ref_sfm_sift, sfm_pairs, num_matched=num_covis) sfm_matches = match_features.main( - matcher_conf, sfm_pairs, feature_conf['output'], outputs) + matcher_conf, sfm_pairs, feature_conf["output"], outputs + ) scale_sfm_images(ref_sfm_sift, ref_sfm_scaled, images) triangulation.main( - ref_sfm, ref_sfm_scaled, - images, - sfm_pairs, - features, - sfm_matches) + ref_sfm, ref_sfm_scaled, images, sfm_pairs, features, sfm_matches + ) loc_matches = match_features.main( - matcher_conf, loc_pairs, feature_conf['output'], outputs) + matcher_conf, loc_pairs, feature_conf["output"], outputs + ) localize_sfm.main( ref_sfm, @@ -73,29 +79,46 @@ def run_scene(images, gt_dir, outputs, results, num_covis, num_loc): loc_matches, results, covisibility_clustering=False, - prepend_camera_name=True) + prepend_camera_name=True, + ) -if __name__ == '__main__': +if __name__ == "__main__": parser = argparse.ArgumentParser() - parser.add_argument('--scenes', default=SCENES, choices=SCENES, nargs='+') - parser.add_argument('--overwrite', action='store_true') - parser.add_argument('--dataset', type=Path, default='datasets/cambridge', - help='Path to the dataset, default: %(default)s') - parser.add_argument('--outputs', type=Path, default='outputs/cambridge', - help='Path to the output directory, default: %(default)s') - parser.add_argument('--num_covis', type=int, default=20, - help='Number of image pairs for SfM, default: %(default)s') - parser.add_argument('--num_loc', type=int, default=10, - help='Number of image pairs for loc, default: %(default)s') + parser.add_argument("--scenes", default=SCENES, choices=SCENES, nargs="+") + parser.add_argument("--overwrite", action="store_true") + parser.add_argument( + "--dataset", + type=Path, + default="datasets/cambridge", + help="Path to the dataset, default: %(default)s", + ) + parser.add_argument( + "--outputs", + type=Path, + default="outputs/cambridge", + help="Path to the output directory, default: %(default)s", + ) + parser.add_argument( + "--num_covis", + type=int, + default=20, + help="Number of image pairs for SfM, default: %(default)s", + ) + parser.add_argument( + "--num_loc", + type=int, + default=10, + help="Number of image pairs for loc, default: %(default)s", + ) args = parser.parse_args() - gt_dirs = args.dataset / 'CambridgeLandmarks_Colmap_Retriangulated_1024px' + gt_dirs = args.dataset / "CambridgeLandmarks_Colmap_Retriangulated_1024px" all_results = {} for scene in args.scenes: logger.info(f'Working on scene "{scene}".') - results = args.outputs / scene / 'results.txt' + results = args.outputs / scene / "results.txt" if args.overwrite or not results.exists(): run_scene( args.dataset / scene, @@ -103,11 +126,15 @@ def run_scene(images, gt_dir, outputs, results, num_covis, num_loc): args.outputs / scene, results, args.num_covis, - args.num_loc) + args.num_loc, + ) all_results[scene] = results for scene in args.scenes: logger.info(f'Evaluate scene "{scene}".') evaluate( - gt_dirs / scene / 'empty_all', all_results[scene], - gt_dirs / scene / 'list_query.txt', ext='.txt') + gt_dirs / scene / "empty_all", + all_results[scene], + gt_dirs / scene / "list_query.txt", + ext=".txt", + ) diff --git a/hloc/pipelines/Cambridge/utils.py b/hloc/pipelines/Cambridge/utils.py index 409daaec..36460f06 100644 --- a/hloc/pipelines/Cambridge/utils.py +++ b/hloc/pipelines/Cambridge/utils.py @@ -1,19 +1,26 @@ -import cv2 import logging + +import cv2 import numpy as np from hloc.utils.read_write_model import ( - read_cameras_binary, read_images_binary, read_model, write_model, - qvec2rotmat, read_images_text, read_cameras_text) + qvec2rotmat, + read_cameras_binary, + read_cameras_text, + read_images_binary, + read_images_text, + read_model, + write_model, +) logger = logging.getLogger(__name__) def scale_sfm_images(full_model, scaled_model, image_dir): - '''Duplicate the provided model and scale the camera intrinsics so that - they match the original image resolution - makes everything easier. - ''' - logger.info('Scaling the COLMAP model to the original image size.') + """Duplicate the provided model and scale the camera intrinsics so that + they match the original image resolution - makes everything easier. + """ + logger.info("Scaling the COLMAP model to the original image size.") scaled_model.mkdir(exist_ok=True) cameras, images, points3D = read_model(full_model) @@ -31,32 +38,34 @@ def scale_sfm_images(full_model, scaled_model, image_dir): continue camera = cameras[cam_id] - assert camera.model == 'SIMPLE_RADIAL' + assert camera.model == "SIMPLE_RADIAL" sx = w / camera.width sy = h / camera.height assert sx == sy, (sx, sy) scaled_cameras[cam_id] = camera._replace( - width=w, height=h, params=camera.params*np.array([sx, sx, sy, 1.])) + width=w, height=h, params=camera.params * np.array([sx, sx, sy, 1.0]) + ) write_model(scaled_cameras, images, points3D, scaled_model) -def create_query_list_with_intrinsics(model, out, list_file=None, ext='.bin', - image_dir=None): - '''Create a list of query images with intrinsics from the colmap model.''' - if ext == '.bin': - images = read_images_binary(model / 'images.bin') - cameras = read_cameras_binary(model / 'cameras.bin') +def create_query_list_with_intrinsics( + model, out, list_file=None, ext=".bin", image_dir=None +): + """Create a list of query images with intrinsics from the colmap model.""" + if ext == ".bin": + images = read_images_binary(model / "images.bin") + cameras = read_cameras_binary(model / "cameras.bin") else: - images = read_images_text(model / 'images.txt') - cameras = read_cameras_text(model / 'cameras.txt') + images = read_images_text(model / "images.txt") + cameras = read_cameras_text(model / "cameras.txt") name2id = {image.name: i for i, image in images.items()} if list_file is None: names = list(name2id) else: - with open(list_file, 'r') as f: - names = f.read().rstrip().split('\n') + with open(list_file, "r") as f: + names = f.read().rstrip().split("\n") data = [] for name in names: image = images[name2id[name]] @@ -68,38 +77,38 @@ def create_query_list_with_intrinsics(model, out, list_file=None, ext='.bin', img = cv2.imread(str(image_dir / name)) assert img is not None, image_dir / name h_orig, w_orig = img.shape[:2] - assert camera.model == 'SIMPLE_RADIAL' + assert camera.model == "SIMPLE_RADIAL" sx = w_orig / w sy = h_orig / h assert sx == sy, (sx, sy) w, h = w_orig, h_orig - params = params * np.array([sx, sx, sy, 1.]) + params = params * np.array([sx, sx, sy, 1.0]) p = [name, camera.model, w, h] + params.tolist() - data.append(' '.join(map(str, p))) - with open(out, 'w') as f: - f.write('\n'.join(data)) + data.append(" ".join(map(str, p))) + with open(out, "w") as f: + f.write("\n".join(data)) -def evaluate(model, results, list_file=None, ext='.bin', only_localized=False): +def evaluate(model, results, list_file=None, ext=".bin", only_localized=False): predictions = {} - with open(results, 'r') as f: - for data in f.read().rstrip().split('\n'): + with open(results, "r") as f: + for data in f.read().rstrip().split("\n"): data = data.split() name = data[0] q, t = np.split(np.array(data[1:], float), [4]) predictions[name] = (qvec2rotmat(q), t) - if ext == '.bin': - images = read_images_binary(model / 'images.bin') + if ext == ".bin": + images = read_images_binary(model / "images.bin") else: - images = read_images_text(model / 'images.txt') + images = read_images_text(model / "images.txt") name2id = {image.name: i for i, image in images.items()} if list_file is None: test_names = list(name2id) else: - with open(list_file, 'r') as f: - test_names = f.read().rstrip().split('\n') + with open(list_file, "r") as f: + test_names = f.read().rstrip().split("\n") errors_t = [] errors_R = [] @@ -108,13 +117,13 @@ def evaluate(model, results, list_file=None, ext='.bin', only_localized=False): if only_localized: continue e_t = np.inf - e_R = 180. + e_R = 180.0 else: image = images[name2id[name]] R_gt, t_gt = image.qvec2rotmat(), image.tvec R, t = predictions[name] e_t = np.linalg.norm(-R_gt.T @ t_gt + R.T @ t, axis=0) - cos = np.clip((np.trace(np.dot(R_gt.T, R)) - 1) / 2, -1., 1.) + cos = np.clip((np.trace(np.dot(R_gt.T, R)) - 1) / 2, -1.0, 1.0) e_R = np.rad2deg(np.abs(np.arccos(cos))) errors_t.append(e_t) errors_R.append(e_R) @@ -124,13 +133,13 @@ def evaluate(model, results, list_file=None, ext='.bin', only_localized=False): med_t = np.median(errors_t) med_R = np.median(errors_R) - out = f'Results for file {results.name}:' - out += f'\nMedian errors: {med_t:.3f}m, {med_R:.3f}deg' + out = f"Results for file {results.name}:" + out += f"\nMedian errors: {med_t:.3f}m, {med_R:.3f}deg" - out += '\nPercentage of test images localized within:' + out += "\nPercentage of test images localized within:" threshs_t = [0.01, 0.02, 0.03, 0.05, 0.25, 0.5, 5.0] threshs_R = [1.0, 2.0, 3.0, 5.0, 2.0, 5.0, 10.0] for th_t, th_R in zip(threshs_t, threshs_R): ratio = np.mean((errors_t < th_t) & (errors_R < th_R)) - out += f'\n\t{th_t*100:.0f}cm, {th_R:.0f}deg : {ratio*100:.2f}%' + out += f"\n\t{th_t*100:.0f}cm, {th_R:.0f}deg : {ratio*100:.2f}%" logger.info(out) diff --git a/hloc/pipelines/RobotCar/colmap_from_nvm.py b/hloc/pipelines/RobotCar/colmap_from_nvm.py index adaeefa1..e90ed72b 100644 --- a/hloc/pipelines/RobotCar/colmap_from_nvm.py +++ b/hloc/pipelines/RobotCar/colmap_from_nvm.py @@ -1,83 +1,96 @@ import argparse +import logging import sqlite3 -from tqdm import tqdm from collections import defaultdict -import numpy as np from pathlib import Path -import logging + +import numpy as np +from tqdm import tqdm from ...colmap_from_nvm import ( - recover_database_images_and_ids, camera_center_to_translation) -from ...utils.read_write_model import Camera, Image, Point3D, CAMERA_MODEL_IDS -from ...utils.read_write_model import write_model + camera_center_to_translation, + recover_database_images_and_ids, +) +from ...utils.read_write_model import ( + CAMERA_MODEL_IDS, + Camera, + Image, + Point3D, + write_model, +) logger = logging.getLogger(__name__) -def read_nvm_model( - nvm_path, database_path, image_ids, camera_ids, skip_points=False): - +def read_nvm_model(nvm_path, database_path, image_ids, camera_ids, skip_points=False): # Extract the intrinsics from the db file instead of the NVM model db = sqlite3.connect(str(database_path)) - ret = db.execute( - 'SELECT camera_id, model, width, height, params FROM cameras;') + ret = db.execute("SELECT camera_id, model, width, height, params FROM cameras;") cameras = {} for camera_id, camera_model, width, height, params in ret: params = np.fromstring(params, dtype=np.double).reshape(-1) camera_model = CAMERA_MODEL_IDS[camera_model] - assert len(params) == camera_model.num_params, (len(params), camera_model.num_params) + assert len(params) == camera_model.num_params, ( + len(params), + camera_model.num_params, + ) camera = Camera( - id=camera_id, model=camera_model.model_name, - width=int(width), height=int(height), params=params) + id=camera_id, + model=camera_model.model_name, + width=int(width), + height=int(height), + params=params, + ) cameras[camera_id] = camera - nvm_f = open(nvm_path, 'r') + nvm_f = open(nvm_path, "r") line = nvm_f.readline() - while line == '\n' or line.startswith('NVM_V3'): + while line == "\n" or line.startswith("NVM_V3"): line = nvm_f.readline() num_images = int(line) # assert num_images == len(cameras), (num_images, len(cameras)) - logger.info(f'Reading {num_images} images...') + logger.info(f"Reading {num_images} images...") image_idx_to_db_image_id = [] image_data = [] i = 0 while i < num_images: line = nvm_f.readline() - if line == '\n': + if line == "\n": continue - data = line.strip('\n').lstrip('./').split(' ') + data = line.strip("\n").lstrip("./").split(" ") image_data.append(data) image_idx_to_db_image_id.append(image_ids[data[0]]) i += 1 line = nvm_f.readline() - while line == '\n': + while line == "\n": line = nvm_f.readline() num_points = int(line) if skip_points: - logger.info(f'Skipping {num_points} points.') + logger.info(f"Skipping {num_points} points.") num_points = 0 else: - logger.info(f'Reading {num_points} points...') + logger.info(f"Reading {num_points} points...") points3D = {} image_idx_to_keypoints = defaultdict(list) i = 0 - pbar = tqdm(total=num_points, unit='pts') + pbar = tqdm(total=num_points, unit="pts") while i < num_points: line = nvm_f.readline() - if line == '\n': + if line == "\n": continue - data = line.strip('\n').split(' ') + data = line.strip("\n").split(" ") x, y, z, r, g, b, num_observations = data[:7] obs_image_ids, point2D_idxs = [], [] for j in range(int(num_observations)): - s = 7 + 4*j - img_index, kp_index, kx, ky = data[s:s+4] + s = 7 + 4 * j + img_index, kp_index, kx, ky = data[s : s + 4] image_idx_to_keypoints[int(img_index)].append( - (int(kp_index), float(kx), float(ky), i)) + (int(kp_index), float(kx), float(ky), i) + ) db_image_id = image_idx_to_db_image_id[int(img_index)] obs_image_ids.append(db_image_id) point2D_idxs.append(kp_index) @@ -86,16 +99,17 @@ def read_nvm_model( id=i, xyz=np.array([x, y, z], float), rgb=np.array([r, g, b], int), - error=1., # fake + error=1.0, # fake image_ids=np.array(obs_image_ids, int), - point2D_idxs=np.array(point2D_idxs, int)) + point2D_idxs=np.array(point2D_idxs, int), + ) points3D[i] = point i += 1 pbar.update(1) pbar.close() - logger.info('Parsing image data...') + logger.info("Parsing image data...") images = {} for i, data in enumerate(image_data): # Skip the focal length. Skip the distortion and terminal 0. @@ -126,9 +140,10 @@ def read_nvm_model( qvec=qvec, tvec=t, camera_id=camera_ids[name], - name=name.replace('png', 'jpg'), # some hack required for RobotCar + name=name.replace("png", "jpg"), # some hack required for RobotCar xys=xys, - point3D_ids=point3D_ids) + point3D_ids=point3D_ids, + ) images[image_id] = image return cameras, images, points3D @@ -140,22 +155,22 @@ def main(nvm, database, output, skip_points=False): image_ids, camera_ids = recover_database_images_and_ids(database) - logger.info('Reading the NVM model...') + logger.info("Reading the NVM model...") model = read_nvm_model( - nvm, database, image_ids, camera_ids, skip_points=skip_points) + nvm, database, image_ids, camera_ids, skip_points=skip_points + ) - logger.info('Writing the COLMAP model...') + logger.info("Writing the COLMAP model...") output.mkdir(exist_ok=True, parents=True) - write_model(*model, path=str(output), ext='.bin') - logger.info('Done.') + write_model(*model, path=str(output), ext=".bin") + logger.info("Done.") -if __name__ == '__main__': +if __name__ == "__main__": parser = argparse.ArgumentParser() - parser.add_argument('--nvm', required=True, type=Path) - parser.add_argument('--database', required=True, type=Path) - parser.add_argument('--output', required=True, type=Path) - parser.add_argument('--skip_points', action='store_true') + parser.add_argument("--nvm", required=True, type=Path) + parser.add_argument("--database", required=True, type=Path) + parser.add_argument("--output", required=True, type=Path) + parser.add_argument("--skip_points", action="store_true") args = parser.parse_args() main(**args.__dict__) - diff --git a/hloc/pipelines/RobotCar/pipeline.py b/hloc/pipelines/RobotCar/pipeline.py index f3f18267..5851df79 100644 --- a/hloc/pipelines/RobotCar/pipeline.py +++ b/hloc/pipelines/RobotCar/pipeline.py @@ -1,106 +1,138 @@ -from pathlib import Path import argparse import glob +from pathlib import Path +from ... import ( + extract_features, + localize_sfm, + match_features, + pairs_from_covisibility, + pairs_from_retrieval, + triangulation, +) from . import colmap_from_nvm -from ... import extract_features, match_features, triangulation -from ... import pairs_from_covisibility, pairs_from_retrieval, localize_sfm - -CONDITIONS = ['dawn', 'dusk', 'night', 'night-rain', 'overcast-summer', - 'overcast-winter', 'rain', 'snow', 'sun'] +CONDITIONS = [ + "dawn", + "dusk", + "night", + "night-rain", + "overcast-summer", + "overcast-winter", + "rain", + "snow", + "sun", +] def generate_query_list(dataset, image_dir, path): h, w = 1024, 1024 - intrinsics_filename = 'intrinsics/{}_intrinsics.txt' + intrinsics_filename = "intrinsics/{}_intrinsics.txt" cameras = {} - for side in ['left', 'right', 'rear']: - with open(dataset / intrinsics_filename.format(side), 'r') as f: + for side in ["left", "right", "rear"]: + with open(dataset / intrinsics_filename.format(side), "r") as f: fx = f.readline().split()[1] fy = f.readline().split()[1] cx = f.readline().split()[1] cy = f.readline().split()[1] assert fx == fy - params = ['SIMPLE_RADIAL', w, h, fx, cx, cy, 0.0] + params = ["SIMPLE_RADIAL", w, h, fx, cx, cy, 0.0] cameras[side] = [str(p) for p in params] - queries = glob.glob((image_dir / '**/*.jpg').as_posix(), recursive=True) - queries = [Path(q).relative_to(image_dir.parents[0]).as_posix() - for q in sorted(queries)] + queries = glob.glob((image_dir / "**/*.jpg").as_posix(), recursive=True) + queries = [ + Path(q).relative_to(image_dir.parents[0]).as_posix() for q in sorted(queries) + ] out = [[q] + cameras[Path(q).parent.name] for q in queries] - with open(path, 'w') as f: - f.write('\n'.join(map(' '.join, out))) + with open(path, "w") as f: + f.write("\n".join(map(" ".join, out))) parser = argparse.ArgumentParser() -parser.add_argument('--dataset', type=Path, default='datasets/robotcar', - help='Path to the dataset, default: %(default)s') -parser.add_argument('--outputs', type=Path, default='outputs/robotcar', - help='Path to the output directory, default: %(default)s') -parser.add_argument('--num_covis', type=int, default=20, - help='Number of image pairs for SfM, default: %(default)s') -parser.add_argument('--num_loc', type=int, default=20, - help='Number of image pairs for loc, default: %(default)s') +parser.add_argument( + "--dataset", + type=Path, + default="datasets/robotcar", + help="Path to the dataset, default: %(default)s", +) +parser.add_argument( + "--outputs", + type=Path, + default="outputs/robotcar", + help="Path to the output directory, default: %(default)s", +) +parser.add_argument( + "--num_covis", + type=int, + default=20, + help="Number of image pairs for SfM, default: %(default)s", +) +parser.add_argument( + "--num_loc", + type=int, + default=20, + help="Number of image pairs for loc, default: %(default)s", +) args = parser.parse_args() # Setup the paths dataset = args.dataset -images = dataset / 'images/' +images = dataset / "images/" outputs = args.outputs # where everything will be saved outputs.mkdir(exist_ok=True, parents=True) -query_list = outputs / '{condition}_queries_with_intrinsics.txt' -sift_sfm = outputs / 'sfm_sift' -reference_sfm = outputs / 'sfm_superpoint+superglue' -sfm_pairs = outputs / f'pairs-db-covis{args.num_covis}.txt' -loc_pairs = outputs / f'pairs-query-netvlad{args.num_loc}.txt' -results = outputs / f'RobotCar_hloc_superpoint+superglue_netvlad{args.num_loc}.txt' +query_list = outputs / "{condition}_queries_with_intrinsics.txt" +sift_sfm = outputs / "sfm_sift" +reference_sfm = outputs / "sfm_superpoint+superglue" +sfm_pairs = outputs / f"pairs-db-covis{args.num_covis}.txt" +loc_pairs = outputs / f"pairs-query-netvlad{args.num_loc}.txt" +results = outputs / f"RobotCar_hloc_superpoint+superglue_netvlad{args.num_loc}.txt" # pick one of the configurations for extraction and matching -retrieval_conf = extract_features.confs['netvlad'] -feature_conf = extract_features.confs['superpoint_aachen'] -matcher_conf = match_features.confs['superglue'] +retrieval_conf = extract_features.confs["netvlad"] +feature_conf = extract_features.confs["superpoint_aachen"] +matcher_conf = match_features.confs["superglue"] for condition in CONDITIONS: generate_query_list( - dataset, images / condition, - str(query_list).format(condition=condition)) + dataset, images / condition, str(query_list).format(condition=condition) + ) features = extract_features.main(feature_conf, images, outputs, as_half=True) colmap_from_nvm.main( - dataset / '3D-models/all-merged/all.nvm', - dataset / '3D-models/overcast-reference.db', - sift_sfm) -pairs_from_covisibility.main( - sift_sfm, sfm_pairs, num_matched=args.num_covis) + dataset / "3D-models/all-merged/all.nvm", + dataset / "3D-models/overcast-reference.db", + sift_sfm, +) +pairs_from_covisibility.main(sift_sfm, sfm_pairs, num_matched=args.num_covis) sfm_matches = match_features.main( - matcher_conf, sfm_pairs, feature_conf['output'], outputs) + matcher_conf, sfm_pairs, feature_conf["output"], outputs +) -triangulation.main( - reference_sfm, - sift_sfm, - images, - sfm_pairs, - features, - sfm_matches) +triangulation.main(reference_sfm, sift_sfm, images, sfm_pairs, features, sfm_matches) global_descriptors = extract_features.main(retrieval_conf, images, outputs) # TODO: do per location and per camera pairs_from_retrieval.main( - global_descriptors, loc_pairs, args.num_loc, - query_prefix=CONDITIONS, db_model=reference_sfm) + global_descriptors, + loc_pairs, + args.num_loc, + query_prefix=CONDITIONS, + db_model=reference_sfm, +) loc_matches = match_features.main( - matcher_conf, loc_pairs, feature_conf['output'], outputs) + matcher_conf, loc_pairs, feature_conf["output"], outputs +) localize_sfm.main( reference_sfm, - Path(str(query_list).format(condition='*')), + Path(str(query_list).format(condition="*")), loc_pairs, features, loc_matches, results, covisibility_clustering=False, - prepend_camera_name=True) + prepend_camera_name=True, +) diff --git a/hloc/reconstruction.py b/hloc/reconstruction.py index 94392054..ea1e7fc0 100644 --- a/hloc/reconstruction.py +++ b/hloc/reconstruction.py @@ -1,43 +1,54 @@ import argparse -import shutil -from typing import Optional, List, Dict, Any import multiprocessing +import shutil from pathlib import Path +from typing import Any, Dict, List, Optional + 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) + OutputCapture, + estimation_and_geometric_verification, + import_features, + import_matches, + parse_option_args, +) +from .utils.database import COLMAPDatabase def create_empty_db(database_path: Path): if database_path.exists(): - logger.warning('The database already exists, deleting it.') + logger.warning("The database already exists, deleting it.") database_path.unlink() - logger.info('Creating an empty database...') + 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...') +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}.') + 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) + 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]: @@ -49,27 +60,29 @@ def get_image_ids(database_path: Path) -> Dict[str, int]: 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' +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...') + logger.info("Running 3D reconstruction...") if options is None: options = {} - options = {'num_threads': min(multiprocessing.cpu_count(), 16), **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) + database_path, image_dir, models_path, options=options + ) if len(reconstructions) == 0: - logger.error('Could not reconstruct any model!') + logger.error("Could not reconstruct any model!") return None - logger.info(f'Reconstructed {len(reconstructions)} model(s).') + logger.info(f"Reconstructed {len(reconstructions)} model(s).") largest_index = None largest_num_images = 0 @@ -79,80 +92,103 @@ def run_reconstruction(sfm_dir: Path, 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.') + logger.info( + f"Largest model is #{largest_index} " f"with {largest_num_images} images." + ) - for filename in ['images.bin', 'cameras.bin', 'points3D.bin']: + 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)) + 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: - +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' + 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) + 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) + 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)}') + logger.info( + f"Reconstruction statistics:\n{reconstruction.summary()}" + + f"\n\tnum_input_images = {len(image_ids)}" + ) return reconstruction -if __name__ == '__main__': +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())) + 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()) + args.pop("image_options"), pycolmap.ImageReaderOptions() + ) mapper_options = parse_option_args( - args.pop("mapper_options"), pycolmap.IncrementalMapperOptions()) + args.pop("mapper_options"), pycolmap.IncrementalMapperOptions() + ) main(**args, image_options=image_options, mapper_options=mapper_options) diff --git a/hloc/triangulation.py b/hloc/triangulation.py index 3dd0179f..103d5158 100644 --- a/hloc/triangulation.py +++ b/hloc/triangulation.py @@ -1,18 +1,19 @@ import argparse import contextlib -from typing import Optional, List, Dict, Any import io import sys from pathlib import Path +from typing import Any, Dict, List, Optional + import numpy as np -from tqdm import tqdm import pycolmap +from tqdm import tqdm from . import logger from .utils.database import COLMAPDatabase +from .utils.geometry import compute_epipolar_errors from .utils.io import get_keypoints, get_matches from .utils.parsers import parse_retrieval -from .utils.geometry import compute_epipolar_errors class OutputCapture: @@ -28,14 +29,15 @@ def __exit__(self, exc_type, *args): if not self.verbose: self.capture.__exit__(exc_type, *args) if exc_type is not None: - logger.error('Failed with output:\n%s', self.out.getvalue()) + logger.error("Failed with output:\n%s", self.out.getvalue()) sys.stdout.flush() -def create_db_from_model(reconstruction: pycolmap.Reconstruction, - database_path: Path) -> Dict[str, int]: +def create_db_from_model( + reconstruction: pycolmap.Reconstruction, database_path: Path +) -> Dict[str, int]: if database_path.exists(): - logger.warning('The database already exists, deleting it.') + logger.warning("The database already exists, deleting it.") database_path.unlink() db = COLMAPDatabase.connect(database_path) @@ -43,8 +45,13 @@ def create_db_from_model(reconstruction: pycolmap.Reconstruction, for i, camera in reconstruction.cameras.items(): db.add_camera( - camera.model_id, camera.width, camera.height, camera.params, - camera_id=i, prior_focal_length=True) + camera.model_id, + camera.width, + camera.height, + camera.params, + camera_id=i, + prior_focal_length=True, + ) for i, image in reconstruction.images.items(): db.add_image(image.name, image.camera_id, image_id=i) @@ -54,10 +61,10 @@ def create_db_from_model(reconstruction: pycolmap.Reconstruction, return {image.name: i for i, image in reconstruction.images.items()} -def import_features(image_ids: Dict[str, int], - database_path: Path, - features_path: Path): - logger.info('Importing features into the database...') +def import_features( + image_ids: Dict[str, int], database_path: Path, features_path: Path +): + logger.info("Importing features into the database...") db = COLMAPDatabase.connect(database_path) for image_name, image_id in tqdm(image_ids.items()): @@ -69,15 +76,17 @@ def import_features(image_ids: Dict[str, int], db.close() -def import_matches(image_ids: Dict[str, int], - database_path: Path, - pairs_path: Path, - matches_path: Path, - min_match_score: Optional[float] = None, - skip_geometric_verification: bool = False): - logger.info('Importing matches into the database...') +def import_matches( + image_ids: Dict[str, int], + database_path: Path, + pairs_path: Path, + matches_path: Path, + min_match_score: Optional[float] = None, + skip_geometric_verification: bool = False, +): + logger.info("Importing matches into the database...") - with open(str(pairs_path), 'r') as f: + with open(str(pairs_path), "r") as f: pairs = [p.split() for p in f.readlines()] db = COLMAPDatabase.connect(database_path) @@ -100,25 +109,27 @@ def import_matches(image_ids: Dict[str, int], db.close() -def estimation_and_geometric_verification(database_path: Path, - pairs_path: Path, - verbose: bool = False): - logger.info('Performing geometric verification of the matches...') +def estimation_and_geometric_verification( + database_path: Path, pairs_path: Path, verbose: bool = False +): + logger.info("Performing geometric verification of the matches...") with OutputCapture(verbose): with pycolmap.ostream(): pycolmap.verify_matches( - database_path, pairs_path, - max_num_trials=20000, min_inlier_ratio=0.1) + database_path, pairs_path, max_num_trials=20000, min_inlier_ratio=0.1 + ) -def geometric_verification(image_ids: Dict[str, int], - reference: pycolmap.Reconstruction, - database_path: Path, - features_path: Path, - pairs_path: Path, - matches_path: Path, - max_error: float = 4.0): - logger.info('Performing geometric verification of the matches...') +def geometric_verification( + image_ids: Dict[str, int], + reference: pycolmap.Reconstruction, + database_path: Path, + features_path: Path, + pairs_path: Path, + matches_path: Path, + max_error: float = 4.0, +): + logger.info("Performing geometric verification of the matches...") pairs = parse_retrieval(pairs_path) db = COLMAPDatabase.connect(database_path) @@ -129,8 +140,7 @@ def geometric_verification(image_ids: Dict[str, int], id0 = image_ids[name0] image0 = reference.images[id0] cam0 = reference.cameras[image0.camera_id] - kps0, noise0 = get_keypoints( - features_path, name0, return_uncertainty=True) + kps0, noise0 = get_keypoints(features_path, name0, return_uncertainty=True) noise0 = 1.0 if noise0 is None else noise0 if len(kps0) > 0: kps0 = np.stack(cam0.image_to_world(kps0)) @@ -141,8 +151,7 @@ def geometric_verification(image_ids: Dict[str, int], id1 = image_ids[name1] image1 = reference.images[id1] cam1 = reference.cameras[image1.camera_id] - kps1, noise1 = get_keypoints( - features_path, name1, return_uncertainty=True) + kps1, noise1 = get_keypoints(features_path, name1, return_uncertainty=True) noise1 = 1.0 if noise1 is None else noise1 if len(kps1) > 0: kps1 = np.stack(cam1.image_to_world(kps1)) @@ -160,118 +169,138 @@ def geometric_verification(image_ids: Dict[str, int], continue qvec_01, tvec_01 = pycolmap.relative_pose( - image0.qvec, image0.tvec, image1.qvec, image1.tvec) + image0.qvec, image0.tvec, image1.qvec, image1.tvec + ) _, errors0, errors1 = compute_epipolar_errors( - qvec_01, tvec_01, kps0[matches[:, 0]], kps1[matches[:, 1]]) + qvec_01, tvec_01, kps0[matches[:, 0]], kps1[matches[:, 1]] + ) valid_matches = np.logical_and( errors0 <= max_error * noise0 / cam0.mean_focal_length(), - errors1 <= max_error * noise1 / cam1.mean_focal_length()) + errors1 <= max_error * noise1 / cam1.mean_focal_length(), + ) # TODO: We could also add E to the database, but we need # to reverse the transformations if id0 > id1 in utils/database.py. db.add_two_view_geometry(id0, id1, matches[valid_matches, :]) inlier_ratios.append(np.mean(valid_matches)) - logger.info('mean/med/min/max valid matches %.2f/%.2f/%.2f/%.2f%%.', - np.mean(inlier_ratios) * 100, np.median(inlier_ratios) * 100, - np.min(inlier_ratios) * 100, np.max(inlier_ratios) * 100) + logger.info( + "mean/med/min/max valid matches %.2f/%.2f/%.2f/%.2f%%.", + np.mean(inlier_ratios) * 100, + np.median(inlier_ratios) * 100, + np.min(inlier_ratios) * 100, + np.max(inlier_ratios) * 100, + ) db.commit() db.close() -def run_triangulation(model_path: Path, - database_path: Path, - image_dir: Path, - reference_model: pycolmap.Reconstruction, - verbose: bool = False, - options: Optional[Dict[str, Any]] = None, - ) -> pycolmap.Reconstruction: +def run_triangulation( + model_path: Path, + database_path: Path, + image_dir: Path, + reference_model: pycolmap.Reconstruction, + verbose: bool = False, + options: Optional[Dict[str, Any]] = None, +) -> pycolmap.Reconstruction: model_path.mkdir(parents=True, exist_ok=True) - logger.info('Running 3D triangulation...') + logger.info("Running 3D triangulation...") if options is None: options = {} with OutputCapture(verbose): with pycolmap.ostream(): reconstruction = pycolmap.triangulate_points( - reference_model, database_path, image_dir, model_path, - options=options) + reference_model, database_path, image_dir, model_path, options=options + ) return reconstruction -def main(sfm_dir: Path, - reference_model: Path, - image_dir: Path, - pairs: Path, - features: Path, - matches: Path, - skip_geometric_verification: bool = False, - estimate_two_view_geometries: bool = False, - min_match_score: Optional[float] = None, - verbose: bool = False, - mapper_options: Optional[Dict[str, Any]] = None, - ) -> pycolmap.Reconstruction: - +def main( + sfm_dir: Path, + reference_model: Path, + image_dir: Path, + pairs: Path, + features: Path, + matches: Path, + skip_geometric_verification: bool = False, + estimate_two_view_geometries: bool = False, + min_match_score: Optional[float] = None, + verbose: bool = False, + mapper_options: Optional[Dict[str, Any]] = None, +) -> pycolmap.Reconstruction: assert reference_model.exists(), reference_model 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' + database = sfm_dir / "database.db" reference = pycolmap.Reconstruction(reference_model) image_ids = create_db_from_model(reference, database) import_features(image_ids, database, features) - import_matches(image_ids, database, pairs, matches, - min_match_score, skip_geometric_verification) + import_matches( + image_ids, + database, + pairs, + matches, + min_match_score, + skip_geometric_verification, + ) if not skip_geometric_verification: if estimate_two_view_geometries: estimation_and_geometric_verification(database, pairs, verbose) else: geometric_verification( - image_ids, reference, database, features, pairs, matches) - reconstruction = run_triangulation(sfm_dir, database, image_dir, reference, - verbose, mapper_options) - logger.info('Finished the triangulation with statistics:\n%s', - reconstruction.summary()) + image_ids, reference, database, features, pairs, matches + ) + reconstruction = run_triangulation( + sfm_dir, database, image_dir, reference, verbose, mapper_options + ) + logger.info( + "Finished the triangulation with statistics:\n%s", reconstruction.summary() + ) return reconstruction def parse_option_args(args: List[str], default_options) -> Dict[str, Any]: options = {} for arg in args: - idx = arg.find('=') + idx = arg.find("=") if idx == -1: - raise ValueError('Options format: key1=value1 key2=value2 etc.') - key, value = arg[:idx], arg[idx+1:] + raise ValueError("Options format: key1=value1 key2=value2 etc.") + key, value = arg[:idx], arg[idx + 1 :] if not hasattr(default_options, key): raise ValueError( f'Unknown option "{key}", allowed options and default values' - f' for {default_options.summary()}') + f" for {default_options.summary()}" + ) value = eval(value) target_type = type(getattr(default_options, key)) if not isinstance(value, target_type): - raise ValueError(f'Incorrect type for option "{key}":' - f' {type(value)} vs {target_type}') + raise ValueError( + f'Incorrect type for option "{key}":' f" {type(value)} vs {target_type}" + ) options[key] = value return options -if __name__ == '__main__': +if __name__ == "__main__": parser = argparse.ArgumentParser() - parser.add_argument('--sfm_dir', type=Path, required=True) - parser.add_argument('--reference_sfm_model', type=Path, required=True) - parser.add_argument('--image_dir', type=Path, required=True) + parser.add_argument("--sfm_dir", type=Path, required=True) + parser.add_argument("--reference_sfm_model", 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("--pairs", type=Path, required=True) + parser.add_argument("--features", type=Path, required=True) + parser.add_argument("--matches", type=Path, required=True) - 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("--skip_geometric_verification", action="store_true") + parser.add_argument("--min_match_score", type=float) + parser.add_argument("--verbose", action="store_true") args = parser.parse_args().__dict__ mapper_options = parse_option_args( - args.pop("mapper_options"), pycolmap.IncrementalMapperOptions()) + args.pop("mapper_options"), pycolmap.IncrementalMapperOptions() + ) main(**args, mapper_options=mapper_options) diff --git a/hloc/utils/base_model.py b/hloc/utils/base_model.py index caf17f05..d16fd4ca 100644 --- a/hloc/utils/base_model.py +++ b/hloc/utils/base_model.py @@ -1,8 +1,9 @@ +import inspect import sys from abc import ABCMeta, abstractmethod -from torch import nn from copy import copy -import inspect + +from torch import nn class BaseModel(nn.Module, metaclass=ABCMeta): @@ -20,7 +21,7 @@ def __init__(self, conf): def forward(self, data): """Check the data and call the _forward method of the child model.""" for key in self.required_inputs: - assert key in data, 'Missing key {} in data'.format(key) + assert key in data, "Missing key {} in data".format(key) return self._forward(data) @abstractmethod @@ -35,8 +36,8 @@ def _forward(self, data): def dynamic_load(root, model): - module_path = f'{root.__name__}.{model}' - module = __import__(module_path, fromlist=['']) + module_path = f"{root.__name__}.{model}" + module = __import__(module_path, fromlist=[""]) classes = inspect.getmembers(module, inspect.isclass) # Filter classes defined in the module classes = [c for c in classes if c[1].__module__ == module_path] diff --git a/hloc/utils/database.py b/hloc/utils/database.py index 870a8c4f..c8a7aa5c 100644 --- a/hloc/utils/database.py +++ b/hloc/utils/database.py @@ -31,10 +31,10 @@ # This script is based on an original implementation by True Price. -import sys import sqlite3 -import numpy as np +import sys +import numpy as np IS_PYTHON3 = sys.version_info[0] >= 3 @@ -68,7 +68,9 @@ prior_tz REAL, CONSTRAINT image_id_check CHECK(image_id >= 0 and image_id < {}), FOREIGN KEY(camera_id) REFERENCES cameras(camera_id)) -""".format(MAX_IMAGE_ID) +""".format( + MAX_IMAGE_ID +) CREATE_TWO_VIEW_GEOMETRIES_TABLE = """ CREATE TABLE IF NOT EXISTS two_view_geometries ( @@ -98,18 +100,19 @@ cols INTEGER NOT NULL, data BLOB)""" -CREATE_NAME_INDEX = \ - "CREATE UNIQUE INDEX IF NOT EXISTS index_name ON images(name)" +CREATE_NAME_INDEX = "CREATE UNIQUE INDEX IF NOT EXISTS index_name ON images(name)" -CREATE_ALL = "; ".join([ - CREATE_CAMERAS_TABLE, - CREATE_IMAGES_TABLE, - CREATE_KEYPOINTS_TABLE, - CREATE_DESCRIPTORS_TABLE, - CREATE_MATCHES_TABLE, - CREATE_TWO_VIEW_GEOMETRIES_TABLE, - CREATE_NAME_INDEX -]) +CREATE_ALL = "; ".join( + [ + CREATE_CAMERAS_TABLE, + CREATE_IMAGES_TABLE, + CREATE_KEYPOINTS_TABLE, + CREATE_DESCRIPTORS_TABLE, + CREATE_MATCHES_TABLE, + CREATE_TWO_VIEW_GEOMETRIES_TABLE, + CREATE_NAME_INDEX, + ] +) def image_ids_to_pair_id(image_id1, image_id2): @@ -139,85 +142,116 @@ def blob_to_array(blob, dtype, shape=(-1,)): class COLMAPDatabase(sqlite3.Connection): - @staticmethod def connect(database_path): return sqlite3.connect(str(database_path), factory=COLMAPDatabase) - def __init__(self, *args, **kwargs): super(COLMAPDatabase, self).__init__(*args, **kwargs) self.create_tables = lambda: self.executescript(CREATE_ALL) - self.create_cameras_table = \ - lambda: self.executescript(CREATE_CAMERAS_TABLE) - self.create_descriptors_table = \ - lambda: self.executescript(CREATE_DESCRIPTORS_TABLE) - self.create_images_table = \ - lambda: self.executescript(CREATE_IMAGES_TABLE) - self.create_two_view_geometries_table = \ - lambda: self.executescript(CREATE_TWO_VIEW_GEOMETRIES_TABLE) - self.create_keypoints_table = \ - lambda: self.executescript(CREATE_KEYPOINTS_TABLE) - self.create_matches_table = \ - lambda: self.executescript(CREATE_MATCHES_TABLE) + self.create_cameras_table = lambda: self.executescript(CREATE_CAMERAS_TABLE) + self.create_descriptors_table = lambda: self.executescript( + CREATE_DESCRIPTORS_TABLE + ) + self.create_images_table = lambda: self.executescript(CREATE_IMAGES_TABLE) + self.create_two_view_geometries_table = lambda: self.executescript( + CREATE_TWO_VIEW_GEOMETRIES_TABLE + ) + self.create_keypoints_table = lambda: self.executescript(CREATE_KEYPOINTS_TABLE) + self.create_matches_table = lambda: self.executescript(CREATE_MATCHES_TABLE) self.create_name_index = lambda: self.executescript(CREATE_NAME_INDEX) - def add_camera(self, model, width, height, params, - prior_focal_length=False, camera_id=None): + def add_camera( + self, model, width, height, params, prior_focal_length=False, camera_id=None + ): params = np.asarray(params, np.float64) cursor = self.execute( "INSERT INTO cameras VALUES (?, ?, ?, ?, ?, ?)", - (camera_id, model, width, height, array_to_blob(params), - prior_focal_length)) + ( + camera_id, + model, + width, + height, + array_to_blob(params), + prior_focal_length, + ), + ) return cursor.lastrowid - def add_image(self, name, camera_id, - prior_q=np.full(4, np.NaN), prior_t=np.full(3, np.NaN), - image_id=None): + def add_image( + self, + name, + camera_id, + prior_q=np.full(4, np.NaN), + prior_t=np.full(3, np.NaN), + image_id=None, + ): cursor = self.execute( "INSERT INTO images VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)", - (image_id, name, camera_id, prior_q[0], prior_q[1], prior_q[2], - prior_q[3], prior_t[0], prior_t[1], prior_t[2])) + ( + image_id, + name, + camera_id, + prior_q[0], + prior_q[1], + prior_q[2], + prior_q[3], + prior_t[0], + prior_t[1], + prior_t[2], + ), + ) return cursor.lastrowid def add_keypoints(self, image_id, keypoints): - assert(len(keypoints.shape) == 2) - assert(keypoints.shape[1] in [2, 4, 6]) + assert len(keypoints.shape) == 2 + assert keypoints.shape[1] in [2, 4, 6] keypoints = np.asarray(keypoints, np.float32) self.execute( "INSERT INTO keypoints VALUES (?, ?, ?, ?)", - (image_id,) + keypoints.shape + (array_to_blob(keypoints),)) + (image_id,) + keypoints.shape + (array_to_blob(keypoints),), + ) def add_descriptors(self, image_id, descriptors): descriptors = np.ascontiguousarray(descriptors, np.uint8) self.execute( "INSERT INTO descriptors VALUES (?, ?, ?, ?)", - (image_id,) + descriptors.shape + (array_to_blob(descriptors),)) + (image_id,) + descriptors.shape + (array_to_blob(descriptors),), + ) def add_matches(self, image_id1, image_id2, matches): - assert(len(matches.shape) == 2) - assert(matches.shape[1] == 2) + assert len(matches.shape) == 2 + assert matches.shape[1] == 2 if image_id1 > image_id2: - matches = matches[:,::-1] + matches = matches[:, ::-1] pair_id = image_ids_to_pair_id(image_id1, image_id2) matches = np.asarray(matches, np.uint32) self.execute( "INSERT INTO matches VALUES (?, ?, ?, ?)", - (pair_id,) + matches.shape + (array_to_blob(matches),)) - - def add_two_view_geometry(self, image_id1, image_id2, matches, - F=np.eye(3), E=np.eye(3), H=np.eye(3), - qvec=np.array([1.0, 0.0, 0.0, 0.0]), - tvec=np.zeros(3), config=2): - assert(len(matches.shape) == 2) - assert(matches.shape[1] == 2) + (pair_id,) + matches.shape + (array_to_blob(matches),), + ) + + def add_two_view_geometry( + self, + image_id1, + image_id2, + matches, + F=np.eye(3), + E=np.eye(3), + H=np.eye(3), + qvec=np.array([1.0, 0.0, 0.0, 0.0]), + tvec=np.zeros(3), + config=2, + ): + assert len(matches.shape) == 2 + assert matches.shape[1] == 2 if image_id1 > image_id2: - matches = matches[:,::-1] + matches = matches[:, ::-1] pair_id = image_ids_to_pair_id(image_id1, image_id2) matches = np.asarray(matches, np.uint32) @@ -228,133 +262,15 @@ def add_two_view_geometry(self, image_id1, image_id2, matches, tvec = np.asarray(tvec, dtype=np.float64) self.execute( "INSERT INTO two_view_geometries VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)", - (pair_id,) + matches.shape + (array_to_blob(matches), config, - array_to_blob(F), array_to_blob(E), array_to_blob(H), - array_to_blob(qvec), array_to_blob(tvec))) - - -def example_usage(): - import os - import argparse - - parser = argparse.ArgumentParser() - parser.add_argument("--database_path", default="database.db") - args = parser.parse_args() - - if os.path.exists(args.database_path): - print("ERROR: database path already exists -- will not modify it.") - return - - # Open the database. - - db = COLMAPDatabase.connect(args.database_path) - - # For convenience, try creating all the tables upfront. - - db.create_tables() - - # Create dummy cameras. - - model1, width1, height1, params1 = \ - 0, 1024, 768, np.array((1024., 512., 384.)) - model2, width2, height2, params2 = \ - 2, 1024, 768, np.array((1024., 512., 384., 0.1)) - - camera_id1 = db.add_camera(model1, width1, height1, params1) - camera_id2 = db.add_camera(model2, width2, height2, params2) - - # Create dummy images. - - image_id1 = db.add_image("image1.png", camera_id1) - image_id2 = db.add_image("image2.png", camera_id1) - image_id3 = db.add_image("image3.png", camera_id2) - image_id4 = db.add_image("image4.png", camera_id2) - - # Create dummy keypoints. - # - # Note that COLMAP supports: - # - 2D keypoints: (x, y) - # - 4D keypoints: (x, y, theta, scale) - # - 6D affine keypoints: (x, y, a_11, a_12, a_21, a_22) - - num_keypoints = 1000 - keypoints1 = np.random.rand(num_keypoints, 2) * (width1, height1) - keypoints2 = np.random.rand(num_keypoints, 2) * (width1, height1) - keypoints3 = np.random.rand(num_keypoints, 2) * (width2, height2) - keypoints4 = np.random.rand(num_keypoints, 2) * (width2, height2) - - db.add_keypoints(image_id1, keypoints1) - db.add_keypoints(image_id2, keypoints2) - db.add_keypoints(image_id3, keypoints3) - db.add_keypoints(image_id4, keypoints4) - - # Create dummy matches. - - M = 50 - matches12 = np.random.randint(num_keypoints, size=(M, 2)) - matches23 = np.random.randint(num_keypoints, size=(M, 2)) - matches34 = np.random.randint(num_keypoints, size=(M, 2)) - - db.add_matches(image_id1, image_id2, matches12) - db.add_matches(image_id2, image_id3, matches23) - db.add_matches(image_id3, image_id4, matches34) - - # Commit the data to the file. - - db.commit() - - # Read and check cameras. - - rows = db.execute("SELECT * FROM cameras") - - camera_id, model, width, height, params, prior = next(rows) - params = blob_to_array(params, np.float64) - assert camera_id == camera_id1 - assert model == model1 and width == width1 and height == height1 - assert np.allclose(params, params1) - - camera_id, model, width, height, params, prior = next(rows) - params = blob_to_array(params, np.float64) - assert camera_id == camera_id2 - assert model == model2 and width == width2 and height == height2 - assert np.allclose(params, params2) - - # Read and check keypoints. - - keypoints = dict( - (image_id, blob_to_array(data, np.float32, (-1, 2))) - for image_id, data in db.execute( - "SELECT image_id, data FROM keypoints")) - - assert np.allclose(keypoints[image_id1], keypoints1) - assert np.allclose(keypoints[image_id2], keypoints2) - assert np.allclose(keypoints[image_id3], keypoints3) - assert np.allclose(keypoints[image_id4], keypoints4) - - # Read and check matches. - - pair_ids = [image_ids_to_pair_id(*pair) for pair in - ((image_id1, image_id2), - (image_id2, image_id3), - (image_id3, image_id4))] - - matches = dict( - (pair_id_to_image_ids(pair_id), - blob_to_array(data, np.uint32, (-1, 2))) - for pair_id, data in db.execute("SELECT pair_id, data FROM matches") - ) - - assert np.all(matches[(image_id1, image_id2)] == matches12) - assert np.all(matches[(image_id2, image_id3)] == matches23) - assert np.all(matches[(image_id3, image_id4)] == matches34) - - # Clean up. - - db.close() - - if os.path.exists(args.database_path): - os.remove(args.database_path) - - -if __name__ == "__main__": - example_usage() + (pair_id,) + + matches.shape + + ( + array_to_blob(matches), + config, + array_to_blob(F), + array_to_blob(E), + array_to_blob(H), + array_to_blob(qvec), + array_to_blob(tvec), + ), + ) diff --git a/hloc/utils/geometry.py b/hloc/utils/geometry.py index 7f5ce101..0fc1bb69 100644 --- a/hloc/utils/geometry.py +++ b/hloc/utils/geometry.py @@ -7,31 +7,27 @@ def to_homogeneous(p): def vector_to_cross_product_matrix(v): - return np.array([ - [0, -v[2], v[1]], - [v[2], 0, -v[0]], - [-v[1], v[0], 0] - ]) + return np.array([[0, -v[2], v[1]], [v[2], 0, -v[0]], [-v[1], v[0], 0]]) def compute_epipolar_errors(qvec_r2t, tvec_r2t, p2d_r, p2d_t): T_r2t = pose_matrix_from_qvec_tvec(qvec_r2t, tvec_r2t) # Compute errors in normalized plane to avoid distortion. - E = vector_to_cross_product_matrix(T_r2t[: 3, -1]) @ T_r2t[: 3, : 3] + E = vector_to_cross_product_matrix(T_r2t[:3, -1]) @ T_r2t[:3, :3] l2d_r2t = (E @ to_homogeneous(p2d_r).T).T l2d_t2r = (E.T @ to_homogeneous(p2d_t).T).T - errors_r = ( - np.abs(np.sum(to_homogeneous(p2d_r) * l2d_t2r, axis=1)) / - np.linalg.norm(l2d_t2r[:, : 2], axis=1)) - errors_t = ( - np.abs(np.sum(to_homogeneous(p2d_t) * l2d_r2t, axis=1)) / - np.linalg.norm(l2d_r2t[:, : 2], axis=1)) + errors_r = np.abs(np.sum(to_homogeneous(p2d_r) * l2d_t2r, axis=1)) / np.linalg.norm( + l2d_t2r[:, :2], axis=1 + ) + errors_t = np.abs(np.sum(to_homogeneous(p2d_t) * l2d_r2t, axis=1)) / np.linalg.norm( + l2d_r2t[:, :2], axis=1 + ) return E, errors_r, errors_t def pose_matrix_from_qvec_tvec(qvec, tvec): pose = np.zeros((4, 4)) - pose[: 3, : 3] = pycolmap.qvec_to_rotmat(qvec) - pose[: 3, -1] = tvec + pose[:3, :3] = pycolmap.qvec_to_rotmat(qvec) + pose[:3, -1] = tvec pose[-1, -1] = 1 return pose diff --git a/hloc/utils/io.py b/hloc/utils/io.py index 92958e96..e7820f9c 100644 --- a/hloc/utils/io.py +++ b/hloc/utils/io.py @@ -1,8 +1,9 @@ -from typing import Tuple from pathlib import Path -import numpy as np +from typing import Tuple + import cv2 import h5py +import numpy as np from .parsers import names_to_pair, names_to_pair_old @@ -14,7 +15,7 @@ def read_image(path, grayscale=False): mode = cv2.IMREAD_COLOR image = cv2.imread(str(path), mode) if image is None: - raise ValueError(f'Cannot read image {path}.') + raise ValueError(f"Cannot read image {path}.") if not grayscale and len(image.shape) == 3: image = image[:, :, ::-1] # BGR to RGB return image @@ -22,20 +23,23 @@ def read_image(path, grayscale=False): def list_h5_names(path): names = [] - with h5py.File(str(path), 'r', libver='latest') as fd: + with h5py.File(str(path), "r", libver="latest") as fd: + def visit_fn(_, obj): if isinstance(obj, h5py.Dataset): - names.append(obj.parent.name.strip('/')) + names.append(obj.parent.name.strip("/")) + fd.visititems(visit_fn) return list(set(names)) -def get_keypoints(path: Path, name: str, - return_uncertainty: bool = False) -> np.ndarray: - with h5py.File(str(path), 'r', libver='latest') as hfile: - dset = hfile[name]['keypoints'] +def get_keypoints( + path: Path, name: str, return_uncertainty: bool = False +) -> np.ndarray: + with h5py.File(str(path), "r", libver="latest") as hfile: + dset = hfile[name]["keypoints"] p = dset.__array__() - uncertainty = dset.attrs.get('uncertainty') + uncertainty = dset.attrs.get("uncertainty") if return_uncertainty: return p, uncertainty return p @@ -56,15 +60,16 @@ def find_pair(hfile: h5py.File, name0: str, name1: str): if pair in hfile: return pair, True raise ValueError( - f'Could not find pair {(name0, name1)}... ' - 'Maybe you matched with a different list of pairs? ') + f"Could not find pair {(name0, name1)}... " + "Maybe you matched with a different list of pairs? " + ) def get_matches(path: Path, name0: str, name1: str) -> Tuple[np.ndarray]: - with h5py.File(str(path), 'r', libver='latest') as hfile: + with h5py.File(str(path), "r", libver="latest") as hfile: pair, reverse = find_pair(hfile, name0, name1) - matches = hfile[pair]['matches0'].__array__() - scores = hfile[pair]['matching_scores0'].__array__() + matches = hfile[pair]["matches0"].__array__() + scores = hfile[pair]["matching_scores0"].__array__() idx = np.where(matches != -1)[0] matches = np.stack([idx, matches[idx]], -1) if reverse: diff --git a/hloc/utils/parsers.py b/hloc/utils/parsers.py index 1f4d9c19..8b5913d8 100644 --- a/hloc/utils/parsers.py +++ b/hloc/utils/parsers.py @@ -1,7 +1,8 @@ -from pathlib import Path import logging -import numpy as np from collections import defaultdict +from pathlib import Path + +import numpy as np import pycolmap logger = logging.getLogger(__name__) @@ -9,10 +10,10 @@ def parse_image_list(path, with_intrinsics=False): images = [] - with open(path, 'r') as f: + with open(path, "r") as f: for line in f: - line = line.strip('\n') - if len(line) == 0 or line[0] == '#': + line = line.strip("\n") + if len(line) == 0 or line[0] == "#": continue name, *data = line.split() if with_intrinsics: @@ -24,7 +25,7 @@ def parse_image_list(path, with_intrinsics=False): images.append(name) assert len(images) > 0 - logger.info(f'Imported {len(images)} images from {path.name}') + logger.info(f"Imported {len(images)} images from {path.name}") return images @@ -39,8 +40,8 @@ def parse_image_lists(paths, with_intrinsics=False): def parse_retrieval(path): retrieval = defaultdict(list) - with open(path, 'r') as f: - for p in f.read().rstrip('\n').split('\n'): + with open(path, "r") as f: + for p in f.read().rstrip("\n").split("\n"): if len(p) == 0: continue q, r = p.split() @@ -48,9 +49,9 @@ def parse_retrieval(path): return dict(retrieval) -def names_to_pair(name0, name1, separator='/'): - return separator.join((name0.replace('/', '-'), name1.replace('/', '-'))) +def names_to_pair(name0, name1, separator="/"): + return separator.join((name0.replace("/", "-"), name1.replace("/", "-"))) def names_to_pair_old(name0, name1): - return names_to_pair(name0, name1, separator='_') + return names_to_pair(name0, name1, separator="_") diff --git a/hloc/utils/read_write_model.py b/hloc/utils/read_write_model.py index df8c9a3a..197921de 100644 --- a/hloc/utils/read_write_model.py +++ b/hloc/utils/read_write_model.py @@ -29,24 +29,27 @@ # # Author: Johannes L. Schoenberger (jsch-at-demuc-dot-de) -import os -import collections -import numpy as np -import struct import argparse +import collections import logging +import os +import struct + +import numpy as np logger = logging.getLogger(__name__) CameraModel = collections.namedtuple( - "CameraModel", ["model_id", "model_name", "num_params"]) -Camera = collections.namedtuple( - "Camera", ["id", "model", "width", "height", "params"]) + "CameraModel", ["model_id", "model_name", "num_params"] +) +Camera = collections.namedtuple("Camera", ["id", "model", "width", "height", "params"]) BaseImage = collections.namedtuple( - "Image", ["id", "qvec", "tvec", "camera_id", "name", "xys", "point3D_ids"]) + "Image", ["id", "qvec", "tvec", "camera_id", "name", "xys", "point3D_ids"] +) Point3D = collections.namedtuple( - "Point3D", ["id", "xyz", "rgb", "error", "image_ids", "point2D_idxs"]) + "Point3D", ["id", "xyz", "rgb", "error", "image_ids", "point2D_idxs"] +) class Image(BaseImage): @@ -65,12 +68,14 @@ def qvec2rotmat(self): CameraModel(model_id=7, model_name="FOV", num_params=5), CameraModel(model_id=8, model_name="SIMPLE_RADIAL_FISHEYE", num_params=4), CameraModel(model_id=9, model_name="RADIAL_FISHEYE", num_params=5), - CameraModel(model_id=10, model_name="THIN_PRISM_FISHEYE", num_params=12) + CameraModel(model_id=10, model_name="THIN_PRISM_FISHEYE", num_params=12), } -CAMERA_MODEL_IDS = dict([(camera_model.model_id, camera_model) - for camera_model in CAMERA_MODELS]) -CAMERA_MODEL_NAMES = dict([(camera_model.model_name, camera_model) - for camera_model in CAMERA_MODELS]) +CAMERA_MODEL_IDS = dict( + [(camera_model.model_id, camera_model) for camera_model in CAMERA_MODELS] +) +CAMERA_MODEL_NAMES = dict( + [(camera_model.model_name, camera_model) for camera_model in CAMERA_MODELS] +) def read_next_bytes(fid, num_bytes, format_char_sequence, endian_character="<"): @@ -121,9 +126,9 @@ def read_cameras_text(path): width = int(elems[2]) height = int(elems[3]) params = np.array(tuple(map(float, elems[4:]))) - cameras[camera_id] = Camera(id=camera_id, model=model, - width=width, height=height, - params=params) + cameras[camera_id] = Camera( + id=camera_id, model=model, width=width, height=height, params=params + ) return cameras @@ -138,20 +143,24 @@ def read_cameras_binary(path_to_model_file): num_cameras = read_next_bytes(fid, 8, "Q")[0] for _ in range(num_cameras): camera_properties = read_next_bytes( - fid, num_bytes=24, format_char_sequence="iiQQ") + fid, num_bytes=24, format_char_sequence="iiQQ" + ) camera_id = camera_properties[0] model_id = camera_properties[1] model_name = CAMERA_MODEL_IDS[camera_properties[1]].model_name width = camera_properties[2] height = camera_properties[3] num_params = CAMERA_MODEL_IDS[model_id].num_params - params = read_next_bytes(fid, num_bytes=8*num_params, - format_char_sequence="d"*num_params) - cameras[camera_id] = Camera(id=camera_id, - model=model_name, - width=width, - height=height, - params=np.array(params)) + params = read_next_bytes( + fid, num_bytes=8 * num_params, format_char_sequence="d" * num_params + ) + cameras[camera_id] = Camera( + id=camera_id, + model=model_name, + width=width, + height=height, + params=np.array(params), + ) assert len(cameras) == num_cameras return cameras @@ -162,9 +171,11 @@ def write_cameras_text(cameras, path): void Reconstruction::WriteCamerasText(const std::string& path) void Reconstruction::ReadCamerasText(const std::string& path) """ - HEADER = "# Camera list with one line of data per camera:\n" + \ - "# CAMERA_ID, MODEL, WIDTH, HEIGHT, PARAMS[]\n" + \ - "# Number of cameras: {}\n".format(len(cameras)) + HEADER = ( + "# Camera list with one line of data per camera:\n" + + "# CAMERA_ID, MODEL, WIDTH, HEIGHT, PARAMS[]\n" + + "# Number of cameras: {}\n".format(len(cameras)) + ) with open(path, "w") as fid: fid.write(HEADER) for _, cam in cameras.items(): @@ -183,10 +194,7 @@ def write_cameras_binary(cameras, path_to_model_file): write_next_bytes(fid, len(cameras), "Q") for _, cam in cameras.items(): model_id = CAMERA_MODEL_NAMES[cam.model].model_id - camera_properties = [cam.id, - model_id, - cam.width, - cam.height] + camera_properties = [cam.id, model_id, cam.width, cam.height] write_next_bytes(fid, camera_properties, "iiQQ") for p in cam.params: write_next_bytes(fid, float(p), "d") @@ -214,13 +222,19 @@ def read_images_text(path): camera_id = int(elems[8]) image_name = elems[9] elems = fid.readline().split() - xys = np.column_stack([tuple(map(float, elems[0::3])), - tuple(map(float, elems[1::3]))]) + xys = np.column_stack( + [tuple(map(float, elems[0::3])), tuple(map(float, elems[1::3]))] + ) point3D_ids = np.array(tuple(map(int, elems[2::3]))) images[image_id] = Image( - id=image_id, qvec=qvec, tvec=tvec, - camera_id=camera_id, name=image_name, - xys=xys, point3D_ids=point3D_ids) + id=image_id, + qvec=qvec, + tvec=tvec, + camera_id=camera_id, + name=image_name, + xys=xys, + point3D_ids=point3D_ids, + ) return images @@ -235,27 +249,38 @@ def read_images_binary(path_to_model_file): num_reg_images = read_next_bytes(fid, 8, "Q")[0] for _ in range(num_reg_images): binary_image_properties = read_next_bytes( - fid, num_bytes=64, format_char_sequence="idddddddi") + fid, num_bytes=64, format_char_sequence="idddddddi" + ) image_id = binary_image_properties[0] qvec = np.array(binary_image_properties[1:5]) tvec = np.array(binary_image_properties[5:8]) camera_id = binary_image_properties[8] image_name = "" current_char = read_next_bytes(fid, 1, "c")[0] - while current_char != b"\x00": # look for the ASCII 0 entry + while current_char != b"\x00": # look for the ASCII 0 entry image_name += current_char.decode("utf-8") current_char = read_next_bytes(fid, 1, "c")[0] - num_points2D = read_next_bytes(fid, num_bytes=8, - format_char_sequence="Q")[0] - x_y_id_s = read_next_bytes(fid, num_bytes=24*num_points2D, - format_char_sequence="ddq"*num_points2D) - xys = np.column_stack([tuple(map(float, x_y_id_s[0::3])), - tuple(map(float, x_y_id_s[1::3]))]) + num_points2D = read_next_bytes(fid, num_bytes=8, format_char_sequence="Q")[ + 0 + ] + x_y_id_s = read_next_bytes( + fid, + num_bytes=24 * num_points2D, + format_char_sequence="ddq" * num_points2D, + ) + xys = np.column_stack( + [tuple(map(float, x_y_id_s[0::3])), tuple(map(float, x_y_id_s[1::3]))] + ) point3D_ids = np.array(tuple(map(int, x_y_id_s[2::3]))) images[image_id] = Image( - id=image_id, qvec=qvec, tvec=tvec, - camera_id=camera_id, name=image_name, - xys=xys, point3D_ids=point3D_ids) + id=image_id, + qvec=qvec, + tvec=tvec, + camera_id=camera_id, + name=image_name, + xys=xys, + point3D_ids=point3D_ids, + ) return images @@ -268,11 +293,17 @@ def write_images_text(images, path): if len(images) == 0: mean_observations = 0 else: - mean_observations = sum((len(img.point3D_ids) for _, img in images.items()))/len(images) - HEADER = "# Image list with two lines of data per image:\n" + \ - "# IMAGE_ID, QW, QX, QY, QZ, TX, TY, TZ, CAMERA_ID, NAME\n" + \ - "# POINTS2D[] as (X, Y, POINT3D_ID)\n" + \ - "# Number of images: {}, mean observations per image: {}\n".format(len(images), mean_observations) + mean_observations = sum( + (len(img.point3D_ids) for _, img in images.items()) + ) / len(images) + HEADER = ( + "# Image list with two lines of data per image:\n" + + "# IMAGE_ID, QW, QX, QY, QZ, TX, TY, TZ, CAMERA_ID, NAME\n" + + "# POINTS2D[] as (X, Y, POINT3D_ID)\n" + + "# Number of images: {}, mean observations per image: {}\n".format( + len(images), mean_observations + ) + ) with open(path, "w") as fid: fid.write(HEADER) @@ -329,9 +360,14 @@ def read_points3D_text(path): error = float(elems[7]) image_ids = np.array(tuple(map(int, elems[8::2]))) point2D_idxs = np.array(tuple(map(int, elems[9::2]))) - points3D[point3D_id] = Point3D(id=point3D_id, xyz=xyz, rgb=rgb, - error=error, image_ids=image_ids, - point2D_idxs=point2D_idxs) + points3D[point3D_id] = Point3D( + id=point3D_id, + xyz=xyz, + rgb=rgb, + error=error, + image_ids=image_ids, + point2D_idxs=point2D_idxs, + ) return points3D @@ -346,22 +382,30 @@ def read_points3D_binary(path_to_model_file): num_points = read_next_bytes(fid, 8, "Q")[0] for _ in range(num_points): binary_point_line_properties = read_next_bytes( - fid, num_bytes=43, format_char_sequence="QdddBBBd") + fid, num_bytes=43, format_char_sequence="QdddBBBd" + ) point3D_id = binary_point_line_properties[0] xyz = np.array(binary_point_line_properties[1:4]) rgb = np.array(binary_point_line_properties[4:7]) error = np.array(binary_point_line_properties[7]) - track_length = read_next_bytes( - fid, num_bytes=8, format_char_sequence="Q")[0] + track_length = read_next_bytes(fid, num_bytes=8, format_char_sequence="Q")[ + 0 + ] track_elems = read_next_bytes( - fid, num_bytes=8*track_length, - format_char_sequence="ii"*track_length) + fid, + num_bytes=8 * track_length, + format_char_sequence="ii" * track_length, + ) image_ids = np.array(tuple(map(int, track_elems[0::2]))) point2D_idxs = np.array(tuple(map(int, track_elems[1::2]))) points3D[point3D_id] = Point3D( - id=point3D_id, xyz=xyz, rgb=rgb, - error=error, image_ids=image_ids, - point2D_idxs=point2D_idxs) + id=point3D_id, + xyz=xyz, + rgb=rgb, + error=error, + image_ids=image_ids, + point2D_idxs=point2D_idxs, + ) return points3D @@ -374,10 +418,16 @@ def write_points3D_text(points3D, path): if len(points3D) == 0: mean_track_length = 0 else: - mean_track_length = sum((len(pt.image_ids) for _, pt in points3D.items()))/len(points3D) - HEADER = "# 3D point list with one line of data per point:\n" + \ - "# POINT3D_ID, X, Y, Z, R, G, B, ERROR, TRACK[] as (IMAGE_ID, POINT2D_IDX)\n" + \ - "# Number of points: {}, mean track length: {}\n".format(len(points3D), mean_track_length) + mean_track_length = sum( + (len(pt.image_ids) for _, pt in points3D.items()) + ) / len(points3D) + HEADER = ( + "# 3D point list with one line of data per point:\n" + + "# POINT3D_ID, X, Y, Z, R, G, B, ERROR, TRACK[] as (IMAGE_ID, POINT2D_IDX)\n" # noqa: E501 + + "# Number of points: {}, mean track length: {}\n".format( + len(points3D), mean_track_length + ) + ) with open(path, "w") as fid: fid.write(HEADER) @@ -410,9 +460,11 @@ def write_points3D_binary(points3D, path_to_model_file): def detect_model_format(path, ext): - if os.path.isfile(os.path.join(path, "cameras" + ext)) and \ - os.path.isfile(os.path.join(path, "images" + ext)) and \ - os.path.isfile(os.path.join(path, "points3D" + ext)): + if ( + os.path.isfile(os.path.join(path, "cameras" + ext)) + and os.path.isfile(os.path.join(path, "images" + ext)) + and os.path.isfile(os.path.join(path, "points3D" + ext)) + ): return True return False @@ -428,12 +480,12 @@ def read_model(path, ext=""): else: try: cameras, images, points3D = read_model(os.path.join(path, "model/")) - logger.warning( - "This SfM file structure was deprecated in hloc v1.1") + logger.warning("This SfM file structure was deprecated in hloc v1.1") return cameras, images, points3D except FileNotFoundError: raise FileNotFoundError( - f"Could not find binary or text COLMAP model at {path}") + f"Could not find binary or text COLMAP model at {path}" + ) if ext == ".txt": cameras = read_cameras_text(os.path.join(path, "cameras" + ext)) @@ -459,25 +511,40 @@ def write_model(cameras, images, points3D, path, ext=".bin"): def qvec2rotmat(qvec): - return np.array([ - [1 - 2 * qvec[2]**2 - 2 * qvec[3]**2, - 2 * qvec[1] * qvec[2] - 2 * qvec[0] * qvec[3], - 2 * qvec[3] * qvec[1] + 2 * qvec[0] * qvec[2]], - [2 * qvec[1] * qvec[2] + 2 * qvec[0] * qvec[3], - 1 - 2 * qvec[1]**2 - 2 * qvec[3]**2, - 2 * qvec[2] * qvec[3] - 2 * qvec[0] * qvec[1]], - [2 * qvec[3] * qvec[1] - 2 * qvec[0] * qvec[2], - 2 * qvec[2] * qvec[3] + 2 * qvec[0] * qvec[1], - 1 - 2 * qvec[1]**2 - 2 * qvec[2]**2]]) + return np.array( + [ + [ + 1 - 2 * qvec[2] ** 2 - 2 * qvec[3] ** 2, + 2 * qvec[1] * qvec[2] - 2 * qvec[0] * qvec[3], + 2 * qvec[3] * qvec[1] + 2 * qvec[0] * qvec[2], + ], + [ + 2 * qvec[1] * qvec[2] + 2 * qvec[0] * qvec[3], + 1 - 2 * qvec[1] ** 2 - 2 * qvec[3] ** 2, + 2 * qvec[2] * qvec[3] - 2 * qvec[0] * qvec[1], + ], + [ + 2 * qvec[3] * qvec[1] - 2 * qvec[0] * qvec[2], + 2 * qvec[2] * qvec[3] + 2 * qvec[0] * qvec[1], + 1 - 2 * qvec[1] ** 2 - 2 * qvec[2] ** 2, + ], + ] + ) def rotmat2qvec(R): Rxx, Ryx, Rzx, Rxy, Ryy, Rzy, Rxz, Ryz, Rzz = R.flat - K = np.array([ - [Rxx - Ryy - Rzz, 0, 0, 0], - [Ryx + Rxy, Ryy - Rxx - Rzz, 0, 0], - [Rzx + Rxz, Rzy + Ryz, Rzz - Rxx - Ryy, 0], - [Ryz - Rzy, Rzx - Rxz, Rxy - Ryx, Rxx + Ryy + Rzz]]) / 3.0 + K = ( + np.array( + [ + [Rxx - Ryy - Rzz, 0, 0, 0], + [Ryx + Rxy, Ryy - Rxx - Rzz, 0, 0], + [Rzx + Rxz, Rzy + Ryz, Rzz - Rxx - Ryy, 0], + [Ryz - Rzy, Rzx - Rxz, Rxy - Ryx, Rxx + Ryy + Rzz], + ] + ) + / 3.0 + ) eigvals, eigvecs = np.linalg.eigh(K) qvec = eigvecs[[3, 0, 1, 2], np.argmax(eigvals)] if qvec[0] < 0: @@ -486,14 +553,23 @@ def rotmat2qvec(R): def main(): - parser = argparse.ArgumentParser(description="Read and write COLMAP binary and text models") + parser = argparse.ArgumentParser( + description="Read and write COLMAP binary and text models" + ) parser.add_argument("--input_model", help="path to input model folder") - parser.add_argument("--input_format", choices=[".bin", ".txt"], - help="input model format", default="") - parser.add_argument("--output_model", - help="path to output model folder") - parser.add_argument("--output_format", choices=[".bin", ".txt"], - help="outut model format", default=".txt") + parser.add_argument( + "--input_format", + choices=[".bin", ".txt"], + help="input model format", + default="", + ) + parser.add_argument("--output_model", help="path to output model folder") + parser.add_argument( + "--output_format", + choices=[".bin", ".txt"], + help="outut model format", + default=".txt", + ) args = parser.parse_args() cameras, images, points3D = read_model(path=args.input_model, ext=args.input_format) @@ -503,7 +579,9 @@ def main(): print("num_points3D:", len(points3D)) if args.output_model is not None: - write_model(cameras, images, points3D, path=args.output_model, ext=args.output_format) + write_model( + cameras, images, points3D, path=args.output_model, ext=args.output_format + ) if __name__ == "__main__": diff --git a/hloc/utils/viz.py b/hloc/utils/viz.py index 24f8c4d2..dde1a1bc 100644 --- a/hloc/utils/viz.py +++ b/hloc/utils/viz.py @@ -7,20 +7,21 @@ """ import matplotlib -import matplotlib.pyplot as plt import matplotlib.patheffects as path_effects +import matplotlib.pyplot as plt import numpy as np def cm_RdGn(x): """Custom colormap: red (0) -> yellow (0.5) -> green (1).""" - x = np.clip(x, 0, 1)[..., None]*2 - c = x*np.array([[0, 1., 0]]) + (2-x)*np.array([[1., 0, 0]]) + x = np.clip(x, 0, 1)[..., None] * 2 + c = x * np.array([[0, 1.0, 0]]) + (2 - x) * np.array([[1.0, 0, 0]]) return np.clip(c, 0, 1) -def plot_images(imgs, titles=None, cmaps='gray', dpi=100, pad=.5, - adaptive=True, figsize=4.5): +def plot_images( + imgs, titles=None, cmaps="gray", dpi=100, pad=0.5, adaptive=True, figsize=4.5 +): """Plot a set of images horizontally. Args: imgs: a list of NumPy or PyTorch images, RGB (H, W, 3) or mono (H, W). @@ -35,10 +36,11 @@ def plot_images(imgs, titles=None, cmaps='gray', dpi=100, pad=.5, if adaptive: ratios = [i.shape[1] / i.shape[0] for i in imgs] # W / H else: - ratios = [4/3] * n - figsize = [sum(ratios)*figsize, figsize] + ratios = [4 / 3] * n + figsize = [sum(ratios) * figsize, figsize] fig, axs = plt.subplots( - 1, n, figsize=figsize, dpi=dpi, gridspec_kw={'width_ratios': ratios}) + 1, n, figsize=figsize, dpi=dpi, gridspec_kw={"width_ratios": ratios} + ) if n == 1: axs = [axs] for i, (img, ax) in enumerate(zip(imgs, axs)): @@ -49,7 +51,7 @@ def plot_images(imgs, titles=None, cmaps='gray', dpi=100, pad=.5, fig.tight_layout(pad=pad) -def plot_keypoints(kpts, colors='lime', ps=4): +def plot_keypoints(kpts, colors="lime", ps=4): """Plot keypoints for existing images. Args: kpts: list of ndarrays of size (N, 2). @@ -63,7 +65,7 @@ def plot_keypoints(kpts, colors='lime', ps=4): a.scatter(k[:, 0], k[:, 1], c=c, s=ps, linewidths=0) -def plot_matches(kpts0, kpts1, color=None, lw=1.5, ps=4, indices=(0, 1), a=1.): +def plot_matches(kpts0, kpts1, color=None, lw=1.5, ps=4, indices=(0, 1), a=1.0): """Plot matches for a pair of existing images. Args: kpts0, kpts1: corresponding keypoints of size (N, 2). @@ -88,11 +90,18 @@ def plot_matches(kpts0, kpts1, color=None, lw=1.5, ps=4, indices=(0, 1), a=1.): if lw > 0: # transform the points into the figure coordinate system for i in range(len(kpts0)): - fig.add_artist(matplotlib.patches.ConnectionPatch( - xyA=(kpts0[i, 0], kpts0[i, 1]), coordsA=ax0.transData, - xyB=(kpts1[i, 0], kpts1[i, 1]), coordsB=ax1.transData, - zorder=1, color=color[i], linewidth=lw, - alpha=a)) + fig.add_artist( + matplotlib.patches.ConnectionPatch( + xyA=(kpts0[i, 0], kpts0[i, 1]), + coordsA=ax0.transData, + xyB=(kpts1[i, 0], kpts1[i, 1]), + coordsB=ax1.transData, + zorder=1, + color=color[i], + linewidth=lw, + alpha=a, + ) + ) # freeze the axes to prevent the transform to change ax0.autoscale(enable=False) @@ -103,17 +112,30 @@ def plot_matches(kpts0, kpts1, color=None, lw=1.5, ps=4, indices=(0, 1), a=1.): ax1.scatter(kpts1[:, 0], kpts1[:, 1], c=color, s=ps) -def add_text(idx, text, pos=(0.01, 0.99), fs=15, color='w', - lcolor='k', lwidth=2, ha='left', va='top'): +def add_text( + idx, + text, + pos=(0.01, 0.99), + fs=15, + color="w", + lcolor="k", + lwidth=2, + ha="left", + va="top", +): ax = plt.gcf().axes[idx] - t = ax.text(*pos, text, fontsize=fs, ha=ha, va=va, - color=color, transform=ax.transAxes) + t = ax.text( + *pos, text, fontsize=fs, ha=ha, va=va, color=color, transform=ax.transAxes + ) if lcolor is not None: - t.set_path_effects([ - path_effects.Stroke(linewidth=lwidth, foreground=lcolor), - path_effects.Normal()]) + t.set_path_effects( + [ + path_effects.Stroke(linewidth=lwidth, foreground=lcolor), + path_effects.Normal(), + ] + ) def save_plot(path, **kw): """Save the current figure without any white margin.""" - plt.savefig(path, bbox_inches='tight', pad_inches=0, **kw) + plt.savefig(path, bbox_inches="tight", pad_inches=0, **kw) diff --git a/hloc/utils/viz_3d.py b/hloc/utils/viz_3d.py index ae14707a..1aecb51c 100644 --- a/hloc/utils/viz_3d.py +++ b/hloc/utils/viz_3d.py @@ -9,13 +9,14 @@ """ from typing import Optional + import numpy as np -import pycolmap import plotly.graph_objects as go +import pycolmap def to_homogeneous(points): - pad = np.ones((points.shape[:-1]+(1,)), dtype=points.dtype) + pad = np.ones((points.shape[:-1] + (1,)), dtype=points.dtype) return np.concatenate([points, pad], axis=-1) @@ -34,57 +35,59 @@ def init_figure(height: int = 800) -> go.Figure: template="plotly_dark", height=height, scene_camera=dict( - eye=dict(x=0., y=-.1, z=-2), - up=dict(x=0, y=-1., z=0), - projection=dict(type="orthographic")), + eye=dict(x=0.0, y=-0.1, z=-2), + up=dict(x=0, y=-1.0, z=0), + projection=dict(type="orthographic"), + ), scene=dict( xaxis=axes, yaxis=axes, zaxis=axes, - aspectmode='data', - dragmode='orbit', + aspectmode="data", + dragmode="orbit", ), margin=dict(l=0, r=0, b=0, t=0, pad=0), - legend=dict( - orientation="h", - yanchor="top", - y=0.99, - xanchor="left", - x=0.1 - ), + legend=dict(orientation="h", yanchor="top", y=0.99, xanchor="left", x=0.1), ) return fig def plot_points( - fig: go.Figure, - pts: np.ndarray, - color: str = 'rgba(255, 0, 0, 1)', - ps: int = 2, - colorscale: Optional[str] = None, - name: Optional[str] = None): + fig: go.Figure, + pts: np.ndarray, + color: str = "rgba(255, 0, 0, 1)", + ps: int = 2, + colorscale: Optional[str] = None, + name: Optional[str] = None, +): """Plot a set of 3D points.""" x, y, z = pts.T tr = go.Scatter3d( - x=x, y=y, z=z, mode='markers', name=name, legendgroup=name, - marker=dict( - size=ps, color=color, line_width=0.0, colorscale=colorscale)) + x=x, + y=y, + z=z, + mode="markers", + name=name, + legendgroup=name, + marker=dict(size=ps, color=color, line_width=0.0, colorscale=colorscale), + ) fig.add_trace(tr) def plot_camera( - fig: go.Figure, - R: np.ndarray, - t: np.ndarray, - K: np.ndarray, - color: str = 'rgb(0, 0, 255)', - name: Optional[str] = None, - legendgroup: Optional[str] = None, - fill: bool = False, - size: float = 1.0, - text: Optional[str] = None): + fig: go.Figure, + R: np.ndarray, + t: np.ndarray, + K: np.ndarray, + color: str = "rgb(0, 0, 255)", + name: Optional[str] = None, + legendgroup: Optional[str] = None, + fill: bool = False, + size: float = 1.0, + text: Optional[str] = None, +): """Plot a camera frustum from pose and intrinsic matrix.""" - W, H = K[0, 2]*2, K[1, 2]*2 + W, H = K[0, 2] * 2, K[1, 2] * 2 corners = np.array([[0, 0], [W, 0], [W, H], [0, H], [0, 0]]) if size is not None: image_extent = max(size * W / 1024.0, size * H / 1024.0) @@ -103,31 +106,46 @@ def plot_camera( if fill: pyramid = go.Mesh3d( - x=x, y=y, z=z, color=color, i=i, j=j, k=k, - legendgroup=legendgroup, name=name, showlegend=False, - hovertemplate=text.replace('\n', '
')) + x=x, + y=y, + z=z, + color=color, + i=i, + j=j, + k=k, + legendgroup=legendgroup, + name=name, + showlegend=False, + hovertemplate=text.replace("\n", "
"), + ) fig.add_trace(pyramid) triangles = np.vstack((i, j, k)).T vertices = np.concatenate(([t], corners)) - tri_points = np.array([ - vertices[i] for i in triangles.reshape(-1) - ]) + tri_points = np.array([vertices[i] for i in triangles.reshape(-1)]) x, y, z = tri_points.T pyramid = go.Scatter3d( - x=x, y=y, z=z, mode='lines', legendgroup=legendgroup, - name=name, line=dict(color=color, width=1), showlegend=False, - hovertemplate=text.replace('\n', '
')) + x=x, + y=y, + z=z, + mode="lines", + legendgroup=legendgroup, + name=name, + line=dict(color=color, width=1), + showlegend=False, + hovertemplate=text.replace("\n", "
"), + ) fig.add_trace(pyramid) def plot_camera_colmap( - fig: go.Figure, - image: pycolmap.Image, - camera: pycolmap.Camera, - name: Optional[str] = None, - **kwargs): + fig: go.Figure, + image: pycolmap.Image, + camera: pycolmap.Camera, + name: Optional[str] = None, + **kwargs +): """Plot a camera frustum from PyCOLMAP objects""" plot_camera( fig, @@ -136,38 +154,43 @@ def plot_camera_colmap( camera.calibration_matrix(), name=name or str(image.image_id), text=image.summary(), - **kwargs) + **kwargs + ) -def plot_cameras( - fig: go.Figure, - reconstruction: pycolmap.Reconstruction, - **kwargs): +def plot_cameras(fig: go.Figure, reconstruction: pycolmap.Reconstruction, **kwargs): """Plot a camera as a cone with camera frustum.""" for image_id, image in reconstruction.images.items(): plot_camera_colmap( - fig, image, reconstruction.cameras[image.camera_id], **kwargs) + fig, image, reconstruction.cameras[image.camera_id], **kwargs + ) def plot_reconstruction( - fig: go.Figure, - rec: pycolmap.Reconstruction, - max_reproj_error: float = 6.0, - color: str = 'rgb(0, 0, 255)', - name: Optional[str] = None, - min_track_length: int = 2, - points: bool = True, - cameras: bool = True, - points_rgb: bool = True, - cs: float = 1.0): + fig: go.Figure, + rec: pycolmap.Reconstruction, + max_reproj_error: float = 6.0, + color: str = "rgb(0, 0, 255)", + name: Optional[str] = None, + min_track_length: int = 2, + points: bool = True, + cameras: bool = True, + points_rgb: bool = True, + cs: float = 1.0, +): # Filter outliers bbs = rec.compute_bounding_box(0.001, 0.999) # Filter points, use original reproj error here - p3Ds = [p3D for _, p3D in rec.points3D.items() if ( - (p3D.xyz >= bbs[0]).all() and - (p3D.xyz <= bbs[1]).all() and - p3D.error <= max_reproj_error and - p3D.track.length() >= min_track_length)] + p3Ds = [ + p3D + for _, p3D in rec.points3D.items() + if ( + (p3D.xyz >= bbs[0]).all() + and (p3D.xyz <= bbs[1]).all() + and p3D.error <= max_reproj_error + and p3D.track.length() >= min_track_length + ) + ] xyzs = [p3D.xyz for p3D in p3Ds] if points_rgb: pcolor = [p3D.color for p3D in p3Ds] diff --git a/hloc/visualization.py b/hloc/visualization.py index 5b0b70d4..7f28c042 100644 --- a/hloc/visualization.py +++ b/hloc/visualization.py @@ -1,67 +1,86 @@ -from matplotlib import cm +import pickle import random + import numpy as np -import pickle import pycolmap +from matplotlib import cm -from .utils.viz import ( - plot_images, plot_keypoints, plot_matches, cm_RdGn, add_text) from .utils.io import read_image +from .utils.viz import add_text, cm_RdGn, plot_images, plot_keypoints, plot_matches -def visualize_sfm_2d(reconstruction, image_dir, color_by='visibility', - selected=[], n=1, seed=0, dpi=75): +def visualize_sfm_2d( + reconstruction, image_dir, color_by="visibility", selected=[], n=1, seed=0, dpi=75 +): assert image_dir.exists() if not isinstance(reconstruction, pycolmap.Reconstruction): reconstruction = pycolmap.Reconstruction(reconstruction) if not selected: image_ids = reconstruction.reg_image_ids() - selected = random.Random(seed).sample( - image_ids, min(n, len(image_ids))) + selected = random.Random(seed).sample(image_ids, min(n, len(image_ids))) for i in selected: image = reconstruction.images[i] keypoints = np.array([p.xy for p in image.points2D]) visible = np.array([p.has_point3D() for p in image.points2D]) - if color_by == 'visibility': + if color_by == "visibility": color = [(0, 0, 1) if v else (1, 0, 0) for v in visible] - text = f'visible: {np.count_nonzero(visible)}/{len(visible)}' - elif color_by == 'track_length': - tl = np.array([reconstruction.points3D[p.point3D_id].track.length() - if p.has_point3D() else 1 for p in image.points2D]) + text = f"visible: {np.count_nonzero(visible)}/{len(visible)}" + elif color_by == "track_length": + tl = np.array( + [ + reconstruction.points3D[p.point3D_id].track.length() + if p.has_point3D() + else 1 + for p in image.points2D + ] + ) max_, med_ = np.max(tl), np.median(tl[tl > 1]) tl = np.log(tl) color = cm.jet(tl / tl.max()).tolist() - text = f'max/median track length: {max_}/{med_}' - elif color_by == 'depth': + text = f"max/median track length: {max_}/{med_}" + elif color_by == "depth": p3ids = [p.point3D_id for p in image.points2D if p.has_point3D()] - z = np.array([image.transform_to_image( - reconstruction.points3D[j].xyz)[-1] for j in p3ids]) + z = np.array( + [ + image.transform_to_image(reconstruction.points3D[j].xyz)[-1] + for j in p3ids + ] + ) z -= z.min() color = cm.jet(z / np.percentile(z, 99.9)) - text = f'visible: {np.count_nonzero(visible)}/{len(visible)}' + text = f"visible: {np.count_nonzero(visible)}/{len(visible)}" keypoints = keypoints[visible] else: - raise NotImplementedError(f'Coloring not implemented: {color_by}.') + raise NotImplementedError(f"Coloring not implemented: {color_by}.") name = image.name plot_images([read_image(image_dir / name)], dpi=dpi) plot_keypoints([keypoints], colors=[color], ps=4) add_text(0, text) - add_text(0, name, pos=(0.01, 0.01), fs=5, lcolor=None, va='bottom') - - -def visualize_loc(results, image_dir, reconstruction=None, db_image_dir=None, - selected=[], n=1, seed=0, prefix=None, **kwargs): + add_text(0, name, pos=(0.01, 0.01), fs=5, lcolor=None, va="bottom") + + +def visualize_loc( + results, + image_dir, + reconstruction=None, + db_image_dir=None, + selected=[], + n=1, + seed=0, + prefix=None, + **kwargs, +): assert image_dir.exists() - with open(str(results)+'_logs.pkl', 'rb') as f: + with open(str(results) + "_logs.pkl", "rb") as f: logs = pickle.load(f) if not selected: - queries = list(logs['loc'].keys()) + queries = list(logs["loc"].keys()) if prefix: queries = [q for q in queries if q.startswith(prefix)] selected = random.Random(seed).sample(queries, min(n, len(queries))) @@ -71,69 +90,74 @@ def visualize_loc(results, image_dir, reconstruction=None, db_image_dir=None, reconstruction = pycolmap.Reconstruction(reconstruction) for qname in selected: - loc = logs['loc'][qname] + loc = logs["loc"][qname] visualize_loc_from_log( - image_dir, qname, loc, reconstruction, db_image_dir, **kwargs) - - -def visualize_loc_from_log(image_dir, query_name, loc, reconstruction=None, - db_image_dir=None, top_k_db=2, dpi=75): - + image_dir, qname, loc, reconstruction, db_image_dir, **kwargs + ) + + +def visualize_loc_from_log( + image_dir, + query_name, + loc, + reconstruction=None, + db_image_dir=None, + top_k_db=2, + dpi=75, +): q_image = read_image(image_dir / query_name) - if loc.get('covisibility_clustering', False): + if loc.get("covisibility_clustering", False): # select the first, largest cluster if the localization failed - loc = loc['log_clusters'][loc['best_cluster'] or 0] + loc = loc["log_clusters"][loc["best_cluster"] or 0] - inliers = np.array(loc['PnP_ret']['inliers']) - mkp_q = loc['keypoints_query'] - n = len(loc['db']) + inliers = np.array(loc["PnP_ret"]["inliers"]) + mkp_q = loc["keypoints_query"] + n = len(loc["db"]) if reconstruction is not None: # for each pair of query keypoint and its matched 3D point, # we need to find its corresponding keypoint in each database image # that observes it. We also count the number of inliers in each. - kp_idxs, kp_to_3D_to_db = loc['keypoint_index_to_db'] + kp_idxs, kp_to_3D_to_db = loc["keypoint_index_to_db"] counts = np.zeros(n) dbs_kp_q_db = [[] for _ in range(n)] inliers_dbs = [[] for _ in range(n)] - for i, (inl, (p3D_id, db_idxs)) in enumerate(zip(inliers, - kp_to_3D_to_db)): + for i, (inl, (p3D_id, db_idxs)) in enumerate(zip(inliers, kp_to_3D_to_db)): track = reconstruction.points3D[p3D_id].track track = {el.image_id: el.point2D_idx for el in track.elements} for db_idx in db_idxs: counts[db_idx] += inl - kp_db = track[loc['db'][db_idx]] + kp_db = track[loc["db"][db_idx]] dbs_kp_q_db[db_idx].append((i, kp_db)) inliers_dbs[db_idx].append(inl) else: # for inloc the database keypoints are already in the logs - assert 'keypoints_db' in loc - assert 'indices_db' in loc - counts = np.array([ - np.sum(loc['indices_db'][inliers] == i) for i in range(n)]) + assert "keypoints_db" in loc + assert "indices_db" in loc + counts = np.array([np.sum(loc["indices_db"][inliers] == i) for i in range(n)]) # display the database images with the most inlier matches db_sort = np.argsort(-counts) for db_idx in db_sort[:top_k_db]: if reconstruction is not None: - db = reconstruction.images[loc['db'][db_idx]] + db = reconstruction.images[loc["db"][db_idx]] db_name = db.name db_kp_q_db = np.array(dbs_kp_q_db[db_idx]) kp_q = mkp_q[db_kp_q_db[:, 0]] kp_db = np.array([db.points2D[i].xy for i in db_kp_q_db[:, 1]]) inliers_db = inliers_dbs[db_idx] else: - db_name = loc['db'][db_idx] - kp_q = mkp_q[loc['indices_db'] == db_idx] - kp_db = loc['keypoints_db'][loc['indices_db'] == db_idx] - inliers_db = inliers[loc['indices_db'] == db_idx] + db_name = loc["db"][db_idx] + kp_q = mkp_q[loc["indices_db"] == db_idx] + kp_db = loc["keypoints_db"][loc["indices_db"] == db_idx] + inliers_db = inliers[loc["indices_db"] == db_idx] db_image = read_image((db_image_dir or image_dir) / db_name) color = cm_RdGn(inliers_db).tolist() - text = f'inliers: {sum(inliers_db)}/{len(inliers_db)}' + text = f"inliers: {sum(inliers_db)}/{len(inliers_db)}" plot_images([q_image, db_image], dpi=dpi) plot_matches(kp_q, kp_db, color, a=0.1) add_text(0, text) - opts = dict(pos=(0.01, 0.01), fs=5, lcolor=None, va='bottom') + opts = dict(pos=(0.01, 0.01), fs=5, lcolor=None, va="bottom") add_text(0, query_name, **opts) add_text(1, db_name, **opts) From 5244632f7a2246d87608a3f72ccfcc1af2425c4c Mon Sep 17 00:00:00 2001 From: Paul-Edouard Sarlin Date: Thu, 11 Jan 2024 16:57:02 +0100 Subject: [PATCH 3/9] Add CI --- .github/workflows/code-quality.yml | 24 ++++++++++++++++++++++++ 1 file changed, 24 insertions(+) create mode 100644 .github/workflows/code-quality.yml diff --git a/.github/workflows/code-quality.yml b/.github/workflows/code-quality.yml new file mode 100644 index 00000000..128ebc06 --- /dev/null +++ b/.github/workflows/code-quality.yml @@ -0,0 +1,24 @@ +name: Format and Lint Checks +on: + push: + branches: + - main + paths: + - '*.py' + pull_request: + types: [ assigned, opened, synchronize, reopened ] +jobs: + check: + name: Format and Lint Checks + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v3 + - uses: actions/setup-python@v4 + with: + python-version: '3.10' + cache: 'pip' + - run: python -m pip install --upgrade pip + - run: python -m pip install black flake8 isort + - run: python -m flake8 hloc + - run: python -m isort hloc --check-only --diff + - run: python -m black hloc --check --diff From 59546687b91cb8f7b91c8ef619be974c963f4241 Mon Sep 17 00:00:00 2001 From: Paul-Edouard Sarlin Date: Thu, 11 Jan 2024 16:57:30 +0100 Subject: [PATCH 4/9] Ignore env file --- .gitignore | 1 + 1 file changed, 1 insertion(+) diff --git a/.gitignore b/.gitignore index 64707625..c363cd22 100644 --- a/.gitignore +++ b/.gitignore @@ -1,3 +1,4 @@ +.venv __pycache__ *.pyc *.egg-info From 449d79f2bd9403fb1fe131365638ecd89d6db593 Mon Sep 17 00:00:00 2001 From: Paul-Edouard Sarlin Date: Thu, 11 Jan 2024 17:03:57 +0100 Subject: [PATCH 5/9] Run CI From 546508f3c4a8676b784f0a2189f616dc1163e010 Mon Sep 17 00:00:00 2001 From: Paul-Edouard Sarlin Date: Thu, 11 Jan 2024 17:06:53 +0100 Subject: [PATCH 6/9] Fix --- .github/workflows/code-quality.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/workflows/code-quality.yml b/.github/workflows/code-quality.yml index 128ebc06..e3327ffb 100644 --- a/.github/workflows/code-quality.yml +++ b/.github/workflows/code-quality.yml @@ -2,7 +2,7 @@ name: Format and Lint Checks on: push: branches: - - main + - master paths: - '*.py' pull_request: From 015c9f5cf80724d0415f2b3e1b79b48aa0a9aceb Mon Sep 17 00:00:00 2001 From: Paul-Edouard Sarlin Date: Thu, 11 Jan 2024 17:11:10 +0100 Subject: [PATCH 7/9] Run CI From e7ead1b65845e925d411a1341e06cd5455bcd0f9 Mon Sep 17 00:00:00 2001 From: Paul-Edouard Sarlin Date: Fri, 12 Jan 2024 12:07:38 +0100 Subject: [PATCH 8/9] Format notebooks as well --- .github/workflows/code-quality.yml | 4 +- demo.ipynb | 67 +++++++++++++++++--------- pipeline_Aachen.ipynb | 77 +++++++++++++++++------------- pipeline_InLoc.ipynb | 26 +++++----- pipeline_SfM.ipynb | 32 ++++++++----- 5 files changed, 123 insertions(+), 83 deletions(-) diff --git a/.github/workflows/code-quality.yml b/.github/workflows/code-quality.yml index e3327ffb..34ad1ee1 100644 --- a/.github/workflows/code-quality.yml +++ b/.github/workflows/code-quality.yml @@ -20,5 +20,5 @@ jobs: - run: python -m pip install --upgrade pip - run: python -m pip install black flake8 isort - run: python -m flake8 hloc - - run: python -m isort hloc --check-only --diff - - run: python -m black hloc --check --diff + - run: python -m isort hloc *.ipynb --check-only --diff + - run: python -m black hloc *.ipynb --check --diff diff --git a/demo.ipynb b/demo.ipynb index 28f30608..c4e50480 100644 --- a/demo.ipynb +++ b/demo.ipynb @@ -18,11 +18,18 @@ "%load_ext autoreload\n", "%autoreload 2\n", "import tqdm, tqdm.notebook\n", + "\n", "tqdm.tqdm = tqdm.notebook.tqdm # notebook-friendly progress bars\n", "from pathlib import Path\n", "import numpy as np\n", "\n", - "from hloc import extract_features, match_features, reconstruction, visualization, pairs_from_exhaustive\n", + "from hloc import (\n", + " extract_features,\n", + " match_features,\n", + " reconstruction,\n", + " visualization,\n", + " pairs_from_exhaustive,\n", + ")\n", "from hloc.visualization import plot_images, read_image\n", "from hloc.utils import viz_3d" ] @@ -43,17 +50,17 @@ "metadata": {}, "outputs": [], "source": [ - "images = Path('datasets/sacre_coeur')\n", - "outputs = Path('outputs/demo/')\n", + "images = Path(\"datasets/sacre_coeur\")\n", + "outputs = Path(\"outputs/demo/\")\n", "!rm -rf $outputs\n", - "sfm_pairs = outputs / 'pairs-sfm.txt'\n", - "loc_pairs = outputs / 'pairs-loc.txt'\n", - "sfm_dir = outputs / 'sfm'\n", - "features = outputs / 'features.h5'\n", - "matches = outputs / 'matches.h5'\n", + "sfm_pairs = outputs / \"pairs-sfm.txt\"\n", + "loc_pairs = outputs / \"pairs-loc.txt\"\n", + "sfm_dir = outputs / \"sfm\"\n", + "features = outputs / \"features.h5\"\n", + "matches = outputs / \"matches.h5\"\n", "\n", - "feature_conf = extract_features.confs['disk']\n", - "matcher_conf = match_features.confs['disk+lightglue']" + "feature_conf = extract_features.confs[\"disk\"]\n", + "matcher_conf = match_features.confs[\"disk+lightglue\"]" ] }, { @@ -90,7 +97,7 @@ } ], "source": [ - "references = [p.relative_to(images).as_posix() for p in (images / 'mapping/').iterdir()]\n", + "references = [p.relative_to(images).as_posix() for p in (images / \"mapping/\").iterdir()]\n", "print(len(references), \"mapping images\")\n", "plot_images([read_image(images / r) for r in references], dpi=25)" ] @@ -142,7 +149,9 @@ } ], "source": [ - "extract_features.main(feature_conf, images, image_list=references, feature_path=features)\n", + "extract_features.main(\n", + " feature_conf, images, image_list=references, feature_path=features\n", + ")\n", "pairs_from_exhaustive.main(sfm_pairs, image_list=references)\n", "match_features.main(matcher_conf, sfm_pairs, features=features, matches=matches);" ] @@ -162,9 +171,13 @@ "metadata": {}, "outputs": [], "source": [ - "model = reconstruction.main(sfm_dir, images, sfm_pairs, features, matches, image_list=references)\n", + "model = reconstruction.main(\n", + " sfm_dir, images, sfm_pairs, features, matches, image_list=references\n", + ")\n", "fig = viz_3d.init_figure()\n", - "viz_3d.plot_reconstruction(fig, model, color='rgba(255,0,0,0.5)', name=\"mapping\", points_rgb=True)\n", + "viz_3d.plot_reconstruction(\n", + " fig, model, color=\"rgba(255,0,0,0.5)\", name=\"mapping\", points_rgb=True\n", + ")\n", "fig.show()" ] }, @@ -204,7 +217,7 @@ } ], "source": [ - "visualization.visualize_sfm_2d(model, images, color_by='visibility', n=2)" + "visualization.visualize_sfm_2d(model, images, color_by=\"visibility\", n=2)" ] }, { @@ -238,7 +251,7 @@ "# try other queries by uncommenting their url\n", "# url = \"https://upload.wikimedia.org/wikipedia/commons/5/59/Basilique_du_Sacr%C3%A9-C%C5%93ur_%285430392880%29.jpg\"\n", "# url = \"https://upload.wikimedia.org/wikipedia/commons/8/8e/Sacr%C3%A9_C%C5%93ur_at_night%21_%285865355326%29.jpg\"\n", - "query = 'query/night.jpg'\n", + "query = \"query/night.jpg\"\n", "!mkdir -p $images/query && wget $url -O $images/$query -q\n", "plot_images([read_image(images / query)], dpi=75)" ] @@ -290,9 +303,13 @@ } ], "source": [ - "extract_features.main(feature_conf, images, image_list=[query], feature_path=features, overwrite=True)\n", + "extract_features.main(\n", + " feature_conf, images, image_list=[query], feature_path=features, overwrite=True\n", + ")\n", "pairs_from_exhaustive.main(loc_pairs, image_list=[query], ref_list=references)\n", - "match_features.main(matcher_conf, loc_pairs, features=features, matches=matches, overwrite=True);" + "match_features.main(\n", + " matcher_conf, loc_pairs, features=features, matches=matches, overwrite=True\n", + ");" ] }, { @@ -344,8 +361,8 @@ "camera = pycolmap.infer_camera_from_image(images / query)\n", "ref_ids = [model.find_image_with_name(r).image_id for r in references]\n", "conf = {\n", - " 'estimation': {'ransac': {'max_error': 12}},\n", - " 'refinement': {'refine_focal_length': True, 'refine_extra_params': True},\n", + " \"estimation\": {\"ransac\": {\"max_error\": 12}},\n", + " \"refinement\": {\"refine_focal_length\": True, \"refine_extra_params\": True},\n", "}\n", "localizer = QueryLocalizer(model, conf)\n", "ret, log = pose_from_cluster(localizer, query, camera, ref_ids, features, matches)\n", @@ -369,10 +386,14 @@ "metadata": {}, "outputs": [], "source": [ - "pose = pycolmap.Image(tvec=ret['tvec'], qvec=ret['qvec'])\n", - "viz_3d.plot_camera_colmap(fig, pose, camera, color='rgba(0,255,0,0.5)', name=query, fill=True)\n", + "pose = pycolmap.Image(tvec=ret[\"tvec\"], qvec=ret[\"qvec\"])\n", + "viz_3d.plot_camera_colmap(\n", + " fig, pose, camera, color=\"rgba(0,255,0,0.5)\", name=query, fill=True\n", + ")\n", "# visualize 2D-3D correspodences\n", - "inl_3d = np.array([model.points3D[pid].xyz for pid in np.array(log['points3D_ids'])[ret['inliers']]])\n", + "inl_3d = np.array(\n", + " [model.points3D[pid].xyz for pid in np.array(log[\"points3D_ids\"])[ret[\"inliers\"]]]\n", + ")\n", "viz_3d.plot_points(fig, inl_3d, color=\"lime\", ps=1, name=query)\n", "fig.show()" ] diff --git a/pipeline_Aachen.ipynb b/pipeline_Aachen.ipynb index d288b1f2..b1f581cf 100644 --- a/pipeline_Aachen.ipynb +++ b/pipeline_Aachen.ipynb @@ -12,7 +12,12 @@ "from pathlib import Path\n", "from pprint import pformat\n", "\n", - "from hloc import extract_features, match_features, pairs_from_covisibility, pairs_from_retrieval\n", + "from hloc import (\n", + " extract_features,\n", + " match_features,\n", + " pairs_from_covisibility,\n", + " pairs_from_retrieval,\n", + ")\n", "from hloc import colmap_from_nvm, triangulation, localize_sfm, visualization" ] }, @@ -93,18 +98,18 @@ } ], "source": [ - "dataset = Path('datasets/aachen/') # change this if your dataset is somewhere else\n", - "images = dataset / 'images_upright/'\n", + "dataset = Path(\"datasets/aachen/\") # change this if your dataset is somewhere else\n", + "images = dataset / \"images_upright/\"\n", "\n", - "outputs = Path('outputs/aachen/') # where everything will be saved\n", - "sfm_pairs = outputs / 'pairs-db-covis20.txt' # top 20 most covisible in SIFT model\n", - "loc_pairs = outputs / 'pairs-query-netvlad20.txt' # top 20 retrieved by NetVLAD\n", - "reference_sfm = outputs / 'sfm_superpoint+superglue' # the SfM model we will build\n", - "results = outputs / 'Aachen_hloc_superpoint+superglue_netvlad20.txt' # the result file\n", + "outputs = Path(\"outputs/aachen/\") # where everything will be saved\n", + "sfm_pairs = outputs / \"pairs-db-covis20.txt\" # top 20 most covisible in SIFT model\n", + "loc_pairs = outputs / \"pairs-query-netvlad20.txt\" # top 20 retrieved by NetVLAD\n", + "reference_sfm = outputs / \"sfm_superpoint+superglue\" # the SfM model we will build\n", + "results = outputs / \"Aachen_hloc_superpoint+superglue_netvlad20.txt\" # the result file\n", "\n", "# list the standard configurations available\n", - "print(f'Configs for feature extractors:\\n{pformat(extract_features.confs)}')\n", - "print(f'Configs for feature matchers:\\n{pformat(match_features.confs)}')" + "print(f\"Configs for feature extractors:\\n{pformat(extract_features.confs)}\")\n", + "print(f\"Configs for feature matchers:\\n{pformat(match_features.confs)}\")" ] }, { @@ -115,9 +120,9 @@ "source": [ "# pick one of the configurations for image retrieval, local feature extraction, and matching\n", "# you can also simply write your own here!\n", - "retrieval_conf = extract_features.confs['netvlad']\n", - "feature_conf = extract_features.confs['superpoint_aachen']\n", - "matcher_conf = match_features.confs['superglue']" + "retrieval_conf = extract_features.confs[\"netvlad\"]\n", + "feature_conf = extract_features.confs[\"superpoint_aachen\"]\n", + "matcher_conf = match_features.confs[\"superglue\"]" ] }, { @@ -158,13 +163,13 @@ "outputs": [], "source": [ "colmap_from_nvm.main(\n", - " dataset / '3D-models/aachen_cvpr2018_db.nvm',\n", - " dataset / '3D-models/database_intrinsics.txt',\n", - " dataset / 'aachen.db',\n", - " outputs / 'sfm_sift')\n", + " dataset / \"3D-models/aachen_cvpr2018_db.nvm\",\n", + " dataset / \"3D-models/database_intrinsics.txt\",\n", + " dataset / \"aachen.db\",\n", + " outputs / \"sfm_sift\",\n", + ")\n", "\n", - "pairs_from_covisibility.main(\n", - " outputs / 'sfm_sift', sfm_pairs, num_matched=20)" + "pairs_from_covisibility.main(outputs / \"sfm_sift\", sfm_pairs, num_matched=20)" ] }, { @@ -180,7 +185,9 @@ "metadata": {}, "outputs": [], "source": [ - "sfm_matches = match_features.main(matcher_conf, sfm_pairs, feature_conf['output'], outputs)" + "sfm_matches = match_features.main(\n", + " matcher_conf, sfm_pairs, feature_conf[\"output\"], outputs\n", + ")" ] }, { @@ -205,12 +212,8 @@ "outputs": [], "source": [ "reconstruction = triangulation.main(\n", - " reference_sfm,\n", - " outputs / 'sfm_sift',\n", - " images,\n", - " sfm_pairs,\n", - " features,\n", - " sfm_matches)" + " reference_sfm, outputs / \"sfm_sift\", images, sfm_pairs, features, sfm_matches\n", + ")" ] }, { @@ -228,7 +231,9 @@ "outputs": [], "source": [ "global_descriptors = extract_features.main(retrieval_conf, images, outputs)\n", - "pairs_from_retrieval.main(global_descriptors, loc_pairs, num_matched=20, db_prefix=\"db\", query_prefix=\"query\")" + "pairs_from_retrieval.main(\n", + " global_descriptors, loc_pairs, num_matched=20, db_prefix=\"db\", query_prefix=\"query\"\n", + ")" ] }, { @@ -244,7 +249,9 @@ "metadata": {}, "outputs": [], "source": [ - "loc_matches = match_features.main(matcher_conf, loc_pairs, feature_conf['output'], outputs)" + "loc_matches = match_features.main(\n", + " matcher_conf, loc_pairs, feature_conf[\"output\"], outputs\n", + ")" ] }, { @@ -263,12 +270,13 @@ "source": [ "localize_sfm.main(\n", " reconstruction,\n", - " dataset / 'queries/*_time_queries_with_intrinsics.txt',\n", + " dataset / \"queries/*_time_queries_with_intrinsics.txt\",\n", " loc_pairs,\n", " features,\n", " loc_matches,\n", " results,\n", - " covisibility_clustering=False) # not required with SuperPoint+SuperGlue" + " covisibility_clustering=False,\n", + ") # not required with SuperPoint+SuperGlue" ] }, { @@ -303,7 +311,7 @@ } ], "source": [ - "visualization.visualize_sfm_2d(reconstruction, images, n=1, color_by='track_length')" + "visualization.visualize_sfm_2d(reconstruction, images, n=1, color_by=\"track_length\")" ] }, { @@ -319,7 +327,7 @@ "metadata": {}, "outputs": [], "source": [ - "visualization.visualize_sfm_2d(reconstruction, images, n=1, color_by='visibility')" + "visualization.visualize_sfm_2d(reconstruction, images, n=1, color_by=\"visibility\")" ] }, { @@ -346,7 +354,7 @@ } ], "source": [ - "visualization.visualize_sfm_2d(reconstruction, images, n=1, color_by='depth')" + "visualization.visualize_sfm_2d(reconstruction, images, n=1, color_by=\"depth\")" ] }, { @@ -375,7 +383,8 @@ ], "source": [ "visualization.visualize_loc(\n", - " results, images, reconstruction, n=1, top_k_db=1, prefix='query/night', seed=2)" + " results, images, reconstruction, n=1, top_k_db=1, prefix=\"query/night\", seed=2\n", + ")" ] } ], diff --git a/pipeline_InLoc.ipynb b/pipeline_InLoc.ipynb index 4c421619..8fc7efde 100644 --- a/pipeline_InLoc.ipynb +++ b/pipeline_InLoc.ipynb @@ -36,13 +36,13 @@ "metadata": {}, "outputs": [], "source": [ - "dataset = Path('datasets/inloc/') # change this if your dataset is somewhere else\n", + "dataset = Path(\"datasets/inloc/\") # change this if your dataset is somewhere else\n", "\n", - "pairs = Path('pairs/inloc/')\n", - "loc_pairs = pairs / 'pairs-query-netvlad40.txt' # top 40 retrieved by NetVLAD\n", + "pairs = Path(\"pairs/inloc/\")\n", + "loc_pairs = pairs / \"pairs-query-netvlad40.txt\" # top 40 retrieved by NetVLAD\n", "\n", - "outputs = Path('outputs/inloc/') # where everything will be saved\n", - "results = outputs / 'InLoc_hloc_superpoint+superglue_netvlad40.txt' # the result file" + "outputs = Path(\"outputs/inloc/\") # where everything will be saved\n", + "results = outputs / \"InLoc_hloc_superpoint+superglue_netvlad40.txt\" # the result file" ] }, { @@ -52,8 +52,8 @@ "outputs": [], "source": [ "# list the standard configurations available\n", - "print(f'Configs for feature extractors:\\n{pformat(extract_features.confs)}')\n", - "print(f'Configs for feature matchers:\\n{pformat(match_features.confs)}')" + "print(f\"Configs for feature extractors:\\n{pformat(extract_features.confs)}\")\n", + "print(f\"Configs for feature matchers:\\n{pformat(match_features.confs)}\")" ] }, { @@ -64,8 +64,8 @@ "source": [ "# pick one of the configurations for extraction and matching\n", "# you can also simply write your own here!\n", - "feature_conf = extract_features.confs['superpoint_inloc']\n", - "matcher_conf = match_features.confs['superglue']" + "feature_conf = extract_features.confs[\"superpoint_inloc\"]\n", + "matcher_conf = match_features.confs[\"superglue\"]" ] }, { @@ -98,7 +98,9 @@ "metadata": {}, "outputs": [], "source": [ - "match_path = match_features.main(matcher_conf, loc_pairs, feature_conf['output'], outputs)" + "match_path = match_features.main(\n", + " matcher_conf, loc_pairs, feature_conf[\"output\"], outputs\n", + ")" ] }, { @@ -116,8 +118,8 @@ "outputs": [], "source": [ "localize_inloc.main(\n", - " dataset, loc_pairs, feature_path, match_path, results,\n", - " skip_matches=20) # skip database images with too few matches" + " dataset, loc_pairs, feature_path, match_path, results, skip_matches=20\n", + ") # skip database images with too few matches" ] }, { diff --git a/pipeline_SfM.ipynb b/pipeline_SfM.ipynb index 473732bb..28c36cb1 100644 --- a/pipeline_SfM.ipynb +++ b/pipeline_SfM.ipynb @@ -11,7 +11,13 @@ "\n", "from pathlib import Path\n", "\n", - "from hloc import extract_features, match_features, reconstruction, visualization, pairs_from_retrieval" + "from hloc import (\n", + " extract_features,\n", + " match_features,\n", + " reconstruction,\n", + " visualization,\n", + " pairs_from_retrieval,\n", + ")" ] }, { @@ -28,15 +34,15 @@ "metadata": {}, "outputs": [], "source": [ - "images = Path('datasets/South-Building/images/')\n", + "images = Path(\"datasets/South-Building/images/\")\n", "\n", - "outputs = Path('outputs/sfm/')\n", - "sfm_pairs = outputs / 'pairs-netvlad.txt'\n", - "sfm_dir = outputs / 'sfm_superpoint+superglue'\n", + "outputs = Path(\"outputs/sfm/\")\n", + "sfm_pairs = outputs / \"pairs-netvlad.txt\"\n", + "sfm_dir = outputs / \"sfm_superpoint+superglue\"\n", "\n", - "retrieval_conf = extract_features.confs['netvlad']\n", - "feature_conf = extract_features.confs['superpoint_aachen']\n", - "matcher_conf = match_features.confs['superglue']" + "retrieval_conf = extract_features.confs[\"netvlad\"]\n", + "feature_conf = extract_features.confs[\"superpoint_aachen\"]\n", + "matcher_conf = match_features.confs[\"superglue\"]" ] }, { @@ -90,7 +96,9 @@ "outputs": [], "source": [ "feature_path = extract_features.main(feature_conf, images, outputs)\n", - "match_path = match_features.main(matcher_conf, sfm_pairs, feature_conf['output'], outputs)" + "match_path = match_features.main(\n", + " matcher_conf, sfm_pairs, feature_conf[\"output\"], outputs\n", + ")" ] }, { @@ -124,7 +132,7 @@ "metadata": {}, "outputs": [], "source": [ - "visualization.visualize_sfm_2d(model, images, color_by='visibility', n=5)" + "visualization.visualize_sfm_2d(model, images, color_by=\"visibility\", n=5)" ] }, { @@ -133,7 +141,7 @@ "metadata": {}, "outputs": [], "source": [ - "visualization.visualize_sfm_2d(model, images, color_by='track_length', n=5)" + "visualization.visualize_sfm_2d(model, images, color_by=\"track_length\", n=5)" ] }, { @@ -142,7 +150,7 @@ "metadata": {}, "outputs": [], "source": [ - "visualization.visualize_sfm_2d(model, images, color_by='depth', n=5)" + "visualization.visualize_sfm_2d(model, images, color_by=\"depth\", n=5)" ] } ], From 5fe9ebae0e7ef9d360166121005da21157f1c043 Mon Sep 17 00:00:00 2001 From: Paul-Edouard Sarlin Date: Fri, 12 Jan 2024 12:07:55 +0100 Subject: [PATCH 9/9] Add missing ignores --- .gitignore | 2 ++ 1 file changed, 2 insertions(+) diff --git a/.gitignore b/.gitignore index c363cd22..1bb733be 100644 --- a/.gitignore +++ b/.gitignore @@ -1,3 +1,4 @@ +build .venv __pycache__ *.pyc @@ -6,3 +7,4 @@ __pycache__ outputs/ datasets/* !datasets/sacre_coeur/ +datasets/sacre_coeur/query