├── .gitignore ├── .pre-commit-config.yaml ├── README.md ├── args.py ├── datasets.py ├── demo.ipynb ├── dump_data.py ├── dump_datasets ├── __init__.py ├── eth3d.py ├── example.py ├── inloc.py ├── megadepth.py └── phototourism.py ├── evaluate_utils.py ├── gluefactory ├── __init__.py ├── configs │ ├── aliked+NN.yaml │ ├── aliked+lightglue-official.yaml │ ├── aliked+lightglue_homography.yaml │ ├── aliked+lightglue_megadepth.yaml │ ├── disk+NN.yaml │ ├── disk+lightglue-official.yaml │ ├── disk+lightglue_homography.yaml │ ├── disk+lightglue_megadepth.yaml │ ├── sift+NN.yaml │ ├── sift+lightglue-official.yaml │ ├── sift+lightglue_homography.yaml │ ├── sift+lightglue_megadepth.yaml │ ├── superpoint+NN.yaml │ ├── superpoint+lightglue-official.yaml │ ├── superpoint+lightglue_homography.yaml │ ├── superpoint+lightglue_megadepth.yaml │ ├── superpoint+lsd+gluestick-homography.yaml │ ├── superpoint+lsd+gluestick-megadepth.yaml │ ├── superpoint+lsd+gluestick.yaml │ ├── superpoint+superglue-official.yaml │ ├── superpoint-open+NN.yaml │ ├── superpoint-open+lightglue_homography.yaml │ └── superpoint-open+lightglue_megadepth.yaml ├── datasets │ ├── __init__.py │ ├── augmentations.py │ ├── base_dataset.py │ ├── eth3d.py │ ├── homographies.py │ ├── hpatches.py │ ├── image_folder.py │ ├── image_pairs.py │ ├── megadepth.py │ ├── megadepth_easy.py │ ├── megadepth_scene_lists │ │ ├── loftr_train_list.txt │ │ ├── loftr_val_list.txt │ │ ├── test_list.txt │ │ ├── test_pairs.txt │ │ ├── test_scenes_clean.txt │ │ ├── train_hard_list.txt │ │ ├── train_pairs.txt │ │ ├── train_scenes.txt │ │ ├── train_scenes_clean.txt │ │ ├── valid_pairs.txt │ │ ├── valid_scenes.txt │ │ └── valid_scenes_clean.txt │ └── utils.py ├── eval │ ├── __init__.py │ ├── eth3d.py │ ├── eval_pipeline.py │ ├── hpatches.py │ ├── inspect.py │ ├── io.py │ ├── megadepth1500.py │ └── utils.py ├── geometry │ ├── __init__.py │ ├── depth.py │ ├── epipolar.py │ ├── gt_generation.py │ ├── homography.py │ ├── utils.py │ └── wrappers.py ├── models │ ├── __init__.py │ ├── backbones │ │ ├── __init__.py │ │ ├── dino.py │ │ └── dinov2.py │ ├── base_model.py │ ├── cache_loader.py │ ├── cosplace.py │ ├── extractors │ │ ├── __init__.py │ │ ├── aliked.py │ │ ├── disk_kornia.py │ │ ├── grid_extractor.py │ │ ├── keynet_affnet_hardnet.py │ │ ├── mixed.py │ │ ├── sift.py │ │ ├── sift_kornia.py │ │ └── superpoint_open.py │ ├── lines │ │ ├── __init__.py │ │ ├── deeplsd.py │ │ ├── lsd.py │ │ └── wireframe.py │ ├── matchers │ │ ├── __init__.py │ │ ├── adalam.py │ │ ├── depth_matcher.py │ │ ├── gluestick.py │ │ ├── homography_matcher.py │ │ ├── kornia_loftr.py │ │ ├── lightglue.py │ │ ├── lightglue_pretrained.py │ │ └── nearest_neighbor_matcher.py │ ├── netvlad.py │ ├── overlap_predictor.py │ ├── triplet_pipeline.py │ ├── two_view_pipeline.py │ └── utils │ │ ├── __init__.py │ │ ├── losses.py │ │ ├── metrics.py │ │ └── misc.py ├── robust_estimators │ ├── __init__.py │ ├── base_estimator.py │ ├── homography │ │ ├── __init__.py │ │ ├── homography_est.py │ │ ├── opencv.py │ │ └── poselib.py │ └── relative_pose │ │ ├── __init__.py │ │ ├── opencv.py │ │ ├── poselib.py │ │ └── pycolmap.py ├── scripts │ ├── __init__.py │ ├── export_local_features.py │ └── export_megadepth.py ├── settings.py ├── train.py ├── utils │ ├── __init__.py │ ├── benchmark.py │ ├── experiments.py │ ├── export_predictions.py │ ├── image.py │ ├── misc.py │ ├── patch_helper.py │ ├── patches.py │ ├── stdout_capturing.py │ ├── tensor.py │ └── tools.py └── visualization │ ├── global_frame.py │ ├── tools.py │ ├── two_view_frame.py │ ├── visualize_batch.py │ └── viz2d.py ├── inloc_localization.py ├── register.py ├── relative_pose.py ├── retrieve.py ├── train_configs ├── best.yaml └── best_easy.yaml └── utils.py /.gitignore: -------------------------------------------------------------------------------- 1 | .venv 2 | /build/ 3 | *.egg-info 4 | *.pyc 5 | /.idea/ 6 | /venv/ 7 | /data/ 8 | /outputs/ 9 | __pycache__ 10 | /logs/ 11 | /configs/ 12 | dumped_data/ 13 | dump_datasets/data_dirs.yaml 14 | *.batch 15 | logs/ 16 | demo/ 17 | out* 18 | -------------------------------------------------------------------------------- /.pre-commit-config.yaml: -------------------------------------------------------------------------------- 1 | ci: 2 | autofix_prs: true 3 | autoupdate_commit_msg: '[pre-commit.ci] pre-commit suggestions' 4 | # submodules: true 5 | 6 | repos: 7 | - repo: https://github.com/pre-commit/pre-commit-hooks 8 | rev: v4.4.0 9 | hooks: 10 | - id: end-of-file-fixer 11 | - id: trailing-whitespace 12 | - id: check-yaml 13 | - id: requirements-txt-fixer 14 | - id: pretty-format-json 15 | exclude: ^(\.devcontainer/|demo\.ipynb)$ 16 | 17 | # - repo: https://github.com/tox-dev/pyproject-fmt 18 | # rev: "1.1.0" 19 | # hooks: 20 | # - id: pyproject-fmt 21 | 22 | # - repo: https://github.com/astral-sh/ruff-pre-commit 23 | # rev: v0.0.286 24 | # hooks: 25 | # - id: ruff 26 | # args: [--fix, --exit-non-zero-on-fix] 27 | 28 | 29 | - repo: https://github.com/codespell-project/codespell 30 | rev: v2.2.5 31 | hooks: 32 | - id: codespell 33 | args: 34 | - --ignore-words-list 35 | - "ans,hist,laf" 36 | - --skip 37 | #- "*.bib,*.ipynb" 38 | exclude: ^.*\.ipynb$ 39 | 40 | 41 | - repo: https://github.com/PyCQA/docformatter 42 | rev: v1.7.5 43 | hooks: 44 | - id: docformatter 45 | args: [--in-place, --wrap-summaries=115, --wrap-descriptions=120] 46 | exclude: ^.*\.ipynb$ 47 | -------------------------------------------------------------------------------- /args.py: -------------------------------------------------------------------------------- 1 | import argparse 2 | import torch 3 | from pathlib import Path 4 | 5 | def create_parser(): 6 | 7 | parser = argparse.ArgumentParser() 8 | # dump data 9 | parser.add_argument('--model', '-m', default='best', 10 | help='The name of the model to be train/test.') 11 | parser.add_argument('--dataset_dir', '-dir', type=Path, default="data/ETH3D_undistorted") 12 | parser.add_argument('--dump_dir', '-dump', type=Path, default="dumped_data", 13 | help='dir to save the features.') 14 | parser.add_argument('--dataset', '-ds', default="eth3d", choices=['all', 'aachen', 'pitts', 'inloc', 'megadepth', 'eth3d', 'phototourism', 'imc2023'], 15 | help='dataset name.') 16 | parser.add_argument('--imsize', '-im', type=int, default=224, 17 | help='The resized image shape.') 18 | parser.add_argument('--overwrite', '-ow', action='store_true', default=False, 19 | help='overwrite the dump data.') 20 | 21 | # type 22 | parser.add_argument('--device', '-device', default='cuda') 23 | parser.add_argument('--dtype', '-dtype', default=torch.float32) 24 | 25 | # retrieval 26 | parser.add_argument('--radius', '-r', type=float, default=-1, 27 | help='radius for radius knn search, default set it as -1 to \ 28 | trigger the procedure of automatically compute the median similarity over 100 samples as radius.') 29 | parser.add_argument('--cls', '-cls', action='store_true', default=True, 30 | help='use CLS tokens as prefilter.') 31 | parser.add_argument('--pre_filter', '-pre', type=int, default=20, 32 | help='how many db images prefiltered for reranking') 33 | parser.add_argument('--weighted', '-w', action='store_true', default=True, 34 | help='use TF-IDF weights for voting scores.') 35 | parser.add_argument('--k', '-k', type=int, default=1, 36 | help='top-k retrievals') 37 | parser.add_argument('--batch_size', '-bs', type=int, default=128) 38 | parser.add_argument('--num_workers', '-nw', type=int, default=8) 39 | parser.add_argument('--output_dir', '-out', type=Path, default='outputs') 40 | parser.add_argument('--radius_round', '-round', default='round', choices=['round, ceil'], 41 | help='vote methods') 42 | 43 | # localization 44 | parser.add_argument('--method', default=None) 45 | parser.add_argument('--num_loc', type=int, default=40, 46 | help='Number of image pairs for loc, default: %(default)s') 47 | parser.add_argument('--loc_pairs', type=Path, default="outputs/inloc/09/cls_100/top40_overlap_pairs_w_auc.txt", 48 | help='Number of image pairs for loc, default: %(default)s') 49 | 50 | return parser.parse_args() 51 | -------------------------------------------------------------------------------- /datasets.py: -------------------------------------------------------------------------------- 1 | import math 2 | import os 3 | import cv2 4 | import h5py 5 | import torch 6 | import numpy as np 7 | import torch.utils.data as data 8 | from pathlib import Path 9 | from lightglue.utils import load_image 10 | from gluefactory.datasets.eth3d import read_cameras, qvec2rotmat 11 | from gluefactory.datasets.utils import scale_intrinsics 12 | from collections import defaultdict 13 | from torch.utils.data import Dataset 14 | 15 | 16 | 17 | class SampleDataset(Dataset): 18 | def __init__(self, data, image_list, dataset='megadepth'): 19 | 20 | self.overlap_features = data 21 | self.image_list = image_list 22 | self.dataset = dataset 23 | 24 | def __len__(self): 25 | return len(self.image_list) 26 | 27 | def __getitem__(self, idx): 28 | """Retrieve a sample from the dataset at the given index. 29 | 30 | Args: 31 | idx (int): Index of the sample to retrieve. 32 | 33 | Returns: 34 | sample (dict): A dictionary containing the input and target data for the sample. 35 | """ 36 | data = defaultdict(list) 37 | 38 | with h5py.File(str(self.overlap_features), 'r') as hfile: 39 | for k in hfile[self.image_list[idx]].keys(): 40 | data[k] = [torch.Tensor(hfile[self.image_list[idx]][k].__array__())] 41 | 42 | input_data0 = {k+'0': torch.concat(data[k]) for k in data.keys()} 43 | input_data1 = {k+'1': torch.concat(data[k]) for k in data.keys()} 44 | return {**input_data0, **input_data1} 45 | -------------------------------------------------------------------------------- /dump_data.py: -------------------------------------------------------------------------------- 1 | import os 2 | import torch 3 | import importlib 4 | import gluefactory 5 | from omegaconf import OmegaConf 6 | from args import * 7 | from utils import * 8 | 9 | opt = create_parser() 10 | # check if gpu device is available 11 | opt.device = torch.device('cuda:0' if torch.cuda.is_available() and opt.device != 'cpu' else 'cpu') 12 | # load the data dirs for all 13 | dataset_dirs = OmegaConf.load('dump_datasets/data_dirs.yaml') 14 | 15 | def dump_data(opt, model, dataset): 16 | print(f"----------Start dumping the dataset: {dataset} on {opt.device}------") 17 | opt.dataset = dataset 18 | 19 | try: 20 | module = importlib.import_module(f"dump_datasets.{dataset}") 21 | module.dump(opt, model) 22 | except ImportError: 23 | print(f"Unknown dataset: {dataset}") 24 | 25 | print(f"----------Finished------") 26 | 27 | if opt.dataset == 'all': 28 | datasets = ['inloc', 'megadepth', 'eth3d', 'phototourism'] 29 | else: 30 | datasets = [opt.dataset] 31 | 32 | for dataset in datasets: 33 | # fetch the data paths from yaml 34 | opt.dataset_dir = dataset_dirs.get('dataset_dirs')[dataset] 35 | if not Path(opt.dataset_dir).exists(): 36 | raise FileNotFoundError('dataset_dir is hardcoded, please adjust it to yours') 37 | if not os.path.exists(f"outputs/training/{opt.model}"): 38 | download_best(opt.model) 39 | modelconf = {} 40 | model = gluefactory.load_experiment(opt.model, conf=modelconf).to(opt.device).eval() 41 | dump_data(opt, model, dataset) 42 | -------------------------------------------------------------------------------- /dump_datasets/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/weitong8591/vop/e964dc779958e691d8511b78dd3c40bf7920e4e3/dump_datasets/__init__.py -------------------------------------------------------------------------------- /dump_datasets/eth3d.py: -------------------------------------------------------------------------------- 1 | 2 | import os 3 | import h5py 4 | import torch 5 | import numpy as np 6 | from tqdm import tqdm 7 | from pathlib import Path 8 | from evaluate_utils import read_cameras 9 | from lightglue.utils import load_image 10 | from gluefactory.datasets.utils import scale_intrinsics 11 | from gluefactory.datasets.eth3d import qvec2rotmat 12 | from gluefactory.datasets.utils import scale_intrinsics 13 | 14 | 15 | def dump(opt, model): 16 | scenes = os.listdir(opt.dataset_dir) 17 | overlap_features = Path(opt.dump_dir) / 'eth3d' / 'overlap_feats.h5' 18 | (Path(opt.dump_dir) / 'eth3d').mkdir(exist_ok=True, parents=True) 19 | 20 | if not os.path.exists(overlap_features) or opt.overwrite: 21 | with h5py.File(str(overlap_features), 'w') as hfile: 22 | for scene in scenes: 23 | image_list = [image for image in os.listdir(Path(opt.dataset_dir) / scene / "images/dslr_images_undistorted")] 24 | cameras = read_cameras(str(Path(opt.dataset_dir, scene, "dslr_calibration_undistorted", "cameras.txt")),1,) 25 | name_to_cam_idx = {name: {} for name in image_list} 26 | with open(str(Path(opt.dataset_dir, scene, "dslr_calibration_jpg", "images.txt")), "r") as f: 27 | raw_data = f.read().rstrip().split("\n")[4::2] 28 | for raw_line in raw_data: 29 | line = raw_line.split(" ") 30 | img_name = os.path.basename(line[-1]) 31 | name_to_cam_idx[img_name]["dist_camera_idx"] = int(line[-2]) 32 | T_world_to_camera = {} 33 | image_visible_points3D = {} 34 | with open(str(Path(opt.dataset_dir, scene, "dslr_calibration_undistorted", "images.txt")), "r") as f: 35 | lines = f.readlines()[4:] # Skip the header 36 | raw_poses = [line.strip("\n").split(" ") for line in lines[::2]] 37 | raw_points = [line.strip("\n").split(" ") for line in lines[1::2]] 38 | for raw_pose, raw_pts in zip(raw_poses, raw_points): 39 | img_name = os.path.basename(raw_pose[-1]) 40 | # Extract the transform from world to camera 41 | target_extrinsics = list(map(float, raw_pose[1:8])) 42 | pose = np.eye(4, dtype=np.float32) 43 | pose[:3, :3] = qvec2rotmat(target_extrinsics[:4]) 44 | pose[:3, 3] = target_extrinsics[4:] 45 | T_world_to_camera[img_name] = pose 46 | name_to_cam_idx[img_name]["undist_camera_idx"] = int(raw_pose[-2]) 47 | # Extract the visible 3D points 48 | point3D_ids = [id for id in map(int, raw_pts[2::3]) if id != -1] 49 | image_visible_points3D[img_name] = set(point3D_ids) 50 | 51 | for image_path in tqdm(image_list): 52 | img_file = scene / Path("images/dslr_images_undistorted") / image_path 53 | image_ = load_image(Path(opt.dataset_dir) / img_file, resize=opt.imsize) 54 | c, h, w = image_.shape 55 | image = torch.zeros((3, opt.imsize, opt.imsize), device=opt.device, dtype=opt.dtype) 56 | image[:, :h, :w] = image_ 57 | feats = model.extractor({'image': image[None]}) 58 | group = hfile.create_group(str(img_file)) 59 | for key in ['keypoints', 'descriptors', 'global_descriptor']: 60 | group.create_dataset(key, data=feats[key][0].detach().cpu().numpy()) 61 | 62 | ori_h, ori_w = load_image(Path(opt.dataset_dir) / scene / Path("images/dslr_images_undistorted") / image_path).shape[-2:] 63 | original_image_size = np.array([ori_w, ori_h]) 64 | scales = opt.imsize/original_image_size 65 | ori_K = cameras[name_to_cam_idx[image_path]["dist_camera_idx"]], 66 | K = scale_intrinsics(ori_K[0], scales) 67 | 68 | group.create_dataset("K", data=K) 69 | group.create_dataset("original_K", data=ori_K[0]) 70 | group.create_dataset("T_w2cam", data=T_world_to_camera[image_path]) 71 | group.create_dataset('image_size', data=[w, h]) 72 | 73 | print(f"{len(image_list)} images in {scene}") 74 | print(f"finished." ) 75 | -------------------------------------------------------------------------------- /dump_datasets/example.py: -------------------------------------------------------------------------------- 1 | 2 | import os 3 | import h5py 4 | import torch 5 | import gluefactory 6 | import numpy as np 7 | from tqdm import tqdm 8 | from pathlib import Path 9 | from args import * 10 | from lightglue.utils import load_image 11 | 12 | def dump(opt): 13 | 14 | # path to save the preprocessed data 15 | overlap_features = Path(opt.dump_dir)/ 'example' / 'overlap_feats.h5' 16 | (Path(opt.dump_dir) / 'example').mkdir(exist_ok=True, parents=True) 17 | modelconf = {} 18 | 19 | # load a model with frozen DINOv2 backbone, for patchfying the imgaes 20 | model = gluefactory.load_experiment(opt.model, conf=modelconf).cuda().eval() 21 | if not os.path.exists(overlap_features) or opt.overwrite: 22 | with h5py.File(str(overlap_features), 'w') as hfile: 23 | for scene in os.listdir(opt.dataset_dir): 24 | # load images and run DINOv2 backbone 25 | dataset_dir = os.path.join(opt.dataset_dir, scene, 'images') 26 | if not os.path.isdir(dataset_dir): 27 | continue 28 | image_list = os.listdir(dataset_dir) 29 | for image_path in tqdm(image_list): 30 | image_ = load_image(dataset_dir +'/' + image_path, resize=opt.imsize) 31 | c, h, w = image_.shape 32 | image = torch.zeros((3, opt.imsize, opt.imsize), device=opt.device, dtype=opt.dtype) 33 | image[:, :h, :w] = image_ 34 | feats = model.extractor({'image': image[None]}) 35 | group = hfile.create_group(scene+'/images/' +image_path) 36 | for key in ['keypoints', 'descriptors', 'global_descriptor']: 37 | group.create_dataset(key, data=feats[key][0].detach().cpu().numpy()) 38 | group.create_dataset('image_size', data=[w, h]) 39 | 40 | # enough for retrieval process, after dumping the data, running register/retrieve.py will return image pair lists 41 | # If pose estimation evaluation, pls load GT pose information 42 | -------------------------------------------------------------------------------- /dump_datasets/inloc.py: -------------------------------------------------------------------------------- 1 | from lightglue.utils import load_image 2 | import gluefactory 3 | from gluefactory.datasets.utils import scale_intrinsics 4 | import os 5 | import h5py 6 | import torch 7 | import numpy as np 8 | from tqdm import tqdm 9 | from pathlib import Path 10 | from args import * 11 | from utils import quaternion_to_rotation_matrix, camera_center_to_translation, Camera 12 | 13 | def dump(opt, model): 14 | 15 | overlap_features = Path(opt.dump_dir)/ 'inloc' / 'overlap_feats.h5' 16 | (Path(opt.dump_dir)/ opt.dataset).mkdir(exist_ok=True, parents=True) 17 | 18 | # save intrinsic and extrinsic parameters for each image, from world to camera 19 | with h5py.File(str(overlap_features), 'w') as hfile: 20 | query_list = [f"query/iphone7/" + q for q in os.listdir(opt.dataset_dir + "/query/iphone7/") if q !='LICENSE.txt'] 21 | db_list = [] 22 | for scene in ['DUC1', 'DUC2']: 23 | if scene != 'LICENSE.txt': 24 | for folder in os.listdir(opt.dataset_dir + "/database/cutouts/" + scene): 25 | db_list += [f"database/cutouts/{scene}/{folder}/{f}" for f in os.listdir(f"{opt.dataset_dir}/database/cutouts/{scene}/{folder}") if f.endswith(".jpg")] 26 | image_list = np.concatenate((query_list, db_list)) 27 | 28 | # save the indices of query and db images 29 | group = hfile.create_group("indices") 30 | group.create_dataset("query", data=query_list, dtype=h5py.string_dtype(encoding='utf-8')) 31 | group.create_dataset("db", data=db_list, dtype=h5py.string_dtype(encoding='utf-8')) 32 | print(f"in total {len(query_list)} queries and {len(db_list)} db images") 33 | 34 | # save the images, embeddings 35 | for image_path in tqdm(image_list): 36 | image_ = load_image(opt.dataset_dir + image_path, resize=opt.imsize) 37 | c, h, w = image_.shape 38 | image = torch.zeros((3, opt.imsize, opt.imsize), device=opt.device, dtype=opt.dtype) 39 | image[:, :h, :w] = image_ 40 | feats = model.extractor({'image': image[None]}) 41 | group = hfile.create_group(image_path) 42 | for key in ['keypoints', 'descriptors', 'global_descriptor']: 43 | group.create_dataset(key, data=feats[key][0].detach().cpu().numpy()) 44 | ori_h, ori_w = load_image(opt.dataset_dir + image_path).shape[-2:] 45 | original_image_size = np.array([ori_w, ori_h]) 46 | scales = opt.imsize/original_image_size 47 | group.create_dataset('image_size', data=[w, h]) 48 | print(f"finished." ) 49 | -------------------------------------------------------------------------------- /dump_datasets/megadepth.py: -------------------------------------------------------------------------------- 1 | import os 2 | import h5py 3 | import torch 4 | import numpy as np 5 | from tqdm import tqdm 6 | from pathlib import Path 7 | from lightglue.utils import load_image 8 | from gluefactory.datasets.utils import scale_intrinsics 9 | from gluefactory.utils.image import ImagePreprocessor 10 | 11 | torch.set_grad_enabled(False) 12 | 13 | def dump(opt, model): 14 | opt.dataset_dir = Path(opt.dataset_dir) 15 | 16 | def get_imagedata(name, suffix): 17 | idx = list(scene_info["image_paths"]).index(name) 18 | depth_process = ImagePreprocessor({'resize': opt.imsize, "side": "long", "square_pad": True}) 19 | dpath = opt.dataset_dir / scene_info["depth_paths"][idx] 20 | with h5py.File(str(dpath), "r") as f: 21 | depth = f["/depth"].__array__().astype(np.float32, copy=False) 22 | depth = torch.Tensor(depth)[None] 23 | ddata = depth_process(depth, interpolation="nearest") 24 | K = scene_info['intrinsics'][idx] 25 | K = scale_intrinsics(K, ddata["scales"]) 26 | return { 27 | 'scales'+suffix: ddata["scales"], 28 | 'original_image_size': torch.from_numpy(ddata["original_image_size"]), 29 | 'original_K': torch.from_numpy(scene_info['intrinsics'][idx]), 30 | 'depth'+suffix: ddata["image"][0], 31 | 'K'+suffix: torch.tensor(K).float(), 32 | 'T_w2cam'+suffix: torch.tensor(scene_info['poses'][idx]).float() 33 | } 34 | 35 | overlap_features = Path(opt.dump_dir) / 'megadepth' / 'overlap_feats.h5' 36 | (Path(opt.dump_dir) / 'megadepth').mkdir(exist_ok=True, parents=True) 37 | 38 | if not os.path.exists(overlap_features) or opt.overwrite: 39 | with h5py.File(str(overlap_features), 'w') as hfile: 40 | for scene in ['0015','0022']: 41 | scene_file = opt.dataset_dir / 'scene_info' / (scene+'.npz') 42 | scene_info = np.load(scene_file, allow_pickle=True) 43 | image_list = [image for image in scene_info['image_paths'] if image is not None] 44 | for image_path in tqdm(image_list): 45 | image_ = load_image(opt.dataset_dir / image_path, resize=opt.imsize) 46 | c, h, w = image_.shape 47 | image = torch.zeros((3, opt.imsize, opt.imsize), device=opt.device, dtype=opt.dtype) 48 | image[:, :h, :w] = image_ 49 | feats = model.extractor({'image': image[None]}) 50 | group = hfile.create_group(image_path) 51 | for key in ['keypoints', 'descriptors', 'global_descriptor']: 52 | group.create_dataset(key, data=feats[key][0].detach().cpu().numpy()) 53 | imdata = get_imagedata(image_path, "") 54 | for key, v in imdata.items(): 55 | group.create_dataset(key, data=v.numpy()) 56 | group.create_dataset('image_size', data=[w, h]) 57 | print(f"Finished." ) 58 | -------------------------------------------------------------------------------- /dump_datasets/phototourism.py: -------------------------------------------------------------------------------- 1 | 2 | import os 3 | import h5py 4 | import torch 5 | import gluefactory 6 | import numpy as np 7 | from tqdm import tqdm 8 | from pathlib import Path 9 | from args import * 10 | from utils import loadh5 11 | from collections import defaultdict 12 | from lightglue.utils import load_image 13 | from gluefactory.datasets.utils import scale_intrinsics 14 | 15 | def dump(opt, model): 16 | 17 | overlap_features = Path(opt.dump_dir)/ 'phototourism' / 'overlap_feats.h5' 18 | (Path(opt.dump_dir) / 'phototourism').mkdir(exist_ok=True, parents=True) 19 | 20 | if not os.path.exists(overlap_features) or opt.overwrite: 21 | with h5py.File(str(overlap_features), 'w') as hfile: 22 | for scene in os.listdir(opt.dataset_dir): 23 | if not os.path.exists(os.path.join(opt.dataset_dir, scene, 'K1_K2.h5')): 24 | continue 25 | K0K1s = loadh5(os.path.join(opt.dataset_dir, scene, 'K1_K2.h5')) 26 | Rs = loadh5(os.path.join(opt.dataset_dir, scene, 'R.h5')) 27 | Ts = loadh5(os.path.join(opt.dataset_dir, scene, 'T.h5')) 28 | 29 | Ks = defaultdict() 30 | for name0_1 in K0K1s.keys(): 31 | name0, name1 = name0_1.split('-') 32 | if name0 not in Ks.keys(): 33 | Ks[name0] = K0K1s[name0_1][0][0] 34 | if name1 not in Ks.keys(): 35 | Ks[name1] = K0K1s[name0_1][0][1] 36 | 37 | dataset_dir = os.path.join(opt.dataset_dir, scene, 'images/') 38 | image_list = [i for i in Rs.keys()] 39 | 40 | # Read intrinsics and extrinsics data 41 | for image_path in tqdm(image_list): 42 | image_ = load_image(dataset_dir +'/' + image_path + str('.jpg'), resize=opt.imsize) 43 | c, h, w = image_.shape 44 | image = torch.zeros((3, opt.imsize, opt.imsize), device=opt.device, dtype=opt.dtype) 45 | image[:, :h, :w] = image_ 46 | feats = model.extractor({'image': image[None]}) 47 | group = hfile.create_group(scene+'/images/' +image_path+str('.jpg')) 48 | for key in ['keypoints', 'descriptors', 'global_descriptor']: 49 | group.create_dataset(key, data=feats[key][0].detach().cpu().numpy()) 50 | ori_h, ori_w = load_image(dataset_dir +'/' + image_path + str('.jpg')).shape[-2:] 51 | original_image_size = np.array([ori_w, ori_h]) 52 | scales = opt.imsize/original_image_size 53 | ori_K = Ks[image_path] 54 | K = scale_intrinsics(ori_K, scales) 55 | T_w2cam = np.zeros((4, 4)) 56 | T_w2cam[-1, -1] = 1 57 | T_w2cam[:3, :3] = Rs[image_path] 58 | T_w2cam[:-1, -1] = Ts[image_path].T[0] 59 | 60 | group.create_dataset("K", data=K) 61 | group.create_dataset("original_K", data=ori_K) 62 | group.create_dataset("T_w2cam", data=T_w2cam) 63 | group.create_dataset('image_size', data=[w, h]) 64 | -------------------------------------------------------------------------------- /gluefactory/__init__.py: -------------------------------------------------------------------------------- 1 | import logging 2 | 3 | from .utils.experiments import load_experiment # noqa: F401 4 | 5 | formatter = logging.Formatter( 6 | fmt="[%(asctime)s %(name)s %(levelname)s] %(message)s", datefmt="%m/%d/%Y %H:%M:%S" 7 | ) 8 | handler = logging.StreamHandler() 9 | handler.setFormatter(formatter) 10 | handler.setLevel(logging.INFO) 11 | 12 | logger = logging.getLogger(__name__) 13 | logger.setLevel(logging.INFO) 14 | logger.addHandler(handler) 15 | logger.propagate = False 16 | 17 | __module_name__ = __name__ 18 | -------------------------------------------------------------------------------- /gluefactory/configs/aliked+NN.yaml: -------------------------------------------------------------------------------- 1 | model: 2 | name: two_view_pipeline 3 | extractor: 4 | name: extractors.aliked 5 | max_num_keypoints: 2048 6 | detection_threshold: 0.0 7 | matcher: 8 | name: matchers.nearest_neighbor_matcher 9 | benchmarks: 10 | megadepth1500: 11 | data: 12 | preprocessing: 13 | side: long 14 | resize: 1600 15 | eval: 16 | estimator: opencv 17 | ransac_th: 0.5 18 | hpatches: 19 | eval: 20 | estimator: opencv 21 | ransac_th: 0.5 22 | model: 23 | extractor: 24 | max_num_keypoints: 1024 # overwrite config above 25 | -------------------------------------------------------------------------------- /gluefactory/configs/aliked+lightglue-official.yaml: -------------------------------------------------------------------------------- 1 | model: 2 | name: two_view_pipeline 3 | extractor: 4 | name: extractors.aliked 5 | max_num_keypoints: 2048 6 | detection_threshold: 0.0 7 | matcher: 8 | name: matchers.lightglue_pretrained 9 | features: aliked 10 | depth_confidence: -1 11 | width_confidence: -1 12 | filter_threshold: 0.1 13 | benchmarks: 14 | megadepth1500: 15 | data: 16 | preprocessing: 17 | side: long 18 | resize: 1600 19 | eval: 20 | estimator: opencv 21 | ransac_th: 0.5 22 | hpatches: 23 | eval: 24 | estimator: opencv 25 | ransac_th: 0.5 26 | model: 27 | extractor: 28 | max_num_keypoints: 1024 # overwrite config above 29 | -------------------------------------------------------------------------------- /gluefactory/configs/aliked+lightglue_homography.yaml: -------------------------------------------------------------------------------- 1 | data: 2 | name: homographies 3 | data_dir: revisitop1m 4 | train_size: 150000 5 | val_size: 2000 6 | batch_size: 128 7 | num_workers: 14 8 | homography: 9 | difficulty: 0.7 10 | max_angle: 45 11 | photometric: 12 | name: lg 13 | model: 14 | name: two_view_pipeline 15 | extractor: 16 | name: extractors.aliked 17 | max_num_keypoints: 512 18 | detection_threshold: 0.0 19 | trainable: False 20 | detector: 21 | name: null 22 | descriptor: 23 | name: null 24 | ground_truth: 25 | name: matchers.homography_matcher 26 | th_positive: 3 27 | th_negative: 3 28 | matcher: 29 | name: matchers.lightglue 30 | filter_threshold: 0.1 31 | flash: false 32 | checkpointed: true 33 | input_dim: 128 34 | train: 35 | seed: 0 36 | epochs: 40 37 | log_every_iter: 100 38 | eval_every_iter: 500 39 | lr: 1e-4 40 | lr_schedule: 41 | start: 20 42 | type: exp 43 | on_epoch: true 44 | exp_div_10: 10 45 | plot: [5, 'gluefactory.visualization.visualize_batch.make_match_figures'] 46 | benchmarks: 47 | hpatches: 48 | eval: 49 | estimator: opencv 50 | ransac_th: 0.5 51 | -------------------------------------------------------------------------------- /gluefactory/configs/aliked+lightglue_megadepth.yaml: -------------------------------------------------------------------------------- 1 | data: 2 | name: megadepth 3 | preprocessing: 4 | resize: 1024 5 | side: long 6 | square_pad: True 7 | train_split: train_scenes_clean.txt 8 | train_num_per_scene: 300 9 | val_split: valid_scenes_clean.txt 10 | val_pairs: valid_pairs.txt 11 | min_overlap: 0.1 12 | max_overlap: 0.7 13 | num_overlap_bins: 3 14 | read_depth: true 15 | read_image: true 16 | batch_size: 32 17 | num_workers: 14 18 | load_features: 19 | do: false # enable this if you have cached predictions 20 | path: exports/megadepth-undist-depth-r1024_ALIKED-k2048-n16/{scene}.h5 21 | padding_length: 2048 22 | padding_fn: pad_local_features 23 | model: 24 | name: two_view_pipeline 25 | extractor: 26 | name: extractors.aliked 27 | max_num_keypoints: 2048 28 | detection_threshold: 0.0 29 | trainable: False 30 | matcher: 31 | name: matchers.lightglue 32 | filter_threshold: 0.1 33 | flash: false 34 | checkpointed: true 35 | input_dim: 128 36 | ground_truth: 37 | name: matchers.depth_matcher 38 | th_positive: 3 39 | th_negative: 5 40 | th_epi: 5 41 | allow_no_extract: True 42 | train: 43 | seed: 0 44 | epochs: 50 45 | log_every_iter: 100 46 | eval_every_iter: 1000 47 | lr: 1e-4 48 | lr_schedule: 49 | start: 30 50 | type: exp 51 | on_epoch: true 52 | exp_div_10: 10 53 | dataset_callback_fn: sample_new_items 54 | plot: [5, 'gluefactory.visualization.visualize_batch.make_match_figures'] 55 | benchmarks: 56 | megadepth1500: 57 | data: 58 | preprocessing: 59 | side: long 60 | resize: 1600 61 | eval: 62 | estimator: opencv 63 | ransac_th: 0.5 64 | hpatches: 65 | eval: 66 | estimator: opencv 67 | ransac_th: 0.5 68 | model: 69 | extractor: 70 | max_num_keypoints: 1024 71 | -------------------------------------------------------------------------------- /gluefactory/configs/disk+NN.yaml: -------------------------------------------------------------------------------- 1 | model: 2 | name: two_view_pipeline 3 | extractor: 4 | name: extractors.disk_kornia 5 | max_num_keypoints: 2048 6 | detection_threshold: 0.0 7 | matcher: 8 | name: matchers.nearest_neighbor_matcher 9 | benchmarks: 10 | megadepth1500: 11 | data: 12 | preprocessing: 13 | side: long 14 | resize: 1600 15 | eval: 16 | estimator: opencv 17 | ransac_th: 0.5 18 | hpatches: 19 | eval: 20 | estimator: opencv 21 | ransac_th: 0.5 22 | model: 23 | extractor: 24 | max_num_keypoints: 1024 # overwrite config above 25 | -------------------------------------------------------------------------------- /gluefactory/configs/disk+lightglue-official.yaml: -------------------------------------------------------------------------------- 1 | model: 2 | name: two_view_pipeline 3 | extractor: 4 | name: extractors.disk_kornia 5 | max_num_keypoints: 2048 6 | detection_threshold: 0.0 7 | matcher: 8 | name: matchers.lightglue_pretrained 9 | features: disk 10 | depth_confidence: -1 11 | width_confidence: -1 12 | filter_threshold: 0.1 13 | benchmarks: 14 | megadepth1500: 15 | data: 16 | preprocessing: 17 | side: long 18 | resize: 1600 19 | eval: 20 | estimator: opencv 21 | ransac_th: 0.5 22 | hpatches: 23 | eval: 24 | estimator: opencv 25 | ransac_th: 0.5 26 | model: 27 | extractor: 28 | max_num_keypoints: 1024 # overwrite config above 29 | -------------------------------------------------------------------------------- /gluefactory/configs/disk+lightglue_homography.yaml: -------------------------------------------------------------------------------- 1 | data: 2 | name: homographies 3 | data_dir: revisitop1m 4 | train_size: 150000 5 | val_size: 2000 6 | batch_size: 128 7 | num_workers: 14 8 | homography: 9 | difficulty: 0.7 10 | max_angle: 45 11 | photometric: 12 | name: lg 13 | model: 14 | name: two_view_pipeline 15 | extractor: 16 | name: extractors.disk_kornia 17 | max_num_keypoints: 512 18 | force_num_keypoints: True 19 | detection_threshold: 0.0 20 | trainable: False 21 | ground_truth: 22 | name: matchers.homography_matcher 23 | th_positive: 3 24 | th_negative: 3 25 | matcher: 26 | name: matchers.lightglue 27 | filter_threshold: 0.1 28 | input_dim: 128 29 | flash: false 30 | checkpointed: true 31 | train: 32 | seed: 0 33 | epochs: 40 34 | log_every_iter: 100 35 | eval_every_iter: 500 36 | lr: 1e-4 37 | lr_schedule: 38 | start: 20 39 | type: exp 40 | on_epoch: true 41 | exp_div_10: 10 42 | plot: [5, 'gluefactory.visualization.visualize_batch.make_match_figures'] 43 | benchmarks: 44 | hpatches: 45 | eval: 46 | estimator: opencv 47 | ransac_th: 0.5 48 | -------------------------------------------------------------------------------- /gluefactory/configs/disk+lightglue_megadepth.yaml: -------------------------------------------------------------------------------- 1 | data: 2 | name: megadepth 3 | preprocessing: 4 | resize: 1024 5 | side: long 6 | square_pad: True 7 | train_split: train_scenes_clean.txt 8 | train_num_per_scene: 300 9 | val_split: valid_scenes_clean.txt 10 | val_pairs: valid_pairs.txt 11 | min_overlap: 0.1 12 | max_overlap: 0.7 13 | num_overlap_bins: 3 14 | read_depth: true 15 | read_image: true 16 | batch_size: 32 17 | num_workers: 14 18 | load_features: 19 | do: false # enable this if you have cached predictions 20 | path: exports/megadepth-undist-depth-r1024_DISK-k2048-nms5/{scene}.h5 21 | padding_length: 2048 22 | padding_fn: pad_local_features 23 | model: 24 | name: two_view_pipeline 25 | extractor: 26 | name: extractors.disk_kornia 27 | max_num_keypoints: 512 28 | force_num_keypoints: True 29 | detection_threshold: 0.0 30 | trainable: False 31 | ground_truth: 32 | name: matchers.homography_matcher 33 | th_positive: 3 34 | th_negative: 3 35 | matcher: 36 | name: matchers.lightglue 37 | filter_threshold: 0.1 38 | input_dim: 128 39 | flash: false 40 | checkpointed: true 41 | allow_no_extract: True 42 | train: 43 | seed: 0 44 | epochs: 50 45 | log_every_iter: 100 46 | eval_every_iter: 1000 47 | lr: 1e-4 48 | lr_schedule: 49 | start: 30 50 | type: exp 51 | on_epoch: true 52 | exp_div_10: 10 53 | dataset_callback_fn: sample_new_items 54 | plot: [5, 'gluefactory.visualization.visualize_batch.make_match_figures'] 55 | benchmarks: 56 | megadepth1500: 57 | data: 58 | preprocessing: 59 | side: long 60 | resize: 1024 61 | eval: 62 | estimator: opencv 63 | ransac_th: 0.5 64 | hpatches: 65 | eval: 66 | estimator: opencv 67 | ransac_th: 0.5 68 | model: 69 | extractor: 70 | max_num_keypoints: 1024 71 | -------------------------------------------------------------------------------- /gluefactory/configs/sift+NN.yaml: -------------------------------------------------------------------------------- 1 | model: 2 | name: two_view_pipeline 3 | extractor: 4 | name: extractors.sift 5 | detector: pycolmap_cuda 6 | max_num_keypoints: 2048 7 | detection_threshold: 0.00666666 8 | nms_radius: -1 9 | pycolmap_options: 10 | first_octave: -1 11 | matcher: 12 | name: matchers.nearest_neighbor_matcher 13 | benchmarks: 14 | megadepth1500: 15 | data: 16 | preprocessing: 17 | side: long 18 | resize: 1600 19 | eval: 20 | estimator: opencv 21 | ransac_th: 0.5 22 | hpatches: 23 | eval: 24 | estimator: opencv 25 | ransac_th: 0.5 26 | model: 27 | extractor: 28 | max_num_keypoints: 1024 # overwrite config above 29 | -------------------------------------------------------------------------------- /gluefactory/configs/sift+lightglue-official.yaml: -------------------------------------------------------------------------------- 1 | model: 2 | name: two_view_pipeline 3 | extractor: 4 | name: extractors.sift 5 | backend: pycolmap_cuda 6 | max_num_keypoints: 4096 7 | matcher: 8 | name: matchers.lightglue_pretrained 9 | features: sift 10 | depth_confidence: -1 11 | width_confidence: -1 12 | filter_threshold: 0.1 13 | benchmarks: 14 | megadepth1500: 15 | data: 16 | preprocessing: 17 | side: long 18 | resize: 1600 19 | eval: 20 | estimator: opencv 21 | ransac_th: 0.5 22 | hpatches: 23 | eval: 24 | estimator: opencv 25 | ransac_th: 0.5 26 | model: 27 | extractor: 28 | max_num_keypoints: 1024 # overwrite config above 29 | -------------------------------------------------------------------------------- /gluefactory/configs/sift+lightglue_homography.yaml: -------------------------------------------------------------------------------- 1 | data: 2 | name: homographies 3 | data_dir: revisitop1m 4 | train_size: 150000 5 | val_size: 2000 6 | batch_size: 64 7 | num_workers: 14 8 | homography: 9 | difficulty: 0.7 10 | max_angle: 45 11 | photometric: 12 | name: lg 13 | model: 14 | name: two_view_pipeline 15 | extractor: 16 | name: extractors.sift 17 | backend: pycolmap_cuda 18 | max_num_keypoints: 1024 19 | force_num_keypoints: True 20 | nms_radius: 3 21 | trainable: False 22 | ground_truth: 23 | name: matchers.homography_matcher 24 | th_positive: 3 25 | th_negative: 3 26 | matcher: 27 | name: matchers.lightglue 28 | filter_threshold: 0.1 29 | flash: false 30 | checkpointed: true 31 | input_dim: 128 32 | train: 33 | seed: 0 34 | epochs: 40 35 | log_every_iter: 100 36 | eval_every_iter: 500 37 | lr: 1e-4 38 | lr_schedule: 39 | start: 20 40 | type: exp 41 | on_epoch: true 42 | exp_div_10: 10 43 | plot: [5, 'gluefactory.visualization.visualize_batch.make_match_figures'] 44 | benchmarks: 45 | hpatches: 46 | eval: 47 | estimator: opencv 48 | ransac_th: 0.5 49 | model: 50 | extractor: 51 | nms_radius: 0 52 | -------------------------------------------------------------------------------- /gluefactory/configs/sift+lightglue_megadepth.yaml: -------------------------------------------------------------------------------- 1 | data: 2 | name: megadepth 3 | preprocessing: 4 | resize: 1024 5 | side: long 6 | square_pad: True 7 | train_split: train_scenes_clean.txt 8 | train_num_per_scene: 300 9 | val_split: valid_scenes_clean.txt 10 | val_pairs: valid_pairs.txt 11 | min_overlap: 0.1 12 | max_overlap: 0.7 13 | num_overlap_bins: 3 14 | read_depth: true 15 | read_image: true 16 | batch_size: 32 17 | num_workers: 14 18 | load_features: 19 | do: false # enable this if you have cached predictions 20 | path: exports/megadepth-undist-depth-r1024_pycolmap_SIFTGPU-nms3-fixed-k2048/{scene}.h5 21 | padding_length: 2048 22 | padding_fn: pad_local_features 23 | data_keys: ["keypoints", "keypoint_scores", "descriptors", "oris", "scales"] 24 | model: 25 | name: two_view_pipeline 26 | extractor: 27 | name: extractors.sift 28 | backend: pycolmap_cuda 29 | max_num_keypoints: 2048 30 | force_num_keypoints: True 31 | nms_radius: 3 32 | trainable: False 33 | matcher: 34 | name: matchers.lightglue 35 | filter_threshold: 0.1 36 | flash: false 37 | checkpointed: true 38 | add_scale_ori: true 39 | input_dim: 128 40 | ground_truth: 41 | name: matchers.depth_matcher 42 | th_positive: 3 43 | th_negative: 5 44 | th_epi: 5 45 | allow_no_extract: True 46 | train: 47 | seed: 0 48 | epochs: 50 49 | log_every_iter: 100 50 | eval_every_iter: 1000 51 | lr: 1e-4 52 | lr_schedule: 53 | start: 30 54 | type: exp 55 | on_epoch: true 56 | exp_div_10: 10 57 | dataset_callback_fn: sample_new_items 58 | plot: [5, 'gluefactory.visualization.visualize_batch.make_match_figures'] 59 | benchmarks: 60 | megadepth1500: 61 | data: 62 | preprocessing: 63 | side: long 64 | resize: 1600 65 | model: 66 | extractor: 67 | nms_radius: 0 68 | eval: 69 | estimator: opencv 70 | ransac_th: 0.5 71 | hpatches: 72 | eval: 73 | estimator: opencv 74 | ransac_th: 0.5 75 | model: 76 | extractor: 77 | max_num_keypoints: 1024 78 | nms_radius: 0 79 | -------------------------------------------------------------------------------- /gluefactory/configs/superpoint+NN.yaml: -------------------------------------------------------------------------------- 1 | model: 2 | name: two_view_pipeline 3 | extractor: 4 | name: gluefactory_nonfree.superpoint 5 | max_num_keypoints: 2048 6 | detection_threshold: 0.0 7 | nms_radius: 3 8 | matcher: 9 | name: matchers.nearest_neighbor_matcher 10 | benchmarks: 11 | megadepth1500: 12 | data: 13 | preprocessing: 14 | side: long 15 | resize: 1600 16 | eval: 17 | estimator: opencv 18 | ransac_th: 1.0 19 | hpatches: 20 | eval: 21 | estimator: opencv 22 | ransac_th: 0.5 23 | model: 24 | extractor: 25 | max_num_keypoints: 1024 # overwrite config above 26 | -------------------------------------------------------------------------------- /gluefactory/configs/superpoint+lightglue-official.yaml: -------------------------------------------------------------------------------- 1 | model: 2 | name: two_view_pipeline 3 | extractor: 4 | name: gluefactory_nonfree.superpoint 5 | max_num_keypoints: 2048 6 | detection_threshold: 0.0 7 | nms_radius: 3 8 | matcher: 9 | name: matchers.lightglue_pretrained 10 | features: superpoint 11 | depth_confidence: -1 12 | width_confidence: -1 13 | filter_threshold: 0.1 14 | benchmarks: 15 | megadepth1500: 16 | data: 17 | preprocessing: 18 | side: long 19 | resize: 1600 20 | eval: 21 | estimator: opencv 22 | ransac_th: 0.5 23 | hpatches: 24 | eval: 25 | estimator: opencv 26 | ransac_th: 0.5 27 | model: 28 | extractor: 29 | max_num_keypoints: 1024 # overwrite config above 30 | -------------------------------------------------------------------------------- /gluefactory/configs/superpoint+lightglue_homography.yaml: -------------------------------------------------------------------------------- 1 | data: 2 | name: homographies 3 | data_dir: revisitop1m 4 | train_size: 150000 5 | val_size: 2000 6 | batch_size: 128 7 | num_workers: 14 8 | homography: 9 | difficulty: 0.7 10 | max_angle: 45 11 | photometric: 12 | name: lg 13 | model: 14 | name: two_view_pipeline 15 | extractor: 16 | name: gluefactory_nonfree.superpoint 17 | max_num_keypoints: 512 18 | force_num_keypoints: True 19 | detection_threshold: 0.0 20 | nms_radius: 3 21 | trainable: False 22 | ground_truth: 23 | name: matchers.homography_matcher 24 | th_positive: 3 25 | th_negative: 3 26 | matcher: 27 | name: matchers.lightglue 28 | filter_threshold: 0.1 29 | flash: false 30 | checkpointed: true 31 | train: 32 | seed: 0 33 | epochs: 40 34 | log_every_iter: 100 35 | eval_every_iter: 500 36 | lr: 1e-4 37 | lr_schedule: 38 | start: 20 39 | type: exp 40 | on_epoch: true 41 | exp_div_10: 10 42 | plot: [5, 'gluefactory.visualization.visualize_batch.make_match_figures'] 43 | benchmarks: 44 | hpatches: 45 | eval: 46 | estimator: opencv 47 | ransac_th: 0.5 48 | -------------------------------------------------------------------------------- /gluefactory/configs/superpoint+lightglue_megadepth.yaml: -------------------------------------------------------------------------------- 1 | data: 2 | name: megadepth 3 | preprocessing: 4 | resize: 1024 5 | side: long 6 | square_pad: True 7 | train_split: train_scenes_clean.txt 8 | train_num_per_scene: 300 9 | val_split: valid_scenes_clean.txt 10 | val_pairs: valid_pairs.txt 11 | min_overlap: 0.1 12 | max_overlap: 0.7 13 | num_overlap_bins: 3 14 | read_depth: true 15 | read_image: true 16 | batch_size: 32 17 | num_workers: 14 18 | load_features: 19 | do: false # enable this if you have cached predictions 20 | path: exports/megadepth-undist-depth-r1024_SP-k2048-nms3/{scene}.h5 21 | padding_length: 2048 22 | padding_fn: pad_local_features 23 | model: 24 | name: two_view_pipeline 25 | extractor: 26 | name: gluefactory_nonfree.superpoint 27 | max_num_keypoints: 2048 28 | force_num_keypoints: True 29 | detection_threshold: 0.0 30 | nms_radius: 3 31 | trainable: False 32 | matcher: 33 | name: matchers.lightglue 34 | filter_threshold: 0.1 35 | flash: false 36 | checkpointed: true 37 | ground_truth: 38 | name: matchers.depth_matcher 39 | th_positive: 3 40 | th_negative: 5 41 | th_epi: 5 42 | allow_no_extract: True 43 | train: 44 | seed: 0 45 | epochs: 50 46 | log_every_iter: 100 47 | eval_every_iter: 1000 48 | lr: 1e-4 49 | lr_schedule: 50 | start: 30 51 | type: exp 52 | on_epoch: true 53 | exp_div_10: 10 54 | dataset_callback_fn: sample_new_items 55 | plot: [5, 'gluefactory.visualization.visualize_batch.make_match_figures'] 56 | benchmarks: 57 | megadepth1500: 58 | data: 59 | preprocessing: 60 | side: long 61 | resize: 1600 62 | eval: 63 | estimator: opencv 64 | ransac_th: 0.5 65 | hpatches: 66 | eval: 67 | estimator: opencv 68 | ransac_th: 0.5 69 | model: 70 | extractor: 71 | max_num_keypoints: 1024 72 | -------------------------------------------------------------------------------- /gluefactory/configs/superpoint+lsd+gluestick-homography.yaml: -------------------------------------------------------------------------------- 1 | data: 2 | name: homographies 3 | homography: 4 | difficulty: 0.7 5 | max_angle: 45 6 | patch_shape: [640, 480] 7 | photometric: 8 | p: 0.75 9 | train_size: 900000 10 | val_size: 1000 11 | batch_size: 160 # 20 per 10GB of GPU mem (12 for triplet) 12 | num_workers: 15 13 | model: 14 | name: gluefactory.models.two_view_pipeline 15 | extractor: 16 | name: gluefactory.models.lines.wireframe 17 | trainable: False 18 | point_extractor: 19 | name: gluefactory.models.extractors.superpoint_open 20 | # name: disk 21 | # chunk: 10 22 | max_num_keypoints: 1000 23 | force_num_keypoints: true 24 | trainable: False 25 | line_extractor: 26 | name: gluefactory.models.lines.lsd 27 | max_num_lines: 250 28 | force_num_lines: True 29 | min_length: 15 30 | trainable: False 31 | wireframe_params: 32 | merge_points: True 33 | merge_line_endpoints: True 34 | nms_radius: 4 35 | detector: 36 | name: null 37 | descriptor: 38 | name: null 39 | ground_truth: 40 | name: gluefactory.models.matchers.homography_matcher 41 | trainable: False 42 | use_points: True 43 | use_lines: True 44 | th_positive: 3 45 | th_negative: 5 46 | matcher: 47 | name: gluefactory.models.matchers.gluestick 48 | input_dim: 256 # 128 for DISK 49 | descriptor_dim: 256 # 128 for DISK 50 | inter_supervision: [2, 5] 51 | GNN_layers: [ 52 | self, cross, self, cross, self, cross, 53 | self, cross, self, cross, self, cross, 54 | self, cross, self, cross, self, cross, 55 | ] 56 | checkpointed: true 57 | train: 58 | seed: 0 59 | epochs: 200 60 | log_every_iter: 400 61 | eval_every_iter: 700 62 | save_every_iter: 1400 63 | lr: 1e-4 64 | lr_schedule: 65 | type: exp # exp or multi_step 66 | start: 200e3 67 | exp_div_10: 200e3 68 | gamma: 0.5 69 | step: 50e3 70 | n_steps: 4 71 | submodules: [] 72 | # clip_grad: 10 # Use only with mixed precision 73 | # load_experiment: -------------------------------------------------------------------------------- /gluefactory/configs/superpoint+lsd+gluestick-megadepth.yaml: -------------------------------------------------------------------------------- 1 | data: 2 | name: gluefactory.datasets.megadepth 3 | train_num_per_scene: 300 4 | val_pairs: valid_pairs.txt 5 | views: 2 6 | min_overlap: 0.1 7 | max_overlap: 0.7 8 | num_overlap_bins: 3 9 | preprocessing: 10 | resize: 640 11 | square_pad: True 12 | batch_size: 160 13 | num_workers: 15 14 | model: 15 | name: gluefactory.models.two_view_pipeline 16 | extractor: 17 | name: gluefactory.models.lines.wireframe 18 | trainable: False 19 | point_extractor: 20 | name: gluefactory.models.extractors.superpoint_open 21 | # name: disk 22 | # chunk: 10 23 | max_num_keypoints: 1000 24 | force_num_keypoints: true 25 | trainable: False 26 | line_extractor: 27 | name: gluefactory.models.lines.lsd 28 | max_num_lines: 250 29 | force_num_lines: True 30 | min_length: 15 31 | trainable: False 32 | wireframe_params: 33 | merge_points: True 34 | merge_line_endpoints: True 35 | nms_radius: 4 36 | detector: 37 | name: null 38 | descriptor: 39 | name: null 40 | ground_truth: 41 | name: gluefactory.models.matchers.depth_matcher 42 | trainable: False 43 | use_points: True 44 | use_lines: True 45 | th_positive: 3 46 | th_negative: 5 47 | matcher: 48 | name: gluefactory.models.matchers.gluestick 49 | input_dim: 256 # 128 for DISK 50 | descriptor_dim: 256 # 128 for DISK 51 | inter_supervision: null 52 | GNN_layers: [ 53 | self, cross, self, cross, self, cross, 54 | self, cross, self, cross, self, cross, 55 | self, cross, self, cross, self, cross, 56 | ] 57 | checkpointed: true 58 | train: 59 | seed: 0 60 | epochs: 200 61 | log_every_iter: 400 62 | eval_every_iter: 700 63 | save_every_iter: 1400 64 | lr: 1e-4 65 | lr_schedule: 66 | type: exp # exp or multi_step 67 | start: 200e3 68 | exp_div_10: 200e3 69 | gamma: 0.5 70 | step: 50e3 71 | n_steps: 4 72 | submodules: [] 73 | # clip_grad: 10 # Use only with mixed precision 74 | load_experiment: gluestick_H -------------------------------------------------------------------------------- /gluefactory/configs/superpoint+lsd+gluestick.yaml: -------------------------------------------------------------------------------- 1 | model: 2 | name: gluefactory.models.two_view_pipeline 3 | extractor: 4 | name: gluefactory.models.lines.wireframe 5 | point_extractor: 6 | name: gluefactory_nonfree.superpoint 7 | trainable: False 8 | dense_outputs: True 9 | max_num_keypoints: 2048 10 | force_num_keypoints: False 11 | detection_threshold: 0 12 | line_extractor: 13 | name: gluefactory.models.lines.lsd 14 | trainable: False 15 | max_num_lines: 512 16 | force_num_lines: False 17 | min_length: 15 18 | wireframe_params: 19 | merge_points: True 20 | merge_line_endpoints: True 21 | nms_radius: 3 22 | matcher: 23 | name: gluefactory.models.matchers.gluestick 24 | weights: checkpoint_GlueStick_MD # This will download weights from internet 25 | 26 | # ground_truth: # for ETH3D, comment otherwise 27 | # name: gluefactory.models.matchers.depth_matcher 28 | # use_lines: True 29 | 30 | benchmarks: 31 | hpatches: 32 | eval: 33 | estimator: homography_est 34 | ransac_th: -1 # [1., 1.5, 2., 2.5, 3.] 35 | megadepth1500: 36 | data: 37 | preprocessing: 38 | side: long 39 | resize: 1600 40 | eval: 41 | estimator: poselib 42 | ransac_th: -1 43 | eth3d: 44 | ground_truth: 45 | name: gluefactory.models.matchers.depth_matcher 46 | use_lines: True 47 | eval: 48 | plot_methods: [ ] # ['sp+NN', 'sp+sg', 'superpoint+lsd+gluestick'] 49 | plot_line_methods: [ ] # ['superpoint+lsd+gluestick', 'sp+deeplsd+gs'] -------------------------------------------------------------------------------- /gluefactory/configs/superpoint+superglue-official.yaml: -------------------------------------------------------------------------------- 1 | model: 2 | name: two_view_pipeline 3 | extractor: 4 | name: gluefactory_nonfree.superpoint 5 | max_num_keypoints: 2048 6 | detection_threshold: 0.0 7 | nms_radius: 3 8 | matcher: 9 | name: gluefactory_nonfree.superglue 10 | benchmarks: 11 | megadepth1500: 12 | data: 13 | preprocessing: 14 | side: long 15 | resize: 1600 16 | eval: 17 | estimator: opencv 18 | ransac_th: 0.5 19 | hpatches: 20 | eval: 21 | estimator: opencv 22 | ransac_th: 0.5 23 | model: 24 | extractor: 25 | max_num_keypoints: 1024 # overwrite config above 26 | 27 | -------------------------------------------------------------------------------- /gluefactory/configs/superpoint-open+NN.yaml: -------------------------------------------------------------------------------- 1 | model: 2 | name: two_view_pipeline 3 | extractor: 4 | name: extractors.superpoint_open 5 | max_num_keypoints: 2048 6 | detection_threshold: 0.0 7 | nms_radius: 3 8 | matcher: 9 | name: matchers.nearest_neighbor_matcher 10 | benchmarks: 11 | megadepth1500: 12 | data: 13 | preprocessing: 14 | side: long 15 | resize: 1600 16 | eval: 17 | estimator: opencv 18 | ransac_th: 1.0 19 | hpatches: 20 | eval: 21 | estimator: opencv 22 | ransac_th: 0.5 23 | model: 24 | extractor: 25 | max_num_keypoints: 1024 # overwrite config above 26 | -------------------------------------------------------------------------------- /gluefactory/configs/superpoint-open+lightglue_homography.yaml: -------------------------------------------------------------------------------- 1 | data: 2 | name: homographies 3 | data_dir: revisitop1m 4 | train_size: 150000 5 | val_size: 2000 6 | batch_size: 128 7 | num_workers: 14 8 | homography: 9 | difficulty: 0.7 10 | max_angle: 45 11 | photometric: 12 | name: lg 13 | model: 14 | name: two_view_pipeline 15 | extractor: 16 | name: extractors.superpoint_open 17 | max_num_keypoints: 512 18 | force_num_keypoints: True 19 | detection_threshold: -1 20 | nms_radius: 3 21 | trainable: False 22 | ground_truth: 23 | name: matchers.homography_matcher 24 | th_positive: 3 25 | th_negative: 3 26 | matcher: 27 | name: matchers.lightglue 28 | filter_threshold: 0.1 29 | flash: false 30 | checkpointed: true 31 | train: 32 | seed: 0 33 | epochs: 40 34 | log_every_iter: 100 35 | eval_every_iter: 500 36 | lr: 1e-4 37 | lr_schedule: 38 | start: 20 39 | type: exp 40 | on_epoch: true 41 | exp_div_10: 10 42 | plot: [5, 'gluefactory.visualization.visualize_batch.make_match_figures'] 43 | benchmarks: 44 | hpatches: 45 | eval: 46 | estimator: opencv 47 | ransac_th: 0.5 48 | -------------------------------------------------------------------------------- /gluefactory/configs/superpoint-open+lightglue_megadepth.yaml: -------------------------------------------------------------------------------- 1 | data: 2 | name: megadepth 3 | preprocessing: 4 | resize: 1024 5 | side: long 6 | square_pad: True 7 | train_split: train_scenes_clean.txt 8 | train_num_per_scene: 300 9 | val_split: valid_scenes_clean.txt 10 | val_pairs: valid_pairs.txt 11 | min_overlap: 0.1 12 | max_overlap: 0.7 13 | num_overlap_bins: 3 14 | read_depth: true 15 | read_image: true 16 | batch_size: 32 17 | num_workers: 14 18 | load_features: 19 | do: false # enable this if you have cached predictions 20 | path: exports/megadepth-undist-depth-r1024_SP-open-k2048-nms3/{scene}.h5 21 | padding_length: 2048 22 | padding_fn: pad_local_features 23 | model: 24 | name: two_view_pipeline 25 | extractor: 26 | name: extractors.superpoint_open 27 | max_num_keypoints: 2048 28 | force_num_keypoints: True 29 | detection_threshold: -1 30 | nms_radius: 3 31 | trainable: False 32 | matcher: 33 | name: matchers.lightglue 34 | filter_threshold: 0.1 35 | flash: false 36 | checkpointed: true 37 | ground_truth: 38 | name: matchers.depth_matcher 39 | th_positive: 3 40 | th_negative: 5 41 | th_epi: 5 42 | allow_no_extract: True 43 | train: 44 | seed: 0 45 | epochs: 50 46 | log_every_iter: 100 47 | eval_every_iter: 1000 48 | lr: 1e-4 49 | lr_schedule: 50 | start: 30 51 | type: exp 52 | on_epoch: true 53 | exp_div_10: 10 54 | dataset_callback_fn: sample_new_items 55 | plot: [5, 'gluefactory.visualization.visualize_batch.make_match_figures'] 56 | benchmarks: 57 | megadepth1500: 58 | data: 59 | preprocessing: 60 | side: long 61 | resize: 1600 62 | eval: 63 | estimator: opencv 64 | ransac_th: 0.5 65 | hpatches: 66 | eval: 67 | estimator: opencv 68 | ransac_th: 0.5 69 | model: 70 | extractor: 71 | max_num_keypoints: 1024 72 | -------------------------------------------------------------------------------- /gluefactory/datasets/__init__.py: -------------------------------------------------------------------------------- 1 | import importlib.util 2 | 3 | from ..utils.tools import get_class 4 | from .base_dataset import BaseDataset 5 | 6 | 7 | def get_dataset(name): 8 | import_paths = [name, f"{__name__}.{name}"] 9 | for path in import_paths: 10 | try: 11 | spec = importlib.util.find_spec(path) 12 | except ModuleNotFoundError: 13 | spec = None 14 | if spec is not None: 15 | try: 16 | return get_class(path, BaseDataset) 17 | except AssertionError: 18 | mod = __import__(path, fromlist=[""]) 19 | try: 20 | return mod.__main_dataset__ 21 | except AttributeError as exc: 22 | print(exc) 23 | continue 24 | 25 | raise RuntimeError(f'Dataset {name} not found in any of [{" ".join(import_paths)}]') 26 | -------------------------------------------------------------------------------- /gluefactory/datasets/hpatches.py: -------------------------------------------------------------------------------- 1 | """ 2 | Simply load images from a folder or nested folders (does not have any split). 3 | """ 4 | import argparse 5 | import logging 6 | import tarfile 7 | 8 | import matplotlib.pyplot as plt 9 | import numpy as np 10 | import torch 11 | from omegaconf import OmegaConf 12 | 13 | from ..settings import DATA_PATH 14 | from ..utils.image import ImagePreprocessor, load_image 15 | from ..utils.tools import fork_rng 16 | from ..visualization.viz2d import plot_image_grid 17 | from .base_dataset import BaseDataset 18 | 19 | logger = logging.getLogger(__name__) 20 | 21 | 22 | def read_homography(path): 23 | with open(path) as f: 24 | result = [] 25 | for line in f.readlines(): 26 | while " " in line: # Remove double spaces 27 | line = line.replace(" ", " ") 28 | line = line.replace(" \n", "").replace("\n", "") 29 | # Split and discard empty strings 30 | elements = list(filter(lambda s: s, line.split(" "))) 31 | if elements: 32 | result.append(elements) 33 | return np.array(result).astype(float) 34 | 35 | 36 | class HPatches(BaseDataset, torch.utils.data.Dataset): 37 | default_conf = { 38 | "preprocessing": ImagePreprocessor.default_conf, 39 | "data_dir": "hpatches-sequences-release", 40 | "subset": None, 41 | "ignore_large_images": True, 42 | "grayscale": False, 43 | } 44 | 45 | # Large images that were ignored in previous papers 46 | ignored_scenes = ( 47 | "i_contruction", 48 | "i_crownnight", 49 | "i_dc", 50 | "i_pencils", 51 | "i_whitebuilding", 52 | "v_artisans", 53 | "v_astronautis", 54 | "v_talent", 55 | ) 56 | url = "http://icvl.ee.ic.ac.uk/vbalnt/hpatches/hpatches-sequences-release.tar.gz" 57 | 58 | def _init(self, conf): 59 | assert conf.batch_size == 1 60 | self.preprocessor = ImagePreprocessor(conf.preprocessing) 61 | 62 | self.root = DATA_PATH / conf.data_dir 63 | if not self.root.exists(): 64 | logger.info("Downloading the HPatches dataset.") 65 | self.download() 66 | self.sequences = sorted([x.name for x in self.root.iterdir()]) 67 | if not self.sequences: 68 | raise ValueError("No image found!") 69 | self.items = [] # (seq, q_idx, is_illu) 70 | for seq in self.sequences: 71 | if conf.ignore_large_images and seq in self.ignored_scenes: 72 | continue 73 | if conf.subset is not None and conf.subset != seq[0]: 74 | continue 75 | for i in range(2, 7): 76 | self.items.append((seq, i, seq[0] == "i")) 77 | 78 | def download(self): 79 | data_dir = self.root.parent 80 | data_dir.mkdir(exist_ok=True, parents=True) 81 | tar_path = data_dir / self.url.rsplit("/", 1)[-1] 82 | torch.hub.download_url_to_file(self.url, tar_path) 83 | with tarfile.open(tar_path) as tar: 84 | tar.extractall(data_dir) 85 | tar_path.unlink() 86 | 87 | def get_dataset(self, split): 88 | assert split in ["val", "test"] 89 | return self 90 | 91 | def _read_image(self, seq: str, idx: int) -> dict: 92 | img = load_image(self.root / seq / f"{idx}.ppm", self.conf.grayscale) 93 | return self.preprocessor(img) 94 | 95 | def __getitem__(self, idx): 96 | seq, q_idx, is_illu = self.items[idx] 97 | data0 = self._read_image(seq, 1) 98 | data1 = self._read_image(seq, q_idx) 99 | H = read_homography(self.root / seq / f"H_1_{q_idx}") 100 | H = data1["transform"] @ H @ np.linalg.inv(data0["transform"]) 101 | return { 102 | "H_0to1": H.astype(np.float32), 103 | "scene": seq, 104 | "idx": idx, 105 | "is_illu": is_illu, 106 | "name": f"{seq}/{idx}.ppm", 107 | "view0": data0, 108 | "view1": data1, 109 | } 110 | 111 | def __len__(self): 112 | return len(self.items) 113 | 114 | 115 | def visualize(args): 116 | conf = { 117 | "batch_size": 1, 118 | "num_workers": 8, 119 | "prefetch_factor": 1, 120 | } 121 | conf = OmegaConf.merge(conf, OmegaConf.from_cli(args.dotlist)) 122 | dataset = HPatches(conf) 123 | loader = dataset.get_data_loader("test") 124 | logger.info("The dataset has %d elements.", len(loader)) 125 | 126 | with fork_rng(seed=dataset.conf.seed): 127 | images = [] 128 | for _, data in zip(range(args.num_items), loader): 129 | images.append( 130 | (data[f"view{i}"]["image"][0].permute(1, 2, 0) for i in range(2)) 131 | ) 132 | plot_image_grid(images, dpi=args.dpi) 133 | plt.tight_layout() 134 | plt.show() 135 | 136 | 137 | if __name__ == "__main__": 138 | from .. import logger # overwrite the logger 139 | 140 | parser = argparse.ArgumentParser() 141 | parser.add_argument("--num_items", type=int, default=8) 142 | parser.add_argument("--dpi", type=int, default=100) 143 | parser.add_argument("dotlist", nargs="*") 144 | args = parser.parse_intermixed_args() 145 | visualize(args) 146 | -------------------------------------------------------------------------------- /gluefactory/datasets/image_folder.py: -------------------------------------------------------------------------------- 1 | """ 2 | Simply load images from a folder or nested folders (does not have any split). 3 | """ 4 | 5 | import logging 6 | from pathlib import Path 7 | 8 | import omegaconf 9 | import torch 10 | 11 | from ..utils.image import ImagePreprocessor, load_image 12 | from .base_dataset import BaseDataset 13 | 14 | 15 | class ImageFolder(BaseDataset, torch.utils.data.Dataset): 16 | default_conf = { 17 | "glob": ["*.jpg", "*.png", "*.jpeg", "*.JPG", "*.PNG"], 18 | "images": "???", 19 | "root_folder": "/", 20 | "preprocessing": ImagePreprocessor.default_conf, 21 | } 22 | 23 | def _init(self, conf): 24 | self.root = conf.root_folder 25 | if isinstance(conf.images, str): 26 | if not Path(conf.images).is_dir(): 27 | with open(conf.images, "r") as f: 28 | self.images = f.read().rstrip("\n").split("\n") 29 | logging.info(f"Found {len(self.images)} images in list file.") 30 | else: 31 | self.images = [] 32 | glob = [conf.glob] if isinstance(conf.glob, str) else conf.glob 33 | for g in glob: 34 | self.images += list(Path(conf.images).glob("**/" + g)) 35 | if len(self.images) == 0: 36 | raise ValueError( 37 | f"Could not find any image in folder: {conf.images}." 38 | ) 39 | self.images = [i.relative_to(conf.images) for i in self.images] 40 | self.root = conf.images 41 | logging.info(f"Found {len(self.images)} images in folder.") 42 | elif isinstance(conf.images, omegaconf.listconfig.ListConfig): 43 | self.images = conf.images.to_container() 44 | else: 45 | raise ValueError(conf.images) 46 | 47 | self.preprocessor = ImagePreprocessor(conf.preprocessing) 48 | 49 | def get_dataset(self, split): 50 | return self 51 | 52 | def __getitem__(self, idx): 53 | path = self.images[idx] 54 | img = load_image(path) 55 | data = {"name": str(path), **self.preprocessor(img)} 56 | return data 57 | 58 | def __len__(self): 59 | return len(self.images) 60 | -------------------------------------------------------------------------------- /gluefactory/datasets/image_pairs.py: -------------------------------------------------------------------------------- 1 | """ 2 | Simply load images from a folder or nested folders (does not have any split). 3 | """ 4 | 5 | from pathlib import Path 6 | 7 | import numpy as np 8 | import torch 9 | 10 | from ..geometry.wrappers import Camera, Pose 11 | from ..settings import DATA_PATH 12 | from ..utils.image import ImagePreprocessor, load_image 13 | from .base_dataset import BaseDataset 14 | 15 | 16 | def names_to_pair(name0, name1, separator="/"): 17 | return separator.join((name0.replace("/", "-"), name1.replace("/", "-"))) 18 | 19 | 20 | def parse_homography(homography_elems) -> Camera: 21 | return ( 22 | np.array([float(x) for x in homography_elems[:9]]) 23 | .reshape(3, 3) 24 | .astype(np.float32) 25 | ) 26 | 27 | 28 | def parse_camera(calib_elems) -> Camera: 29 | # assert len(calib_list) == 9 30 | K = np.array([float(x) for x in calib_elems[:9]]).reshape(3, 3).astype(np.float32) 31 | return Camera.from_calibration_matrix(K) 32 | 33 | 34 | def parse_relative_pose(pose_elems) -> Pose: 35 | # assert len(calib_list) == 9 36 | R, t = pose_elems[:9], pose_elems[9:12] 37 | R = np.array([float(x) for x in R]).reshape(3, 3).astype(np.float32) 38 | t = np.array([float(x) for x in t]).astype(np.float32) 39 | return Pose.from_Rt(R, t) 40 | 41 | 42 | class ImagePairs(BaseDataset, torch.utils.data.Dataset): 43 | default_conf = { 44 | "pairs": "???", # ToDo: add image folder interface 45 | "root": "???", 46 | "preprocessing": ImagePreprocessor.default_conf, 47 | "extra_data": None, # relative_pose, homography 48 | } 49 | 50 | def _init(self, conf): 51 | pair_f = ( 52 | Path(conf.pairs) if Path(conf.pairs).exists() else DATA_PATH / conf.pairs 53 | ) 54 | with open(str(pair_f), "r") as f: 55 | self.items = [line.rstrip() for line in f] 56 | self.preprocessor = ImagePreprocessor(conf.preprocessing) 57 | 58 | def get_dataset(self, split): 59 | return self 60 | 61 | def _read_view(self, name): 62 | path = DATA_PATH / self.conf.root / name 63 | img = load_image(path) 64 | return self.preprocessor(img) 65 | 66 | def __getitem__(self, idx): 67 | line = self.items[idx] 68 | pair_data = line.split(" ") 69 | name0, name1 = pair_data[:2] 70 | data0 = self._read_view(name0) 71 | data1 = self._read_view(name1) 72 | data = { 73 | "view0": data0, 74 | "view1": data1, 75 | } 76 | if self.conf.extra_data == "relative_pose": 77 | data["view0"]["camera"] = parse_camera(pair_data[2:11]).scale( 78 | data0["scales"] 79 | ) 80 | data["view1"]["camera"] = parse_camera(pair_data[11:20]).scale( 81 | data1["scales"] 82 | ) 83 | data["T_0to1"] = parse_relative_pose(pair_data[20:32]) 84 | elif self.conf.extra_data == "homography": 85 | data["H_0to1"] = ( 86 | data1["transform"] 87 | @ parse_homography(pair_data[2:11]) 88 | @ np.linalg.inv(data0["transform"]) 89 | ) 90 | else: 91 | assert ( 92 | self.conf.extra_data is None 93 | ), f"Unknown extra data format {self.conf.extra_data}" 94 | 95 | data["name"] = names_to_pair(name0, name1) 96 | return data 97 | 98 | def __len__(self): 99 | return len(self.items) 100 | -------------------------------------------------------------------------------- /gluefactory/datasets/megadepth_scene_lists/loftr_val_list.txt: -------------------------------------------------------------------------------- 1 | 0022_0.1_0.3 2 | 0015_0.1_0.3 3 | 0015_0.3_0.5 4 | 0022_0.3_0.5 5 | 0022_0.5_0.7 6 | -------------------------------------------------------------------------------- /gluefactory/datasets/megadepth_scene_lists/test_list.txt: -------------------------------------------------------------------------------- 1 | 0063_0.1_0.3 2 | 0008_0.1_0.3 3 | 0021_0.3_0.5 4 | 0021_0.1_0.3 5 | 0063_0.5_0.7 6 | 0019_0.5_0.7 7 | 0032_0.3_0.5 8 | 0063_0.3_0.5 9 | 0024_0.5_0.7 10 | 1589_0.5_0.7 11 | 1589_0.3_0.5 12 | 0021_0.5_0.7 13 | 0019_0.1_0.3 14 | 1589_0.1_0.3 15 | 0019_0.3_0.5 16 | 0032_0.5_0.7 17 | 0024_0.3_0.5 18 | 0025_0.1_0.3 19 | 0024_0.1_0.3 20 | 0032_0.1_0.3 21 | 0025_0.5_0.7 22 | 0008_0.5_0.7 23 | 0025_0.3_0.5 24 | 0008_0.3_0.5 25 | -------------------------------------------------------------------------------- /gluefactory/datasets/megadepth_scene_lists/test_scenes_clean.txt: -------------------------------------------------------------------------------- 1 | 0008 2 | 0019 3 | 0021 4 | 0024 5 | 0025 6 | 0032 7 | 0063 8 | 1589 -------------------------------------------------------------------------------- /gluefactory/datasets/megadepth_scene_lists/train_hard_list.txt: -------------------------------------------------------------------------------- 1 | 0041_0.1_0.3 2 | 0099_0.1_0.3 3 | 0387_0.1_0.3 4 | 0185_0.1_0.3 5 | 5006_0.1_0.3 6 | 0044_0.1_0.3 7 | 0331_0.1_0.3 8 | 0141_0.1_0.3 9 | 0360_0.1_0.3 10 | 0326_0.1_0.3 11 | 5018_0.1_0.3 12 | 0117_0.1_0.3 13 | 0023_0.1_0.3 14 | 5008_0.1_0.3 15 | 1017_0.1_0.3 16 | 0243_0.1_0.3 17 | 0275_0.1_0.3 18 | 0238_0.1_0.3 19 | 0406_0.1_0.3 20 | 0252_0.1_0.3 21 | 0130_0.1_0.3 22 | 0307_0.1_0.3 23 | 0237_0.1_0.3 24 | 0768_0.1_0.3 25 | 0306_0.1_0.3 26 | 0043_0.1_0.3 27 | 0042_0.1_0.3 28 | 0186_0.1_0.3 29 | 0107_0.1_0.3 30 | 0017_0.1_0.3 31 | 0223_0.1_0.3 32 | 0476_0.1_0.3 33 | 0003_0.1_0.3 34 | 0214_0.1_0.3 35 | 0062_0.1_0.3 36 | 0049_0.1_0.3 37 | 0013_0.1_0.3 38 | 0407_0.1_0.3 39 | 0482_0.1_0.3 40 | 0100_0.1_0.3 41 | 0026_0.1_0.3 42 | 0115_0.1_0.3 43 | 0183_0.1_0.3 44 | 0377_0.1_0.3 45 | 0286_0.1_0.3 46 | 0076_0.1_0.3 47 | 0733_0.1_0.3 48 | 0098_0.1_0.3 49 | 0102_0.1_0.3 50 | 0151_0.1_0.3 51 | 0007_0.1_0.3 52 | 0150_0.1_0.3 53 | 0446_0.1_0.3 54 | 0197_0.1_0.3 55 | 0505_0.1_0.3 56 | 0039_0.1_0.3 57 | 0231_0.1_0.3 58 | 5013_0.1_0.3 59 | 5002_0.1_0.3 60 | 0057_0.1_0.3 61 | 0402_0.1_0.3 62 | 0257_0.1_0.3 63 | 3346_0.1_0.3 64 | 0258_0.1_0.3 65 | 0156_0.1_0.3 66 | 0559_0.1_0.3 67 | 5004_0.1_0.3 68 | 0299_0.1_0.3 69 | 0071_0.1_0.3 70 | 0411_0.1_0.3 71 | 0056_0.1_0.3 72 | 0160_0.1_0.3 73 | 0064_0.1_0.3 74 | 0083_0.1_0.3 75 | 0290_0.1_0.3 76 | 0205_0.1_0.3 77 | 5009_0.1_0.3 78 | 0294_0.1_0.3 79 | 0027_0.1_0.3 80 | 0327_0.1_0.3 81 | 0065_0.1_0.3 82 | 5007_0.1_0.3 83 | 0149_0.1_0.3 84 | 0493_0.1_0.3 85 | 0175_0.1_0.3 86 | 5005_0.1_0.3 87 | 0034_0.1_0.3 88 | 0001_0.1_0.3 89 | 0067_0.1_0.3 90 | 0496_0.1_0.3 91 | 0189_0.1_0.3 92 | 5003_0.1_0.3 93 | 0269_0.1_0.3 94 | 0137_0.1_0.3 95 | 0277_0.1_0.3 96 | 0235_0.1_0.3 97 | 0101_0.1_0.3 98 | 0217_0.1_0.3 99 | 0078_0.1_0.3 100 | 0005_0.1_0.3 101 | 0312_0.1_0.3 102 | 0455_0.1_0.3 103 | 0162_0.1_0.3 104 | 0335_0.1_0.3 105 | 0037_0.1_0.3 106 | 0303_0.1_0.3 107 | 0323_0.1_0.3 108 | 0200_0.1_0.3 109 | 0046_0.1_0.3 110 | 0285_0.1_0.3 111 | 0094_0.1_0.3 112 | 0048_0.1_0.3 113 | 0086_0.1_0.3 114 | 0004_0.1_0.3 115 | 0148_0.1_0.3 116 | 5011_0.1_0.3 117 | 0129_0.1_0.3 118 | 0281_0.1_0.3 119 | 0478_0.1_0.3 120 | 0190_0.1_0.3 121 | 0389_0.1_0.3 122 | 0060_0.1_0.3 123 | 0012_0.1_0.3 124 | 0047_0.1_0.3 125 | 0122_0.1_0.3 126 | 0147_0.1_0.3 127 | 0080_0.1_0.3 128 | 0181_0.1_0.3 129 | 0087_0.1_0.3 130 | 5000_0.1_0.3 131 | 0472_0.1_0.3 132 | 0341_0.1_0.3 133 | 0058_0.1_0.3 134 | 0204_0.1_0.3 135 | 0036_0.1_0.3 136 | 0104_0.1_0.3 137 | 0380_0.1_0.3 138 | 5012_0.1_0.3 139 | 0061_0.1_0.3 140 | 0240_0.1_0.3 141 | 5017_0.1_0.3 142 | 5010_0.1_0.3 143 | 0090_0.1_0.3 144 | 0095_0.1_0.3 145 | 0394_0.1_0.3 146 | 0016_0.1_0.3 147 | 0224_0.1_0.3 148 | 0035_0.1_0.3 149 | 0070_0.1_0.3 150 | 5001_0.1_0.3 151 | 0271_0.1_0.3 152 | 0212_0.1_0.3 153 | 0348_0.1_0.3 154 | 0493_0.3_0.5 155 | 0240_0.3_0.5 156 | 0137_0.3_0.5 157 | 0181_0.3_0.5 158 | 0326_0.3_0.5 159 | 0377_0.3_0.5 160 | 0285_0.3_0.5 161 | 0205_0.3_0.5 162 | 0476_0.3_0.5 163 | 0286_0.3_0.5 164 | 0212_0.3_0.5 165 | 0269_0.3_0.5 166 | 0281_0.3_0.5 167 | 0048_0.3_0.5 168 | 0102_0.3_0.5 169 | 0026_0.3_0.5 170 | 0095_0.3_0.5 171 | 0086_0.3_0.5 172 | 0012_0.3_0.5 173 | 0478_0.3_0.5 174 | 0389_0.3_0.5 175 | 1017_0.3_0.5 176 | 0446_0.3_0.5 177 | 5010_0.3_0.5 178 | 0100_0.3_0.5 179 | 0004_0.3_0.5 180 | 0360_0.3_0.5 181 | 0005_0.3_0.5 182 | 0060_0.3_0.5 183 | 0007_0.3_0.5 184 | 0115_0.3_0.5 185 | 0335_0.3_0.5 186 | 0214_0.3_0.5 187 | 0034_0.3_0.5 188 | 0098_0.3_0.5 189 | 0122_0.3_0.5 190 | 0071_0.3_0.5 191 | 0037_0.3_0.5 192 | 0348_0.3_0.5 193 | 0312_0.3_0.5 194 | 0090_0.3_0.5 195 | 5012_0.3_0.5 196 | 0001_0.3_0.5 197 | 0257_0.3_0.5 198 | 0331_0.3_0.5 199 | 0307_0.3_0.5 200 | 0197_0.3_0.5 201 | 5009_0.3_0.5 202 | 0056_0.3_0.5 203 | 0341_0.3_0.5 204 | 0559_0.3_0.5 205 | 0496_0.3_0.5 206 | 0327_0.3_0.5 207 | 0041_0.3_0.5 208 | 0027_0.3_0.5 209 | 0149_0.3_0.5 210 | 0277_0.3_0.5 211 | 5008_0.3_0.5 212 | 0190_0.3_0.5 213 | 0183_0.3_0.5 214 | 0394_0.3_0.5 215 | 0150_0.3_0.5 216 | 0299_0.3_0.5 217 | 0482_0.3_0.5 218 | 0147_0.3_0.5 219 | 0058_0.3_0.5 220 | 5011_0.3_0.5 221 | 0061_0.3_0.5 222 | 0200_0.3_0.5 223 | 0189_0.3_0.5 224 | 0323_0.3_0.5 225 | 0062_0.3_0.5 226 | 0023_0.3_0.5 227 | 0042_0.3_0.5 228 | 0235_0.3_0.5 229 | 5006_0.3_0.5 230 | 0117_0.3_0.5 231 | 5013_0.3_0.5 232 | 0407_0.3_0.5 233 | 0080_0.3_0.5 234 | 0231_0.3_0.5 235 | 0243_0.3_0.5 236 | 0306_0.3_0.5 237 | 0217_0.3_0.5 238 | 5017_0.3_0.5 239 | 0017_0.3_0.5 240 | 0505_0.3_0.5 241 | 0472_0.3_0.5 242 | 0076_0.3_0.5 243 | 0402_0.3_0.5 244 | 0003_0.3_0.5 245 | 0064_0.3_0.5 246 | 0046_0.3_0.5 247 | 0290_0.3_0.5 248 | 0252_0.3_0.5 249 | 0258_0.3_0.5 250 | 0411_0.3_0.5 251 | 0067_0.3_0.5 252 | 0237_0.3_0.5 253 | 0099_0.3_0.5 254 | 0141_0.3_0.5 255 | 0406_0.3_0.5 256 | 0065_0.3_0.5 257 | 0039_0.3_0.5 258 | 5002_0.3_0.5 259 | 0104_0.3_0.5 260 | 0130_0.3_0.5 261 | 0455_0.3_0.5 262 | 0380_0.3_0.5 263 | 0107_0.3_0.5 264 | 5004_0.3_0.5 265 | 0036_0.3_0.5 266 | 5003_0.3_0.5 267 | 0013_0.3_0.5 268 | 0083_0.3_0.5 269 | 0070_0.3_0.5 270 | 0043_0.3_0.5 271 | 5000_0.3_0.5 272 | 0224_0.3_0.5 273 | 0156_0.3_0.5 274 | 0238_0.3_0.5 275 | 0094_0.3_0.5 276 | 0078_0.3_0.5 277 | 0101_0.3_0.5 278 | 0160_0.3_0.5 279 | 0275_0.3_0.5 280 | 5007_0.3_0.5 281 | 5005_0.3_0.5 282 | 0223_0.3_0.5 283 | 0057_0.3_0.5 284 | 0044_0.3_0.5 285 | 0148_0.3_0.5 286 | 0035_0.3_0.5 287 | 5001_0.3_0.5 288 | 0162_0.3_0.5 289 | -------------------------------------------------------------------------------- /gluefactory/datasets/megadepth_scene_lists/train_scenes.txt: -------------------------------------------------------------------------------- 1 | 0000 2 | 0001 3 | 0002 4 | 0003 5 | 0004 6 | 0005 7 | 0007 8 | 0008 9 | 0011 10 | 0012 11 | 0013 12 | 0015 13 | 0017 14 | 0019 15 | 0020 16 | 0021 17 | 0022 18 | 0023 19 | 0024 20 | 0025 21 | 0026 22 | 0027 23 | 0032 24 | 0035 25 | 0036 26 | 0037 27 | 0039 28 | 0042 29 | 0043 30 | 0046 31 | 0048 32 | 0050 33 | 0056 34 | 0057 35 | 0060 36 | 0061 37 | 0063 38 | 0065 39 | 0070 40 | 0080 41 | 0083 42 | 0086 43 | 0087 44 | 0092 45 | 0095 46 | 0098 47 | 0100 48 | 0101 49 | 0103 50 | 0104 51 | 0105 52 | 0107 53 | 0115 54 | 0117 55 | 0122 56 | 0130 57 | 0137 58 | 0143 59 | 0147 60 | 0148 61 | 0149 62 | 0150 63 | 0156 64 | 0160 65 | 0176 66 | 0183 67 | 0189 68 | 0190 69 | 0200 70 | 0214 71 | 0224 72 | 0235 73 | 0237 74 | 0240 75 | 0243 76 | 0258 77 | 0265 78 | 0269 79 | 0299 80 | 0312 81 | 0326 82 | 0327 83 | 0331 84 | 0335 85 | 0341 86 | 0348 87 | 0366 88 | 0377 89 | 0380 90 | 0394 91 | 0407 92 | 0411 93 | 0430 94 | 0446 95 | 0455 96 | 0472 97 | 0474 98 | 0476 99 | 0478 100 | 0493 101 | 0494 102 | 0496 103 | 0505 104 | 0559 105 | 0733 106 | 0860 107 | 1017 108 | 1589 109 | 4541 110 | 5004 111 | 5005 112 | 5006 113 | 5007 114 | 5009 115 | 5010 116 | 5012 117 | 5013 118 | 5017 119 | -------------------------------------------------------------------------------- /gluefactory/datasets/megadepth_scene_lists/train_scenes_clean.txt: -------------------------------------------------------------------------------- 1 | 0001 2 | 0003 3 | 0004 4 | 0005 5 | 0007 6 | 0012 7 | 0013 8 | 0016 9 | 0017 10 | 0023 11 | 0026 12 | 0027 13 | 0034 14 | 0035 15 | 0036 16 | 0037 17 | 0039 18 | 0041 19 | 0042 20 | 0043 21 | 0044 22 | 0046 23 | 0047 24 | 0048 25 | 0049 26 | 0056 27 | 0057 28 | 0058 29 | 0060 30 | 0061 31 | 0062 32 | 0064 33 | 0065 34 | 0067 35 | 0070 36 | 0071 37 | 0076 38 | 0078 39 | 0080 40 | 0083 41 | 0086 42 | 0087 43 | 0090 44 | 0094 45 | 0095 46 | 0098 47 | 0099 48 | 0100 49 | 0101 50 | 0102 51 | 0104 52 | 0107 53 | 0115 54 | 0117 55 | 0122 56 | 0129 57 | 0130 58 | 0137 59 | 0141 60 | 0147 61 | 0148 62 | 0149 63 | 0150 64 | 0151 65 | 0156 66 | 0160 67 | 0162 68 | 0175 69 | 0181 70 | 0183 71 | 0185 72 | 0186 73 | 0189 74 | 0190 75 | 0197 76 | 0200 77 | 0204 78 | 0205 79 | 0212 80 | 0214 81 | 0217 82 | 0223 83 | 0224 84 | 0231 85 | 0235 86 | 0237 87 | 0238 88 | 0240 89 | 0243 90 | 0252 91 | 0257 92 | 0258 93 | 0269 94 | 0271 95 | 0275 96 | 0277 97 | 0281 98 | 0285 99 | 0286 100 | 0290 101 | 0294 102 | 0299 103 | 0303 104 | 0306 105 | 0307 106 | 0312 107 | 0323 108 | 0326 109 | 0327 110 | 0331 111 | 0335 112 | 0341 113 | 0348 114 | 0360 115 | 0377 116 | 0380 117 | 0387 118 | 0389 119 | 0394 120 | 0402 121 | 0406 122 | 0407 123 | 0411 124 | 0446 125 | 0455 126 | 0472 127 | 0476 128 | 0478 129 | 0482 130 | 0493 131 | 0496 132 | 0505 133 | 0559 134 | 0733 135 | 0768 136 | 1017 137 | 3346 138 | 5000 139 | 5001 140 | 5002 141 | 5003 142 | 5004 143 | 5005 144 | 5006 145 | 5007 146 | 5008 147 | 5009 148 | 5010 149 | 5011 150 | 5012 151 | 5013 152 | 5017 153 | 5018 -------------------------------------------------------------------------------- /gluefactory/datasets/megadepth_scene_lists/valid_scenes.txt: -------------------------------------------------------------------------------- 1 | 0016 2 | 0033 3 | 0034 4 | 0041 5 | 0044 6 | 0047 7 | 0049 8 | 0058 9 | 0062 10 | 0064 11 | 0067 12 | 0071 13 | 0076 14 | 0078 15 | 0090 16 | 0094 17 | 0099 18 | 0102 19 | 0121 20 | 0129 21 | 0133 22 | 0141 23 | 0151 24 | 0162 25 | 0168 26 | 0175 27 | 0177 28 | 0178 29 | 0181 30 | 0185 31 | 0186 32 | 0197 33 | 0204 34 | 0205 35 | 0209 36 | 0212 37 | 0217 38 | 0223 39 | 0229 40 | 0231 41 | 0238 42 | 0252 43 | 0257 44 | 0271 45 | 0275 46 | 0277 47 | 0281 48 | 0285 49 | 0286 50 | 0290 51 | 0294 52 | 0303 53 | 0306 54 | 0307 55 | 0323 56 | 0349 57 | 0360 58 | 0387 59 | 0389 60 | 0402 61 | 0406 62 | 0412 63 | 0443 64 | 0482 65 | 0768 66 | 1001 67 | 3346 68 | 5000 69 | 5001 70 | 5002 71 | 5003 72 | 5008 73 | 5011 74 | 5014 75 | 5015 76 | 5016 77 | 5018 78 | -------------------------------------------------------------------------------- /gluefactory/datasets/megadepth_scene_lists/valid_scenes_clean.txt: -------------------------------------------------------------------------------- 1 | 0015 2 | 0022 -------------------------------------------------------------------------------- /gluefactory/datasets/utils.py: -------------------------------------------------------------------------------- 1 | import cv2 2 | import numpy as np 3 | import torch 4 | 5 | 6 | def read_image(path, grayscale=False): 7 | """Read an image from path as RGB or grayscale""" 8 | mode = cv2.IMREAD_GRAYSCALE if grayscale else cv2.IMREAD_COLOR 9 | image = cv2.imread(str(path), mode) 10 | if image is None: 11 | raise IOError(f"Could not read image at {path}.") 12 | if not grayscale: 13 | image = image[..., ::-1] 14 | return image 15 | 16 | 17 | def numpy_image_to_torch(image): 18 | """Normalize the image tensor and reorder the dimensions.""" 19 | if image.ndim == 3: 20 | image = image.transpose((2, 0, 1)) # HxWxC to CxHxW 21 | elif image.ndim == 2: 22 | image = image[None] # add channel axis 23 | else: 24 | raise ValueError(f"Not an image: {image.shape}") 25 | return torch.tensor(image / 255.0, dtype=torch.float) 26 | 27 | 28 | def rotate_intrinsics(K, image_shape, rot): 29 | """image_shape is the shape of the image after rotation""" 30 | assert rot <= 3 31 | h, w = image_shape[:2][:: -1 if (rot % 2) else 1] 32 | fx, fy, cx, cy = K[0, 0], K[1, 1], K[0, 2], K[1, 2] 33 | rot = rot % 4 34 | if rot == 1: 35 | return np.array( 36 | [[fy, 0.0, cy], [0.0, fx, w - cx], [0.0, 0.0, 1.0]], dtype=K.dtype 37 | ) 38 | elif rot == 2: 39 | return np.array( 40 | [[fx, 0.0, w - cx], [0.0, fy, h - cy], [0.0, 0.0, 1.0]], 41 | dtype=K.dtype, 42 | ) 43 | else: # if rot == 3: 44 | return np.array( 45 | [[fy, 0.0, h - cy], [0.0, fx, cx], [0.0, 0.0, 1.0]], dtype=K.dtype 46 | ) 47 | 48 | 49 | def rotate_pose_inplane(i_T_w, rot): 50 | rotation_matrices = [ 51 | np.array( 52 | [ 53 | [np.cos(r), -np.sin(r), 0.0, 0.0], 54 | [np.sin(r), np.cos(r), 0.0, 0.0], 55 | [0.0, 0.0, 1.0, 0.0], 56 | [0.0, 0.0, 0.0, 1.0], 57 | ], 58 | dtype=np.float32, 59 | ) 60 | for r in [np.deg2rad(d) for d in (0, 270, 180, 90)] 61 | ] 62 | return np.dot(rotation_matrices[rot], i_T_w) 63 | 64 | 65 | def scale_intrinsics(K, scales): 66 | """Scale intrinsics after resizing the corresponding image.""" 67 | scales = np.diag(np.concatenate([scales, [1.0]])) 68 | return np.dot(scales.astype(K.dtype, copy=False), K) 69 | 70 | 71 | def get_divisible_wh(w, h, df=None): 72 | if df is not None: 73 | w_new, h_new = map(lambda x: int(x // df * df), [w, h]) 74 | else: 75 | w_new, h_new = w, h 76 | return w_new, h_new 77 | 78 | 79 | def resize(image, size, fn=None, interp="linear", df=None): 80 | """Resize an image to a fixed size, or according to max or min edge.""" 81 | h, w = image.shape[:2] 82 | if isinstance(size, int): 83 | scale = size / fn(h, w) 84 | h_new, w_new = int(round(h * scale)), int(round(w * scale)) 85 | w_new, h_new = get_divisible_wh(w_new, h_new, df) 86 | scale = (w_new / w, h_new / h) 87 | elif isinstance(size, (tuple, list)): 88 | h_new, w_new = size 89 | scale = (w_new / w, h_new / h) 90 | else: 91 | raise ValueError(f"Incorrect new size: {size}") 92 | mode = { 93 | "linear": cv2.INTER_LINEAR, 94 | "cubic": cv2.INTER_CUBIC, 95 | "nearest": cv2.INTER_NEAREST, 96 | "area": cv2.INTER_AREA, 97 | }[interp] 98 | return cv2.resize(image, (w_new, h_new), interpolation=mode), scale 99 | 100 | 101 | def crop(image, size, random=True, other=None, K=None, return_bbox=False): 102 | """Random or deterministic crop of an image, adjust depth and intrinsics.""" 103 | h, w = image.shape[:2] 104 | h_new, w_new = (size, size) if isinstance(size, int) else size 105 | top = np.random.randint(0, h - h_new + 1) if random else 0 106 | left = np.random.randint(0, w - w_new + 1) if random else 0 107 | image = image[top : top + h_new, left : left + w_new] 108 | ret = [image] 109 | if other is not None: 110 | ret += [other[top : top + h_new, left : left + w_new]] 111 | if K is not None: 112 | K[0, 2] -= left 113 | K[1, 2] -= top 114 | ret += [K] 115 | if return_bbox: 116 | ret += [(top, top + h_new, left, left + w_new)] 117 | return ret 118 | 119 | 120 | def zero_pad(size, *images): 121 | """zero pad images to size x size""" 122 | ret = [] 123 | for image in images: 124 | if image is None: 125 | ret.append(None) 126 | continue 127 | h, w = image.shape[:2] 128 | padded = np.zeros((size, size) + image.shape[2:], dtype=image.dtype) 129 | padded[:h, :w] = image 130 | ret.append(padded) 131 | return ret 132 | -------------------------------------------------------------------------------- /gluefactory/eval/__init__.py: -------------------------------------------------------------------------------- 1 | import torch 2 | 3 | from ..utils.tools import get_class 4 | from .eval_pipeline import EvalPipeline 5 | 6 | 7 | def get_benchmark(benchmark): 8 | return get_class(f"{__name__}.{benchmark}", EvalPipeline) 9 | 10 | 11 | @torch.no_grad() 12 | def run_benchmark(benchmark, eval_conf, experiment_dir, model=None): 13 | """This overwrites existing benchmarks""" 14 | experiment_dir.mkdir(exist_ok=True, parents=True) 15 | bm = get_benchmark(benchmark) 16 | 17 | pipeline = bm(eval_conf) 18 | return pipeline.run( 19 | experiment_dir, model=model, overwrite=True, overwrite_eval=True 20 | ) 21 | -------------------------------------------------------------------------------- /gluefactory/eval/eval_pipeline.py: -------------------------------------------------------------------------------- 1 | import json 2 | 3 | import h5py 4 | import numpy as np 5 | from omegaconf import OmegaConf 6 | 7 | 8 | def load_eval(dir): 9 | summaries, results = {}, {} 10 | with h5py.File(str(dir / "results.h5"), "r") as hfile: 11 | for k in hfile.keys(): 12 | r = np.array(hfile[k]) 13 | if len(r.shape) < 3: 14 | results[k] = r 15 | for k, v in hfile.attrs.items(): 16 | summaries[k] = v 17 | with open(dir / "summaries.json", "r") as f: 18 | s = json.load(f) 19 | summaries = {k: v if v is not None else np.nan for k, v in s.items()} 20 | return summaries, results 21 | 22 | 23 | def save_eval(dir, summaries, figures, results): 24 | with h5py.File(str(dir / "results.h5"), "w") as hfile: 25 | for k, v in results.items(): 26 | arr = np.array(v) 27 | if not np.issubdtype(arr.dtype, np.number): 28 | arr = arr.astype("object") 29 | hfile.create_dataset(k, data=arr) 30 | # just to be safe, not used in practice 31 | for k, v in summaries.items(): 32 | hfile.attrs[k] = v 33 | s = { 34 | k: float(v) if np.isfinite(v) else None 35 | for k, v in summaries.items() 36 | if not isinstance(v, list) 37 | } 38 | s = {**s, **{k: v for k, v in summaries.items() if isinstance(v, list)}} 39 | with open(dir / "summaries.json", "w") as f: 40 | json.dump(s, f, indent=4) 41 | 42 | for fig_name, fig in figures.items(): 43 | fig.savefig(dir / f"{fig_name}.png") 44 | 45 | 46 | def exists_eval(dir): 47 | return (dir / "results.h5").exists() and (dir / "summaries.json").exists() 48 | 49 | 50 | class EvalPipeline: 51 | default_conf = {} 52 | 53 | export_keys = [] 54 | optional_export_keys = [] 55 | 56 | def __init__(self, conf): 57 | """Assumes""" 58 | self.default_conf = OmegaConf.create(self.default_conf) 59 | self.conf = OmegaConf.merge(self.default_conf, conf) 60 | self._init(self.conf) 61 | 62 | def _init(self, conf): 63 | pass 64 | 65 | @classmethod 66 | def get_dataloader(self, data_conf=None): 67 | """Returns a data loader with samples for each eval datapoint""" 68 | raise NotImplementedError 69 | 70 | def get_predictions(self, experiment_dir, model=None, overwrite=False): 71 | """Export a prediction file for each eval datapoint""" 72 | raise NotImplementedError 73 | 74 | def run_eval(self, loader, pred_file): 75 | """Run the eval on cached predictions""" 76 | raise NotImplementedError 77 | 78 | def run(self, experiment_dir, model=None, overwrite=False, overwrite_eval=False): 79 | """Run export+eval loop""" 80 | self.save_conf( 81 | experiment_dir, overwrite=overwrite, overwrite_eval=overwrite_eval 82 | ) 83 | pred_file = self.get_predictions( 84 | experiment_dir, model=model, overwrite=overwrite 85 | ) 86 | 87 | f = {} 88 | if not exists_eval(experiment_dir) or overwrite_eval or overwrite: 89 | s, f, r = self.run_eval(self.get_dataloader(), pred_file) 90 | save_eval(experiment_dir, s, f, r) 91 | s, r = load_eval(experiment_dir) 92 | return s, f, r 93 | 94 | def save_conf(self, experiment_dir, overwrite=False, overwrite_eval=False): 95 | # store config 96 | conf_output_path = experiment_dir / "conf.yaml" 97 | if conf_output_path.exists(): 98 | saved_conf = OmegaConf.load(conf_output_path) 99 | if (saved_conf.data != self.conf.data) or ( 100 | saved_conf.model != self.conf.model 101 | ): 102 | assert ( 103 | overwrite 104 | ), "configs changed, add --overwrite to rerun experiment with new conf" 105 | if saved_conf.eval != self.conf.eval: 106 | assert ( 107 | overwrite or overwrite_eval 108 | ), "eval configs changed, add --overwrite_eval to rerun evaluation" 109 | OmegaConf.save(self.conf, experiment_dir / "conf.yaml") 110 | -------------------------------------------------------------------------------- /gluefactory/eval/inspect.py: -------------------------------------------------------------------------------- 1 | import argparse 2 | from collections import defaultdict 3 | from pathlib import Path 4 | from pprint import pprint 5 | 6 | import matplotlib 7 | import matplotlib.pyplot as plt 8 | 9 | from ..settings import EVAL_PATH 10 | from ..visualization.global_frame import GlobalFrame 11 | from ..visualization.two_view_frame import TwoViewFrame 12 | from . import get_benchmark 13 | from .eval_pipeline import load_eval 14 | 15 | if __name__ == "__main__": 16 | parser = argparse.ArgumentParser() 17 | parser.add_argument("benchmark", type=str) 18 | parser.add_argument("--x", type=str, default=None) 19 | parser.add_argument("--y", type=str, default=None) 20 | parser.add_argument("--backend", type=str, default=None) 21 | parser.add_argument( 22 | "--default_plot", type=str, default=TwoViewFrame.default_conf["default"] 23 | ) 24 | 25 | parser.add_argument("dotlist", nargs="*") 26 | args = parser.parse_intermixed_args() 27 | 28 | output_dir = Path(EVAL_PATH, args.benchmark) 29 | 30 | results = {} 31 | summaries = defaultdict(dict) 32 | 33 | predictions = {} 34 | 35 | if args.backend: 36 | matplotlib.use(args.backend) 37 | 38 | bm = get_benchmark(args.benchmark) 39 | loader = bm.get_dataloader() 40 | 41 | for name in args.dotlist: 42 | experiment_dir = output_dir / name 43 | pred_file = experiment_dir / "predictions.h5" 44 | s, results[name] = load_eval(experiment_dir) 45 | predictions[name] = pred_file 46 | for k, v in s.items(): 47 | summaries[k][name] = v 48 | 49 | pprint(summaries) 50 | 51 | plt.close("all") 52 | 53 | frame = GlobalFrame( 54 | {"child": {"default": args.default_plot}, **vars(args)}, 55 | results, 56 | loader, 57 | predictions, 58 | child_frame=TwoViewFrame, 59 | ) 60 | frame.draw() 61 | plt.show() 62 | -------------------------------------------------------------------------------- /gluefactory/eval/io.py: -------------------------------------------------------------------------------- 1 | import argparse 2 | from pathlib import Path 3 | from pprint import pprint 4 | from typing import Optional 5 | 6 | import pkg_resources 7 | from omegaconf import OmegaConf 8 | 9 | from ..models import get_model 10 | from ..settings import TRAINING_PATH 11 | from ..utils.experiments import load_experiment 12 | 13 | 14 | def parse_config_path(name_or_path: Optional[str], defaults: str) -> Path: 15 | default_configs = {} 16 | for c in pkg_resources.resource_listdir("gluefactory", str(defaults)): 17 | if c.endswith(".yaml"): 18 | default_configs[Path(c).stem] = Path( 19 | pkg_resources.resource_filename("gluefactory", defaults + c) 20 | ) 21 | if name_or_path is None: 22 | return None 23 | if name_or_path in default_configs: 24 | return default_configs[name_or_path] 25 | path = Path(name_or_path) 26 | if not path.exists(): 27 | raise FileNotFoundError( 28 | f"Cannot find the config file: {name_or_path}. " 29 | f"Not in the default configs {list(default_configs.keys())} " 30 | "and not an existing path." 31 | ) 32 | return Path(path) 33 | 34 | 35 | def extract_benchmark_conf(conf, benchmark): 36 | mconf = OmegaConf.create( 37 | { 38 | "model": conf.get("model", {}), 39 | } 40 | ) 41 | if "benchmarks" in conf.keys(): 42 | return OmegaConf.merge(mconf, conf.benchmarks.get(benchmark, {})) 43 | else: 44 | return mconf 45 | 46 | 47 | def parse_eval_args(benchmark, args, configs_path, default=None): 48 | conf = {"data": {}, "model": {}, "eval": {}} 49 | if args.conf: 50 | conf_path = parse_config_path(args.conf, configs_path) 51 | custom_conf = OmegaConf.load(conf_path) 52 | conf = extract_benchmark_conf(OmegaConf.merge(conf, custom_conf), benchmark) 53 | args.tag = ( 54 | args.tag if args.tag is not None else conf_path.name.replace(".yaml", "") 55 | ) 56 | 57 | cli_conf = OmegaConf.from_cli(args.dotlist) 58 | conf = OmegaConf.merge(conf, cli_conf) 59 | conf.checkpoint = args.checkpoint if args.checkpoint else conf.get("checkpoint") 60 | 61 | if conf.checkpoint and not conf.checkpoint.endswith(".tar"): 62 | checkpoint_conf = OmegaConf.load( 63 | TRAINING_PATH / conf.checkpoint / "config.yaml" 64 | ) 65 | conf = OmegaConf.merge(extract_benchmark_conf(checkpoint_conf, benchmark), conf) 66 | 67 | if default: 68 | conf = OmegaConf.merge(default, conf) 69 | 70 | if args.tag is not None: 71 | name = args.tag 72 | elif args.conf and conf.checkpoint: 73 | name = f"{args.conf}_{conf.checkpoint}" 74 | elif args.conf: 75 | name = args.conf 76 | elif conf.checkpoint: 77 | name = conf.checkpoint 78 | if len(args.dotlist) > 0 and not args.tag: 79 | name = name + "_" + ":".join(args.dotlist) 80 | print("Running benchmark:", benchmark) 81 | print("Experiment tag:", name) 82 | print("Config:") 83 | pprint(OmegaConf.to_container(conf)) 84 | return name, conf 85 | 86 | 87 | def load_model(model_conf, checkpoint): 88 | if checkpoint: 89 | model = load_experiment(checkpoint, conf=model_conf).eval() 90 | else: 91 | model = get_model("two_view_pipeline")(model_conf).eval() 92 | if not model.is_initialized(): 93 | raise ValueError( 94 | "The provided model has non-initialized parameters. " 95 | + "Try to load a checkpoint instead." 96 | ) 97 | return model 98 | 99 | 100 | def get_eval_parser(): 101 | parser = argparse.ArgumentParser() 102 | parser.add_argument("--tag", type=str, default=None) 103 | parser.add_argument("--checkpoint", type=str, default=None) 104 | parser.add_argument("--conf", type=str, default=None) 105 | parser.add_argument("--overwrite", action="store_true") 106 | parser.add_argument("--overwrite_eval", action="store_true") 107 | parser.add_argument("--plot", action="store_true") 108 | parser.add_argument("dotlist", nargs="*") 109 | return parser 110 | -------------------------------------------------------------------------------- /gluefactory/geometry/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/weitong8591/vop/e964dc779958e691d8511b78dd3c40bf7920e4e3/gluefactory/geometry/__init__.py -------------------------------------------------------------------------------- /gluefactory/models/__init__.py: -------------------------------------------------------------------------------- 1 | import importlib.util 2 | 3 | from ..utils.tools import get_class 4 | from .base_model import BaseModel 5 | 6 | 7 | def get_model(name): 8 | import_paths = [ 9 | name, 10 | f"{__name__}.{name}", 11 | f"{__name__}.extractors.{name}", # backward compatibility 12 | f"{__name__}.matchers.{name}", # backward compatibility 13 | ] 14 | for path in import_paths: 15 | try: 16 | spec = importlib.util.find_spec(path) 17 | except ModuleNotFoundError: 18 | spec = None 19 | if spec is not None: 20 | try: 21 | return get_class(path, BaseModel) 22 | except AssertionError: 23 | mod = __import__(path, fromlist=[""]) 24 | try: 25 | return mod.__main_model__ 26 | except AttributeError as exc: 27 | print(exc) 28 | continue 29 | 30 | raise RuntimeError(f'Model {name} not found in any of [{" ".join(import_paths)}]') 31 | -------------------------------------------------------------------------------- /gluefactory/models/backbones/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/weitong8591/vop/e964dc779958e691d8511b78dd3c40bf7920e4e3/gluefactory/models/backbones/__init__.py -------------------------------------------------------------------------------- /gluefactory/models/backbones/dino.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn.functional as F 3 | 4 | from ..base_model import BaseModel 5 | 6 | 7 | class Dino(BaseModel): 8 | default_conf = {"weights": "dino_vits8", "allow_resize": False} 9 | required_data_keys = ["image"] 10 | 11 | def _init(self, conf): 12 | self.net = torch.hub.load("facebookresearch/dino:main", conf.weights) 13 | self.set_initialized() 14 | 15 | def _forward(self, data): 16 | img = data["image"] 17 | import pdb; pdb.set_trace() 18 | if self.conf.allow_resize: 19 | img = F.upsample(img, [int(x // 14 * 14) for x in img.shape[-2:]]) 20 | desc, cls_token = self.net.get_intermediate_layers( 21 | img, n=1, return_class_token=True, reshape=True 22 | )[0] 23 | 24 | return { 25 | "features": desc, 26 | "global_descriptor": cls_token, 27 | "descriptors": desc.flatten(-2).transpose(-2, -1), 28 | } 29 | 30 | def loss(self, pred, data): 31 | raise NotImplementedError 32 | -------------------------------------------------------------------------------- /gluefactory/models/backbones/dinov2.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn.functional as F 3 | 4 | from ..base_model import BaseModel 5 | 6 | 7 | class DinoV2(BaseModel): 8 | default_conf = {"weights": "dinov2_vits14", "allow_resize": False} 9 | required_data_keys = ["image"] 10 | 11 | def _init(self, conf): 12 | self.net = torch.hub.load("facebookresearch/dinov2", conf.weights) 13 | # https://github.com/facebookresearch/dinov2/blob/main/dinov2/models/vision_transformer.py#L44 14 | self.set_initialized() 15 | 16 | def _forward(self, data): 17 | img = data["image"] 18 | if self.conf.allow_resize: 19 | img = F.upsample(img, [int(x // 14 * 14) for x in img.shape[-2:]]) 20 | desc, cls_token = self.net.get_intermediate_layers( 21 | img, n=1, return_class_token=True, reshape=True 22 | )[0] 23 | 24 | return { 25 | "features": desc, 26 | "global_descriptor": cls_token, 27 | "descriptors": desc.flatten(-2).transpose(-2, -1) 28 | } 29 | 30 | def loss(self, pred, data): 31 | raise NotImplementedError 32 | -------------------------------------------------------------------------------- /gluefactory/models/cache_loader.py: -------------------------------------------------------------------------------- 1 | import string 2 | 3 | import h5py 4 | import torch 5 | 6 | from ..datasets.base_dataset import collate 7 | from ..settings import DATA_PATH 8 | from ..utils.tensor import batch_to_device 9 | from .base_model import BaseModel 10 | from .utils.misc import pad_to_length 11 | 12 | 13 | def pad_local_features(pred: dict, seq_l: int): 14 | pred["keypoints"] = pad_to_length( 15 | pred["keypoints"], 16 | seq_l, 17 | -2, 18 | mode="random_c", 19 | ) 20 | if "keypoint_scores" in pred.keys(): 21 | pred["keypoint_scores"] = pad_to_length( 22 | pred["keypoint_scores"], seq_l, -1, mode="zeros" 23 | ) 24 | if "descriptors" in pred.keys(): 25 | pred["descriptors"] = pad_to_length( 26 | pred["descriptors"], seq_l, -2, mode="random" 27 | ) 28 | if "scales" in pred.keys(): 29 | pred["scales"] = pad_to_length(pred["scales"], seq_l, -1, mode="zeros") 30 | if "oris" in pred.keys(): 31 | pred["oris"] = pad_to_length(pred["oris"], seq_l, -1, mode="zeros") 32 | 33 | if "depth_keypoints" in pred.keys(): 34 | pred["depth_keypoints"] = pad_to_length( 35 | pred["depth_keypoints"], seq_l, -1, mode="zeros" 36 | ) 37 | if "valid_depth_keypoints" in pred.keys(): 38 | pred["valid_depth_keypoints"] = pad_to_length( 39 | pred["valid_depth_keypoints"], seq_l, -1, mode="zeros" 40 | ) 41 | return pred 42 | 43 | 44 | def pad_line_features(pred, seq_l: int = None): 45 | raise NotImplementedError 46 | 47 | 48 | def recursive_load(grp, pkeys): 49 | return { 50 | k: torch.from_numpy(grp[k].__array__()) 51 | if isinstance(grp[k], h5py.Dataset) 52 | else recursive_load(grp[k], list(grp.keys())) 53 | for k in pkeys 54 | } 55 | 56 | 57 | class CacheLoader(BaseModel): 58 | default_conf = { 59 | "path": "???", # can be a format string like exports/{scene}/ 60 | "data_keys": None, # load all keys 61 | "device": None, # load to same device as data 62 | "trainable": False, 63 | "add_data_path": True, 64 | "collate": True, 65 | "scale": ["keypoints", "lines", "orig_lines"], 66 | "padding_fn": None, 67 | "padding_length": None, # required for batching! 68 | "numeric_type": "float32", # [None, "float16", "float32", "float64"] 69 | } 70 | 71 | required_data_keys = ["name"] # we need an identifier 72 | 73 | def _init(self, conf): 74 | self.hfiles = {} 75 | self.padding_fn = conf.padding_fn 76 | if self.padding_fn is not None: 77 | self.padding_fn = eval(self.padding_fn) 78 | self.numeric_dtype = { 79 | None: None, 80 | "float16": torch.float16, 81 | "float32": torch.float32, 82 | "float64": torch.float64, 83 | }[conf.numeric_type] 84 | 85 | def _forward(self, data): 86 | preds = [] 87 | device = self.conf.device 88 | if not device: 89 | devices = set( 90 | [v.device for v in data.values() if isinstance(v, torch.Tensor)] 91 | ) 92 | if len(devices) == 0: 93 | device = "cpu" 94 | else: 95 | assert len(devices) == 1 96 | device = devices.pop() 97 | 98 | var_names = [x[1] for x in string.Formatter().parse(self.conf.path) if x[1]] 99 | for i, name in enumerate(data["name"]): 100 | fpath = self.conf.path.format(**{k: data[k][i] for k in var_names}) 101 | if self.conf.add_data_path: 102 | fpath = DATA_PATH / fpath 103 | hfile = h5py.File(str(fpath), "r") 104 | grp = hfile[name] 105 | pkeys = ( 106 | self.conf.data_keys if self.conf.data_keys is not None else grp.keys() 107 | ) 108 | pred = recursive_load(grp, pkeys) 109 | if self.numeric_dtype is not None: 110 | pred = { 111 | k: v 112 | if not isinstance(v, torch.Tensor) or not torch.is_floating_point(v) 113 | else v.to(dtype=self.numeric_dtype) 114 | for k, v in pred.items() 115 | } 116 | pred = batch_to_device(pred, device) 117 | for k, v in pred.items(): 118 | for pattern in self.conf.scale: 119 | if k.startswith(pattern): 120 | view_idx = k.replace(pattern, "") 121 | scales = ( 122 | data["scales"] 123 | if len(view_idx) == 0 124 | else data[f"view{view_idx}"]["scales"] 125 | ) 126 | pred[k] = pred[k] * scales[i] 127 | # use this function to fix number of keypoints etc. 128 | if self.padding_fn is not None: 129 | pred = self.padding_fn(pred, self.conf.padding_length) 130 | preds.append(pred) 131 | hfile.close() 132 | if self.conf.collate: 133 | return batch_to_device(collate(preds), device) 134 | else: 135 | assert len(preds) == 1 136 | return batch_to_device(preds[0], device) 137 | 138 | def loss(self, pred, data): 139 | raise NotImplementedError 140 | -------------------------------------------------------------------------------- /gluefactory/models/cosplace.py: -------------------------------------------------------------------------------- 1 | ''' 2 | Code for loading models trained with CosPlace as a global features extractor 3 | for geolocalization through image retrieval. 4 | Multiple models are available with different backbones. Below is a summary of 5 | models available (backbone : list of available output descriptors 6 | dimensionality). For example you can use a model based on a ResNet50 with 7 | descriptors dimensionality 1024. 8 | ResNet18: [32, 64, 128, 256, 512] 9 | ResNet50: [32, 64, 128, 256, 512, 1024, 2048] 10 | ResNet101: [32, 64, 128, 256, 512, 1024, 2048] 11 | ResNet152: [32, 64, 128, 256, 512, 1024, 2048] 12 | VGG16: [ 64, 128, 256, 512] 13 | 14 | CosPlace paper: https://arxiv.org/abs/2204.02287 15 | ''' 16 | 17 | import torch 18 | import torchvision.transforms as tvf 19 | 20 | from ..utils.base_model import BaseModel 21 | 22 | 23 | class CosPlace(BaseModel): 24 | default_conf = { 25 | 'backbone': 'ResNet50', 26 | 'fc_output_dim' : 2048 27 | } 28 | required_inputs = ['image'] 29 | def _init(self, conf): 30 | self.net = torch.hub.load( 31 | 'gmberton/CosPlace', 32 | 'get_trained_model', 33 | backbone=conf['backbone'], 34 | fc_output_dim=conf['fc_output_dim'] 35 | ).eval() 36 | 37 | mean = [0.485, 0.456, 0.406] 38 | std = [0.229, 0.224, 0.225] 39 | self.norm_rgb = tvf.Normalize(mean=mean, std=std) 40 | 41 | def _forward(self, data): 42 | image = self.norm_rgb(data['image']) 43 | desc = self.net(image) 44 | return { 45 | 'global_descriptor': desc, 46 | } 47 | -------------------------------------------------------------------------------- /gluefactory/models/extractors/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/weitong8591/vop/e964dc779958e691d8511b78dd3c40bf7920e4e3/gluefactory/models/extractors/__init__.py -------------------------------------------------------------------------------- /gluefactory/models/extractors/disk_kornia.py: -------------------------------------------------------------------------------- 1 | import kornia 2 | import torch 3 | 4 | from ..base_model import BaseModel 5 | from ..utils.misc import pad_and_stack 6 | 7 | 8 | class DISK(BaseModel): 9 | default_conf = { 10 | "weights": "depth", 11 | "dense_outputs": False, 12 | "max_num_keypoints": None, 13 | "desc_dim": 128, 14 | "nms_window_size": 5, 15 | "detection_threshold": 0.0, 16 | "force_num_keypoints": False, 17 | "pad_if_not_divisible": True, 18 | "chunk": 4, # for reduced VRAM in training 19 | } 20 | required_data_keys = ["image"] 21 | 22 | def _init(self, conf): 23 | self.model = kornia.feature.DISK.from_pretrained(conf.weights) 24 | self.set_initialized() 25 | 26 | def _get_dense_outputs(self, images): 27 | B = images.shape[0] 28 | if self.conf.pad_if_not_divisible: 29 | h, w = images.shape[2:] 30 | pd_h = 16 - h % 16 if h % 16 > 0 else 0 31 | pd_w = 16 - w % 16 if w % 16 > 0 else 0 32 | images = torch.nn.functional.pad(images, (0, pd_w, 0, pd_h), value=0.0) 33 | 34 | heatmaps, descriptors = self.model.heatmap_and_dense_descriptors(images) 35 | if self.conf.pad_if_not_divisible: 36 | heatmaps = heatmaps[..., :h, :w] 37 | descriptors = descriptors[..., :h, :w] 38 | 39 | keypoints = kornia.feature.disk.detector.heatmap_to_keypoints( 40 | heatmaps, 41 | n=self.conf.max_num_keypoints, 42 | window_size=self.conf.nms_window_size, 43 | score_threshold=self.conf.detection_threshold, 44 | ) 45 | 46 | features = [] 47 | for i in range(B): 48 | features.append(keypoints[i].merge_with_descriptors(descriptors[i])) 49 | 50 | return features, descriptors 51 | 52 | def _forward(self, data): 53 | image = data["image"] 54 | 55 | keypoints, scores, descriptors = [], [], [] 56 | if self.conf.dense_outputs: 57 | dense_descriptors = [] 58 | chunk = self.conf.chunk 59 | for i in range(0, image.shape[0], chunk): 60 | if self.conf.dense_outputs: 61 | features, d_descriptors = self._get_dense_outputs( 62 | image[: min(image.shape[0], i + chunk)] 63 | ) 64 | dense_descriptors.append(d_descriptors) 65 | else: 66 | features = self.model( 67 | image[: min(image.shape[0], i + chunk)], 68 | n=self.conf.max_num_keypoints, 69 | window_size=self.conf.nms_window_size, 70 | score_threshold=self.conf.detection_threshold, 71 | pad_if_not_divisible=self.conf.pad_if_not_divisible, 72 | ) 73 | keypoints += [f.keypoints for f in features] 74 | scores += [f.detection_scores for f in features] 75 | descriptors += [f.descriptors for f in features] 76 | del features 77 | 78 | if self.conf.force_num_keypoints: 79 | # pad to target_length 80 | target_length = self.conf.max_num_keypoints 81 | keypoints = pad_and_stack( 82 | keypoints, 83 | target_length, 84 | -2, 85 | mode="random_c", 86 | bounds=( 87 | 0, 88 | data.get("image_size", torch.tensor(image.shape[-2:])).min().item(), 89 | ), 90 | ) 91 | scores = pad_and_stack(scores, target_length, -1, mode="zeros") 92 | descriptors = pad_and_stack(descriptors, target_length, -2, mode="zeros") 93 | else: 94 | keypoints = torch.stack(keypoints, 0) 95 | scores = torch.stack(scores, 0) 96 | descriptors = torch.stack(descriptors, 0) 97 | 98 | pred = { 99 | "keypoints": keypoints.to(image) + 0.5, 100 | "keypoint_scores": scores.to(image), 101 | "descriptors": descriptors.to(image), 102 | } 103 | if self.conf.dense_outputs: 104 | pred["dense_descriptors"] = torch.cat(dense_descriptors, 0) 105 | return pred 106 | 107 | def loss(self, pred, data): 108 | raise NotImplementedError 109 | -------------------------------------------------------------------------------- /gluefactory/models/extractors/grid_extractor.py: -------------------------------------------------------------------------------- 1 | import math 2 | 3 | import torch 4 | 5 | from ..base_model import BaseModel 6 | 7 | 8 | def to_sequence(map): 9 | return map.flatten(-2).transpose(-1, -2) 10 | 11 | 12 | def to_map(sequence): 13 | n = sequence.shape[-2] 14 | e = math.isqrt(n) 15 | assert e * e == n 16 | assert e * e == n 17 | sequence.transpose(-1, -2).unflatten(-1, [e, e]) 18 | 19 | 20 | class GridExtractor(BaseModel): 21 | default_conf = {"cell_size": 14} 22 | required_data_keys = ["image"] 23 | 24 | def _init(self, conf): 25 | pass 26 | 27 | def _forward(self, data): 28 | b, c, h, w = data["image"].shape 29 | 30 | cgrid = ( 31 | torch.stack( 32 | torch.meshgrid( 33 | torch.arange( 34 | h // self.conf.cell_size, 35 | dtype=torch.float32, 36 | device=data["image"].device, 37 | ), 38 | torch.arange( 39 | w // self.conf.cell_size, 40 | dtype=torch.float32, 41 | device=data["image"].device, 42 | ), 43 | indexing="ij", 44 | )[::-1], 45 | dim=0, 46 | ) 47 | .unsqueeze(0) 48 | .repeat([b, 1, 1, 1]) 49 | * self.conf.cell_size 50 | + self.conf.cell_size / 2 51 | ) 52 | pred = { 53 | "grid": cgrid + 0.5, 54 | "keypoints": to_sequence(cgrid) + 0.5, 55 | } 56 | 57 | return pred 58 | 59 | def loss(self, pred, data): 60 | raise NotImplementedError 61 | -------------------------------------------------------------------------------- /gluefactory/models/extractors/keynet_affnet_hardnet.py: -------------------------------------------------------------------------------- 1 | import kornia 2 | import torch 3 | 4 | from ..base_model import BaseModel 5 | from ..utils.misc import pad_to_length 6 | 7 | 8 | class KeyNetAffNetHardNet(BaseModel): 9 | default_conf = { 10 | "max_num_keypoints": None, 11 | "desc_dim": 128, 12 | "upright": False, 13 | "scale_laf": 1.0, 14 | "chunk": 4, # for reduced VRAM in training 15 | } 16 | required_data_keys = ["image"] 17 | 18 | def _init(self, conf): 19 | self.model = kornia.feature.KeyNetHardNet( 20 | num_features=conf.max_num_keypoints, 21 | upright=conf.upright, 22 | scale_laf=conf.scale_laf, 23 | ) 24 | self.set_initialized() 25 | 26 | def _forward(self, data): 27 | image = data["image"] 28 | if image.shape[1] == 3: # RGB 29 | scale = image.new_tensor([0.299, 0.587, 0.114]).view(1, 3, 1, 1) 30 | image = (image * scale).sum(1, keepdim=True) 31 | lafs, scores, descs = [], [], [] 32 | im_size = data.get("image_size") 33 | for i in range(image.shape[0]): 34 | img_i = image[i : i + 1, :1] 35 | if im_size is not None: 36 | img_i = img_i[:, :, : im_size[i, 1], : im_size[i, 0]] 37 | laf, score, desc = self.model(img_i) 38 | xn = pad_to_length( 39 | kornia.feature.get_laf_center(laf), 40 | self.conf.max_num_keypoints, 41 | pad_dim=-2, 42 | mode="random_c", 43 | bounds=(0, min(img_i.shape[-2:])), 44 | ) 45 | laf = torch.cat( 46 | [ 47 | laf, 48 | kornia.feature.laf_from_center_scale_ori(xn[:, score.shape[-1] :]), 49 | ], 50 | -3, 51 | ) 52 | lafs.append(laf) 53 | scores.append(pad_to_length(score, self.conf.max_num_keypoints, -1)) 54 | descs.append(pad_to_length(desc, self.conf.max_num_keypoints, -2)) 55 | 56 | lafs = torch.cat(lafs, 0) 57 | scores = torch.cat(scores, 0) 58 | descs = torch.cat(descs, 0) 59 | keypoints = kornia.feature.get_laf_center(lafs) 60 | scales = kornia.feature.get_laf_scale(lafs)[..., 0] 61 | oris = kornia.feature.get_laf_orientation(lafs) 62 | pred = { 63 | "keypoints": keypoints, 64 | "scales": scales.squeeze(-1), 65 | "oris": oris.squeeze(-1), 66 | "lafs": lafs, 67 | "keypoint_scores": scores, 68 | "descriptors": descs, 69 | } 70 | 71 | return pred 72 | 73 | def loss(self, pred, data): 74 | raise NotImplementedError 75 | -------------------------------------------------------------------------------- /gluefactory/models/extractors/mixed.py: -------------------------------------------------------------------------------- 1 | import torch.nn.functional as F 2 | from omegaconf import OmegaConf 3 | 4 | from .. import get_model 5 | from ..base_model import BaseModel 6 | 7 | to_ctr = OmegaConf.to_container # convert DictConfig to dict 8 | 9 | 10 | class MixedExtractor(BaseModel): 11 | default_conf = { 12 | "detector": {"name": None}, 13 | "descriptor": {"name": None}, 14 | "interpolate_descriptors_from": None, # field name 15 | } 16 | 17 | required_data_keys = ["image"] 18 | required_cache_keys = [] 19 | 20 | def _init(self, conf): 21 | if conf.detector.name: 22 | self.detector = get_model(conf.detector.name)(to_ctr(conf.detector)) 23 | else: 24 | self.required_data_keys += ["cache"] 25 | self.required_cache_keys += ["keypoints"] 26 | 27 | if conf.descriptor.name: 28 | self.descriptor = get_model(conf.descriptor.name)(to_ctr(conf.descriptor)) 29 | else: 30 | self.required_data_keys += ["cache"] 31 | self.required_cache_keys += ["descriptors"] 32 | 33 | def _forward(self, data): 34 | if self.conf.detector.name: 35 | pred = self.detector(data) 36 | else: 37 | pred = data["cache"] 38 | if self.conf.detector.name: 39 | pred = {**pred, **self.descriptor({**pred, **data})} 40 | 41 | if self.conf.interpolate_descriptors_from: 42 | h, w = data["image"].shape[-2:] 43 | kpts = pred["keypoints"] 44 | pts = (kpts / kpts.new_tensor([[w, h]]) * 2 - 1)[:, None] 45 | pred["descriptors"] = ( 46 | F.grid_sample( 47 | pred[self.conf.interpolate_descriptors_from], 48 | pts, 49 | align_corners=False, 50 | mode="bilinear", 51 | ) 52 | .squeeze(-2) 53 | .transpose(-2, -1) 54 | .contiguous() 55 | ) 56 | 57 | return pred 58 | 59 | def loss(self, pred, data): 60 | losses = {} 61 | metrics = {} 62 | total = 0 63 | 64 | for k in ["detector", "descriptor"]: 65 | apply = True 66 | if "apply_loss" in self.conf[k].keys(): 67 | apply = self.conf[k].apply_loss 68 | if self.conf[k].name and apply: 69 | try: 70 | losses_, metrics_ = getattr(self, k).loss(pred, {**pred, **data}) 71 | except NotImplementedError: 72 | continue 73 | losses = {**losses, **losses_} 74 | metrics = {**metrics, **metrics_} 75 | total = losses_["total"] + total 76 | return {**losses, "total": total}, metrics 77 | -------------------------------------------------------------------------------- /gluefactory/models/extractors/sift_kornia.py: -------------------------------------------------------------------------------- 1 | import kornia 2 | import torch 3 | 4 | from ..base_model import BaseModel 5 | 6 | 7 | class KorniaSIFT(BaseModel): 8 | default_conf = { 9 | "has_detector": True, 10 | "has_descriptor": True, 11 | "max_num_keypoints": -1, 12 | "detection_threshold": None, 13 | "rootsift": True, 14 | } 15 | 16 | required_data_keys = ["image"] 17 | 18 | def _init(self, conf): 19 | self.sift = kornia.feature.SIFTFeature( 20 | num_features=self.conf.max_num_keypoints, rootsift=self.conf.rootsift 21 | ) 22 | self.set_initialized() 23 | 24 | def _forward(self, data): 25 | lafs, scores, descriptors = self.sift(data["image"]) 26 | keypoints = kornia.feature.get_laf_center(lafs) 27 | scales = kornia.feature.get_laf_scale(lafs).squeeze(-1).squeeze(-1) 28 | oris = kornia.feature.get_laf_orientation(lafs).squeeze(-1) 29 | pred = { 30 | "keypoints": keypoints, # @TODO: confirm keypoints are in corner convention 31 | "scales": scales, 32 | "oris": oris, 33 | "keypoint_scores": scores, 34 | } 35 | 36 | if self.conf.has_descriptor: 37 | pred["descriptors"] = descriptors 38 | 39 | pred = {k: pred[k].to(device=data["image"].device) for k in pred.keys()} 40 | 41 | pred["scales"] = pred["scales"] 42 | pred["oris"] = torch.deg2rad(pred["oris"]) 43 | return pred 44 | 45 | def loss(self, pred, data): 46 | raise NotImplementedError 47 | -------------------------------------------------------------------------------- /gluefactory/models/lines/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/weitong8591/vop/e964dc779958e691d8511b78dd3c40bf7920e4e3/gluefactory/models/lines/__init__.py -------------------------------------------------------------------------------- /gluefactory/models/lines/deeplsd.py: -------------------------------------------------------------------------------- 1 | import deeplsd.models.deeplsd_inference as deeplsd_inference 2 | import numpy as np 3 | import torch 4 | 5 | from ...settings import DATA_PATH 6 | from ..base_model import BaseModel 7 | 8 | 9 | class DeepLSD(BaseModel): 10 | default_conf = { 11 | "min_length": 15, 12 | "max_num_lines": None, 13 | "force_num_lines": False, 14 | "model_conf": { 15 | "detect_lines": True, 16 | "line_detection_params": { 17 | "merge": False, 18 | "grad_nfa": True, 19 | "filtering": "normal", 20 | "grad_thresh": 3, 21 | }, 22 | }, 23 | } 24 | required_data_keys = ["image"] 25 | 26 | def _init(self, conf): 27 | if self.conf.force_num_lines: 28 | assert ( 29 | self.conf.max_num_lines is not None 30 | ), "Missing max_num_lines parameter" 31 | ckpt = DATA_PATH / "weights/deeplsd_md.tar" 32 | if not ckpt.is_file(): 33 | self.download_model(ckpt) 34 | ckpt = torch.load(ckpt, map_location="cpu") 35 | self.net = deeplsd_inference.DeepLSD(conf.model_conf).eval() 36 | self.net.load_state_dict(ckpt["model"]) 37 | self.set_initialized() 38 | 39 | def download_model(self, path): 40 | import subprocess 41 | 42 | if not path.parent.is_dir(): 43 | path.parent.mkdir(parents=True, exist_ok=True) 44 | link = "https://cvg-data.inf.ethz.ch/DeepLSD/deeplsd_md.tar" 45 | cmd = ["wget", link, "-O", path] 46 | print("Downloading DeepLSD model...") 47 | subprocess.run(cmd, check=True) 48 | 49 | def _forward(self, data): 50 | image = data["image"] 51 | lines, line_scores, valid_lines = [], [], [] 52 | if image.shape[1] == 3: 53 | # Convert to grayscale 54 | scale = image.new_tensor([0.299, 0.587, 0.114]).view(1, 3, 1, 1) 55 | image = (image * scale).sum(1, keepdim=True) 56 | 57 | # Forward pass 58 | with torch.no_grad(): 59 | segs = self.net({"image": image})["lines"] 60 | 61 | # Line scores are the sqrt of the length 62 | for seg in segs: 63 | lengths = np.linalg.norm(seg[:, 0] - seg[:, 1], axis=1) 64 | segs = seg[lengths >= self.conf.min_length] 65 | scores = np.sqrt(lengths[lengths >= self.conf.min_length]) 66 | 67 | # Keep the best lines 68 | indices = np.argsort(-scores) 69 | if self.conf.max_num_lines is not None: 70 | indices = indices[: self.conf.max_num_lines] 71 | segs = segs[indices] 72 | scores = scores[indices] 73 | 74 | # Pad if necessary 75 | n = len(segs) 76 | valid_mask = np.ones(n, dtype=bool) 77 | if self.conf.force_num_lines: 78 | pad = self.conf.max_num_lines - n 79 | segs = np.concatenate( 80 | [segs, np.zeros((pad, 2, 2), dtype=np.float32)], axis=0 81 | ) 82 | scores = np.concatenate( 83 | [scores, np.zeros(pad, dtype=np.float32)], axis=0 84 | ) 85 | valid_mask = np.concatenate( 86 | [valid_mask, np.zeros(pad, dtype=bool)], axis=0 87 | ) 88 | 89 | lines.append(segs) 90 | line_scores.append(scores) 91 | valid_lines.append(valid_mask) 92 | 93 | # Batch if possible 94 | if len(image) == 1 or self.conf.force_num_lines: 95 | lines = torch.tensor(lines, dtype=torch.float, device=image.device) 96 | line_scores = torch.tensor( 97 | line_scores, dtype=torch.float, device=image.device 98 | ) 99 | valid_lines = torch.tensor( 100 | valid_lines, dtype=torch.bool, device=image.device 101 | ) 102 | 103 | return {"lines": lines, "line_scores": line_scores, "valid_lines": valid_lines} 104 | 105 | def loss(self, pred, data): 106 | raise NotImplementedError 107 | -------------------------------------------------------------------------------- /gluefactory/models/lines/lsd.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import torch 3 | from joblib import Parallel, delayed 4 | from pytlsd import lsd 5 | 6 | from ..base_model import BaseModel 7 | 8 | 9 | class LSD(BaseModel): 10 | default_conf = { 11 | "min_length": 15, 12 | "max_num_lines": None, 13 | "force_num_lines": False, 14 | "n_jobs": 4, 15 | } 16 | required_data_keys = ["image"] 17 | 18 | def _init(self, conf): 19 | if self.conf.force_num_lines: 20 | assert ( 21 | self.conf.max_num_lines is not None 22 | ), "Missing max_num_lines parameter" 23 | 24 | def detect_lines(self, img): 25 | # Run LSD 26 | segs = lsd(img) 27 | 28 | # Filter out keylines that do not meet the minimum length criteria 29 | lengths = np.linalg.norm(segs[:, 2:4] - segs[:, 0:2], axis=1) 30 | to_keep = lengths >= self.conf.min_length 31 | segs, lengths = segs[to_keep], lengths[to_keep] 32 | 33 | # Keep the best lines 34 | scores = segs[:, -1] * np.sqrt(lengths) 35 | segs = segs[:, :4].reshape(-1, 2, 2) 36 | indices = np.argsort(-scores) 37 | if self.conf.max_num_lines is not None: 38 | indices = indices[: self.conf.max_num_lines] 39 | segs = segs[indices] 40 | scores = scores[indices] 41 | 42 | # Pad if necessary 43 | n = len(segs) 44 | valid_mask = np.ones(n, dtype=bool) 45 | if self.conf.force_num_lines: 46 | pad = self.conf.max_num_lines - n 47 | segs = np.concatenate( 48 | [segs, np.zeros((pad, 2, 2), dtype=np.float32)], axis=0 49 | ) 50 | scores = np.concatenate([scores, np.zeros(pad, dtype=np.float32)], axis=0) 51 | valid_mask = np.concatenate([valid_mask, np.zeros(pad, dtype=bool)], axis=0) 52 | 53 | return segs, scores, valid_mask 54 | 55 | def _forward(self, data): 56 | # Convert to the right data format 57 | image = data["image"] 58 | if image.shape[1] == 3: 59 | # Convert to grayscale 60 | scale = image.new_tensor([0.299, 0.587, 0.114]).view(1, 3, 1, 1) 61 | image = (image * scale).sum(1, keepdim=True) 62 | device = image.device 63 | b_size = len(image) 64 | image = np.uint8(image.squeeze(1).cpu().numpy() * 255) 65 | 66 | # LSD detection in parallel 67 | if b_size == 1: 68 | lines, line_scores, valid_lines = self.detect_lines(image[0]) 69 | lines = [lines] 70 | line_scores = [line_scores] 71 | valid_lines = [valid_lines] 72 | else: 73 | lines, line_scores, valid_lines = zip( 74 | *Parallel(n_jobs=self.conf.n_jobs)( 75 | delayed(self.detect_lines)(img) for img in image 76 | ) 77 | ) 78 | 79 | # Batch if possible 80 | if b_size == 1 or self.conf.force_num_lines: 81 | lines = torch.tensor(lines, dtype=torch.float, device=device) 82 | line_scores = torch.tensor(line_scores, dtype=torch.float, device=device) 83 | valid_lines = torch.tensor(valid_lines, dtype=torch.bool, device=device) 84 | 85 | return {"lines": lines, "line_scores": line_scores, "valid_lines": valid_lines} 86 | 87 | def loss(self, pred, data): 88 | raise NotImplementedError 89 | -------------------------------------------------------------------------------- /gluefactory/models/matchers/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/weitong8591/vop/e964dc779958e691d8511b78dd3c40bf7920e4e3/gluefactory/models/matchers/__init__.py -------------------------------------------------------------------------------- /gluefactory/models/matchers/adalam.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/weitong8591/vop/e964dc779958e691d8511b78dd3c40bf7920e4e3/gluefactory/models/matchers/adalam.py -------------------------------------------------------------------------------- /gluefactory/models/matchers/depth_matcher.py: -------------------------------------------------------------------------------- 1 | import torch 2 | 3 | from ...geometry.gt_generation import ( 4 | gt_line_matches_from_pose_depth, 5 | gt_matches_from_pose_depth, 6 | ) 7 | from ..base_model import BaseModel 8 | 9 | 10 | class DepthMatcher(BaseModel): 11 | default_conf = { 12 | # GT parameters for points 13 | "use_points": True, 14 | "th_positive": 3.0, 15 | "th_negative": 5.0, 16 | "th_epi": None, # add some more epi outliers 17 | "th_consistency": None, # check for projection consistency in px 18 | # GT parameters for lines 19 | "use_lines": False, 20 | "n_line_sampled_pts": 50, 21 | "line_perp_dist_th": 5, 22 | "overlap_th": 0.2, 23 | "min_visibility_th": 0.5, 24 | } 25 | 26 | required_data_keys = ["view0", "view1", "T_0to1", "T_1to0"] 27 | 28 | def _init(self, conf): 29 | # TODO (iago): Is this just boilerplate code? 30 | if self.conf.use_points: 31 | self.required_data_keys += ["keypoints0", "keypoints1"] 32 | if self.conf.use_lines: 33 | self.required_data_keys += [ 34 | "lines0", 35 | "lines1", 36 | "valid_lines0", 37 | "valid_lines1", 38 | ] 39 | 40 | @torch.cuda.amp.custom_fwd(cast_inputs=torch.float32) 41 | def _forward(self, data): 42 | result = {} 43 | if self.conf.use_points: 44 | if "depth_keypoints0" in data: 45 | keys = [ 46 | "depth_keypoints0", 47 | "valid_depth_keypoints0", 48 | "depth_keypoints1", 49 | "valid_depth_keypoints1", 50 | ] 51 | kw = {k: data[k] for k in keys} 52 | else: 53 | kw = {} 54 | result = gt_matches_from_pose_depth( 55 | data["keypoints0"], 56 | data["keypoints1"], 57 | data, 58 | pos_th=self.conf.th_positive, 59 | neg_th=self.conf.th_negative, 60 | epi_th=self.conf.th_epi, 61 | cc_th=self.conf.th_consistency, 62 | **kw, 63 | ) 64 | if self.conf.use_lines: 65 | line_assignment, line_m0, line_m1 = gt_line_matches_from_pose_depth( 66 | data["lines0"], 67 | data["lines1"], 68 | data["valid_lines0"], 69 | data["valid_lines1"], 70 | data, 71 | self.conf.n_line_sampled_pts, 72 | self.conf.line_perp_dist_th, 73 | self.conf.overlap_th, 74 | self.conf.min_visibility_th, 75 | ) 76 | result["line_matches0"] = line_m0 77 | result["line_matches1"] = line_m1 78 | result["line_assignment"] = line_assignment 79 | return result 80 | 81 | def loss(self, pred, data): 82 | raise NotImplementedError 83 | -------------------------------------------------------------------------------- /gluefactory/models/matchers/homography_matcher.py: -------------------------------------------------------------------------------- 1 | from ...geometry.gt_generation import ( 2 | gt_line_matches_from_homography, 3 | gt_matches_from_homography, 4 | ) 5 | from ..base_model import BaseModel 6 | 7 | 8 | class HomographyMatcher(BaseModel): 9 | default_conf = { 10 | # GT parameters for points 11 | "use_points": True, 12 | "th_positive": 3.0, 13 | "th_negative": 3.0, 14 | # GT parameters for lines 15 | "use_lines": False, 16 | "n_line_sampled_pts": 50, 17 | "line_perp_dist_th": 5, 18 | "overlap_th": 0.2, 19 | "min_visibility_th": 0.5, 20 | } 21 | 22 | required_data_keys = ["H_0to1"] 23 | 24 | def _init(self, conf): 25 | # TODO (iago): Is this just boilerplate code? 26 | if self.conf.use_points: 27 | self.required_data_keys += ["keypoints0", "keypoints1"] 28 | if self.conf.use_lines: 29 | self.required_data_keys += [ 30 | "lines0", 31 | "lines1", 32 | "valid_lines0", 33 | "valid_lines1", 34 | ] 35 | 36 | def _forward(self, data): 37 | result = {} 38 | if self.conf.use_points: 39 | result = gt_matches_from_homography( 40 | data["keypoints0"], 41 | data["keypoints1"], 42 | data["H_0to1"], 43 | pos_th=self.conf.th_positive, 44 | neg_th=self.conf.th_negative, 45 | ) 46 | if self.conf.use_lines: 47 | line_assignment, line_m0, line_m1 = gt_line_matches_from_homography( 48 | data["lines0"], 49 | data["lines1"], 50 | data["valid_lines0"], 51 | data["valid_lines1"], 52 | data["view0"]["image"].shape, 53 | data["view1"]["image"].shape, 54 | data["H_0to1"], 55 | self.conf.n_line_sampled_pts, 56 | self.conf.line_perp_dist_th, 57 | self.conf.overlap_th, 58 | self.conf.min_visibility_th, 59 | ) 60 | result["line_matches0"] = line_m0 61 | result["line_matches1"] = line_m1 62 | result["line_assignment"] = line_assignment 63 | return result 64 | 65 | def loss(self, pred, data): 66 | raise NotImplementedError 67 | -------------------------------------------------------------------------------- /gluefactory/models/matchers/kornia_loftr.py: -------------------------------------------------------------------------------- 1 | import kornia 2 | import torch 3 | 4 | from ...models import BaseModel 5 | 6 | 7 | class LoFTRModule(BaseModel): 8 | default_conf = { 9 | "topk": None, 10 | "zero_pad": False, 11 | } 12 | required_data_keys = ["view0", "view1"] 13 | 14 | def _init(self, conf): 15 | self.net = kornia.feature.LoFTR(pretrained="outdoor") 16 | self.set_initialized() 17 | 18 | def _forward(self, data): 19 | image0 = data["view0"]["image"] 20 | image1 = data["view1"]["image"] 21 | if self.conf.zero_pad: 22 | image0, mask0 = self.zero_pad(image0) 23 | image1, mask1 = self.zero_pad(image1) 24 | res = self.net( 25 | {"image0": image0, "image1": image1, "mask0": mask0, "mask1": mask1} 26 | ) 27 | res = self.net({"image0": image0, "image1": image1}) 28 | else: 29 | res = self.net({"image0": image0, "image1": image1}) 30 | topk = self.conf.topk 31 | if topk is not None and res["confidence"].shape[-1] > topk: 32 | _, top = torch.topk(res["confidence"], topk, -1) 33 | m_kpts0 = res["keypoints0"][None][:, top] 34 | m_kpts1 = res["keypoints1"][None][:, top] 35 | scores = res["confidence"][None][:, top] 36 | else: 37 | m_kpts0 = res["keypoints0"][None] 38 | m_kpts1 = res["keypoints1"][None] 39 | scores = res["confidence"][None] 40 | 41 | m0 = torch.arange(0, scores.shape[-1]).to(scores.device)[None] 42 | m1 = torch.arange(0, scores.shape[-1]).to(scores.device)[None] 43 | return { 44 | "matches0": m0, 45 | "matches1": m1, 46 | "matching_scores0": scores, 47 | "keypoints0": m_kpts0, 48 | "keypoints1": m_kpts1, 49 | "keypoint_scores0": scores, 50 | "keypoint_scores1": scores, 51 | "matching_scores1": scores, 52 | } 53 | 54 | def zero_pad(self, img): 55 | b, c, h, w = img.shape 56 | if h == w: 57 | return img 58 | s = max(h, w) 59 | image = torch.zeros((b, c, s, s)).to(img) 60 | image[:, :, :h, :w] = img 61 | mask = torch.zeros_like(image) 62 | mask[:, :, :h, :w] = 1.0 63 | return image, mask.squeeze(0).float() 64 | 65 | def loss(self, pred, data): 66 | return NotImplementedError 67 | -------------------------------------------------------------------------------- /gluefactory/models/matchers/lightglue_pretrained.py: -------------------------------------------------------------------------------- 1 | from lightglue import LightGlue as LightGlue_ 2 | from omegaconf import OmegaConf 3 | 4 | from ..base_model import BaseModel 5 | 6 | 7 | class LightGlue(BaseModel): 8 | default_conf = {"features": "superpoint", **LightGlue_.default_conf} 9 | required_data_keys = [ 10 | "view0", 11 | "keypoints0", 12 | "descriptors0", 13 | "view1", 14 | "keypoints1", 15 | "descriptors1", 16 | ] 17 | 18 | def _init(self, conf): 19 | dconf = OmegaConf.to_container(conf) 20 | self.net = LightGlue_(dconf.pop("features"), **dconf) 21 | self.set_initialized() 22 | 23 | def _forward(self, data): 24 | required_keys = ["keypoints", "descriptors", "scales", "oris"] 25 | view0 = { 26 | **data["view0"], 27 | **{k: data[k + "0"] for k in required_keys if (k + "0") in data}, 28 | } 29 | view1 = { 30 | **data["view1"], 31 | **{k: data[k + "1"] for k in required_keys if (k + "1") in data}, 32 | } 33 | return self.net({"image0": view0, "image1": view1}) 34 | 35 | def loss(pred, data): 36 | raise NotImplementedError 37 | -------------------------------------------------------------------------------- /gluefactory/models/matchers/nearest_neighbor_matcher.py: -------------------------------------------------------------------------------- 1 | """ 2 | Nearest neighbor matcher for normalized descriptors. 3 | Optionally apply the mutual check and threshold the distance or ratio. 4 | """ 5 | 6 | import logging 7 | 8 | import torch 9 | import torch.nn.functional as F 10 | 11 | from ..base_model import BaseModel 12 | from ..utils.metrics import matcher_metrics 13 | 14 | 15 | @torch.no_grad() 16 | def find_nn(sim, ratio_thresh, distance_thresh): 17 | sim_nn, ind_nn = sim.topk(2 if ratio_thresh else 1, dim=-1, largest=True) 18 | dist_nn = 2 * (1 - sim_nn) 19 | mask = torch.ones(ind_nn.shape[:-1], dtype=torch.bool, device=sim.device) 20 | if ratio_thresh: 21 | mask = mask & (dist_nn[..., 0] <= (ratio_thresh**2) * dist_nn[..., 1]) 22 | if distance_thresh: 23 | mask = mask & (dist_nn[..., 0] <= distance_thresh**2) 24 | matches = torch.where(mask, ind_nn[..., 0], ind_nn.new_tensor(-1)) 25 | return matches 26 | 27 | 28 | def mutual_check(m0, m1): 29 | inds0 = torch.arange(m0.shape[-1], device=m0.device) 30 | inds1 = torch.arange(m1.shape[-1], device=m1.device) 31 | loop0 = torch.gather(m1, -1, torch.where(m0 > -1, m0, m0.new_tensor(0))) 32 | loop1 = torch.gather(m0, -1, torch.where(m1 > -1, m1, m1.new_tensor(0))) 33 | m0_new = torch.where((m0 > -1) & (inds0 == loop0), m0, m0.new_tensor(-1)) 34 | m1_new = torch.where((m1 > -1) & (inds1 == loop1), m1, m1.new_tensor(-1)) 35 | return m0_new, m1_new 36 | 37 | 38 | class NearestNeighborMatcher(BaseModel): 39 | default_conf = { 40 | "ratio_thresh": None, 41 | "distance_thresh": None, 42 | "mutual_check": True, 43 | "loss": None, 44 | } 45 | required_data_keys = ["descriptors0", "descriptors1"] 46 | 47 | def _init(self, conf): 48 | if conf.loss == "N_pair": 49 | temperature = torch.nn.Parameter(torch.tensor(1.0)) 50 | self.register_parameter("temperature", temperature) 51 | 52 | def _forward(self, data): 53 | sim = torch.einsum("bnd,bmd->bnm", data["descriptors0"], data["descriptors1"]) 54 | matches0 = find_nn(sim, self.conf.ratio_thresh, self.conf.distance_thresh) 55 | matches1 = find_nn( 56 | sim.transpose(1, 2), self.conf.ratio_thresh, self.conf.distance_thresh 57 | ) 58 | if self.conf.mutual_check: 59 | matches0, matches1 = mutual_check(matches0, matches1) 60 | b, m, n = sim.shape 61 | la = sim.new_zeros(b, m + 1, n + 1) 62 | la[:, :-1, :-1] = F.log_softmax(sim, -1) + F.log_softmax(sim, -2) 63 | mscores0 = (matches0 > -1).float() 64 | mscores1 = (matches1 > -1).float() 65 | return { 66 | "matches0": matches0, 67 | "matches1": matches1, 68 | "matching_scores0": mscores0, 69 | "matching_scores1": mscores1, 70 | "similarity": sim, 71 | "log_assignment": la, 72 | } 73 | 74 | def loss(self, pred, data): 75 | losses = {} 76 | if self.conf.loss == "N_pair": 77 | sim = pred["similarity"] 78 | if torch.any(sim > (1.0 + 1e-6)): 79 | logging.warning(f"Similarity larger than 1, max={sim.max()}") 80 | scores = torch.sqrt(torch.clamp(2 * (1 - sim), min=1e-6)) 81 | scores = self.temperature * (2 - scores) 82 | assert not torch.any(torch.isnan(scores)), torch.any(torch.isnan(sim)) 83 | prob0 = torch.nn.functional.log_softmax(scores, 2) 84 | prob1 = torch.nn.functional.log_softmax(scores, 1) 85 | 86 | assignment = data["gt_assignment"].float() 87 | num = torch.max(assignment.sum((1, 2)), assignment.new_tensor(1)) 88 | nll0 = (prob0 * assignment).sum((1, 2)) / num 89 | nll1 = (prob1 * assignment).sum((1, 2)) / num 90 | nll = -(nll0 + nll1) / 2 91 | losses["n_pair_nll"] = losses["total"] = nll 92 | losses["num_matchable"] = num 93 | losses["n_pair_temperature"] = self.temperature[None] 94 | else: 95 | raise NotImplementedError 96 | metrics = {} if self.training else matcher_metrics(pred, data) 97 | return losses, metrics 98 | -------------------------------------------------------------------------------- /gluefactory/models/triplet_pipeline.py: -------------------------------------------------------------------------------- 1 | """ 2 | A two-view sparse feature matching pipeline on triplets. 3 | 4 | If a triplet is found, runs the extractor on three images and 5 | then runs matcher/filter/solver for all three pairs. 6 | 7 | Losses and metrics get accumulated accordingly. 8 | 9 | If no triplet is found, this falls back to two_view_pipeline.py 10 | """ 11 | 12 | import torch 13 | 14 | from ..utils.misc import get_twoview, stack_twoviews, unstack_twoviews 15 | from .two_view_pipeline import TwoViewPipeline 16 | 17 | 18 | def has_triplet(data): 19 | # we already check for image0 and image1 in required_keys 20 | return "view2" in data.keys() 21 | 22 | 23 | class TripletPipeline(TwoViewPipeline): 24 | default_conf = {"batch_triplets": True, **TwoViewPipeline.default_conf} 25 | 26 | def _forward(self, data): 27 | if not has_triplet(data): 28 | return super()._forward(data) 29 | # the two-view outputs are stored in 30 | # pred['0to1'],pred['0to2'], pred['1to2'] 31 | 32 | assert not self.conf.run_gt_in_forward 33 | pred0 = self.extract_view(data, "0") 34 | pred1 = self.extract_view(data, "1") 35 | pred2 = self.extract_view(data, "2") 36 | 37 | pred = {} 38 | pred = { 39 | **{k + "0": v for k, v in pred0.items()}, 40 | **{k + "1": v for k, v in pred1.items()}, 41 | **{k + "2": v for k, v in pred2.items()}, 42 | } 43 | 44 | def predict_twoview(pred, data): 45 | # forward pass 46 | if self.conf.matcher.name: 47 | pred = {**pred, **self.matcher({**data, **pred})} 48 | 49 | if self.conf.filter.name: 50 | pred = {**pred, **self.filter({**m_data, **pred})} 51 | 52 | if self.conf.solver.name: 53 | pred = {**pred, **self.solver({**m_data, **pred})} 54 | return pred 55 | 56 | if self.conf.batch_triplets: 57 | B = data["image1"].shape[0] 58 | # stack on batch dimension 59 | m_data = stack_twoviews(data) 60 | m_pred = stack_twoviews(pred) 61 | 62 | # forward pass 63 | m_pred = predict_twoview(m_pred, m_data) 64 | 65 | # unstack 66 | pred = {**pred, **unstack_twoviews(m_pred, B)} 67 | else: 68 | for idx in ["0to1", "0to2", "1to2"]: 69 | m_data = get_twoview(data, idx) 70 | m_pred = get_twoview(pred, idx) 71 | pred[idx] = predict_twoview(m_pred, m_data) 72 | return pred 73 | 74 | def loss(self, pred, data): 75 | if not has_triplet(data): 76 | return super().loss(pred, data) 77 | if self.conf.batch_triplets: 78 | m_data = stack_twoviews(data) 79 | m_pred = stack_twoviews(pred) 80 | losses, metrics = super().loss(m_pred, m_data) 81 | else: 82 | losses = {} 83 | metrics = {} 84 | for idx in ["0to1", "0to2", "1to2"]: 85 | data_i = get_twoview(data, idx) 86 | pred_i = pred[idx] 87 | losses_i, metrics_i = super().loss(pred_i, data_i) 88 | for k, v in losses_i.items(): 89 | if k in losses.keys(): 90 | losses[k] = losses[k] + v 91 | else: 92 | losses[k] = v 93 | for k, v in metrics_i.items(): 94 | if k in metrics.keys(): 95 | metrics[k] = torch.cat([metrics[k], v], 0) 96 | else: 97 | metrics[k] = v 98 | 99 | return losses, metrics 100 | -------------------------------------------------------------------------------- /gluefactory/models/two_view_pipeline.py: -------------------------------------------------------------------------------- 1 | """ 2 | A two-view sparse feature matching pipeline. 3 | 4 | This model contains sub-models for each step: 5 | feature extraction, feature matching, outlier filtering, pose estimation. 6 | Each step is optional, and the features or matches can be provided as input. 7 | Default: SuperPoint with nearest neighbor matching. 8 | 9 | Convention for the matches: m0[i] is the index of the keypoint in image 1 10 | that corresponds to the keypoint i in image 0. m0[i] = -1 if i is unmatched. 11 | """ 12 | 13 | from omegaconf import OmegaConf 14 | 15 | from . import get_model 16 | from .base_model import BaseModel 17 | to_ctr = OmegaConf.to_container # convert DictConfig to dict 18 | 19 | 20 | class TwoViewPipeline(BaseModel): 21 | default_conf = { 22 | "extractor": { 23 | "name": None, 24 | "trainable": False, 25 | }, 26 | "matcher": {"name": None}, 27 | "filter": {"name": None}, 28 | "solver": {"name": None}, 29 | "ground_truth": {"name": None}, 30 | "allow_no_extract": False, 31 | "run_gt_in_forward": False, 32 | } 33 | required_data_keys = ["view0", "view1"] 34 | strict_conf = False # need to pass new confs to children models 35 | components = [ 36 | "extractor", 37 | "matcher", 38 | "filter", 39 | "solver", 40 | "ground_truth", 41 | ] 42 | 43 | def _init(self, conf): 44 | if conf.extractor.name: 45 | self.extractor = get_model(conf.extractor.name)(to_ctr(conf.extractor)) 46 | 47 | if conf.matcher.name: 48 | self.matcher = get_model(conf.matcher.name)(to_ctr(conf.matcher)) 49 | 50 | if conf.filter.name: 51 | self.filter = get_model(conf.filter.name)(to_ctr(conf.filter)) 52 | 53 | if conf.solver.name: 54 | self.solver = get_model(conf.solver.name)(to_ctr(conf.solver)) 55 | 56 | if conf.ground_truth.name: 57 | self.ground_truth = get_model(conf.ground_truth.name)( 58 | to_ctr(conf.ground_truth) 59 | ) 60 | 61 | def extract_view(self, data, i): 62 | data_i = data[f"view{i}"] 63 | pred_i = data_i.get("cache", {}) 64 | skip_extract = len(pred_i) > 0 and self.conf.allow_no_extract 65 | if self.conf.extractor.name and not skip_extract: 66 | pred_i = {**pred_i, **self.extractor(data_i)} 67 | elif self.conf.extractor.name and not self.conf.allow_no_extract: 68 | pred_i = {**pred_i, **self.extractor({**data_i, **pred_i})} 69 | return pred_i 70 | 71 | def _forward(self, data): 72 | pred0 = self.extract_view(data, "0") 73 | pred1 = self.extract_view(data, "1") 74 | pred = { 75 | **{k + "0": v for k, v in pred0.items()}, 76 | **{k + "1": v for k, v in pred1.items()}, 77 | } 78 | 79 | if self.conf.matcher.name: 80 | pred = {**pred, **self.matcher({**data, **pred})} 81 | if self.conf.filter.name: 82 | pred = {**pred, **self.filter({**data, **pred})} 83 | if self.conf.solver.name: 84 | pred = {**pred, **self.solver({**data, **pred})} 85 | 86 | if self.conf.ground_truth.name and self.conf.run_gt_in_forward: 87 | gt_pred = self.ground_truth({**data, **pred}) 88 | pred.update({f"gt_{k}": v for k, v in gt_pred.items()}) 89 | return pred 90 | 91 | def loss(self, pred, data): 92 | losses = {} 93 | metrics = {} 94 | total = 0 95 | 96 | # get labels 97 | if self.conf.ground_truth.name and not self.conf.run_gt_in_forward: 98 | gt_pred = self.ground_truth({**data, **pred}) 99 | pred.update({f"gt_{k}": v for k, v in gt_pred.items()}) 100 | 101 | for k in self.components: 102 | apply = True 103 | if "apply_loss" in self.conf[k].keys(): 104 | apply = self.conf[k].apply_loss 105 | if self.conf[k].name and apply: 106 | try: 107 | losses_, metrics_ = getattr(self, k).loss(pred, {**pred, **data}) 108 | except NotImplementedError: 109 | continue 110 | losses = {**losses, **losses_} 111 | metrics = {**metrics, **metrics_} 112 | total = losses_["total"] + total 113 | return {**losses, "total": total}, metrics 114 | -------------------------------------------------------------------------------- /gluefactory/models/utils/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/weitong8591/vop/e964dc779958e691d8511b78dd3c40bf7920e4e3/gluefactory/models/utils/__init__.py -------------------------------------------------------------------------------- /gluefactory/models/utils/losses.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | from omegaconf import OmegaConf 4 | from pytorch_metric_learning import losses, miners 5 | from pytorch_metric_learning.distances import CosineSimilarity, DotProductSimilarity 6 | from pytorch_metric_learning.reducers import PerAnchorReducer 7 | 8 | 9 | def weight_loss(log_assignment, weights, gamma=0.0): 10 | b, m, n = log_assignment.shape 11 | m -= 1 12 | n -= 1 13 | 14 | loss_sc = log_assignment * weights 15 | 16 | num_neg0 = weights[:, :m, -1].sum(-1).clamp(min=1.0) 17 | num_neg1 = weights[:, -1, :n].sum(-1).clamp(min=1.0) 18 | num_pos = weights[:, :m, :n].sum((-1, -2)).clamp(min=1.0) 19 | 20 | nll_pos = -loss_sc[:, :m, :n].sum((-1, -2)) 21 | nll_pos /= num_pos.clamp(min=1.0) 22 | 23 | nll_neg0 = -loss_sc[:, :m, -1].sum(-1) 24 | nll_neg1 = -loss_sc[:, -1, :n].sum(-1) 25 | 26 | nll_neg = (nll_neg0 + nll_neg1) / (num_neg0 + num_neg1) 27 | 28 | return nll_pos, nll_neg, num_pos, (num_neg0 + num_neg1) / 2.0 29 | 30 | 31 | class NLLLoss(nn.Module): 32 | default_conf = { 33 | "nll_balancing": 0.5, 34 | "gamma_f": 0.0, # focal loss 35 | } 36 | 37 | def __init__(self, conf): 38 | super().__init__() 39 | self.conf = OmegaConf.merge(self.default_conf, conf) 40 | self.loss_fn = self.nll_loss 41 | 42 | def forward(self, pred, data, weights=None): 43 | log_assignment = pred["log_assignment"] 44 | if weights is None: 45 | weights = self.loss_fn(log_assignment, data) 46 | nll_pos, nll_neg, num_pos, num_neg = weight_loss( 47 | log_assignment, weights, gamma=self.conf.gamma_f 48 | ) 49 | nll = ( 50 | self.conf.nll_balancing * nll_pos + (1 - self.conf.nll_balancing) * nll_neg 51 | ) 52 | 53 | return ( 54 | nll, 55 | weights, 56 | { 57 | "assignment_nll": nll, 58 | "nll_pos": nll_pos, 59 | "nll_neg": nll_neg, 60 | "num_matchable": num_pos, 61 | "num_unmatchable": num_neg, 62 | }, 63 | ) 64 | 65 | def nll_loss(self, log_assignment, data): 66 | m, n = data["gt_matches0"].size(-1), data["gt_matches1"].size(-1) 67 | positive = data["gt_assignment"].float() 68 | neg0 = (data["gt_matches0"] == -1).float() 69 | neg1 = (data["gt_matches1"] == -1).float() 70 | 71 | weights = torch.zeros_like(log_assignment) 72 | weights[:, :m, :n] = positive 73 | 74 | weights[:, :m, -1] = neg0 75 | weights[:, -1, :m] = neg1 76 | return weights 77 | 78 | def get_loss(loss_name): 79 | if loss_name == 'SupConLoss': return losses.SupConLoss(temperature=0.07) 80 | if loss_name == 'CircleLoss': return losses.CircleLoss(m=0.4, gamma=80) #these are params for image retrieval 81 | if loss_name == 'MultiSimilarityLoss': return losses.MultiSimilarityLoss(alpha=1.0, beta=50, base=0.0, distance=DotProductSimilarity()) 82 | if loss_name == 'ContrastiveLoss': return losses.ContrastiveLoss(pos_margin=0, neg_margin=1, reducer=PerAnchorReducer()) 83 | if loss_name == 'Lifted': return losses.GeneralizedLiftedStructureLoss(neg_margin=0, pos_margin=1, distance=DotProductSimilarity()) 84 | if loss_name == 'FastAPLoss': return losses.FastAPLoss(num_bins=30) 85 | if loss_name == 'NTXentLoss': return losses.NTXentLoss(temperature=0.07) #The MoCo paper uses 0.07, while SimCLR uses 0.5. 86 | if loss_name == 'TripletMarginLoss': return losses.TripletMarginLoss(margin=0.1, swap=False, smooth_loss=False, triplets_per_anchor='all') #or an int, for example 100 87 | if loss_name == 'CentroidTripletLoss': return losses.CentroidTripletLoss(margin=0.05, 88 | swap=False, 89 | smooth_loss=False, 90 | triplets_per_anchor="all",) 91 | raise NotImplementedError(f'Sorry, <{loss_name}> loss function is not implemented!') 92 | 93 | def get_miner(miner_name, margin=0.1): 94 | if miner_name == 'TripletMarginMiner' : return miners.TripletMarginMiner(margin=margin, type_of_triplets="semihard") # all, hard, semihard, easy 95 | if miner_name == 'MultiSimilarityMiner' : return miners.MultiSimilarityMiner(epsilon=margin, distance=CosineSimilarity()) 96 | if miner_name == 'PairMarginMiner' : return miners.PairMarginMiner(pos_margin=0.7, neg_margin=0.3, distance=DotProductSimilarity()) 97 | return None 98 | -------------------------------------------------------------------------------- /gluefactory/models/utils/metrics.py: -------------------------------------------------------------------------------- 1 | import torch 2 | 3 | 4 | @torch.no_grad() 5 | def matcher_metrics(pred, data, prefix="", prefix_gt=None): 6 | def recall(m, gt_m): 7 | mask = (gt_m > -1).float() 8 | return ((m == gt_m) * mask).sum(1) / (1e-8 + mask.sum(1)) 9 | 10 | def accuracy(m, gt_m): 11 | mask = (gt_m >= -1).float() 12 | return ((m == gt_m) * mask).sum(1) / (1e-8 + mask.sum(1)) 13 | 14 | def precision(m, gt_m): 15 | mask = ((m > -1) & (gt_m >= -1)).float() 16 | return ((m == gt_m) * mask).sum(1) / (1e-8 + mask.sum(1)) 17 | 18 | def ranking_ap(m, gt_m, scores): 19 | p_mask = ((m > -1) & (gt_m >= -1)).float() 20 | r_mask = (gt_m > -1).float() 21 | sort_ind = torch.argsort(-scores) 22 | sorted_p_mask = torch.gather(p_mask, -1, sort_ind) 23 | sorted_r_mask = torch.gather(r_mask, -1, sort_ind) 24 | sorted_tp = torch.gather(m == gt_m, -1, sort_ind) 25 | p_pts = torch.cumsum(sorted_tp * sorted_p_mask, -1) / ( 26 | 1e-8 + torch.cumsum(sorted_p_mask, -1) 27 | ) 28 | r_pts = torch.cumsum(sorted_tp * sorted_r_mask, -1) / ( 29 | 1e-8 + sorted_r_mask.sum(-1)[:, None] 30 | ) 31 | r_pts_diff = r_pts[..., 1:] - r_pts[..., :-1] 32 | return torch.sum(r_pts_diff * p_pts[:, None, -1], dim=-1) 33 | 34 | if prefix_gt is None: 35 | prefix_gt = prefix 36 | rec = recall(pred[f"{prefix}matches0"], data[f"gt_{prefix_gt}matches0"]) 37 | prec = precision(pred[f"{prefix}matches0"], data[f"gt_{prefix_gt}matches0"]) 38 | acc = accuracy(pred[f"{prefix}matches0"], data[f"gt_{prefix_gt}matches0"]) 39 | ap = ranking_ap( 40 | pred[f"{prefix}matches0"], 41 | data[f"gt_{prefix_gt}matches0"], 42 | pred[f"{prefix}matching_scores0"], 43 | ) 44 | metrics = { 45 | f"{prefix}match_recall": rec, 46 | f"{prefix}match_precision": prec, 47 | f"{prefix}accuracy": acc, 48 | f"{prefix}average_precision": ap, 49 | } 50 | return metrics 51 | -------------------------------------------------------------------------------- /gluefactory/models/utils/misc.py: -------------------------------------------------------------------------------- 1 | import math 2 | from typing import List, Optional, Tuple 3 | 4 | import torch 5 | 6 | 7 | def to_sequence(map): 8 | return map.flatten(-2).transpose(-1, -2) 9 | 10 | 11 | def to_map(sequence): 12 | n = sequence.shape[-2] 13 | e = math.isqrt(n) 14 | assert e * e == n 15 | assert e * e == n 16 | sequence.transpose(-1, -2).unflatten(-1, [e, e]) 17 | 18 | 19 | def pad_to_length( 20 | x, 21 | length: int, 22 | pad_dim: int = -2, 23 | mode: str = "zeros", # zeros, ones, random, random_c 24 | bounds: Tuple[int] = (None, None), 25 | ): 26 | shape = list(x.shape) 27 | d = x.shape[pad_dim] 28 | assert d <= length 29 | if d == length: 30 | return x 31 | shape[pad_dim] = length - d 32 | 33 | low, high = bounds 34 | 35 | if mode == "zeros": 36 | xn = torch.zeros(*shape, device=x.device, dtype=x.dtype) 37 | elif mode == "ones": 38 | xn = torch.ones(*shape, device=x.device, dtype=x.dtype) 39 | elif mode == "random": 40 | low = low if low is not None else x.min() 41 | high = high if high is not None else x.max() 42 | xn = torch.empty(*shape, device=x.device).uniform_(low, high) 43 | elif mode == "random_c": 44 | low, high = bounds # we use the bounds as fallback for empty seq. 45 | xn = torch.cat( 46 | [ 47 | torch.empty(*shape[:-1], 1, device=x.device).uniform_( 48 | x[..., i].min() if d > 0 else low, 49 | x[..., i].max() if d > 0 else high, 50 | ) 51 | for i in range(shape[-1]) 52 | ], 53 | dim=-1, 54 | ) 55 | else: 56 | raise ValueError(mode) 57 | return torch.cat([x, xn], dim=pad_dim) 58 | 59 | 60 | def pad_and_stack( 61 | sequences: List[torch.Tensor], 62 | length: Optional[int] = None, 63 | pad_dim: int = -2, 64 | **kwargs, 65 | ): 66 | if length is None: 67 | length = max([x.shape[pad_dim] for x in sequences]) 68 | 69 | y = torch.stack([pad_to_length(x, length, pad_dim, **kwargs) for x in sequences], 0) 70 | return y 71 | -------------------------------------------------------------------------------- /gluefactory/robust_estimators/__init__.py: -------------------------------------------------------------------------------- 1 | import inspect 2 | 3 | from .base_estimator import BaseEstimator 4 | 5 | 6 | def load_estimator(type, estimator): 7 | module_path = f"{__name__}.{type}.{estimator}" 8 | module = __import__(module_path, fromlist=[""]) 9 | classes = inspect.getmembers(module, inspect.isclass) 10 | # Filter classes defined in the module 11 | classes = [c for c in classes if c[1].__module__ == module_path] 12 | # Filter classes inherited from BaseModel 13 | classes = [c for c in classes if issubclass(c[1], BaseEstimator)] 14 | assert len(classes) == 1, classes 15 | return classes[0][1] 16 | -------------------------------------------------------------------------------- /gluefactory/robust_estimators/base_estimator.py: -------------------------------------------------------------------------------- 1 | from copy import copy 2 | 3 | from omegaconf import OmegaConf 4 | 5 | 6 | class BaseEstimator: 7 | base_default_conf = { 8 | "name": "???", 9 | "ransac_th": "???", 10 | } 11 | test_thresholds = [1.0] 12 | required_data_keys = [] 13 | 14 | strict_conf = False 15 | 16 | def __init__(self, conf): 17 | """Perform some logic and call the _init method of the child model.""" 18 | default_conf = OmegaConf.merge( 19 | self.base_default_conf, OmegaConf.create(self.default_conf) 20 | ) 21 | if self.strict_conf: 22 | OmegaConf.set_struct(default_conf, True) 23 | 24 | if isinstance(conf, dict): 25 | conf = OmegaConf.create(conf) 26 | self.conf = conf = OmegaConf.merge(default_conf, conf) 27 | OmegaConf.set_readonly(conf, True) 28 | OmegaConf.set_struct(conf, True) 29 | self.required_data_keys = copy(self.required_data_keys) 30 | self._init(conf) 31 | 32 | def __call__(self, data): 33 | return self._forward(data) 34 | -------------------------------------------------------------------------------- /gluefactory/robust_estimators/homography/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/weitong8591/vop/e964dc779958e691d8511b78dd3c40bf7920e4e3/gluefactory/robust_estimators/homography/__init__.py -------------------------------------------------------------------------------- /gluefactory/robust_estimators/homography/homography_est.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import torch 3 | from homography_est import ( 4 | LineSegment, 5 | ransac_line_homography, 6 | ransac_point_homography, 7 | ransac_point_line_homography, 8 | ) 9 | 10 | from ...utils.tensor import batch_to_numpy 11 | from ..base_estimator import BaseEstimator 12 | 13 | 14 | def H_estimation_hybrid(kpts0=None, kpts1=None, lines0=None, lines1=None, tol_px=5): 15 | """Estimate a homography from points and lines with hybrid RANSAC. 16 | All features are expected in x-y convention 17 | """ 18 | # Check that we have at least 4 features 19 | n_features = 0 20 | if kpts0 is not None: 21 | n_features += len(kpts0) + len(kpts1) 22 | if lines0 is not None: 23 | n_features += len(lines0) + len(lines1) 24 | if n_features < 4: 25 | return None 26 | 27 | if lines0 is None: 28 | # Point-only RANSAC 29 | H = ransac_point_homography(kpts0, kpts1, tol_px, False, []) 30 | elif kpts0 is None: 31 | # Line-only RANSAC 32 | ls0 = [LineSegment(line[0], line[1]) for line in lines0] 33 | ls1 = [LineSegment(line[0], line[1]) for line in lines1] 34 | H = ransac_line_homography(ls0, ls1, tol_px, False, []) 35 | else: 36 | # Point-lines RANSAC 37 | ls0 = [LineSegment(line[0], line[1]) for line in lines0] 38 | ls1 = [LineSegment(line[0], line[1]) for line in lines1] 39 | H = ransac_point_line_homography(kpts0, kpts1, ls0, ls1, tol_px, False, [], []) 40 | if np.abs(H[-1, -1]) > 1e-8: 41 | H /= H[-1, -1] 42 | return H 43 | 44 | 45 | class PointLineHomographyEstimator(BaseEstimator): 46 | default_conf = {"ransac_th": 2.0, "options": {}} 47 | 48 | required_data_keys = ["m_kpts0", "m_kpts1", "m_lines0", "m_lines1"] 49 | 50 | def _init(self, conf): 51 | pass 52 | 53 | def _forward(self, data): 54 | feat = data["m_kpts0"] if "m_kpts0" in data else data["m_lines0"] 55 | data = batch_to_numpy(data) 56 | m_features = { 57 | "kpts0": data["m_kpts1"] if "m_kpts1" in data else None, 58 | "kpts1": data["m_kpts0"] if "m_kpts0" in data else None, 59 | "lines0": data["m_lines1"] if "m_lines1" in data else None, 60 | "lines1": data["m_lines0"] if "m_lines0" in data else None, 61 | } 62 | M = H_estimation_hybrid(**m_features, tol_px=self.conf.ransac_th) 63 | success = M is not None 64 | if not success: 65 | M = torch.eye(3, device=feat.device, dtype=feat.dtype) 66 | else: 67 | M = torch.from_numpy(M).to(feat) 68 | 69 | estimation = { 70 | "success": success, 71 | "M_0to1": M, 72 | } 73 | 74 | return estimation 75 | -------------------------------------------------------------------------------- /gluefactory/robust_estimators/homography/opencv.py: -------------------------------------------------------------------------------- 1 | import cv2 2 | import torch 3 | 4 | from ..base_estimator import BaseEstimator 5 | 6 | 7 | class OpenCVHomographyEstimator(BaseEstimator): 8 | default_conf = { 9 | "ransac_th": 3.0, 10 | "options": {"method": "ransac", "max_iters": 3000, "confidence": 0.995}, 11 | } 12 | 13 | required_data_keys = ["m_kpts0", "m_kpts1"] 14 | 15 | def _init(self, conf): 16 | self.solver = { 17 | "ransac": cv2.RANSAC, 18 | "lmeds": cv2.LMEDS, 19 | "rho": cv2.RHO, 20 | "usac": cv2.USAC_DEFAULT, 21 | "usac_fast": cv2.USAC_FAST, 22 | "usac_accurate": cv2.USAC_ACCURATE, 23 | "usac_prosac": cv2.USAC_PROSAC, 24 | "usac_magsac": cv2.USAC_MAGSAC, 25 | }[conf.options.method] 26 | 27 | def _forward(self, data): 28 | pts0, pts1 = data["m_kpts0"], data["m_kpts1"] 29 | 30 | try: 31 | M, mask = cv2.findHomography( 32 | pts0.numpy(), 33 | pts1.numpy(), 34 | self.solver, 35 | self.conf.ransac_th, 36 | maxIters=self.conf.options.max_iters, 37 | confidence=self.conf.options.confidence, 38 | ) 39 | success = M is not None 40 | except cv2.error: 41 | success = False 42 | if not success: 43 | M = torch.eye(3, device=pts0.device, dtype=pts0.dtype) 44 | inl = torch.zeros_like(pts0[:, 0]).bool() 45 | else: 46 | M = torch.tensor(M).to(pts0) 47 | inl = torch.tensor(mask).bool().to(pts0.device) 48 | 49 | return { 50 | "success": success, 51 | "M_0to1": M, 52 | "inliers": inl, 53 | } 54 | -------------------------------------------------------------------------------- /gluefactory/robust_estimators/homography/poselib.py: -------------------------------------------------------------------------------- 1 | import poselib 2 | import torch 3 | from omegaconf import OmegaConf 4 | 5 | from ..base_estimator import BaseEstimator 6 | 7 | 8 | class PoseLibHomographyEstimator(BaseEstimator): 9 | default_conf = {"ransac_th": 2.0, "options": {}} 10 | 11 | required_data_keys = ["m_kpts0", "m_kpts1"] 12 | 13 | def _init(self, conf): 14 | pass 15 | 16 | def _forward(self, data): 17 | pts0, pts1 = data["m_kpts0"], data["m_kpts1"] 18 | M, info = poselib.estimate_homography( 19 | pts0.detach().cpu().numpy(), 20 | pts1.detach().cpu().numpy(), 21 | { 22 | "max_reproj_error": self.conf.ransac_th, 23 | **OmegaConf.to_container(self.conf.options), 24 | }, 25 | ) 26 | success = M is not None 27 | if not success: 28 | M = torch.eye(3, device=pts0.device, dtype=pts0.dtype) 29 | inl = torch.zeros_like(pts0[:, 0]).bool() 30 | else: 31 | M = torch.tensor(M).to(pts0) 32 | inl = torch.tensor(info["inliers"]).bool().to(pts0.device) 33 | 34 | estimation = { 35 | "success": success, 36 | "M_0to1": M, 37 | "inliers": inl, 38 | } 39 | 40 | return estimation 41 | -------------------------------------------------------------------------------- /gluefactory/robust_estimators/relative_pose/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/weitong8591/vop/e964dc779958e691d8511b78dd3c40bf7920e4e3/gluefactory/robust_estimators/relative_pose/__init__.py -------------------------------------------------------------------------------- /gluefactory/robust_estimators/relative_pose/opencv.py: -------------------------------------------------------------------------------- 1 | import cv2 2 | import numpy as np 3 | import torch 4 | 5 | from ...geometry.utils import from_homogeneous 6 | from ...geometry.wrappers import Pose 7 | from ..base_estimator import BaseEstimator 8 | 9 | 10 | class OpenCVRelativePoseEstimator(BaseEstimator): 11 | default_conf = { 12 | "ransac_th": 0.5, 13 | "options": {"confidence": 0.99999, "method": "ransac"}, 14 | } 15 | 16 | required_data_keys = ["m_kpts0", "m_kpts1", "camera0", "camera1"] 17 | 18 | def _init(self, conf): 19 | self.solver = {"ransac": cv2.RANSAC, "usac_magsac": cv2.USAC_MAGSAC}[ 20 | self.conf.options.method 21 | ] 22 | 23 | def _forward(self, data): 24 | # import pdb; pdb.set_trace() 25 | kpts0, kpts1 = data["m_kpts0"], data["m_kpts1"] 26 | camera0 = data["camera0"] 27 | camera1 = data["camera1"] 28 | M, inl = None, torch.zeros_like(kpts0[:, 0]).bool() 29 | 30 | if len(kpts0) >= 5: 31 | f_mean = torch.cat([camera0.f, camera1.f]).mean().item() 32 | norm_thresh = self.conf.ransac_th / f_mean 33 | 34 | pts0 = from_homogeneous(camera0.image2cam(kpts0)).cpu().detach().numpy() 35 | pts1 = from_homogeneous(camera1.image2cam(kpts1)).cpu().detach().numpy() 36 | 37 | E, mask = cv2.findEssentialMat( 38 | pts0, 39 | pts1, 40 | np.eye(3), 41 | threshold=norm_thresh, 42 | prob=self.conf.options.confidence, 43 | method=self.solver, 44 | ) 45 | # import pdb; pdb.set_trace() 46 | if E is not None: 47 | best_num_inliers = 0 48 | for _E in np.split(E, len(E) / 3): 49 | n, R, t, _ = cv2.recoverPose( 50 | _E, pts0, pts1, np.eye(3), 1e9, mask=mask 51 | ) 52 | if n > best_num_inliers: 53 | best_num_inliers = n 54 | inl = torch.tensor(mask.ravel() > 0) 55 | M = Pose.from_Rt( 56 | torch.tensor(R).to(kpts0), torch.tensor(t[:, 0]).to(kpts0) 57 | ) 58 | 59 | estimation = { 60 | "success": M is not None, 61 | "M_0to1": M if M is not None else Pose.from_4x4mat(torch.eye(4).to(kpts0)), 62 | "inliers": inl.to(device=kpts0.device), 63 | } 64 | 65 | return estimation 66 | -------------------------------------------------------------------------------- /gluefactory/robust_estimators/relative_pose/poselib.py: -------------------------------------------------------------------------------- 1 | import poselib 2 | import torch 3 | from omegaconf import OmegaConf 4 | 5 | from ...geometry.wrappers import Pose 6 | from ..base_estimator import BaseEstimator 7 | 8 | 9 | class PoseLibRelativePoseEstimator(BaseEstimator): 10 | default_conf = {"ransac_th": 2.0, "options": {}} 11 | 12 | required_data_keys = ["m_kpts0", "m_kpts1", "camera0", "camera1"] 13 | 14 | def _init(self, conf): 15 | pass 16 | 17 | def _forward(self, data): 18 | pts0, pts1 = data["m_kpts0"], data["m_kpts1"] 19 | camera0 = data["camera0"] 20 | camera1 = data["camera1"] 21 | M, info = poselib.estimate_relative_pose( 22 | pts0.numpy(), 23 | pts1.numpy(), 24 | camera0.to_cameradict(), 25 | camera1.to_cameradict(), 26 | { 27 | "max_epipolar_error": self.conf.ransac_th, 28 | **OmegaConf.to_container(self.conf.options), 29 | }, 30 | ) 31 | success = M is not None 32 | if success: 33 | M = Pose.from_Rt(torch.tensor(M.R), torch.tensor(M.t)).to(pts0) 34 | else: 35 | M = Pose.from_4x4mat(torch.eye(4)).to(pts0) 36 | 37 | estimation = { 38 | "success": success, 39 | "M_0to1": M, 40 | "inliers": torch.tensor(info.pop("inliers")).to(pts0), 41 | **info, 42 | } 43 | 44 | return estimation 45 | -------------------------------------------------------------------------------- /gluefactory/robust_estimators/relative_pose/pycolmap.py: -------------------------------------------------------------------------------- 1 | import pycolmap 2 | import torch 3 | from omegaconf import OmegaConf 4 | 5 | from ...geometry.wrappers import Pose 6 | from ..base_estimator import BaseEstimator 7 | 8 | 9 | class PycolmapTwoViewEstimator(BaseEstimator): 10 | default_conf = { 11 | "ransac_th": 4.0, 12 | "options": {**pycolmap.TwoViewGeometryOptions().todict()}, 13 | } 14 | 15 | required_data_keys = ["m_kpts0", "m_kpts1", "camera0", "camera1"] 16 | 17 | def _init(self, conf): 18 | opts = OmegaConf.to_container(conf.options) 19 | self.options = pycolmap.TwoViewGeometryOptions(opts) 20 | self.options.ransac.max_error = conf.ransac_th 21 | 22 | def _forward(self, data): 23 | pts0, pts1 = data["m_kpts0"], data["m_kpts1"] 24 | camera0 = data["camera0"] 25 | camera1 = data["camera1"] 26 | info = pycolmap.two_view_geometry_estimation( 27 | pts0.numpy(), 28 | pts1.numpy(), 29 | camera0.to_cameradict(), 30 | camera1.to_cameradict(), 31 | self.options, 32 | ) 33 | success = info["success"] 34 | if success: 35 | R = pycolmap.qvec_to_rotmat(info["qvec"]) 36 | t = info["tvec"] 37 | M = Pose.from_Rt(torch.tensor(R), torch.tensor(t)).to(pts0) 38 | inl = torch.tensor(info.pop("inliers")).to(pts0) 39 | else: 40 | M = Pose.from_4x4mat(torch.eye(4)).to(pts0) 41 | inl = torch.zeros_like(pts0[:, 0]).bool() 42 | 43 | estimation = { 44 | "success": success, 45 | "M_0to1": M, 46 | "inliers": inl, 47 | "type": str( 48 | info.get("configuration_type", pycolmap.TwoViewGeometry.UNDEFINED) 49 | ), 50 | } 51 | 52 | return estimation 53 | -------------------------------------------------------------------------------- /gluefactory/scripts/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/weitong8591/vop/e964dc779958e691d8511b78dd3c40bf7920e4e3/gluefactory/scripts/__init__.py -------------------------------------------------------------------------------- /gluefactory/scripts/export_local_features.py: -------------------------------------------------------------------------------- 1 | import argparse 2 | import logging 3 | from pathlib import Path 4 | 5 | import torch 6 | from omegaconf import OmegaConf 7 | 8 | from ..datasets import get_dataset 9 | from ..models import get_model 10 | from ..settings import DATA_PATH 11 | from ..utils.export_predictions import export_predictions 12 | 13 | resize = 1600 14 | 15 | sp_keys = ["keypoints", "descriptors", "keypoint_scores"] 16 | 17 | # SuperPoint 18 | n_kpts = 2048 19 | configs = { 20 | "sp": { 21 | "name": f"r{resize}_SP-k{n_kpts}-nms3", 22 | "keys": ["keypoints", "descriptors", "keypoint_scores"], 23 | "gray": True, 24 | "conf": { 25 | "name": "gluefactory_nonfree.superpoint", 26 | "nms_radius": 3, 27 | "max_num_keypoints": n_kpts, 28 | "detection_threshold": 0.000, 29 | }, 30 | }, 31 | "sift": { 32 | "name": f"r{resize}_SIFT-k{n_kpts}", 33 | "keys": ["keypoints", "descriptors", "keypoint_scores", "oris", "scales"], 34 | "gray": True, 35 | "conf": { 36 | "name": "sift", 37 | "max_num_keypoints": n_kpts, 38 | "options": { 39 | "peak_threshold": 0.001, 40 | }, 41 | "peak_threshold": 0.001, 42 | "device": "cpu", 43 | }, 44 | }, 45 | "disk": { 46 | "name": f"r{resize}_DISK-k{n_kpts}-nms6", 47 | "keys": ["keypoints", "descriptors", "keypoint_scores"], 48 | "gray": False, 49 | "conf": { 50 | "name": "disk", 51 | "max_num_keypoints": n_kpts, 52 | }, 53 | }, 54 | } 55 | 56 | 57 | def run_export(feature_file, images, args): 58 | conf = { 59 | "data": { 60 | "name": "image_folder", 61 | "grayscale": configs[args.method]["gray"], 62 | "preprocessing": { 63 | "resize": resize, 64 | }, 65 | "images": str(images), 66 | "batch_size": 1, 67 | "num_workers": args.num_workers, 68 | }, 69 | "split": "train", 70 | "model": configs[args.method]["conf"], 71 | } 72 | 73 | conf = OmegaConf.create(conf) 74 | 75 | keys = configs[args.method]["keys"] 76 | dataset = get_dataset(conf.data.name)(conf.data) 77 | loader = dataset.get_data_loader(conf.split or "test") 78 | 79 | device = "cuda" if torch.cuda.is_available() else "cpu" 80 | model = get_model(conf.model.name)(conf.model).eval().to(device) 81 | 82 | export_predictions(loader, model, feature_file, as_half=True, keys=keys) 83 | 84 | 85 | if __name__ == "__main__": 86 | parser = argparse.ArgumentParser() 87 | parser.add_argument("dataset", type=str) 88 | parser.add_argument("--export_prefix", type=str, default="") 89 | parser.add_argument("--method", type=str, default="sp") 90 | parser.add_argument("--scenes", type=str, default=None) 91 | parser.add_argument("--num_workers", type=int, default=0) 92 | args = parser.parse_args() 93 | 94 | export_name = configs[args.method]["name"] 95 | 96 | if args.dataset == "megadepth": 97 | data_root = Path(DATA_PATH, "megadepth/Undistorted_SfM") 98 | export_root = Path(DATA_PATH, "exports", "megadepth-undist-" + export_name) 99 | export_root.mkdir(parents=True, exist_ok=True) 100 | 101 | if args.scenes is None: 102 | scenes = [p.name for p in data_root.iterdir() if p.is_dir()] 103 | else: 104 | with open(DATA_PATH / "megadepth" / args.scenes, "r") as f: 105 | scenes = f.read().split() 106 | for i, scene in enumerate(scenes): 107 | # print(f'{i} / {len(scenes)}', scene) 108 | print(scene) 109 | feature_file = export_root / (scene + ".h5") 110 | if feature_file.exists(): 111 | continue 112 | if not (data_root / scene / "images").exists(): 113 | logging.info("Skip " + scene) 114 | continue 115 | logging.info(f"Export local features for scene {scene}") 116 | run_export(feature_file, data_root / scene / "images", args) 117 | else: 118 | data_root = Path(DATA_PATH, args.dataset) 119 | feature_file = Path( 120 | DATA_PATH, "exports", args.export_prefix + export_name + ".h5" 121 | ) 122 | feature_file.parent.mkdir(exist_ok=True, parents=True) 123 | logging.info( 124 | f"Export local features for dataset {args.dataset} " 125 | f"to file {feature_file}" 126 | ) 127 | run_export(feature_file, data_root) 128 | -------------------------------------------------------------------------------- /gluefactory/settings.py: -------------------------------------------------------------------------------- 1 | from pathlib import Path 2 | 3 | root = Path(__file__).parent.parent # top-level directory 4 | DATA_PATH = root / "data/" # datasets and pretrained weights 5 | TRAINING_PATH = root / "outputs/training/" # training checkpoints 6 | TRAINING_PATH1 = "/home/weitong/code/VisualOverlap/glue-factory/outputs/training/" # training checkpoints 7 | 8 | EVAL_PATH = root / "outputs/results/" # evaluation results 9 | -------------------------------------------------------------------------------- /gluefactory/utils/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/weitong8591/vop/e964dc779958e691d8511b78dd3c40bf7920e4e3/gluefactory/utils/__init__.py -------------------------------------------------------------------------------- /gluefactory/utils/benchmark.py: -------------------------------------------------------------------------------- 1 | import time 2 | 3 | import numpy as np 4 | import torch 5 | 6 | 7 | def benchmark(model, data, device, r=100): 8 | timings = np.zeros((r, 1)) 9 | if device.type == "cuda": 10 | starter = torch.cuda.Event(enable_timing=True) 11 | ender = torch.cuda.Event(enable_timing=True) 12 | # warmup 13 | for _ in range(10): 14 | _ = model(data) 15 | # measurements 16 | with torch.no_grad(): 17 | for rep in range(r): 18 | if device.type == "cuda": 19 | starter.record() 20 | _ = model(data) 21 | ender.record() 22 | # sync gpu 23 | torch.cuda.synchronize() 24 | curr_time = starter.elapsed_time(ender) 25 | else: 26 | start = time.perf_counter() 27 | _ = model(data) 28 | curr_time = (time.perf_counter() - start) * 1e3 29 | timings[rep] = curr_time 30 | 31 | mean_syn = np.sum(timings) / r 32 | std_syn = np.std(timings) 33 | return {"mean": mean_syn, "std": std_syn} 34 | -------------------------------------------------------------------------------- /gluefactory/utils/experiments.py: -------------------------------------------------------------------------------- 1 | """A set of utilities to manage and load checkpoints of training experiments. 2 | 3 | Author: Paul-Edouard Sarlin (skydes) 4 | """ 5 | 6 | import logging 7 | import os 8 | import re 9 | import shutil 10 | from pathlib import Path 11 | 12 | import torch 13 | from omegaconf import OmegaConf 14 | 15 | from ..models import get_model 16 | from ..settings import TRAINING_PATH, TRAINING_PATH1 17 | 18 | logger = logging.getLogger(__name__) 19 | 20 | 21 | def list_checkpoints(dir_): 22 | """List all valid checkpoints in a given directory.""" 23 | checkpoints = [] 24 | for p in dir_.glob("checkpoint_*.tar"): 25 | numbers = re.findall(r"(\d+)", p.name) 26 | assert len(numbers) <= 2 27 | if len(numbers) == 0: 28 | continue 29 | if len(numbers) == 1: 30 | checkpoints.append((int(numbers[0]), p)) 31 | else: 32 | checkpoints.append((int(numbers[1]), p)) 33 | return checkpoints 34 | 35 | 36 | def get_last_checkpoint(exper, allow_interrupted=True): 37 | """Get the last saved checkpoint for a given experiment name.""" 38 | # ckpts = list_checkpoints(Path(TRAINING_PATH, exper)) 39 | try: 40 | ckpts = list_checkpoints(Path(TRAINING_PATH, exper)) 41 | except: 42 | ckpts = list_checkpoints(Path(TRAINING_PATH1, exper)) 43 | 44 | if not allow_interrupted: 45 | ckpts = [(n, p) for (n, p) in ckpts if "_interrupted" not in p.name] 46 | assert len(ckpts) > 0 47 | return sorted(ckpts)[-1][1] 48 | 49 | 50 | def get_best_checkpoint(exper): 51 | """Get the checkpoint with the best loss, for a given experiment name.""" 52 | if os.path.exists(TRAINING_PATH/ exper/ "checkpoint_best.tar"): 53 | p = Path(TRAINING_PATH, exper, "checkpoint_best.tar") 54 | else: 55 | p = Path(TRAINING_PATH1, exper, "checkpoint_best.tar") 56 | return p 57 | 58 | 59 | def delete_old_checkpoints(dir_, num_keep): 60 | """Delete all but the num_keep last saved checkpoints.""" 61 | ckpts = list_checkpoints(dir_) 62 | ckpts = sorted(ckpts)[::-1] 63 | kept = 0 64 | for ckpt in ckpts: 65 | if ("_interrupted" in str(ckpt[1]) and kept > 0) or kept >= num_keep: 66 | logger.info(f"Deleting checkpoint {ckpt[1].name}") 67 | ckpt[1].unlink() 68 | else: 69 | kept += 1 70 | 71 | 72 | def load_experiment(exper, conf={}, get_last=False, ckpt=None): 73 | """Load and return the model of a given experiment.""" 74 | exper = Path(exper) 75 | if exper.suffix != ".tar": 76 | if get_last: 77 | ckpt = get_last_checkpoint(exper) 78 | else: 79 | ckpt = get_best_checkpoint(exper) 80 | else: 81 | ckpt = exper 82 | logger.info(f"Loading checkpoint {ckpt.name}") 83 | ckpt = torch.load(str(ckpt), map_location="cpu") 84 | 85 | loaded_conf = OmegaConf.create(ckpt["conf"]) 86 | OmegaConf.set_struct(loaded_conf, False) 87 | conf = OmegaConf.merge(loaded_conf.model, OmegaConf.create(conf)) 88 | model = get_model(conf.name)(conf).eval() 89 | 90 | state_dict = ckpt["model"] 91 | dict_params = set(state_dict.keys()) 92 | model_params = set(map(lambda n: n[0], model.named_parameters())) 93 | diff = model_params - dict_params 94 | if len(diff) > 0: 95 | subs = os.path.commonprefix(list(diff)).rstrip(".") 96 | logger.warning(f"Missing {len(diff)} parameters in {subs}") 97 | model.load_state_dict(state_dict, strict=False) 98 | return model 99 | 100 | 101 | # @TODO: also copy the respective module scripts (i.e. the code) 102 | def save_experiment( 103 | model, 104 | optimizer, 105 | lr_scheduler, 106 | conf, 107 | losses, 108 | results, 109 | best_eval, 110 | epoch, 111 | iter_i, 112 | output_dir, 113 | stop=False, 114 | distributed=False, 115 | cp_name=None, 116 | ): 117 | """Save the current model to a checkpoint and return the best result so far.""" 118 | state = (model.module if distributed else model).state_dict() 119 | checkpoint = { 120 | "model": state, 121 | "optimizer": optimizer.state_dict(), 122 | "lr_scheduler": lr_scheduler.state_dict(), 123 | "conf": OmegaConf.to_container(conf, resolve=True), 124 | "epoch": epoch, 125 | "losses": losses, 126 | "eval": results, 127 | } 128 | if cp_name is None: 129 | cp_name = ( 130 | f"checkpoint_{epoch}_{iter_i}" + ("_interrupted" if stop else "") + ".tar" 131 | ) 132 | logger.info(f"Saving checkpoint {cp_name}") 133 | cp_path = str(output_dir / cp_name) 134 | torch.save(checkpoint, cp_path) 135 | if cp_name != "checkpoint_best.tar" and results[conf.train.best_key] < best_eval: 136 | best_eval = results[conf.train.best_key] 137 | logger.info(f"New best val: {conf.train.best_key}={best_eval}") 138 | shutil.copy(cp_path, str(output_dir / "checkpoint_best.tar")) 139 | delete_old_checkpoints(output_dir, conf.train.keep_last_checkpoints) 140 | return best_eval 141 | -------------------------------------------------------------------------------- /gluefactory/utils/export_predictions.py: -------------------------------------------------------------------------------- 1 | """ 2 | Export the predictions of a model for a given dataloader (e.g. ImageFolder). 3 | Use a standalone script with `python3 -m dsfm.scipts.export_predictions dir` 4 | or call from another script. 5 | """ 6 | 7 | from pathlib import Path 8 | 9 | import h5py 10 | import numpy as np 11 | import torch 12 | from tqdm import tqdm 13 | 14 | from .tensor import batch_to_device 15 | 16 | 17 | @torch.no_grad() 18 | def export_predictions( 19 | loader, 20 | model, 21 | output_file, 22 | as_half=False, 23 | keys="*", 24 | callback_fn=None, 25 | optional_keys=[], 26 | ): 27 | import pdb; pdb.set_trace() 28 | assert keys == "*" or isinstance(keys, (tuple, list)) 29 | Path(output_file).parent.mkdir(exist_ok=True, parents=True) 30 | hfile = h5py.File(str(output_file), "w") 31 | device = "cuda" if torch.cuda.is_available() else "cpu" 32 | model = model.to(device).eval() 33 | for data_ in tqdm(loader): 34 | data = batch_to_device(data_, device, non_blocking=True) 35 | pred = model(data) 36 | if callback_fn is not None: 37 | pred = {**callback_fn(pred, data), **pred} 38 | if keys != "*": 39 | if len(set(keys) - set(pred.keys())) > 0: 40 | raise ValueError(f"Missing key {set(keys) - set(pred.keys())}") 41 | pred = {k: v for k, v in pred.items() if k in keys + optional_keys} 42 | assert len(pred) > 0 43 | 44 | # renormalization 45 | for k in pred.keys(): 46 | if k.startswith("keypoints"): 47 | idx = k.replace("keypoints", "") 48 | scales = 1.0 / ( 49 | data["scales"] if len(idx) == 0 else data[f"view{idx}"]["scales"] 50 | ) 51 | pred[k] = pred[k] * scales[None] 52 | if k.startswith("lines"): 53 | idx = k.replace("lines", "") 54 | scales = 1.0 / ( 55 | data["scales"] if len(idx) == 0 else data[f"view{idx}"]["scales"] 56 | ) 57 | pred[k] = pred[k] * scales[None] 58 | if k.startswith("orig_lines"): 59 | idx = k.replace("orig_lines", "") 60 | scales = 1.0 / ( 61 | data["scales"] if len(idx) == 0 else data[f"view{idx}"]["scales"] 62 | ) 63 | pred[k] = pred[k] * scales[None] 64 | 65 | pred = {k: v[0].cpu().numpy() for k, v in pred.items()} 66 | 67 | if as_half: 68 | for k in pred: 69 | dt = pred[k].dtype 70 | if (dt == np.float32) and (dt != np.float16): 71 | pred[k] = pred[k].astype(np.float16) 72 | try: 73 | name = data["name"][0] 74 | grp = hfile.create_group(name) 75 | for k, v in pred.items(): 76 | grp.create_dataset(k, data=v) 77 | except RuntimeError: 78 | continue 79 | 80 | del pred 81 | hfile.close() 82 | return output_file 83 | -------------------------------------------------------------------------------- /gluefactory/utils/image.py: -------------------------------------------------------------------------------- 1 | import collections.abc as collections 2 | from pathlib import Path 3 | from typing import Optional, Tuple 4 | 5 | import cv2 6 | import kornia 7 | import numpy as np 8 | import torch 9 | from omegaconf import OmegaConf 10 | 11 | 12 | class ImagePreprocessor: 13 | default_conf = { 14 | "resize": None, # target edge length, None for no resizing 15 | "edge_divisible_by": None, 16 | "side": "long", 17 | "interpolation": "bilinear", 18 | "align_corners": None, 19 | "antialias": True, 20 | "square_pad": False, 21 | "add_padding_mask": False, 22 | } 23 | 24 | def __init__(self, conf) -> None: 25 | super().__init__() 26 | default_conf = OmegaConf.create(self.default_conf) 27 | OmegaConf.set_struct(default_conf, True) 28 | self.conf = OmegaConf.merge(default_conf, conf) 29 | 30 | def __call__(self, img: torch.Tensor, interpolation: Optional[str] = None) -> dict: 31 | """Resize and preprocess an image, return image and resize scale""" 32 | h, w = img.shape[-2:] 33 | size = h, w 34 | if self.conf.resize is not None: 35 | if interpolation is None: 36 | interpolation = self.conf.interpolation 37 | size = self.get_new_image_size(h, w) 38 | img = kornia.geometry.transform.resize( 39 | img, 40 | size, 41 | side=self.conf.side, 42 | antialias=self.conf.antialias, 43 | align_corners=self.conf.align_corners, 44 | interpolation=interpolation, 45 | ) 46 | scale = torch.Tensor([img.shape[-1] / w, img.shape[-2] / h]).to(img) 47 | T = np.diag([scale[0], scale[1], 1]) 48 | 49 | data = { 50 | "scales": scale, 51 | "image_size": np.array(size[::-1]), 52 | "transform": T, 53 | "original_image_size": np.array([w, h]), 54 | } 55 | if self.conf.square_pad: 56 | sl = max(img.shape[-2:]) 57 | data["image"] = torch.zeros( 58 | *img.shape[:-2], sl, sl, device=img.device, dtype=img.dtype 59 | ) 60 | data["image"][:, : img.shape[-2], : img.shape[-1]] = img 61 | if self.conf.add_padding_mask: 62 | data["padding_mask"] = torch.zeros( 63 | *img.shape[:-3], 1, sl, sl, device=img.device, dtype=torch.bool 64 | ) 65 | data["padding_mask"][:, : img.shape[-2], : img.shape[-1]] = True 66 | 67 | else: 68 | data["image"] = img 69 | return data 70 | 71 | def load_image(self, image_path: Path) -> dict: 72 | return self(load_image(image_path)) 73 | 74 | def get_new_image_size( 75 | self, 76 | h: int, 77 | w: int, 78 | ) -> Tuple[int, int]: 79 | side = self.conf.side 80 | if isinstance(self.conf.resize, collections.Iterable): 81 | assert len(self.conf.resize) == 2 82 | return tuple(self.conf.resize) 83 | side_size = self.conf.resize 84 | aspect_ratio = w / h 85 | if side not in ("short", "long", "vert", "horz"): 86 | raise ValueError( 87 | f"side can be one of 'short', 'long', 'vert', and 'horz'. Got '{side}'" 88 | ) 89 | if side == "vert": 90 | size = side_size, int(side_size * aspect_ratio) 91 | elif side == "horz": 92 | size = int(side_size / aspect_ratio), side_size 93 | elif (side == "short") ^ (aspect_ratio < 1.0): 94 | size = side_size, int(side_size * aspect_ratio) 95 | else: 96 | size = int(side_size / aspect_ratio), side_size 97 | 98 | if self.conf.edge_divisible_by is not None: 99 | df = self.conf.edge_divisible_by 100 | size = list(map(lambda x: int(x // df * df), size)) 101 | return size 102 | 103 | 104 | def read_image(path: Path, grayscale: bool = False) -> np.ndarray: 105 | """Read an image from path as RGB or grayscale""" 106 | if not Path(path).exists(): 107 | raise FileNotFoundError(f"No image at path {path}.") 108 | mode = cv2.IMREAD_GRAYSCALE if grayscale else cv2.IMREAD_COLOR 109 | image = cv2.imread(str(path), mode) # B G R 110 | if image is None: 111 | raise IOError(f"Could not read image at {path}.") 112 | if not grayscale: 113 | image = image[..., ::-1]# reverses the order of elements along the last axis 114 | return image # H W C, rows, columns, channels 115 | 116 | 117 | def numpy_image_to_torch(image: np.ndarray) -> torch.Tensor: 118 | """Normalize the image tensor and reorder the dimensions.""" 119 | if image.ndim == 3: 120 | image = image.transpose((2, 0, 1)) # HxWxC to CxHxW 121 | elif image.ndim == 2: 122 | image = image[None] # add channel axis 123 | else: 124 | raise ValueError(f"Not an image: {image.shape}") 125 | return torch.tensor(image / 255.0, dtype=torch.float) 126 | 127 | 128 | def load_image(path: Path, grayscale=False) -> torch.Tensor: 129 | image = read_image(path, grayscale=grayscale) 130 | return numpy_image_to_torch(image) 131 | -------------------------------------------------------------------------------- /gluefactory/utils/misc.py: -------------------------------------------------------------------------------- 1 | import torch 2 | 3 | 4 | def to_view(data, i): 5 | return {k + i: v for k, v in data.items()} 6 | 7 | 8 | def get_view(data, i): 9 | data_g = {k: v for k, v in data.items() if not k[-1].isnumeric()} 10 | data_i = {k[:-1]: v for k, v in data.items() if k[-1] == i} 11 | return {**data_g, **data_i} 12 | 13 | 14 | def get_twoview(data, idx): 15 | li = idx[0] 16 | ri = idx[-1] 17 | assert idx == f"{li}to{ri}" 18 | data_lr = {k[:-4] + "0to1": v for k, v in data.items() if k[-4:] == f"{li}to{ri}"} 19 | data_rl = {k[:-4] + "1to0": v for k, v in data.items() if k[-4:] == f"{ri}ito{li}"} 20 | data_l = { 21 | k[:-1] + "0": v for k, v in data.items() if k[-1:] == li and k[-3:-1] != "to" 22 | } 23 | data_r = { 24 | k[:-1] + "1": v for k, v in data.items() if k[-1:] == ri and k[-3:-1] != "to" 25 | } 26 | return {**data_lr, **data_rl, **data_l, **data_r} 27 | 28 | 29 | def stack_twoviews(data, indices=["0to1", "0to2", "1to2"]): 30 | idx0 = indices[0] 31 | m_data = data[idx0] if idx0 in data else get_twoview(data, idx0) 32 | # stack on dim=0 33 | for idx in indices[1:]: 34 | data_i = data[idx] if idx in data else get_twoview(data, idx) 35 | for k, v in data_i.items(): 36 | m_data[k] = torch.cat([m_data[k], v], dim=0) 37 | return m_data 38 | 39 | 40 | def unstack_twoviews(data, B, indices=["0to1", "0to2", "1to2"]): 41 | out = {} 42 | for i, idx in enumerate(indices): 43 | out[idx] = {k: v[i * B : (i + 1) * B] for k, v in data.items()} 44 | return out 45 | -------------------------------------------------------------------------------- /gluefactory/utils/patches.py: -------------------------------------------------------------------------------- 1 | import torch 2 | 3 | 4 | def extract_patches( 5 | tensor: torch.Tensor, 6 | required_corners: torch.Tensor, 7 | ps: int, 8 | ) -> torch.Tensor: 9 | c, h, w = tensor.shape 10 | corner = required_corners.long() 11 | corner[:, 0] = corner[:, 0].clamp(min=0, max=w - 1 - ps) 12 | corner[:, 1] = corner[:, 1].clamp(min=0, max=h - 1 - ps) 13 | offset = torch.arange(0, ps) 14 | 15 | kw = {"indexing": "ij"} if torch.__version__ >= "1.10" else {} 16 | x, y = torch.meshgrid(offset, offset, **kw) 17 | patches = torch.stack((x, y)).permute(2, 1, 0).unsqueeze(2) 18 | patches = patches.to(corner) + corner[None, None] 19 | pts = patches.reshape(-1, 2) 20 | sampled = tensor.permute(1, 2, 0)[tuple(pts.T)[::-1]] 21 | sampled = sampled.reshape(ps, ps, -1, c) 22 | assert sampled.shape[:3] == patches.shape[:3] 23 | return sampled.permute(2, 3, 0, 1), corner.float() 24 | 25 | 26 | def batch_extract_patches(tensor: torch.Tensor, kpts: torch.Tensor, ps: int): 27 | b, c, h, w = tensor.shape 28 | b, n, _ = kpts.shape 29 | out = torch.zeros((b, n, c, ps, ps), dtype=tensor.dtype, device=tensor.device) 30 | corners = torch.zeros((b, n, 2), dtype=tensor.dtype, device=tensor.device) 31 | for i in range(b): 32 | out[i], corners[i] = extract_patches(tensor[i], kpts[i] - ps / 2 - 1, ps) 33 | return out, corners 34 | 35 | 36 | def draw_image_patches(img, patches, corners): 37 | b, c, h, w = img.shape 38 | b, n, c, p, p = patches.shape 39 | b, n, _ = corners.shape 40 | for i in range(b): 41 | for k in range(n): 42 | y, x = corners[i, k] 43 | img[i, :, x : x + p, y : y + p] = patches[i, k] 44 | 45 | 46 | def build_heatmap(img, patches, corners): 47 | hmap = torch.zeros_like(img) 48 | draw_image_patches(hmap, patches, corners.long()) 49 | hmap = hmap.squeeze(1) 50 | return hmap, (hmap > 0.0).float() # bxhxw 51 | -------------------------------------------------------------------------------- /gluefactory/utils/stdout_capturing.py: -------------------------------------------------------------------------------- 1 | """ 2 | Based on sacred/stdout_capturing.py in project Sacred 3 | https://github.com/IDSIA/sacred 4 | 5 | Author: Paul-Edouard Sarlin (skydes) 6 | """ 7 | 8 | from __future__ import division, print_function, unicode_literals 9 | 10 | import os 11 | import subprocess 12 | import sys 13 | from contextlib import contextmanager 14 | from threading import Timer 15 | 16 | 17 | def apply_backspaces_and_linefeeds(text): 18 | """ 19 | Interpret backspaces and linefeeds in text like a terminal would. 20 | Interpret text like a terminal by removing backspace and linefeed 21 | characters and applying them line by line. 22 | If final line ends with a carriage it keeps it to be concatenable with next 23 | output chunk. 24 | """ 25 | orig_lines = text.split("\n") 26 | orig_lines_len = len(orig_lines) 27 | new_lines = [] 28 | for orig_line_idx, orig_line in enumerate(orig_lines): 29 | chars, cursor = [], 0 30 | orig_line_len = len(orig_line) 31 | for orig_char_idx, orig_char in enumerate(orig_line): 32 | if orig_char == "\r" and ( 33 | orig_char_idx != orig_line_len - 1 34 | or orig_line_idx != orig_lines_len - 1 35 | ): 36 | cursor = 0 37 | elif orig_char == "\b": 38 | cursor = max(0, cursor - 1) 39 | else: 40 | if ( 41 | orig_char == "\r" 42 | and orig_char_idx == orig_line_len - 1 43 | and orig_line_idx == orig_lines_len - 1 44 | ): 45 | cursor = len(chars) 46 | if cursor == len(chars): 47 | chars.append(orig_char) 48 | else: 49 | chars[cursor] = orig_char 50 | cursor += 1 51 | new_lines.append("".join(chars)) 52 | return "\n".join(new_lines) 53 | 54 | 55 | def flush(): 56 | """Try to flush all stdio buffers, both from python and from C.""" 57 | try: 58 | sys.stdout.flush() 59 | sys.stderr.flush() 60 | except (AttributeError, ValueError, IOError): 61 | pass # unsupported 62 | 63 | 64 | # Duplicate stdout and stderr to a file. Inspired by: 65 | # http://eli.thegreenplace.net/2015/redirecting-all-kinds-of-stdout-in-python/ 66 | # http://stackoverflow.com/a/651718/1388435 67 | # http://stackoverflow.com/a/22434262/1388435 68 | @contextmanager 69 | def capture_outputs(filename): 70 | """Duplicate stdout and stderr to a file on the file descriptor level.""" 71 | with open(str(filename), "a+") as target: 72 | original_stdout_fd = 1 73 | original_stderr_fd = 2 74 | target_fd = target.fileno() 75 | 76 | # Save a copy of the original stdout and stderr file descriptors 77 | saved_stdout_fd = os.dup(original_stdout_fd) 78 | saved_stderr_fd = os.dup(original_stderr_fd) 79 | 80 | tee_stdout = subprocess.Popen( 81 | ["tee", "-a", "-i", "/dev/stderr"], 82 | start_new_session=True, 83 | stdin=subprocess.PIPE, 84 | stderr=target_fd, 85 | stdout=1, 86 | ) 87 | tee_stderr = subprocess.Popen( 88 | ["tee", "-a", "-i", "/dev/stderr"], 89 | start_new_session=True, 90 | stdin=subprocess.PIPE, 91 | stderr=target_fd, 92 | stdout=2, 93 | ) 94 | 95 | flush() 96 | os.dup2(tee_stdout.stdin.fileno(), original_stdout_fd) 97 | os.dup2(tee_stderr.stdin.fileno(), original_stderr_fd) 98 | 99 | try: 100 | yield 101 | finally: 102 | flush() 103 | 104 | # then redirect stdout back to the saved fd 105 | tee_stdout.stdin.close() 106 | tee_stderr.stdin.close() 107 | 108 | # restore original fds 109 | os.dup2(saved_stdout_fd, original_stdout_fd) 110 | os.dup2(saved_stderr_fd, original_stderr_fd) 111 | 112 | # wait for completion of the tee processes with timeout 113 | # implemented using a timer because timeout support is py3 only 114 | def kill_tees(): 115 | tee_stdout.kill() 116 | tee_stderr.kill() 117 | 118 | tee_timer = Timer(1, kill_tees) 119 | try: 120 | tee_timer.start() 121 | tee_stdout.wait() 122 | tee_stderr.wait() 123 | finally: 124 | tee_timer.cancel() 125 | 126 | os.close(saved_stdout_fd) 127 | os.close(saved_stderr_fd) 128 | 129 | # Cleanup log file 130 | with open(str(filename), "r") as target: 131 | text = target.read() 132 | text = apply_backspaces_and_linefeeds(text) 133 | with open(str(filename), "w") as target: 134 | target.write(text) 135 | -------------------------------------------------------------------------------- /gluefactory/utils/tensor.py: -------------------------------------------------------------------------------- 1 | """ 2 | Author: Paul-Edouard Sarlin (skydes) 3 | """ 4 | 5 | import collections.abc as collections 6 | 7 | import numpy as np 8 | import torch 9 | 10 | string_classes = (str, bytes) 11 | 12 | 13 | def map_tensor(input_, func): 14 | if isinstance(input_, string_classes): 15 | return input_ 16 | elif isinstance(input_, collections.Mapping): 17 | return {k: map_tensor(sample, func) for k, sample in input_.items()} 18 | elif isinstance(input_, collections.Sequence): 19 | return [map_tensor(sample, func) for sample in input_] 20 | elif input_ is None: 21 | return None 22 | else: 23 | return func(input_) 24 | 25 | 26 | def batch_to_numpy(batch): 27 | return map_tensor(batch, lambda tensor: tensor.cpu().numpy()) 28 | 29 | 30 | def batch_to_device(batch, device, non_blocking=True): 31 | def _func(tensor): 32 | return tensor.to(device=device, non_blocking=non_blocking) 33 | 34 | return map_tensor(batch, _func) 35 | 36 | 37 | def rbd(data: dict) -> dict: 38 | """Remove batch dimension from elements in data""" 39 | return { 40 | k: v[0] if isinstance(v, (torch.Tensor, np.ndarray, list)) else v 41 | for k, v in data.items() 42 | } 43 | 44 | 45 | def index_batch(tensor_dict): 46 | batch_size = len(next(iter(tensor_dict.values()))) 47 | for i in range(batch_size): 48 | yield map_tensor(tensor_dict, lambda t: t[i]) 49 | -------------------------------------------------------------------------------- /gluefactory/visualization/two_view_frame.py: -------------------------------------------------------------------------------- 1 | import pprint 2 | 3 | import numpy as np 4 | 5 | from . import viz2d 6 | from .tools import RadioHideTool, ToggleTool, __plot_dict__ 7 | 8 | 9 | class FormatPrinter(pprint.PrettyPrinter): 10 | def __init__(self, formats): 11 | super(FormatPrinter, self).__init__() 12 | self.formats = formats 13 | 14 | def format(self, obj, ctx, maxlvl, lvl): 15 | if type(obj) in self.formats: 16 | return self.formats[type(obj)] % obj, 1, 0 17 | return pprint.PrettyPrinter.format(self, obj, ctx, maxlvl, lvl) 18 | 19 | 20 | class TwoViewFrame: 21 | default_conf = { 22 | "default": "matches", 23 | "summary_visible": False, 24 | } 25 | 26 | plot_dict = __plot_dict__ 27 | 28 | childs = [] 29 | 30 | event_to_image = [None, "color", "depth", "color+depth"] 31 | 32 | def __init__(self, conf, data, preds, title=None, event=1, summaries=None): 33 | self.conf = conf 34 | self.data = data 35 | self.preds = preds 36 | self.names = list(preds.keys()) 37 | self.plot = self.event_to_image[event] 38 | self.summaries = summaries 39 | self.fig, self.axes, self.summary_arts = self.init_frame() 40 | if title is not None: 41 | self.fig.canvas.manager.set_window_title(title) 42 | 43 | keys = None 44 | for _, pred in preds.items(): 45 | if keys is None: 46 | keys = set(pred.keys()) 47 | else: 48 | keys = keys.intersection(pred.keys()) 49 | keys = keys.union(data.keys()) 50 | 51 | self.options = [ 52 | k for k, v in self.plot_dict.items() if set(v.required_keys).issubset(keys) 53 | ] 54 | self.handle = None 55 | self.radios = self.fig.canvas.manager.toolmanager.add_tool( 56 | "switch plot", 57 | RadioHideTool, 58 | options=self.options, 59 | callback_fn=self.draw, 60 | active=conf.default, 61 | keymap="R", 62 | ) 63 | 64 | self.toggle_summary = self.fig.canvas.manager.toolmanager.add_tool( 65 | "toggle summary", 66 | ToggleTool, 67 | toggled=self.conf.summary_visible, 68 | callback_fn=self.set_summary_visible, 69 | keymap="t", 70 | ) 71 | 72 | if self.fig.canvas.manager.toolbar is not None: 73 | self.fig.canvas.manager.toolbar.add_tool("switch plot", "navigation") 74 | self.draw(conf.default) 75 | 76 | def init_frame(self): 77 | """initialize frame""" 78 | view0, view1 = self.data["view0"], self.data["view1"] 79 | if self.plot == "color" or self.plot == "color+depth": 80 | imgs = [ 81 | view0["image"][0].permute(1, 2, 0), 82 | view1["image"][0].permute(1, 2, 0), 83 | ] 84 | elif self.plot == "depth": 85 | imgs = [view0["depth"][0], view1["depth"][0]] 86 | else: 87 | raise ValueError(self.plot) 88 | imgs = [imgs for _ in self.names] # repeat for each model 89 | 90 | fig, axes = viz2d.plot_image_grid(imgs, return_fig=True, titles=None, figs=5) 91 | [viz2d.add_text(0, n, axes=axes[i]) for i, n in enumerate(self.names)] 92 | 93 | if ( 94 | self.plot == "color+depth" 95 | and "depth" in view0.keys() 96 | and view0["depth"] is not None 97 | ): 98 | hmaps = [[view0["depth"][0], view1["depth"][0]] for _ in self.names] 99 | [ 100 | viz2d.plot_heatmaps(hmaps[i], axes=axes[i], cmap="Spectral") 101 | for i, _ in enumerate(hmaps) 102 | ] 103 | 104 | fig.canvas.mpl_connect("pick_event", self.click_artist) 105 | if self.summaries is not None: 106 | formatter = FormatPrinter({np.float32: "%.4f", np.float64: "%.4f"}) 107 | toggle_artists = [ 108 | viz2d.add_text( 109 | 0, 110 | formatter.pformat(self.summaries[n]), 111 | axes=axes[i], 112 | pos=(0.01, 0.01), 113 | va="bottom", 114 | backgroundcolor=(0, 0, 0, 0.5), 115 | visible=self.conf.summary_visible, 116 | ) 117 | for i, n in enumerate(self.names) 118 | ] 119 | else: 120 | toggle_artists = [] 121 | return fig, axes, toggle_artists 122 | 123 | def draw(self, value): 124 | """redraw content in frame""" 125 | self.clear() 126 | self.conf.default = value 127 | self.handle = self.plot_dict[value](self.fig, self.axes, self.data, self.preds) 128 | return self.handle 129 | 130 | def clear(self): 131 | if self.handle is not None: 132 | try: 133 | self.handle.clear() 134 | except AttributeError: 135 | pass 136 | self.handle = None 137 | for row in self.axes: 138 | for ax in row: 139 | [li.remove() for li in ax.lines] 140 | [c.remove() for c in ax.collections] 141 | self.fig.artists.clear() 142 | self.fig.canvas.draw_idle() 143 | self.handle = None 144 | 145 | def click_artist(self, event): 146 | art = event.artist 147 | select = art.get_arrowstyle().arrow == "-" 148 | art.set_arrowstyle("<|-|>" if select else "-") 149 | if select: 150 | art.set_zorder(1) 151 | if hasattr(self.handle, "click_artist"): 152 | self.handle.click_artist(event) 153 | self.fig.canvas.draw_idle() 154 | 155 | def set_summary_visible(self, visible): 156 | self.conf.summary_visible = visible 157 | [s.set_visible(visible) for s in self.summary_arts] 158 | self.fig.canvas.draw_idle() 159 | -------------------------------------------------------------------------------- /gluefactory/visualization/visualize_batch.py: -------------------------------------------------------------------------------- 1 | import torch 2 | 3 | from ..utils.tensor import batch_to_device 4 | from .viz2d import cm_RdGn, plot_heatmaps, plot_image_grid, plot_keypoints, plot_matches 5 | 6 | 7 | def make_match_figures(pred_, data_, n_pairs=2): 8 | # print first n pairs in batch 9 | if "0to1" in pred_.keys(): 10 | pred_ = pred_["0to1"] 11 | images, kpts, matches, mcolors = [], [], [], [] 12 | heatmaps = [] 13 | pred = batch_to_device(pred_, "cpu", non_blocking=False) 14 | data = batch_to_device(data_, "cpu", non_blocking=False) 15 | 16 | view0, view1 = data["view0"], data["view1"] 17 | 18 | n_pairs = min(n_pairs, view0["image"].shape[0]) 19 | assert view0["image"].shape[0] >= n_pairs 20 | 21 | kp0, kp1 = pred["keypoints0"], pred["keypoints1"] 22 | m0 = pred["matches0"] 23 | gtm0 = pred["gt_matches0"] 24 | 25 | for i in range(n_pairs): 26 | valid = (m0[i] > -1) & (gtm0[i] >= -1) 27 | kpm0, kpm1 = kp0[i][valid].numpy(), kp1[i][m0[i][valid]].numpy() 28 | images.append( 29 | [view0["image"][i].permute(1, 2, 0), view1["image"][i].permute(1, 2, 0)] 30 | ) 31 | kpts.append([kp0[i], kp1[i]]) 32 | matches.append((kpm0, kpm1)) 33 | 34 | correct = gtm0[i][valid] == m0[i][valid] 35 | 36 | if "heatmap0" in pred.keys(): 37 | heatmaps.append( 38 | [ 39 | torch.sigmoid(pred["heatmap0"][i, 0]), 40 | torch.sigmoid(pred["heatmap1"][i, 0]), 41 | ] 42 | ) 43 | elif "depth" in view0.keys() and view0["depth"] is not None: 44 | heatmaps.append([view0["depth"][i], view1["depth"][i]]) 45 | 46 | mcolors.append(cm_RdGn(correct).tolist()) 47 | 48 | fig, axes = plot_image_grid(images, return_fig=True, set_lim=True) 49 | if len(heatmaps) > 0: 50 | [plot_heatmaps(heatmaps[i], axes=axes[i], a=1.0) for i in range(n_pairs)] 51 | [plot_keypoints(kpts[i], axes=axes[i], colors="royalblue") for i in range(n_pairs)] 52 | [ 53 | plot_matches(*matches[i], color=mcolors[i], axes=axes[i], a=0.5, lw=1.0, ps=0.0) 54 | for i in range(n_pairs) 55 | ] 56 | 57 | return {"matching": fig} 58 | -------------------------------------------------------------------------------- /inloc_localization.py: -------------------------------------------------------------------------------- 1 | from pathlib import Path 2 | from pprint import pformat 3 | from omegaconf import OmegaConf 4 | from args import * 5 | import sys 6 | sys.path.append('/mnt/personal/weitong/cache/Hierarchical-Localization/') 7 | from hloc import extract_features, match_features, localize_inloc 8 | 9 | opt = create_parser() 10 | dataset_dirs = OmegaConf.load('dump_datasets/data_dirs.yaml') 11 | dataset = Path(dataset_dirs.get('dataset_dirs')[opt.dataset]) 12 | 13 | loc_pairs = opt.loc_pairs # top 40 retrieved 14 | outputs = Path(f"{opt.output_dir}/{opt.dataset}/{opt.method}{opt.num_loc}") if opt.method is not None else Path(f"{opt.output_dir}/{opt.dataset}/{opt.model}{opt.num_loc}")# where everything will be saved 15 | outputs.mkdir(exist_ok=True, parents=True) 16 | 17 | results = outputs / f'InLoc_hloc_superpoint+superglue.txt' 18 | 19 | # list the standard configurations available 20 | # print(f'Configs for feature extractors:\n{pformat(extract_features.confs)}') 21 | # print(f'Configs for feature matchers:\n{pformat(match_features.confs)}') 22 | 23 | # pick one of the configurations for extraction and matching 24 | # you can also simply write your own here! 25 | feature_conf = extract_features.confs['superpoint_inloc'] 26 | matcher_conf = match_features.confs['superglue'] 27 | 28 | # ## Extract local features for database and query images 29 | feature_path = extract_features.main(feature_conf, dataset, outputs) 30 | 31 | # Here we assume that the localization pairs are already computed using image retrieval. To generate new pairs from your own global descriptors, have a look at `hloc/pairs_from_retrieval.py`. These pairs are also used for the localization - see below. 32 | match_path = match_features.main(matcher_conf, loc_pairs, feature_conf['output'], outputs) 33 | 34 | # ## Localize! 35 | # Perform hierarchical localization using the precomputed retrieval and matches. Different from when localizing with Aachen, here we do not need a 3D SfM model here: the dataset already has 3D lidar scans. The file `InLoc_hloc_superpoint+superglue_netvlad40.txt` will contain the estimated query poses. 36 | localize_inloc.main( 37 | dataset, loc_pairs, feature_path, match_path, results, 38 | skip_matches=10) # skip database images with too few matches 39 | 40 | print(f"Done, now you can submit {results} to LONG-TERM VISUAL LOCALIZATION benchmark!") 41 | -------------------------------------------------------------------------------- /relative_pose.py: -------------------------------------------------------------------------------- 1 | import os 2 | import torch 3 | import numpy as np 4 | import gluefactory 5 | from pathlib import Path 6 | from omegaconf import OmegaConf 7 | from hloc import extract_features, match_features 8 | from evaluate_utils import * 9 | from args import * 10 | 11 | opt = create_parser() 12 | 13 | # check if gpu device is available 14 | opt.device = torch.device('cuda:0' if torch.cuda.is_available() and opt.device != 'cpu' else 'cpu') 15 | torch.set_grad_enabled(False) 16 | 17 | # load the data dirs for all 18 | dataset_dirs = OmegaConf.load('dump_datasets/data_dirs.yaml') 19 | opt.dataset_dir = Path(dataset_dirs.get('dataset_dirs')[opt.dataset]) 20 | # load the query and db lists in the dumped data 21 | overlap_features = Path(opt.dump_dir) / opt.dataset / "overlap_feats.h5" 22 | avg_results = Path(opt.output_dir) / opt.model / "avg_results.txt" 23 | assert os.path.exists(overlap_features) 24 | 25 | scenes = ['Undistorted_SfM/0015/images','Undistorted_SfM/0022/images'] if opt.dataset == 'megadepth' else os.listdir(opt.dataset_dir) 26 | 27 | all = [] 28 | all_num_matches = [] 29 | 30 | for scene in scenes: 31 | print("Start testing on scene:", scene) 32 | if opt.dataset == 'eth3d': scene = f"{scene}/images/dslr_images_undistorted" 33 | if opt.dataset == 'phototourism': scene = f"{scene}/images/" 34 | if not (opt.dataset_dir/scene).is_dir(): 35 | print(f"skipped scene:{scene}") 36 | continue 37 | 38 | with h5py.File(str(overlap_features), 'r') as hfile: 39 | image_list = [f"{scene}/{f}" for f in hfile[scene].keys()] 40 | 41 | output_dir = Path(opt.output_dir) / opt.model / opt.dataset / scene 42 | output_dir.mkdir(exist_ok=True, parents=True) 43 | 44 | conf = extract_features.confs['superpoint_aachen'] 45 | features = output_dir / 'features.h5' 46 | if not os.path.exists(features) or opt.overwrite: 47 | extract_features.main(conf, opt.dataset_dir, feature_path=features, image_list=image_list) 48 | 49 | if opt.cls: 50 | output_dir = output_dir / ('cls_' + str(opt.pre_filter)) 51 | output_dir.mkdir(exist_ok=True) 52 | 53 | baseline_helper = BaselineHelper(image_list) 54 | radius = np.loadtxt(output_dir / "median_radius.txt").item() if opt.radius == -1 else opt.radius 55 | 56 | matches = output_dir / (f"top{opt.k}_{radius}_overlap_matches.h5") 57 | overlap_pairs = output_dir / Path(f"top{opt.k}_{radius}_overlap_pairs.txt") 58 | overlap_results = output_dir / "overlap_results.txt" 59 | 60 | if os.path.exists(overlap_pairs): 61 | if not os.path.exists(matches) or opt.overwrite: 62 | # local feature matching 63 | match_conf = match_features.confs['superpoint+lightglue'] 64 | match_features.main(match_conf, overlap_pairs, features, matches=matches, overwrite=opt.overwrite, unique_pairs=False) 65 | # check the number of matches in the retrieved image pairs 66 | pred_matches, _ = baseline_helper.match_after_retrival(matches, overlap_pairs) 67 | # relative pose estimation on the retrieved image pairs 68 | pred_pose_results, best_thre = baseline_helper.relative_pose_after_retrieval(matches, overlap_pairs, features, overlap_features) 69 | pose_results = [pred_pose_results[key] for key in ['rel_pose_error@10°', 'mrel_pose_error', 'mransac_inl']]#pred_pose_results.keys()] 70 | print("Finished test on scene:", scene) 71 | 72 | with open(overlap_results, "a") as doc: 73 | doc.write(f"{scene} recall@{opt.k} {radius} {[pred_matches[key] for key in pred_matches.keys()][0]} " + " ".join(map(str, pose_results[:16])) + " \n") 74 | all += [pose_results] 75 | all_num_matches += [pred_matches[key] for key in pred_matches.keys()] 76 | else: 77 | print(f"Skipped scene:{scene}, as no pairs provided in {overlap_pairs}!") 78 | continue 79 | 80 | with open(avg_results, "a") as doc: 81 | doc.write(f"{opt.dataset} {len(scenes)} {len(all)} recall@{opt.k} {opt.radius} {np.mean(all_num_matches)} " + " ".join(map(str, np.vstack(all).mean(0))) + " \n") 82 | -------------------------------------------------------------------------------- /train_configs/best.yaml: -------------------------------------------------------------------------------- 1 | data: 2 | name: megadepth 3 | preprocessing: 4 | resize: 224 5 | side: long 6 | square_pad: true 7 | data_dir: "/mnt/personal/weitong/olap_predictor_mlp_large/data/megadepth/" 8 | info_dir: "/mnt/personal/weitong/olap_predictor_mlp_large/data/megadepth/scene_info/" 9 | train_split: train_scenes_clean.txt 10 | val_split: valid_scenes_clean.txt 11 | train_num_per_scene: [200, 200] 12 | val_num_per_scene: [150, 150] 13 | num_overlap_bins: 3 14 | min_overlap: 0.1 15 | max_overlap: 0.7 16 | batch_size: 32 17 | num_workers: 8 18 | one_to_one: true 19 | aug: true 20 | photometric: { 21 | "name": "flip", 22 | "p": 0.95, 23 | } 24 | 25 | model: 26 | name: two_view_pipeline 27 | extractor: 28 | name: extractors.mixed 29 | detector: 30 | name: extractors.grid_extractor 31 | cell_size: 14 32 | descriptor: 33 | name: backbones.dinov2 34 | trainable: false 35 | weights: dinov2_vitl14 36 | trainable: false 37 | ground_truth: 38 | name: null 39 | matcher: 40 | name: overlap_predictor 41 | input_dim: 1024 42 | descriptor_dim: 256 43 | dropout_prob: 0.5 44 | train: 45 | seed: 0 46 | epochs: 30 47 | log_every_iter: 500 48 | eval_every_iter: 500 49 | optimizer: adamw 50 | lr: 0.0001 51 | lr_schedule: 52 | start: 10 53 | type: exp 54 | on_epoch: true 55 | exp_div_10: 10 56 | median_metrics: 57 | - match_recall 58 | submodules: [] 59 | plot: 60 | # - 10 61 | # - gluefactory.visualization.visualize_batch.make_overlap_figures 62 | -------------------------------------------------------------------------------- /train_configs/best_easy.yaml: -------------------------------------------------------------------------------- 1 | data: 2 | name: megadepth_easy 3 | preprocessing: 4 | resize: 224 5 | side: long 6 | square_pad: true 7 | data_dir: "/mnt/personal/weitong/olap_predictor_mlp_large/data/megadepth/" 8 | info_dir: "/mnt/personal/weitong/olap_predictor_mlp_large/data/megadepth/scene_info/" 9 | train_split: train_scenes_clean.txt 10 | val_split: test_scenes_clean.txt 11 | train_pairs: train_pairs.txt 12 | val_pairs: test_pairs.txt 13 | train_num_per_scene: null 14 | val_num_per_scene: null 15 | num_overlap_bins: 3 16 | min_overlap: 0.1 17 | max_overlap: 0.7 18 | batch_size: 32 19 | num_workers: 8 20 | one_to_one: true 21 | aug: true 22 | photometric: { 23 | "name": "flip", 24 | "p": 0.95, 25 | } 26 | gt_label_path: "/mnt/personal/weitong/cache/glue-factory/train_prepare/" 27 | 28 | model: 29 | name: two_view_pipeline 30 | extractor: 31 | name: extractors.mixed 32 | detector: 33 | name: extractors.grid_extractor 34 | cell_size: 14 35 | descriptor: 36 | name: backbones.dinov2 37 | trainable: false 38 | weights: dinov2_vitl14 39 | trainable: false 40 | ground_truth: 41 | name: null 42 | matcher: 43 | name: overlap_predictor 44 | input_dim: 1024 45 | descriptor_dim: 256 46 | dropout_prob: 0.5 47 | train: 48 | seed: 0 49 | epochs: 30 50 | log_every_iter: 500 51 | eval_every_iter: 500 52 | optimizer: adamw 53 | lr: 0.0001 54 | lr_schedule: 55 | start: 10 56 | type: exp 57 | on_epoch: true 58 | exp_div_10: 10 59 | median_metrics: 60 | - match_recall 61 | submodules: [] 62 | plot: 63 | # - 10 64 | # - gluefactory.visualization.visualize_batch.make_overlap_figures 65 | --------------------------------------------------------------------------------