add: superglue and dedode

third_party/DeDoDe/.gitignore

@@ -0,0 +1,162 @@
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+
+# C extensions
+*.so
+
+# Distribution / packaging
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+share/python-wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+MANIFEST
+
+# PyInstaller
+#  Usually these files are written by a python script from a template
+#  before PyInstaller builds the exe, so as to inject date/other infos into it.
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.nox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*.cover
+*.py,cover
+.hypothesis/
+.pytest_cache/
+cover/
+
+# Translations
+*.mo
+*.pot
+
+# Django stuff:
+*.log
+local_settings.py
+db.sqlite3
+db.sqlite3-journal
+
+# Flask stuff:
+instance/
+.webassets-cache
+
+# Scrapy stuff:
+.scrapy
+
+# Sphinx documentation
+docs/_build/
+
+# PyBuilder
+.pybuilder/
+target/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# IPython
+profile_default/
+ipython_config.py
+
+# pyenv
+#   For a library or package, you might want to ignore these files since the code is
+#   intended to run in multiple environments; otherwise, check them in:
+# .python-version
+
+# pipenv
+#   According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
+#   However, in case of collaboration, if having platform-specific dependencies or dependencies
+#   having no cross-platform support, pipenv may install dependencies that don't work, or not
+#   install all needed dependencies.
+#Pipfile.lock
+
+# poetry
+#   Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control.
+#   This is especially recommended for binary packages to ensure reproducibility, and is more
+#   commonly ignored for libraries.
+#   https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control
+#poetry.lock
+
+# pdm
+#   Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control.
+#pdm.lock
+#   pdm stores project-wide configurations in .pdm.toml, but it is recommended to not include it
+#   in version control.
+#   https://pdm.fming.dev/#use-with-ide
+.pdm.toml
+
+# PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm
+__pypackages__/
+
+# Celery stuff
+celerybeat-schedule
+celerybeat.pid
+
+# SageMath parsed files
+*.sage.py
+
+# Environments
+.env
+.venv
+env/
+venv/
+ENV/
+env.bak/
+venv.bak/
+
+# Spyder project settings
+.spyderproject
+.spyproject
+
+# Rope project settings
+.ropeproject
+
+# mkdocs documentation
+/site
+
+# mypy
+.mypy_cache/
+.dmypy.json
+dmypy.json
+
+# Pyre type checker
+.pyre/
+
+# pytype static type analyzer
+.pytype/
+
+# Cython debug symbols
+cython_debug/
+
+# PyCharm
+#  JetBrains specific template is maintained in a separate JetBrains.gitignore that can
+#  be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore
+#  and can be added to the global gitignore or merged into this file.  For a more nuclear
+#  option (not recommended) you can uncomment the following to ignore the entire idea folder.
+#.idea/
+
+.vscode*

third_party/DeDoDe/DeDoDe/__init__.py

@@ -0,0 +1 @@
+from .model_zoo import dedode_detector_B, dedode_detector_L, dedode_descriptor_B

third_party/DeDoDe/DeDoDe/benchmarks/__init__.py

@@ -0,0 +1,3 @@
+from .num_inliers import NumInliersBenchmark
+from .mega_pose_est import MegaDepthPoseEstimationBenchmark
+from .mega_pose_est_mnn import MegaDepthPoseMNNBenchmark

third_party/DeDoDe/DeDoDe/benchmarks/mega_pose_est.py

@@ -0,0 +1,114 @@
+import numpy as np
+import torch
+from DeDoDe.utils import *
+from PIL import Image
+from tqdm import tqdm
+import torch.nn.functional as F
+
+class MegaDepthPoseEstimationBenchmark:
+    def __init__(self, data_root="data/megadepth", scene_names = None) -> None:
+        if scene_names is None:
+            self.scene_names = [
+                "0015_0.1_0.3.npz",
+                "0015_0.3_0.5.npz",
+                "0022_0.1_0.3.npz",
+                "0022_0.3_0.5.npz",
+                "0022_0.5_0.7.npz",
+            ]
+        else:
+            self.scene_names = scene_names
+        self.scenes = [
+            np.load(f"{data_root}/{scene}", allow_pickle=True)
+            for scene in self.scene_names
+        ]
+        self.data_root = data_root
+
+    def benchmark(self, keypoint_model, matching_model, model_name = None, resolution = None, scale_intrinsics = True, calibrated = True):
+        H,W = matching_model.get_output_resolution()
+        with torch.no_grad():
+            data_root = self.data_root
+            tot_e_t, tot_e_R, tot_e_pose = [], [], []
+            thresholds = [5, 10, 20]
+            for scene_ind in range(len(self.scenes)):
+                import os
+                scene_name = os.path.splitext(self.scene_names[scene_ind])[0]
+                scene = self.scenes[scene_ind]
+                pairs = scene["pair_infos"]
+                intrinsics = scene["intrinsics"]
+                poses = scene["poses"]
+                im_paths = scene["image_paths"]
+                pair_inds = range(len(pairs))
+                for pairind in tqdm(pair_inds):
+                    idx1, idx2 = pairs[pairind][0]
+                    K1 = intrinsics[idx1].copy()
+                    T1 = poses[idx1].copy()
+                    R1, t1 = T1[:3, :3], T1[:3, 3]
+                    K2 = intrinsics[idx2].copy()
+                    T2 = poses[idx2].copy()
+                    R2, t2 = T2[:3, :3], T2[:3, 3]
+                    R, t = compute_relative_pose(R1, t1, R2, t2)
+                    T1_to_2 = np.concatenate((R,t[:,None]), axis=-1)
+                    im_A_path = f"{data_root}/{im_paths[idx1]}"
+                    im_B_path = f"{data_root}/{im_paths[idx2]}"
+                    
+                    keypoints_A = keypoint_model.detect_from_path(im_A_path, num_keypoints = 20_000)["keypoints"][0]
+                    keypoints_B = keypoint_model.detect_from_path(im_B_path, num_keypoints = 20_000)["keypoints"][0]
+                    warp, certainty = matching_model.match(im_A_path, im_B_path)
+                    matches = matching_model.match_keypoints(keypoints_A, keypoints_B, warp, certainty, return_tuple = False)                    
+                    im_A = Image.open(im_A_path)
+                    w1, h1 = im_A.size
+                    im_B = Image.open(im_B_path)
+                    w2, h2 = im_B.size
+                    if scale_intrinsics:
+                        scale1 = 1200 / max(w1, h1)
+                        scale2 = 1200 / max(w2, h2)
+                        w1, h1 = scale1 * w1, scale1 * h1
+                        w2, h2 = scale2 * w2, scale2 * h2
+                        K1, K2 = K1.copy(), K2.copy()
+                        K1[:2] = K1[:2] * scale1
+                        K2[:2] = K2[:2] * scale2
+                    kpts1, kpts2 = matching_model.to_pixel_coordinates(matches, h1, w1, h2, w2)
+                    for _ in range(1):
+                        shuffling = np.random.permutation(np.arange(len(kpts1)))
+                        kpts1 = kpts1[shuffling]
+                        kpts2 = kpts2[shuffling]
+                        try:
+                            threshold = 0.5 
+                            if calibrated:
+                                norm_threshold = threshold / (np.mean(np.abs(K1[:2, :2])) + np.mean(np.abs(K2[:2, :2])))
+                                R_est, t_est, mask = estimate_pose(
+                                    kpts1.cpu().numpy(),
+                                    kpts2.cpu().numpy(),
+                                    K1,
+                                    K2,
+                                    norm_threshold,
+                                    conf=0.99999,
+                                )
+                            T1_to_2_est = np.concatenate((R_est, t_est), axis=-1)  #
+                            e_t, e_R = compute_pose_error(T1_to_2_est, R, t)
+                            e_pose = max(e_t, e_R)
+                        except Exception as e:
+                            print(repr(e))
+                            e_t, e_R = 90, 90
+                            e_pose = max(e_t, e_R)
+                        tot_e_t.append(e_t)
+                        tot_e_R.append(e_R)
+                        tot_e_pose.append(e_pose)
+            tot_e_pose = np.array(tot_e_pose)
+            auc = pose_auc(tot_e_pose, thresholds)
+            acc_5 = (tot_e_pose < 5).mean()
+            acc_10 = (tot_e_pose < 10).mean()
+            acc_15 = (tot_e_pose < 15).mean()
+            acc_20 = (tot_e_pose < 20).mean()
+            map_5 = acc_5
+            map_10 = np.mean([acc_5, acc_10])
+            map_20 = np.mean([acc_5, acc_10, acc_15, acc_20])
+            print(f"{model_name} auc: {auc}")
+            return {
+                "auc_5": auc[0],
+                "auc_10": auc[1],
+                "auc_20": auc[2],
+                "map_5": map_5,
+                "map_10": map_10,
+                "map_20": map_20,
+            }

third_party/DeDoDe/DeDoDe/benchmarks/mega_pose_est_mnn.py

@@ -0,0 +1,119 @@
+import numpy as np
+import torch
+from DeDoDe.utils import *
+from PIL import Image
+from tqdm import tqdm
+import torch.nn.functional as F
+
+class MegaDepthPoseMNNBenchmark:
+    def __init__(self, data_root="data/megadepth", scene_names = None) -> None:
+        if scene_names is None:
+            self.scene_names = [
+                "0015_0.1_0.3.npz",
+                "0015_0.3_0.5.npz",
+                "0022_0.1_0.3.npz",
+                "0022_0.3_0.5.npz",
+                "0022_0.5_0.7.npz",
+            ]
+        else:
+            self.scene_names = scene_names
+        self.scenes = [
+            np.load(f"{data_root}/{scene}", allow_pickle=True)
+            for scene in self.scene_names
+        ]
+        self.data_root = data_root
+
+    def benchmark(self, detector_model, descriptor_model, matcher_model, model_name = None, resolution = None, scale_intrinsics = True, calibrated = True):
+        with torch.no_grad():
+            data_root = self.data_root
+            tot_e_t, tot_e_R, tot_e_pose = [], [], []
+            thresholds = [5, 10, 20]
+            for scene_ind in range(len(self.scenes)):
+                import os
+                scene_name = os.path.splitext(self.scene_names[scene_ind])[0]
+                scene = self.scenes[scene_ind]
+                pairs = scene["pair_infos"]
+                intrinsics = scene["intrinsics"]
+                poses = scene["poses"]
+                im_paths = scene["image_paths"]
+                pair_inds = range(len(pairs))
+                for pairind in tqdm(pair_inds):
+                    idx1, idx2 = pairs[pairind][0]
+                    K1 = intrinsics[idx1].copy()
+                    T1 = poses[idx1].copy()
+                    R1, t1 = T1[:3, :3], T1[:3, 3]
+                    K2 = intrinsics[idx2].copy()
+                    T2 = poses[idx2].copy()
+                    R2, t2 = T2[:3, :3], T2[:3, 3]
+                    R, t = compute_relative_pose(R1, t1, R2, t2)
+                    T1_to_2 = np.concatenate((R,t[:,None]), axis=-1)
+                    im_A_path = f"{data_root}/{im_paths[idx1]}"
+                    im_B_path = f"{data_root}/{im_paths[idx2]}"
+                    detections_A = detector_model.detect_from_path(im_A_path)
+                    keypoints_A, P_A = detections_A["keypoints"], detections_A["confidence"]
+                    detections_B = detector_model.detect_from_path(im_B_path)
+                    keypoints_B, P_B = detections_B["keypoints"], detections_B["confidence"]
+                    description_A = descriptor_model.describe_keypoints_from_path(im_A_path, keypoints_A)["descriptions"]
+                    description_B = descriptor_model.describe_keypoints_from_path(im_B_path, keypoints_B)["descriptions"]
+                    matches_A, matches_B, batch_ids = matcher_model.match(keypoints_A, description_A,
+                        keypoints_B, description_B,
+                        P_A = P_A, P_B = P_B,
+                        normalize = True, inv_temp=20, threshold = 0.01)
+
+                    im_A = Image.open(im_A_path)
+                    w1, h1 = im_A.size
+                    im_B = Image.open(im_B_path)
+                    w2, h2 = im_B.size
+                    if scale_intrinsics:
+                        scale1 = 1200 / max(w1, h1)
+                        scale2 = 1200 / max(w2, h2)
+                        w1, h1 = scale1 * w1, scale1 * h1
+                        w2, h2 = scale2 * w2, scale2 * h2
+                        K1, K2 = K1.copy(), K2.copy()
+                        K1[:2] = K1[:2] * scale1
+                        K2[:2] = K2[:2] * scale2
+                    kpts1, kpts2 = matcher_model.to_pixel_coords(matches_A, matches_B, h1, w1, h2, w2)
+                    for _ in range(1):
+                        shuffling = np.random.permutation(np.arange(len(kpts1)))
+                        kpts1 = kpts1[shuffling]
+                        kpts2 = kpts2[shuffling]
+                        try:
+                            threshold = 0.5 
+                            if calibrated:
+                                norm_threshold = threshold / (np.mean(np.abs(K1[:2, :2])) + np.mean(np.abs(K2[:2, :2])))
+                                R_est, t_est, mask = estimate_pose(
+                                    kpts1.cpu().numpy(),
+                                    kpts2.cpu().numpy(),
+                                    K1,
+                                    K2,
+                                    norm_threshold,
+                                    conf=0.99999,
+                                )
+                            T1_to_2_est = np.concatenate((R_est, t_est), axis=-1)  #
+                            e_t, e_R = compute_pose_error(T1_to_2_est, R, t)
+                            e_pose = max(e_t, e_R)
+                        except Exception as e:
+                            print(repr(e))
+                            e_t, e_R = 90, 90
+                            e_pose = max(e_t, e_R)
+                        tot_e_t.append(e_t)
+                        tot_e_R.append(e_R)
+                        tot_e_pose.append(e_pose)
+            tot_e_pose = np.array(tot_e_pose)
+            auc = pose_auc(tot_e_pose, thresholds)
+            acc_5 = (tot_e_pose < 5).mean()
+            acc_10 = (tot_e_pose < 10).mean()
+            acc_15 = (tot_e_pose < 15).mean()
+            acc_20 = (tot_e_pose < 20).mean()
+            map_5 = acc_5
+            map_10 = np.mean([acc_5, acc_10])
+            map_20 = np.mean([acc_5, acc_10, acc_15, acc_20])
+            print(f"{model_name} auc: {auc}")
+            return {
+                "auc_5": auc[0],
+                "auc_10": auc[1],
+                "auc_20": auc[2],
+                "map_5": map_5,
+                "map_10": map_10,
+                "map_20": map_20,
+            }

third_party/DeDoDe/DeDoDe/benchmarks/num_inliers.py

@@ -0,0 +1,76 @@
+import torch
+import torch.nn as nn
+from DeDoDe.utils import *
+import DeDoDe
+
+class NumInliersBenchmark(nn.Module):
+    
+    def __init__(self, dataset, num_samples = 1000, batch_size = 8, num_keypoints = 10_000, device = "cuda") -> None:
+        super().__init__()
+        sampler = torch.utils.data.WeightedRandomSampler(
+                torch.ones(len(dataset)), replacement=False, num_samples=num_samples
+            )
+        dataloader = torch.utils.data.DataLoader(
+                dataset, batch_size=batch_size, num_workers=batch_size, sampler=sampler
+            )
+        self.dataloader = dataloader
+        self.tracked_metrics = {}
+        self.batch_size = batch_size
+        self.N = len(dataloader)
+        self.num_keypoints = num_keypoints
+    
+    def compute_batch_metrics(self, outputs, batch, device = "cuda"):
+        kpts_A, kpts_B = outputs["keypoints_A"], outputs["keypoints_B"]
+        B, K, H, W = batch["im_A"].shape
+        gt_warp_A_to_B, valid_mask_A_to_B = get_gt_warp(                
+                    batch["im_A_depth"],
+                    batch["im_B_depth"],
+                    batch["T_1to2"],
+                    batch["K1"],
+                    batch["K2"],
+                    H=H,
+                    W=W,
+                )
+        kpts_A_to_B = F.grid_sample(gt_warp_A_to_B[...,2:].float().permute(0,3,1,2), kpts_A[...,None,:], 
+                                    align_corners=False, mode = 'bilinear')[...,0].mT
+        legit_A_to_B = F.grid_sample(valid_mask_A_to_B.reshape(B,1,H,W), kpts_A[...,None,:], 
+                                    align_corners=False, mode = 'bilinear')[...,0,:,0]
+        dists = (torch.cdist(kpts_A_to_B, kpts_B).min(dim=-1).values[legit_A_to_B > 0.]).float()
+        if legit_A_to_B.sum() == 0:
+            return
+        percent_inliers_at_1 = (dists < 0.02).float().mean()
+        percent_inliers_at_05 = (dists < 0.01).float().mean()
+        percent_inliers_at_025 = (dists < 0.005).float().mean()
+        percent_inliers_at_01 = (dists < 0.002).float().mean()
+        percent_inliers_at_005 = (dists < 0.001).float().mean()
+
+        inlier_bins = torch.linspace(0, 0.002, steps = 100, device = device)[None]
+        inlier_counts = (dists[...,None] < inlier_bins).float().mean(dim=0)
+        self.tracked_metrics["inlier_counts"] = self.tracked_metrics.get("inlier_counts", 0) + 1/self.N * inlier_counts
+        self.tracked_metrics["percent_inliers_at_1"] = self.tracked_metrics.get("percent_inliers_at_1", 0) + 1/self.N * percent_inliers_at_1
+        self.tracked_metrics["percent_inliers_at_05"] = self.tracked_metrics.get("percent_inliers_at_05", 0) + 1/self.N * percent_inliers_at_05
+        self.tracked_metrics["percent_inliers_at_025"] = self.tracked_metrics.get("percent_inliers_at_025", 0) + 1/self.N * percent_inliers_at_025
+        self.tracked_metrics["percent_inliers_at_01"] = self.tracked_metrics.get("percent_inliers_at_01", 0) + 1/self.N * percent_inliers_at_01
+        self.tracked_metrics["percent_inliers_at_005"] = self.tracked_metrics.get("percent_inliers_at_005", 0) + 1/self.N * percent_inliers_at_005
+
+    def benchmark(self, detector):
+        self.tracked_metrics = {}
+        from tqdm import tqdm
+        print("Evaluating percent inliers...")
+        for idx, batch in tqdm(enumerate(self.dataloader), mininterval = 10.):
+            batch = to_cuda(batch)
+            outputs = detector.detect(batch, num_keypoints = self.num_keypoints)
+            keypoints_A, keypoints_B = outputs["keypoints"][:self.batch_size], outputs["keypoints"][self.batch_size:] 
+            if isinstance(outputs["keypoints"], (tuple, list)):
+                keypoints_A, keypoints_B = torch.stack(keypoints_A), torch.stack(keypoints_B)
+            outputs = {"keypoints_A": keypoints_A, "keypoints_B": keypoints_B}
+            self.compute_batch_metrics(outputs, batch)
+        import matplotlib.pyplot as plt
+        plt.plot(torch.linspace(0, 0.002, steps = 100), self.tracked_metrics["inlier_counts"].cpu())
+        import numpy as np
+        x = np.linspace(0,0.002, 100)
+        sigma = 0.52 * 2 / 512
+        F = 1 - np.exp(-x**2 / (2*sigma**2))
+        plt.plot(x, F)
+        plt.savefig("vis/inlier_counts")
+        [print(name, metric.item() * self.N / (idx+1)) for name, metric in self.tracked_metrics.items() if "percent" in name]

third_party/DeDoDe/DeDoDe/checkpoint.py

@@ -0,0 +1,59 @@
+import os
+import torch
+from torch.nn.parallel.data_parallel import DataParallel
+from torch.nn.parallel.distributed import DistributedDataParallel
+import gc
+
+import DeDoDe
+
+class CheckPoint:
+    def __init__(self, dir=None, name="tmp"):
+        self.name = name
+        self.dir = dir
+        os.makedirs(self.dir, exist_ok=True)
+
+    def save(
+        self,
+        model,
+        optimizer,
+        lr_scheduler,
+        n,
+        ):
+        if DeDoDe.RANK == 0:
+            assert model is not None
+            if isinstance(model, (DataParallel, DistributedDataParallel)):
+                model = model.module
+            states = {
+                "model": model.state_dict(),
+                "n": n,
+                "optimizer": optimizer.state_dict(),
+                "lr_scheduler": lr_scheduler.state_dict(),
+            }
+            torch.save(states, self.dir + self.name + f"_latest.pth")
+            print(f"Saved states {list(states.keys())}, at step {n}")
+    
+    def load(
+        self,
+        model,
+        optimizer,
+        lr_scheduler,
+        n,
+        ):
+        if os.path.exists(self.dir + self.name + f"_latest.pth") and DeDoDe.RANK == 0:
+            states = torch.load(self.dir + self.name + f"_latest.pth")
+            if "model" in states:
+                model.load_state_dict(states["model"])
+            if "n" in states:
+                n = states["n"] if states["n"] else n
+            if "optimizer" in states:
+                try:
+                    optimizer.load_state_dict(states["optimizer"])
+                except Exception as e:
+                    print(f"Failed to load states for optimizer, with error {e}")
+            if "lr_scheduler" in states:
+                lr_scheduler.load_state_dict(states["lr_scheduler"])
+            print(f"Loaded states {list(states.keys())}, at step {n}")
+            del states
+            gc.collect()
+            torch.cuda.empty_cache()
+        return model, optimizer, lr_scheduler, n

third_party/DeDoDe/DeDoDe/datasets/megadepth.py

@@ -0,0 +1,269 @@
+import os
+from PIL import Image
+import h5py
+import numpy as np
+import torch
+import torchvision.transforms.functional as tvf
+from tqdm import tqdm
+
+from DeDoDe.utils import get_depth_tuple_transform_ops, get_tuple_transform_ops
+import DeDoDe
+from DeDoDe.utils import *
+
+class MegadepthScene:
+    def __init__(
+        self,
+        data_root,
+        scene_info,
+        ht=512,
+        wt=512,
+        min_overlap=0.0,
+        max_overlap=1.0,
+        shake_t=0,
+        scene_info_detections=None,
+        scene_info_detections3D=None,
+        normalize=True,
+        max_num_pairs = 100_000,
+        scene_name = None,
+        use_horizontal_flip_aug = False,
+        grayscale = False,
+        clahe = False,
+    ) -> None:
+        self.data_root = data_root
+        self.scene_name = os.path.splitext(scene_name)[0]+f"_{min_overlap}_{max_overlap}"
+        self.image_paths = scene_info["image_paths"]
+        self.depth_paths = scene_info["depth_paths"]
+        self.intrinsics = scene_info["intrinsics"]
+        self.poses = scene_info["poses"]
+        self.pairs = scene_info["pairs"]
+        self.overlaps = scene_info["overlaps"]
+        threshold = (self.overlaps > min_overlap) & (self.overlaps < max_overlap)
+        self.pairs = self.pairs[threshold]
+        self.overlaps = self.overlaps[threshold]
+        self.detections = scene_info_detections
+        self.tracks3D = scene_info_detections3D
+        if len(self.pairs) > max_num_pairs:
+            pairinds = np.random.choice(
+                np.arange(0, len(self.pairs)), max_num_pairs, replace=False
+            )
+            self.pairs = self.pairs[pairinds]
+            self.overlaps = self.overlaps[pairinds]
+        self.im_transform_ops = get_tuple_transform_ops(
+            resize=(ht, wt), normalize=normalize, clahe = clahe,
+        )
+        self.depth_transform_ops = get_depth_tuple_transform_ops(
+            resize=(ht, wt), normalize=False
+        )
+        self.wt, self.ht = wt, ht
+        self.shake_t = shake_t
+        self.use_horizontal_flip_aug = use_horizontal_flip_aug
+        self.grayscale = grayscale
+
+    def load_im(self, im_B, crop=None):
+        im = Image.open(im_B)
+        return im
+    
+    def horizontal_flip(self, im_A, im_B, depth_A, depth_B,  K_A, K_B):
+        im_A = im_A.flip(-1)
+        im_B = im_B.flip(-1)
+        depth_A, depth_B = depth_A.flip(-1), depth_B.flip(-1) 
+        flip_mat = torch.tensor([[-1, 0, self.wt],[0,1,0],[0,0,1.]]).to(K_A.device)
+        K_A = flip_mat@K_A  
+        K_B = flip_mat@K_B  
+        
+        return im_A, im_B, depth_A, depth_B, K_A, K_B
+        
+    def load_depth(self, depth_ref, crop=None):
+        depth = np.array(h5py.File(depth_ref, "r")["depth"])
+        return torch.from_numpy(depth)
+
+    def __len__(self):
+        return len(self.pairs)
+
+    def scale_intrinsic(self, K, wi, hi):
+        sx, sy = self.wt / wi, self.ht / hi
+        sK = torch.tensor([[sx, 0, 0], [0, sy, 0], [0, 0, 1]])
+        return sK @ K
+
+    def scale_detections(self, detections, wi, hi):
+        sx, sy = self.wt / wi, self.ht / hi
+        return detections * torch.tensor([[sx,sy]])
+    
+    def rand_shake(self, *things):
+        t = np.random.choice(range(-self.shake_t, self.shake_t + 1), size=(2))
+        return [
+            tvf.affine(thing, angle=0.0, translate=list(t), scale=1.0, shear=[0.0, 0.0])
+            for thing in things
+        ], t
+
+    def tracks_to_detections(self, tracks3D, pose, intrinsics, H, W):
+        tracks3D = tracks3D.double()
+        intrinsics = intrinsics.double()
+        bearing_vectors = pose[...,:3,:3] @ tracks3D.mT + pose[...,:3,3:]        
+        hom_pixel_coords = (intrinsics @ bearing_vectors).mT
+        pixel_coords = hom_pixel_coords[...,:2] / (hom_pixel_coords[...,2:]+1e-12)
+        legit_detections = (pixel_coords > 0).prod(dim = -1) * (pixel_coords[...,0] < W - 1) * (pixel_coords[...,1] < H - 1) * (tracks3D != 0).prod(dim=-1)
+        return pixel_coords.float(), legit_detections.bool()
+
+    def __getitem__(self, pair_idx):
+        try:
+            # read intrinsics of original size
+            idx1, idx2 = self.pairs[pair_idx]
+            K1 = torch.tensor(self.intrinsics[idx1].copy(), dtype=torch.float).reshape(3, 3)
+            K2 = torch.tensor(self.intrinsics[idx2].copy(), dtype=torch.float).reshape(3, 3)
+
+            # read and compute relative poses
+            T1 = self.poses[idx1]
+            T2 = self.poses[idx2]
+            T_1to2 = torch.tensor(np.matmul(T2, np.linalg.inv(T1)), dtype=torch.float)[
+                :4, :4
+            ]  # (4, 4)
+
+            # Load positive pair data
+            im_A, im_B = self.image_paths[idx1], self.image_paths[idx2]
+            depth1, depth2 = self.depth_paths[idx1], self.depth_paths[idx2]
+            im_A_ref = os.path.join(self.data_root, im_A)
+            im_B_ref = os.path.join(self.data_root, im_B)
+            depth_A_ref = os.path.join(self.data_root, depth1)
+            depth_B_ref = os.path.join(self.data_root, depth2)
+            # return torch.randn((1000,1000))
+            im_A = self.load_im(im_A_ref)
+            im_B = self.load_im(im_B_ref)
+            depth_A = self.load_depth(depth_A_ref)
+            depth_B = self.load_depth(depth_B_ref)
+
+            # Recompute camera intrinsic matrix due to the resize
+            W_A, H_A = im_A.width, im_A.height
+            W_B, H_B = im_B.width, im_B.height
+
+            detections2D_A = self.detections[idx1]
+            detections2D_B = self.detections[idx2]
+            
+            K = 10000
+            tracks3D_A = torch.zeros(K,3)
+            tracks3D_B = torch.zeros(K,3)
+            tracks3D_A[:len(detections2D_A)] = torch.tensor(self.tracks3D[detections2D_A[:K,-1].astype(np.int32)])
+            tracks3D_B[:len(detections2D_B)] = torch.tensor(self.tracks3D[detections2D_B[:K,-1].astype(np.int32)])
+            
+            #projs_A, _ = self.tracks_to_detections(tracks3D_A, T1, K1, W_A, H_A)
+            #tracks3D_B = torch.zeros(K,2)
+
+            K1 = self.scale_intrinsic(K1, W_A, H_A)
+            K2 = self.scale_intrinsic(K2, W_B, H_B)
+            
+            # Process images
+            im_A, im_B = self.im_transform_ops((im_A, im_B))
+            depth_A, depth_B = self.depth_transform_ops(
+                (depth_A[None, None], depth_B[None, None])
+            )
+            [im_A, depth_A], t_A = self.rand_shake(im_A, depth_A)
+            [im_B, depth_B], t_B = self.rand_shake(im_B, depth_B)
+
+            detections_A = -torch.ones(K,2)
+            detections_B = -torch.ones(K,2)
+            detections_A[:len(self.detections[idx1])] = self.scale_detections(torch.tensor(detections2D_A[:K,:2]), W_A, H_A) + t_A
+            detections_B[:len(self.detections[idx2])] = self.scale_detections(torch.tensor(detections2D_B[:K,:2]), W_B, H_B) + t_B
+
+            
+            K1[:2, 2] += t_A
+            K2[:2, 2] += t_B
+                    
+            if self.use_horizontal_flip_aug:
+                if np.random.rand() > 0.5:
+                    im_A, im_B, depth_A, depth_B, K1, K2 = self.horizontal_flip(im_A, im_B, depth_A, depth_B, K1, K2)
+                    detections_A[:,0] = W-detections_A
+                    detections_B[:,0] = W-detections_B
+                    
+            if DeDoDe.DEBUG_MODE:
+                tensor_to_pil(im_A[0], unnormalize=True).save(
+                                f"vis/im_A.jpg")
+                tensor_to_pil(im_B[0], unnormalize=True).save(
+                                f"vis/im_B.jpg")
+            if self.grayscale:
+                im_A = im_A.mean(dim=-3,keepdim=True)
+                im_B = im_B.mean(dim=-3,keepdim=True)
+            data_dict = {
+                "im_A": im_A,
+                "im_A_identifier": self.image_paths[idx1].split("/")[-1].split(".jpg")[0],
+                "im_B": im_B,
+                "im_B_identifier": self.image_paths[idx2].split("/")[-1].split(".jpg")[0],
+                "im_A_depth": depth_A[0, 0],
+                "im_B_depth": depth_B[0, 0],
+                "pose_A": T1,
+                "pose_B": T2,
+                "detections_A": detections_A,
+                "detections_B": detections_B,
+                "tracks3D_A": tracks3D_A,
+                "tracks3D_B": tracks3D_B,
+                "K1": K1,
+                "K2": K2,
+                "T_1to2": T_1to2,
+                "im_A_path": im_A_ref,
+                "im_B_path": im_B_ref,
+            }
+        except Exception as e:
+            print(e)
+            print(f"Failed to load image pair {self.pairs[pair_idx]}")
+            print("Loading a random pair in scene instead")
+            rand_ind = np.random.choice(range(len(self)))
+            return self[rand_ind]
+        return data_dict
+
+
+class MegadepthBuilder:
+    def __init__(self, data_root="data/megadepth", loftr_ignore=True, imc21_ignore = True) -> None:
+        self.data_root = data_root
+        self.scene_info_root = os.path.join(data_root, "prep_scene_info")
+        self.all_scenes = os.listdir(self.scene_info_root)
+        self.test_scenes = ["0017.npy", "0004.npy", "0048.npy", "0013.npy"]
+        # LoFTR did the D2-net preprocessing differently than we did and got more ignore scenes, can optionially ignore those
+        self.loftr_ignore_scenes = set(['0121.npy', '0133.npy', '0168.npy', '0178.npy', '0229.npy', '0349.npy', '0412.npy', '0430.npy', '0443.npy', '1001.npy', '5014.npy', '5015.npy', '5016.npy'])
+        self.imc21_scenes = set(['0008.npy', '0019.npy', '0021.npy', '0024.npy', '0025.npy', '0032.npy', '0063.npy', '1589.npy'])
+        self.test_scenes_loftr = ["0015.npy", "0022.npy"]
+        self.loftr_ignore = loftr_ignore
+        self.imc21_ignore = imc21_ignore
+
+    def build_scenes(self, split="train", min_overlap=0.0, scene_names = None, **kwargs):
+        if split == "train":
+            scene_names = set(self.all_scenes) - set(self.test_scenes)
+        elif split == "train_loftr":
+            scene_names = set(self.all_scenes) - set(self.test_scenes_loftr)
+        elif split == "test":
+            scene_names = self.test_scenes
+        elif split == "test_loftr":
+            scene_names = self.test_scenes_loftr
+        elif split == "custom":
+            scene_names = scene_names
+        else:
+            raise ValueError(f"Split {split} not available")
+        scenes = []
+        for scene_name in tqdm(scene_names):
+            if self.loftr_ignore and scene_name in self.loftr_ignore_scenes:
+                continue
+            if self.imc21_ignore and scene_name in self.imc21_scenes:
+                continue
+            if ".npy" not in scene_name:
+                continue
+            scene_info = np.load(
+                os.path.join(self.scene_info_root, scene_name), allow_pickle=True
+            ).item()
+            scene_info_detections = np.load(
+                os.path.join(self.scene_info_root, "detections", f"detections_{scene_name}"), allow_pickle=True
+            ).item()
+            scene_info_detections3D = np.load(
+                os.path.join(self.scene_info_root, "detections3D", f"detections3D_{scene_name}"), allow_pickle=True
+            )
+
+            scenes.append(
+                MegadepthScene(
+                    self.data_root, scene_info, scene_info_detections = scene_info_detections, scene_info_detections3D = scene_info_detections3D, min_overlap=min_overlap,scene_name = scene_name, **kwargs
+                )
+            )
+        return scenes
+
+    def weight_scenes(self, concat_dataset, alpha=0.5):
+        ns = []
+        for d in concat_dataset.datasets:
+            ns.append(len(d))
+        ws = torch.cat([torch.ones(n) / n**alpha for n in ns])
+        return ws

third_party/DeDoDe/DeDoDe/decoder.py

@@ -0,0 +1,90 @@
+import torch
+import torch.nn as nn
+import torchvision.models as tvm
+
+
+class Decoder(nn.Module):
+    def __init__(self, layers, *args, super_resolution = False, num_prototypes = 1, **kwargs) -> None:
+        super().__init__(*args, **kwargs)
+        self.layers = layers
+        self.scales = self.layers.keys()
+        self.super_resolution = super_resolution
+        self.num_prototypes = num_prototypes
+    def forward(self, features, context = None, scale = None):
+        if context is not None:
+            features = torch.cat((features, context), dim = 1)
+        stuff = self.layers[scale](features)
+        logits, context = stuff[:,:self.num_prototypes], stuff[:,self.num_prototypes:]
+        return logits, context
+
+class ConvRefiner(nn.Module):
+    def __init__(
+        self,
+        in_dim=6,
+        hidden_dim=16,
+        out_dim=2,
+        dw=True,
+        kernel_size=5,
+        hidden_blocks=5,
+        amp = True,
+        residual = False,
+        amp_dtype = torch.float16,
+    ):
+        super().__init__()
+        self.block1 = self.create_block(
+            in_dim, hidden_dim, dw=False, kernel_size=1,
+        )
+        self.hidden_blocks = nn.Sequential(
+            *[
+                self.create_block(
+                    hidden_dim,
+                    hidden_dim,
+                    dw=dw,
+                    kernel_size=kernel_size,
+                )
+                for hb in range(hidden_blocks)
+            ]
+        )
+        self.hidden_blocks = self.hidden_blocks
+        self.out_conv = nn.Conv2d(hidden_dim, out_dim, 1, 1, 0)
+        self.amp = amp
+        self.amp_dtype = amp_dtype
+        self.residual = residual
+        
+    def create_block(
+        self,
+        in_dim,
+        out_dim,
+        dw=True,
+        kernel_size=5,
+        bias = True,
+        norm_type = nn.BatchNorm2d,
+    ):
+        num_groups = 1 if not dw else in_dim
+        if dw:
+            assert (
+                out_dim % in_dim == 0
+            ), "outdim must be divisible by indim for depthwise"
+        conv1 = nn.Conv2d(
+            in_dim,
+            out_dim,
+            kernel_size=kernel_size,
+            stride=1,
+            padding=kernel_size // 2,
+            groups=num_groups,
+            bias=bias,
+        )
+        norm = norm_type(out_dim) if norm_type is nn.BatchNorm2d else norm_type(num_channels = out_dim)
+        relu = nn.ReLU(inplace=True)
+        conv2 = nn.Conv2d(out_dim, out_dim, 1, 1, 0)
+        return nn.Sequential(conv1, norm, relu, conv2)
+        
+    def forward(self, feats):
+        b,c,hs,ws = feats.shape
+        with torch.autocast("cuda", enabled=self.amp, dtype = self.amp_dtype):
+            x0 = self.block1(feats)
+            x = self.hidden_blocks(x0)
+            if self.residual:
+                x = (x + x0)/1.4
+            x = self.out_conv(x)
+            return x

third_party/DeDoDe/DeDoDe/descriptors/dedode_descriptor.py

@@ -0,0 +1,49 @@
+import torch
+from PIL import Image
+import torch.nn as nn
+import torchvision.models as tvm
+import torch.nn.functional as F
+import numpy as np
+
+class DeDoDeDescriptor(nn.Module):
+    def __init__(self, encoder, decoder, *args, **kwargs) -> None:
+        super().__init__(*args, **kwargs)
+        self.encoder = encoder
+        self.decoder = decoder
+        import torchvision.transforms as transforms
+        self.normalizer = transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
+        
+    def forward(
+        self,
+        batch,
+    ):
+        if "im_A" in batch:
+            images = torch.cat((batch["im_A"], batch["im_B"]))
+        else:
+            images = batch["image"]
+        features, sizes = self.encoder(images)
+        descriptor = 0
+        context = None
+        scales = self.decoder.scales
+        for idx, (feature_map, scale) in enumerate(zip(reversed(features), scales)):
+            delta_descriptor, context = self.decoder(feature_map, scale = scale, context = context)
+            descriptor = descriptor + delta_descriptor
+            if idx < len(scales) - 1:
+                size = sizes[-(idx+2)]
+                descriptor = F.interpolate(descriptor, size = size, mode = "bilinear", align_corners = False)
+                context = F.interpolate(context, size = size, mode = "bilinear", align_corners = False)
+        return {"description_grid" : descriptor}
+    
+    @torch.inference_mode()
+    def describe_keypoints(self, batch, keypoints):
+        self.train(False)
+        description_grid = self.forward(batch)["description_grid"]
+        described_keypoints = F.grid_sample(description_grid.float(), keypoints[:,None], mode = "bilinear", align_corners = False)[:,:,0].mT
+        return {"descriptions": described_keypoints}
+    
+    def read_image(self, im_path, H = 560, W = 560):
+        return self.normalizer(torch.from_numpy(np.array(Image.open(im_path).resize((W,H)))/255.).permute(2,0,1)).cuda().float()[None]
+
+    def describe_keypoints_from_path(self, im_path, keypoints, H = 768, W = 768):
+        batch = {"image": self.read_image(im_path, H = H, W = W)}
+        return self.describe_keypoints(batch, keypoints)

third_party/DeDoDe/DeDoDe/descriptors/descriptor_loss.py

@@ -0,0 +1,75 @@
+import torch
+import torch.nn as nn
+import math
+import torch.nn.functional as F
+
+from DeDoDe.utils import *
+import DeDoDe
+
+class DescriptorLoss(nn.Module):
+    
+    def __init__(self, detector, num_keypoints = 5000, normalize_descriptions = False, inv_temp = 1, device = "cuda") -> None:
+        super().__init__()
+        self.detector = detector
+        self.tracked_metrics = {}
+        self.num_keypoints = num_keypoints
+        self.normalize_descriptions = normalize_descriptions
+        self.inv_temp = inv_temp
+    
+    def warp_from_depth(self, batch, kpts_A, kpts_B):
+        mask_A_to_B, kpts_A_to_B = warp_kpts(kpts_A, 
+                    batch["im_A_depth"],
+                    batch["im_B_depth"],
+                    batch["T_1to2"],
+                    batch["K1"],
+                    batch["K2"],)
+        mask_B_to_A, kpts_B_to_A = warp_kpts(kpts_B, 
+                    batch["im_B_depth"],
+                    batch["im_A_depth"],
+                    batch["T_1to2"].inverse(),
+                    batch["K2"],
+                    batch["K1"],)
+        return (mask_A_to_B, kpts_A_to_B), (mask_B_to_A, kpts_B_to_A)
+    
+    def warp_from_homog(self, batch, kpts_A, kpts_B):
+        kpts_A_to_B = homog_transform(batch["Homog_A_to_B"], kpts_A)
+        kpts_B_to_A = homog_transform(batch["Homog_A_to_B"].inverse(), kpts_B)
+        return (None, kpts_A_to_B), (None, kpts_B_to_A)
+
+    def supervised_loss(self, outputs, batch):
+        kpts_A, kpts_B = self.detector.detect(batch, num_keypoints = self.num_keypoints)['keypoints'].clone().chunk(2)
+        desc_grid_A, desc_grid_B = outputs["description_grid"].chunk(2)
+        desc_A = F.grid_sample(desc_grid_A.float(), kpts_A[:,None], mode = "bilinear", align_corners = False)[:,:,0].mT
+        desc_B = F.grid_sample(desc_grid_B.float(), kpts_B[:,None], mode = "bilinear", align_corners = False)[:,:,0].mT
+        if "im_A_depth" in batch:
+            (mask_A_to_B, kpts_A_to_B), (mask_B_to_A, kpts_B_to_A) = self.warp_from_depth(batch, kpts_A, kpts_B)
+        elif "Homog_A_to_B" in batch:
+            (mask_A_to_B, kpts_A_to_B), (mask_B_to_A, kpts_B_to_A) = self.warp_from_homog(batch, kpts_A, kpts_B)
+            
+        with torch.no_grad():
+            D_B = torch.cdist(kpts_A_to_B, kpts_B)
+            D_A = torch.cdist(kpts_A, kpts_B_to_A)
+            inds = torch.nonzero((D_B == D_B.min(dim=-1, keepdim = True).values) 
+                                 * (D_A == D_A.min(dim=-2, keepdim = True).values)
+                                 * (D_B < 0.01)
+                                 * (D_A < 0.01))
+            
+        logP_A_B = dual_log_softmax_matcher(desc_A, desc_B, 
+                                            normalize = self.normalize_descriptions,
+                                            inv_temperature = self.inv_temp)
+        neg_log_likelihood = -logP_A_B[inds[:,0], inds[:,1], inds[:,2]].mean()
+        if False:
+            import matplotlib.pyplot as plt
+            inds0 = inds[inds[:,0] == 0]
+            mnn_A = kpts_A[0,inds0[:,1]].detach().cpu()
+            mnn_B = kpts_B[0,inds0[:,2]].detach().cpu()
+            plt.scatter(mnn_A[:,0], -mnn_A[:,1], s = 0.5)
+            plt.savefig("vis/mnn_A.jpg")
+        self.tracked_metrics["neg_log_likelihood"] = (0.99 * self.tracked_metrics.get("neg_log_likelihood", neg_log_likelihood.detach().item()) + 0.01 * neg_log_likelihood.detach().item())
+        if np.random.rand() > 0.99:
+            print(self.tracked_metrics["neg_log_likelihood"])
+        return neg_log_likelihood
+    
+    def forward(self, outputs, batch):
+        losses = self.supervised_loss(outputs, batch)
+        return losses

third_party/DeDoDe/DeDoDe/detectors/dedode_detector.py

@@ -0,0 +1,75 @@
+import torch
+from PIL import Image
+import torch.nn as nn
+import torchvision.models as tvm
+import torch.nn.functional as F
+import numpy as np
+
+from DeDoDe.utils import sample_keypoints, to_pixel_coords, to_normalized_coords
+
+
+
+class DeDoDeDetector(nn.Module):
+    def __init__(self, encoder, decoder, *args, **kwargs) -> None:
+        super().__init__(*args, **kwargs)
+        self.encoder = encoder
+        self.decoder = decoder
+        import torchvision.transforms as transforms
+        self.normalizer = transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
+        
+    def forward(
+        self,
+        batch,
+    ):
+        if "im_A" in batch:
+            images = torch.cat((batch["im_A"], batch["im_B"]))
+        else:
+            images = batch["image"]
+        features, sizes = self.encoder(images)
+        logits = 0
+        context = None
+        scales = ["8", "4", "2", "1"]
+        for idx, (feature_map, scale) in enumerate(zip(reversed(features), scales)):
+            delta_logits, context = self.decoder(feature_map, context = context, scale = scale)
+            logits = logits + delta_logits.float() # ensure float (need bf16 doesnt have f.interpolate)
+            if idx < len(scales) - 1:
+                size = sizes[-(idx+2)]
+                logits = F.interpolate(logits, size = size, mode = "bicubic", align_corners = False)
+                context = F.interpolate(context.float(), size = size, mode = "bilinear", align_corners = False)
+        return {"keypoint_logits" : logits.float()}
+    
+    @torch.inference_mode()
+    def detect(self, batch, num_keypoints = 10_000):
+        self.train(False)
+        keypoint_logits = self.forward(batch)["keypoint_logits"]
+        B,K,H,W = keypoint_logits.shape
+        keypoint_p = keypoint_logits.reshape(B, K*H*W).softmax(dim=-1).reshape(B, K, H*W).sum(dim=1)
+        keypoints, confidence = sample_keypoints(keypoint_p.reshape(B,H,W), 
+                                  use_nms = False, sample_topk = True, num_samples = num_keypoints, 
+                                  return_scoremap=True, sharpen = False, upsample = False,
+                                  increase_coverage=True)
+        return {"keypoints": keypoints, "confidence": confidence}
+
+    @torch.inference_mode()
+    def detect_dense(self, batch):
+        self.train(False)
+        keypoint_logits = self.forward(batch)["keypoint_logits"]
+        return {"dense_keypoint_logits": keypoint_logits}
+
+    def read_image(self, im_path, H = 560, W = 560):
+        pil_im = Image.open(im_path).resize((W, H))
+        standard_im = np.array(pil_im)/255.
+        return self.normalizer(torch.from_numpy(standard_im).permute(2,0,1)).cuda().float()[None]
+
+    def detect_from_path(self, im_path, num_keypoints = 30_000, H = 768, W = 768, dense = False):
+        batch = {"image": self.read_image(im_path, H = H, W = W)}
+        if dense:
+            return self.detect_dense(batch)
+        else:
+            return self.detect(batch, num_keypoints = num_keypoints)
+
+    def to_pixel_coords(self, x, H, W):
+        return to_pixel_coords(x, H, W)
+    
+    def to_normalized_coords(self, x, H, W):
+        return to_normalized_coords(x, H, W)

third_party/DeDoDe/DeDoDe/detectors/loss.py

@@ -0,0 +1,275 @@
+import torch
+import torch.nn as nn
+import math
+
+from DeDoDe.utils import *
+import DeDoDe
+
+class KeyPointLoss(nn.Module):
+    
+    def __init__(self, smoothing_size = 1, use_max_logit = False, entropy_target = 80, 
+                 num_matches = 1024, jacobian_density_adjustment = False,
+                 matchability_weight = 1, device = "cuda") -> None:
+        super().__init__()
+        X = torch.linspace(-1,1,smoothing_size, device = device)
+        G = (-X**2 / (2 *1/2**2)).exp()
+        G = G/G.sum()
+        self.use_max_logit = use_max_logit
+        self.entropy_target = entropy_target
+        self.smoothing_kernel = G[None, None, None,:]
+        self.smoothing_size = smoothing_size
+        self.tracked_metrics = {}
+        self.center = None
+        self.num_matches = num_matches
+        self.jacobian_density_adjustment = jacobian_density_adjustment
+        self.matchability_weight = matchability_weight
+        
+    def compute_consistency(self, logits_A, logits_B_to_A, mask = None):
+        
+        masked_logits_A = torch.full_like(logits_A, -torch.inf)
+        masked_logits_A[mask] = logits_A[mask]
+
+        masked_logits_B_to_A = torch.full_like(logits_B_to_A, -torch.inf)
+        masked_logits_B_to_A[mask] = logits_B_to_A[mask]
+
+        log_p_A = masked_logits_A.log_softmax(dim=-1)[mask]
+        log_p_B_to_A = masked_logits_B_to_A.log_softmax(dim=-1)[mask]
+
+        return self.compute_jensen_shannon_div(log_p_A, log_p_B_to_A)
+    
+    def compute_joint_neg_log_likelihood(self, logits_A, logits_B_to_A, detections_A = None, detections_B_to_A = None, mask = None, device = "cuda", dtype = torch.float32, num_matches = None):
+        B, K, HW = logits_A.shape
+        logits_A, logits_B_to_A = logits_A.to(dtype), logits_B_to_A.to(dtype)
+        mask = mask[:,None].expand(B, K, HW).reshape(B, K*HW)
+        log_p_B_to_A = self.masked_log_softmax(logits_B_to_A.reshape(B,K*HW), mask = mask)
+        log_p_A = self.masked_log_softmax(logits_A.reshape(B,K*HW), mask = mask)
+        log_p = log_p_A + log_p_B_to_A
+        if detections_A is None:
+            detections_A = torch.zeros_like(log_p_A)
+        if detections_B_to_A is None:
+            detections_B_to_A = torch.zeros_like(log_p_B_to_A)
+        detections_A = detections_A.reshape(B, HW)
+        detections_A[~mask] = 0         
+        detections_B_to_A = detections_B_to_A.reshape(B, HW)
+        detections_B_to_A[~mask] = 0
+        log_p_target = log_p.detach() + 50*detections_A + 50*detections_B_to_A
+        num_matches = self.num_matches if num_matches is None else num_matches
+        best_k = -(-log_p_target).flatten().kthvalue(k = B * num_matches, dim=-1).values
+        p_target = (log_p_target > best_k[..., None]).float().reshape(B,K*HW)/num_matches
+        return self.compute_cross_entropy(log_p_A[mask], p_target[mask]) + self.compute_cross_entropy(log_p_B_to_A[mask], p_target[mask])
+                
+    def compute_jensen_shannon_div(self, log_p, log_q):
+        return 1/2 * (self.compute_kl_div(log_p, log_q) + self.compute_kl_div(log_q, log_p))
+    
+    def compute_kl_div(self, log_p, log_q):
+        return (log_p.exp()*(log_p-log_q)).sum(dim=-1)
+    
+    def masked_log_softmax(self, logits, mask):
+        masked_logits = torch.full_like(logits, -torch.inf)
+        masked_logits[mask] = logits[mask]
+        log_p = masked_logits.log_softmax(dim=-1)
+        return log_p
+    
+    def masked_softmax(self, logits, mask):
+        masked_logits = torch.full_like(logits, -torch.inf)
+        masked_logits[mask] = logits[mask]
+        log_p = masked_logits.softmax(dim=-1)
+        return log_p
+    
+    def compute_entropy(self, logits, mask = None):
+        p = self.masked_softmax(logits, mask)[mask]
+        log_p = self.masked_log_softmax(logits, mask)[mask]
+        return -(log_p * p).sum(dim=-1) 
+
+    def compute_detection_img(self, detections, mask, B, H, W, device = "cuda"):
+        kernel_size = 5
+        X = torch.linspace(-2,2,kernel_size, device = device)
+        G = (-X**2 / (2 * (1/2)**2)).exp() # half pixel std
+        G = G/G.sum()
+        det_smoothing_kernel = G[None, None, None,:]
+        det_img = torch.zeros((B,1,H,W), device = device) # add small epsilon for later logstuff
+        for b in range(B):
+            valid_detections = (detections[b][mask[b]]).int()
+            det_img[b,0][valid_detections[:,1], valid_detections[:,0]] = 1
+        det_img = F.conv2d(det_img, weight = det_smoothing_kernel, padding = (kernel_size//2, 0))
+        det_img = F.conv2d(det_img, weight = det_smoothing_kernel.mT, padding = (0, kernel_size//2))
+        return det_img
+
+    def compute_cross_entropy(self, log_p_hat, p):
+        return -(log_p_hat * p).sum(dim=-1)
+
+    def compute_matchability(self, keypoint_p, has_depth, B, K, H, W, device = "cuda"):
+        smooth_keypoint_p = F.conv2d(keypoint_p.reshape(B,1,H,W), weight = self.smoothing_kernel, padding = (self.smoothing_size//2,0))
+        smooth_keypoint_p = F.conv2d(smooth_keypoint_p, weight = self.smoothing_kernel.mT, padding = (0,self.smoothing_size//2))
+        log_p_hat = (smooth_keypoint_p+1e-8).log().reshape(B,H*W).log_softmax(dim=-1)
+        smooth_has_depth = F.conv2d(has_depth.reshape(B,1,H,W), weight = self.smoothing_kernel, padding = (0,self.smoothing_size//2))
+        smooth_has_depth = F.conv2d(smooth_has_depth, weight = self.smoothing_kernel.mT, padding = (self.smoothing_size//2,0)).reshape(B,H*W)
+        p = smooth_has_depth/smooth_has_depth.sum(dim=-1,keepdim=True)
+        return self.compute_cross_entropy(log_p_hat, p) - self.compute_cross_entropy((p+1e-12).log(), p)
+
+    def tracks_to_detections(self, tracks3D, pose, intrinsics, H, W):
+        tracks3D = tracks3D.double()
+        intrinsics = intrinsics.double()
+        bearing_vectors = pose[:,:3,:3] @ tracks3D.mT + pose[:,:3,3:]        
+        hom_pixel_coords = (intrinsics @ bearing_vectors).mT
+        pixel_coords = hom_pixel_coords[...,:2] / (hom_pixel_coords[...,2:]+1e-12)
+        legit_detections = (pixel_coords > 0).prod(dim = -1) * (pixel_coords[...,0] < W - 1) * (pixel_coords[...,1] < H - 1) * (tracks3D != 0).prod(dim=-1)
+        return pixel_coords.float(), legit_detections.bool()
+    
+    def self_supervised_loss(self, outputs, batch):
+        keypoint_logits_A, keypoint_logits_B = outputs["keypoint_logits"].chunk(2)
+        B, K, H, W = keypoint_logits_A.shape
+        keypoint_logits_A = keypoint_logits_A.reshape(B, K, H*W)
+        keypoint_logits_B = keypoint_logits_B.reshape(B, K, H*W)
+        keypoint_logits = torch.cat((keypoint_logits_A, keypoint_logits_B))
+
+        warp_A_to_B, mask_A_to_B = get_homog_warp(
+            batch["Homog_A_to_B"], H, W
+        )
+        warp_B_to_A, mask_B_to_A = get_homog_warp(
+            torch.linalg.inv(batch["Homog_A_to_B"]), H, W
+        )
+        B = 2*B
+        
+        warp = torch.cat((warp_A_to_B, warp_B_to_A)).reshape(B, H*W, 4)
+        mask = torch.cat((mask_A_to_B, mask_B_to_A)).reshape(B,H*W)
+        
+        keypoint_logits_backwarped = F.grid_sample(torch.cat((keypoint_logits_B, keypoint_logits_A)).reshape(B,K,H,W), 
+            warp[...,-2:].reshape(B,H,W,2).float(), align_corners = False, mode = "bicubic")
+        
+        keypoint_logits_backwarped = keypoint_logits_backwarped.reshape(B,K,H*W)
+        joint_log_likelihood_loss = self.compute_joint_neg_log_likelihood(keypoint_logits, keypoint_logits_backwarped, 
+                                                                          mask = mask.bool(), num_matches = 5_000).mean()
+        return joint_log_likelihood_loss
+    
+    def supervised_loss(self, outputs, batch):
+        keypoint_logits_A, keypoint_logits_B = outputs["keypoint_logits"].chunk(2)
+        B, K, H, W = keypoint_logits_A.shape
+
+        detections_A, detections_B = batch["detections_A"], batch["detections_B"]
+        
+        tracks3D_A, tracks3D_B = batch["tracks3D_A"], batch["tracks3D_B"]
+        gt_warp_A_to_B, valid_mask_A_to_B = get_gt_warp(                
+                    batch["im_A_depth"],
+                    batch["im_B_depth"],
+                    batch["T_1to2"],
+                    batch["K1"],
+                    batch["K2"],
+                    H=H,
+                    W=W,
+                )
+        gt_warp_B_to_A, valid_mask_B_to_A = get_gt_warp(                
+            batch["im_B_depth"],
+            batch["im_A_depth"],
+            batch["T_1to2"].inverse(),
+            batch["K2"],
+            batch["K1"],
+            H=H,
+            W=W,
+        )
+        keypoint_logits_A = keypoint_logits_A.reshape(B, K, H*W)
+        keypoint_logits_B = keypoint_logits_B.reshape(B, K, H*W)
+        keypoint_logits = torch.cat((keypoint_logits_A, keypoint_logits_B))
+
+        B = 2*B
+        gt_warp = torch.cat((gt_warp_A_to_B, gt_warp_B_to_A))
+        valid_mask = torch.cat((valid_mask_A_to_B, valid_mask_B_to_A))
+        valid_mask = valid_mask.reshape(B,H*W)
+        binary_mask = valid_mask == 1
+        if self.jacobian_density_adjustment:
+            j_logdet = jacobi_determinant(gt_warp.reshape(B,H,W,4), valid_mask.reshape(B,H,W).float())[:,None]
+        else:
+            j_logdet = 0
+        tracks3D = torch.cat((tracks3D_A, tracks3D_B))
+        
+        #detections, legit_detections = self.tracks_to_detections(tracks3D, torch.cat((batch["pose_A"],batch["pose_B"])), torch.cat((batch["K1"],batch["K2"])), H, W)
+        #detections_backwarped, legit_backwarped_detections = self.tracks_to_detections(torch.cat((tracks3D_B, tracks3D_A)), torch.cat((batch["pose_A"],batch["pose_B"])), torch.cat((batch["K1"],batch["K2"])), H, W)
+        detections = torch.cat((detections_A, detections_B))
+        legit_detections = ((detections > 0).prod(dim = -1) * (detections[...,0] < W) * (detections[...,1] < H)).bool()
+        det_imgs_A, det_imgs_B = self.compute_detection_img(detections, legit_detections, B, H, W).chunk(2)
+        det_imgs = torch.cat((det_imgs_A, det_imgs_B))
+        #det_imgs_backwarped = self.compute_detection_img(detections_backwarped, legit_backwarped_detections, B, H, W)
+        det_imgs_backwarped = F.grid_sample(torch.cat((det_imgs_B, det_imgs_A)).reshape(B,1,H,W), 
+            gt_warp[...,-2:].reshape(B,H,W,2).float(), align_corners = False, mode = "bicubic")
+
+        keypoint_logits_backwarped = F.grid_sample(torch.cat((keypoint_logits_B, keypoint_logits_A)).reshape(B,K,H,W), 
+            gt_warp[...,-2:].reshape(B,H,W,2).float(), align_corners = False, mode = "bicubic")
+        
+        # Note: Below step should be taken, but seems difficult to get it to work well.
+        #keypoint_logits_B_to_A = keypoint_logits_B_to_A + j_logdet_A_to_B # adjust for the viewpoint by log jacobian of warp
+        keypoint_logits_backwarped = (keypoint_logits_backwarped + j_logdet).reshape(B,K,H*W)
+
+
+        depth = F.interpolate(torch.cat((batch["im_A_depth"][:,None],batch["im_B_depth"][:,None]),dim=0), size = (H,W), mode = "bilinear", align_corners=False)
+        has_depth = (depth > 0).float().reshape(B,H*W)
+        
+        joint_log_likelihood_loss = self.compute_joint_neg_log_likelihood(keypoint_logits, keypoint_logits_backwarped, 
+                                                                          mask = binary_mask, detections_A = det_imgs, 
+                                                                          detections_B_to_A = det_imgs_backwarped).mean()
+        keypoint_p = keypoint_logits.reshape(B, K*H*W).softmax(dim=-1).reshape(B, K, H*W).sum(dim=1)
+        matchability_loss = self.compute_matchability(keypoint_p, has_depth, B, K, H, W).mean()
+        
+        #peakiness_loss = self.compute_negative_peakiness(keypoint_logits.reshape(B,H,W), mask = binary_mask)
+        #mnn_loss = self.compute_mnn_loss(keypoint_logits_A, keypoint_logits_B, gt_warp_A_to_B, valid_mask_A_to_B, B, H, W)
+        B = B//2
+        import matplotlib.pyplot as plt
+        kpts_A = sample_keypoints(keypoint_p[:B].reshape(B,H,W), 
+                                use_nms = False, sample_topk = True, num_samples = 4*2048)
+        kpts_B = sample_keypoints(keypoint_p[B:].reshape(B,H,W),
+                                use_nms = False, sample_topk = True, num_samples = 4*2048)
+        kpts_A_to_B = F.grid_sample(gt_warp_A_to_B[...,2:].float().permute(0,3,1,2), kpts_A[...,None,:], 
+                                    align_corners=False, mode = 'bilinear')[...,0].mT
+        legit_A_to_B = F.grid_sample(valid_mask_A_to_B.reshape(B,1,H,W), kpts_A[...,None,:], 
+                                    align_corners=False, mode = 'bilinear')[...,0,:,0]
+        percent_inliers = (torch.cdist(kpts_A_to_B, kpts_B).min(dim=-1).values[legit_A_to_B > 0] < 0.01).float().mean()
+        self.tracked_metrics["mega_percent_inliers"] = (0.9 * self.tracked_metrics.get("mega_percent_inliers", percent_inliers) + 0.1 * percent_inliers)
+
+        if torch.rand(1) > 0.995:
+            keypoint_logits_A_to_B = keypoint_logits_backwarped[:B]
+            import matplotlib.pyplot as plt
+            import os
+            os.makedirs("vis",exist_ok = True)
+            for b in range(0, B, 2):
+                #import cv2
+                plt.scatter(kpts_A_to_B[b,:,0].cpu(),-kpts_A_to_B[b,:,1].cpu(), s = 1)
+                plt.scatter(kpts_B[b,:,0].cpu(),-kpts_B[b,:,1].cpu(), s = 1)
+                plt.xlim(-1,1)
+                plt.ylim(-1,1)
+                plt.savefig(f"vis/keypoints_A_to_B_vs_B_{b}.png")
+                plt.close()
+                tensor_to_pil(keypoint_logits_A[b].reshape(1,H,W).expand(3,H,W).detach().cpu(), 
+                            autoscale = True).save(f"vis/logits_A_{b}.png")
+                tensor_to_pil(keypoint_logits_B[b].reshape(1,H,W).expand(3,H,W).detach().cpu(), 
+                            autoscale = True).save(f"vis/logits_B_{b}.png")
+                tensor_to_pil(keypoint_logits_A_to_B[b].reshape(1,H,W).expand(3,H,W).detach().cpu(), 
+                            autoscale = True).save(f"vis/logits_A_to_B{b}.png")
+                tensor_to_pil(keypoint_logits_A[b].softmax(dim=-1).reshape(1,H,W).expand(3,H,W).detach().cpu(), 
+                            autoscale = True).save(f"vis/keypoint_p_A_{b}.png")
+                tensor_to_pil(keypoint_logits_B[b].softmax(dim=-1).reshape(1,H,W).expand(3,H,W).detach().cpu(), 
+                            autoscale = True).save(f"vis/keypoint_p_B_{b}.png")
+                tensor_to_pil(has_depth[b].reshape(1,H,W).expand(3,H,W).detach().cpu(), autoscale=True).save(f"vis/has_depth_A_{b}.png")                            
+                tensor_to_pil(valid_mask_A_to_B[b].reshape(1,H,W).expand(3,H,W).detach().cpu(), autoscale=True).save(f"vis/valid_mask_A_to_B_{b}.png")                            
+                tensor_to_pil(batch['im_A'][b], unnormalize=True).save(
+                                    f"vis/im_A_{b}.jpg")
+                tensor_to_pil(batch['im_B'][b], unnormalize=True).save(
+                                    f"vis/im_B_{b}.jpg")
+            plt.close()
+        tot_loss = joint_log_likelihood_loss + self.matchability_weight * matchability_loss# 
+        #tot_loss = tot_loss + (-2*consistency_loss).detach().exp()*compression_loss
+        if torch.rand(1) > 1:
+            print(f"Precent Inlier: {self.tracked_metrics.get('mega_percent_inliers', 0)}")
+            print(f"{joint_log_likelihood_loss=} {matchability_loss=}")
+            print(f"Total Loss: {tot_loss.item()}")
+        return  tot_loss
+    
+    def forward(self, outputs, batch):
+        
+        if not isinstance(outputs, list):
+            outputs = [outputs]
+        losses = 0
+        for output in outputs:
+            if "Homog_A_to_B" in batch:
+                losses = losses + self.self_supervised_loss(output, batch)
+            else:
+                losses = losses + self.supervised_loss(output, batch)
+        return losses

third_party/DeDoDe/DeDoDe/encoder.py

@@ -0,0 +1,47 @@
+import torch
+import torch.nn as nn
+import torchvision.models as tvm
+
+
+class VGG19(nn.Module):
+    def __init__(self, pretrained=False, amp = False, amp_dtype = torch.float16) -> None:
+        super().__init__()
+        self.layers = nn.ModuleList(tvm.vgg19_bn(pretrained=pretrained).features[:40])
+        # Maxpool layers: 6, 13, 26, 39
+        self.amp = amp
+        self.amp_dtype = amp_dtype
+
+    def forward(self, x, **kwargs):
+        with torch.autocast("cuda", enabled=self.amp, dtype = self.amp_dtype):
+            feats = []
+            sizes = []
+            for layer in self.layers:
+                if isinstance(layer, nn.MaxPool2d):
+                    feats.append(x)
+                    sizes.append(x.shape[-2:])
+                x = layer(x)
+            return feats, sizes
+
+class VGG(nn.Module):
+    def __init__(self, size = "19", pretrained=False, amp = False, amp_dtype = torch.float16) -> None:
+        super().__init__()
+        if size == "11":
+            self.layers = nn.ModuleList(tvm.vgg11_bn(pretrained=pretrained).features[:22])
+        elif size == "13": 
+            self.layers = nn.ModuleList(tvm.vgg13_bn(pretrained=pretrained).features[:28])
+        elif size == "19": 
+            self.layers = nn.ModuleList(tvm.vgg19_bn(pretrained=pretrained).features[:40])
+        # Maxpool layers: 6, 13, 26, 39
+        self.amp = amp
+        self.amp_dtype = amp_dtype
+
+    def forward(self, x, **kwargs):
+        with torch.autocast("cuda", enabled=self.amp, dtype = self.amp_dtype):
+            feats = []
+            sizes = []
+            for layer in self.layers:
+                if isinstance(layer, nn.MaxPool2d):
+                    feats.append(x)
+                    sizes.append(x.shape[-2:])
+                x = layer(x)
+            return feats, sizes

third_party/DeDoDe/DeDoDe/matchers/dual_softmax_matcher.py

@@ -0,0 +1,38 @@
+import torch
+from PIL import Image
+import torch.nn as nn
+import torchvision.models as tvm
+import torch.nn.functional as F
+import numpy as np
+from DeDoDe.utils import dual_softmax_matcher, to_pixel_coords, to_normalized_coords
+
+class DualSoftMaxMatcher(nn.Module):        
+    @torch.inference_mode()
+    def match(self, keypoints_A, descriptions_A, 
+              keypoints_B, descriptions_B, P_A = None, P_B = None, 
+              normalize = False, inv_temp = 1, threshold = 0.0):
+        if isinstance(descriptions_A, list):
+            matches = [self.match(k_A[None], d_A[None], k_B[None], d_B[None], normalize = normalize,
+                               inv_temp = inv_temp, threshold = threshold) 
+                    for k_A,d_A,k_B,d_B in
+                    zip(keypoints_A, descriptions_A, keypoints_B, descriptions_B)]
+            matches_A = torch.cat([m[0] for m in matches])
+            matches_B = torch.cat([m[1] for m in matches])
+            inds = torch.cat([m[2] + b for b, m in enumerate(matches)])
+            return matches_A, matches_B, inds
+        
+        P = dual_softmax_matcher(descriptions_A, descriptions_B, 
+                                 normalize = normalize, inv_temperature=inv_temp,
+                                 )
+        inds = torch.nonzero((P == P.max(dim=-1, keepdim = True).values) 
+                        * (P == P.max(dim=-2, keepdim = True).values) * (P > threshold))
+        batch_inds = inds[:,0]
+        matches_A = keypoints_A[batch_inds, inds[:,1]]
+        matches_B = keypoints_B[batch_inds, inds[:,2]]
+        return matches_A, matches_B, batch_inds
+
+    def to_pixel_coords(self, x_A, x_B, H_A, W_A, H_B, W_B):
+        return to_pixel_coords(x_A, H_A, W_A), to_pixel_coords(x_B, H_B, W_B)
+    
+    def to_normalized_coords(self, x_A, x_B, H_A, W_A, H_B, W_B):
+        return to_normalized_coords(x_A, H_A, W_A), to_normalized_coords(x_B, H_B, W_B)

third_party/DeDoDe/DeDoDe/model_zoo/__init__.py

@@ -0,0 +1,3 @@
+from .dedode_models import dedode_detector_B, dedode_detector_L, dedode_descriptor_B
+    
+    

third_party/DeDoDe/DeDoDe/model_zoo/dedode_models.py

@@ -0,0 +1,177 @@
+import torch
+import torch.nn as nn
+
+from DeDoDe.detectors.dedode_detector import DeDoDeDetector
+from DeDoDe.descriptors.dedode_descriptor import DeDoDeDescriptor
+from DeDoDe.decoder import ConvRefiner, Decoder
+from DeDoDe.encoder import VGG19, VGG
+
+
+
+def dedode_detector_B(device = "cuda", weights = None):
+    residual = True
+    hidden_blocks = 5
+    amp_dtype = torch.float16
+    amp = True
+    NUM_PROTOTYPES = 1
+    conv_refiner = nn.ModuleDict(
+        {
+            "8": ConvRefiner(
+                512,
+                512,
+                256 + NUM_PROTOTYPES,
+                hidden_blocks = hidden_blocks,
+                residual = residual,
+                amp = amp,
+                amp_dtype = amp_dtype,
+            ),
+            "4": ConvRefiner(
+                256+256,
+                256,
+                128 + NUM_PROTOTYPES,
+                hidden_blocks = hidden_blocks,
+                residual = residual,
+                amp = amp,
+                amp_dtype = amp_dtype,
+
+            ),
+            "2": ConvRefiner(
+                128+128,
+                64,
+                32 + NUM_PROTOTYPES,
+                hidden_blocks = hidden_blocks,
+                residual = residual,
+                amp = amp,
+                amp_dtype = amp_dtype,
+
+            ),
+            "1": ConvRefiner(
+                64 + 32,
+                32,
+                1 + NUM_PROTOTYPES,
+                hidden_blocks = hidden_blocks,
+                residual = residual,
+                amp = amp,
+                amp_dtype = amp_dtype,
+            ),
+        }
+    )
+    encoder = VGG19(pretrained = False, amp = amp, amp_dtype = amp_dtype)
+    decoder = Decoder(conv_refiner)
+    model = DeDoDeDetector(encoder = encoder, decoder = decoder).to(device)
+    if weights is not None:
+        model.load_state_dict(weights)
+    return model
+
+
+def dedode_detector_L(device = "cuda", weights = None):
+    NUM_PROTOTYPES = 1
+    residual = True
+    hidden_blocks = 8
+    amp_dtype = torch.float16#torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16
+    amp = True
+    conv_refiner = nn.ModuleDict(
+        {
+            "8": ConvRefiner(
+                512,
+                512,
+                256 + NUM_PROTOTYPES,
+                hidden_blocks = hidden_blocks,
+                residual = residual,
+                amp = amp,
+                amp_dtype = amp_dtype,
+            ),
+            "4": ConvRefiner(
+                256+256,
+                256,
+                128 + NUM_PROTOTYPES,
+                hidden_blocks = hidden_blocks,
+                residual = residual,
+                amp = amp,
+                amp_dtype = amp_dtype,
+
+            ),
+            "2": ConvRefiner(
+                128+128,
+                128,
+                64 + NUM_PROTOTYPES,
+                hidden_blocks = hidden_blocks,
+                residual = residual,
+                amp = amp,
+                amp_dtype = amp_dtype,
+
+            ),
+            "1": ConvRefiner(
+                64 + 64,
+                64,
+                1 + NUM_PROTOTYPES,
+                hidden_blocks = hidden_blocks,
+                residual = residual,
+                amp = amp,
+                amp_dtype = amp_dtype,
+            ),
+        }
+    )
+    encoder = VGG19(pretrained = False, amp = amp, amp_dtype = amp_dtype)
+    decoder = Decoder(conv_refiner)
+    model = DeDoDeDetector(encoder = encoder, decoder = decoder).to(device)
+    if weights is not None:
+        model.load_state_dict(weights)
+    return model
+
+
+
+def dedode_descriptor_B(device = "cuda", weights = None):
+    NUM_PROTOTYPES = 256 # == descriptor size
+    residual = True
+    hidden_blocks = 5
+    amp_dtype = torch.float16#torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16
+    amp = True
+    conv_refiner = nn.ModuleDict(
+        {
+            "8": ConvRefiner(
+                512,
+                512,
+                256 + NUM_PROTOTYPES,
+                hidden_blocks = hidden_blocks,
+                residual = residual,
+                amp = amp,
+                amp_dtype = amp_dtype,
+            ),
+            "4": ConvRefiner(
+                256+256,
+                256,
+                128 + NUM_PROTOTYPES,
+                hidden_blocks = hidden_blocks,
+                residual = residual,
+                amp = amp,
+                amp_dtype = amp_dtype,
+
+            ),
+            "2": ConvRefiner(
+                128+128,
+                64,
+                32 + NUM_PROTOTYPES,
+                hidden_blocks = hidden_blocks,
+                residual = residual,
+                amp = amp,
+                amp_dtype = amp_dtype,
+
+            ),
+            "1": ConvRefiner(
+                64 + 32,
+                32,
+                1 + NUM_PROTOTYPES,
+                hidden_blocks = hidden_blocks,
+                residual = residual,
+                amp = amp,
+                amp_dtype = amp_dtype,
+            ),
+        }
+    )
+    encoder = VGG(size = "19", pretrained = False, amp = amp, amp_dtype = amp_dtype)
+    decoder = Decoder(conv_refiner, num_prototypes=NUM_PROTOTYPES)
+    model = DeDoDeDescriptor(encoder = encoder, decoder = decoder).to(device)    
+    if weights is not None:
+        model.load_state_dict(weights)
+    return model

third_party/DeDoDe/DeDoDe/train.py

@@ -0,0 +1,76 @@
+import torch
+from tqdm import tqdm
+from DeDoDe.utils import to_cuda
+
+
+def train_step(train_batch, model, objective, optimizer, grad_scaler = None,**kwargs):
+    optimizer.zero_grad()
+    out = model(train_batch)
+    l = objective(out, train_batch)
+    if grad_scaler is not None:
+        grad_scaler.scale(l).backward()
+        grad_scaler.unscale_(optimizer)
+        torch.nn.utils.clip_grad_norm_(model.parameters(), 0.01)
+        grad_scaler.step(optimizer)
+        grad_scaler.update()
+    else:
+        l.backward()
+        optimizer.step()
+    return {"train_out": out, "train_loss": l.item()}
+
+
+def train_k_steps(
+    n_0, k, dataloader, model, objective, optimizer, lr_scheduler, grad_scaler = None, progress_bar=True
+):
+    for n in tqdm(range(n_0, n_0 + k), disable=not progress_bar, mininterval = 10.):
+        batch = next(dataloader)
+        model.train(True)
+        batch = to_cuda(batch)
+        train_step(
+            train_batch=batch,
+            model=model,
+            objective=objective,
+            optimizer=optimizer,
+            lr_scheduler=lr_scheduler,
+            n=n,
+            grad_scaler = grad_scaler,
+        )
+        lr_scheduler.step()
+
+
+def train_epoch(
+    dataloader=None,
+    model=None,
+    objective=None,
+    optimizer=None,
+    lr_scheduler=None,
+    epoch=None,
+):
+    model.train(True)
+    print(f"At epoch {epoch}")
+    for batch in tqdm(dataloader, mininterval=5.0):
+        batch = to_cuda(batch)
+        train_step(
+            train_batch=batch, model=model, objective=objective, optimizer=optimizer
+        )
+    lr_scheduler.step()
+    return {
+        "model": model,
+        "optimizer": optimizer,
+        "lr_scheduler": lr_scheduler,
+        "epoch": epoch,
+    }
+
+
+def train_k_epochs(
+    start_epoch, end_epoch, dataloader, model, objective, optimizer, lr_scheduler
+):
+    for epoch in range(start_epoch, end_epoch + 1):
+        train_epoch(
+            dataloader=dataloader,
+            model=model,
+            objective=objective,
+            optimizer=optimizer,
+            lr_scheduler=lr_scheduler,
+            epoch=epoch,
+        )

third_party/DeDoDe/DeDoDe/utils.py

@@ -0,0 +1,759 @@
+import warnings
+import numpy as np
+import math
+import cv2
+import torch
+from torchvision import transforms
+from torchvision.transforms.functional import InterpolationMode
+import torch.nn.functional as F
+from PIL import Image
+from einops import rearrange
+import torch
+from time import perf_counter
+
+def recover_pose(E, kpts0, kpts1, K0, K1, mask):
+    best_num_inliers = 0
+    K0inv = np.linalg.inv(K0[:2,:2])
+    K1inv = np.linalg.inv(K1[:2,:2])
+
+    kpts0_n = (K0inv @ (kpts0-K0[None,:2,2]).T).T 
+    kpts1_n = (K1inv @ (kpts1-K1[None,:2,2]).T).T
+
+    for _E in np.split(E, len(E) / 3):
+        n, R, t, _ = cv2.recoverPose(_E, kpts0_n, kpts1_n, np.eye(3), 1e9, mask=mask)
+        if n > best_num_inliers:
+            best_num_inliers = n
+            ret = (R, t, mask.ravel() > 0)
+    return ret
+
+
+
+# Code taken from https://github.com/PruneTruong/DenseMatching/blob/40c29a6b5c35e86b9509e65ab0cd12553d998e5f/validation/utils_pose_estimation.py
+# --- GEOMETRY ---
+def estimate_pose(kpts0, kpts1, K0, K1, norm_thresh, conf=0.99999):
+    if len(kpts0) < 5:
+        return None
+    K0inv = np.linalg.inv(K0[:2,:2])
+    K1inv = np.linalg.inv(K1[:2,:2])
+
+    kpts0 = (K0inv @ (kpts0-K0[None,:2,2]).T).T 
+    kpts1 = (K1inv @ (kpts1-K1[None,:2,2]).T).T
+    E, mask = cv2.findEssentialMat(
+        kpts0, kpts1, np.eye(3), threshold=norm_thresh, prob=conf
+    )
+
+    ret = None
+    if E is not None:
+        best_num_inliers = 0
+
+        for _E in np.split(E, len(E) / 3):
+            n, R, t, _ = cv2.recoverPose(_E, kpts0, kpts1, np.eye(3), 1e9, mask=mask)
+            if n > best_num_inliers:
+                best_num_inliers = n
+                ret = (R, t, mask.ravel() > 0)
+    return ret
+
+
+def get_grid(B,H,W, device = "cuda"):
+    x1_n = torch.meshgrid(
+    *[
+        torch.linspace(
+            -1 + 1 / n, 1 - 1 / n, n, device=device
+        )
+        for n in (B, H, W)
+    ]
+    )
+    x1_n = torch.stack((x1_n[2], x1_n[1]), dim=-1).reshape(B, H * W, 2)
+    return x1_n
+
+@torch.no_grad()
+def finite_diff_hessian(f: tuple(["B", "H", "W"]), device = "cuda"):
+    dxx = torch.tensor([[0,0,0],[1,-2,1],[0,0,0]], device = device)[None,None]/2
+    dxy = torch.tensor([[1,0,-1],[0,0,0],[-1,0,1]], device = device)[None,None]/4
+    dyy = dxx.mT
+    Hxx = F.conv2d(f[:,None], dxx, padding = 1)[:,0]
+    Hxy = F.conv2d(f[:,None], dxy, padding = 1)[:,0]
+    Hyy = F.conv2d(f[:,None], dyy, padding = 1)[:,0]
+    H = torch.stack((Hxx, Hxy, Hxy, Hyy), dim = -1).reshape(*f.shape,2,2)
+    return H
+
+def finite_diff_grad(f: tuple(["B", "H", "W"]), device = "cuda"):
+    dx = torch.tensor([[0,0,0],[-1,0,1],[0,0,0]],device = device)[None,None]/2
+    dy = dx.mT
+    gx = F.conv2d(f[:,None], dx, padding = 1)
+    gy = F.conv2d(f[:,None], dy, padding = 1)
+    g = torch.cat((gx, gy), dim = 1)
+    return g
+
+def fast_inv_2x2(matrix: tuple[...,2,2], eps = 1e-10):
+    return 1/(torch.linalg.det(matrix)[...,None,None]+eps) * torch.stack((matrix[...,1,1],-matrix[...,0,1],
+                                                     -matrix[...,1,0],matrix[...,0,0]),dim=-1).reshape(*matrix.shape)
+
+def newton_step(f:tuple["B","H","W"], inds, device = "cuda"):
+    B,H,W = f.shape
+    Hess = finite_diff_hessian(f).reshape(B,H*W,2,2)
+    Hess = torch.gather(Hess, dim = 1, index = inds[...,None].expand(B,-1,2,2))
+    grad = finite_diff_grad(f).reshape(B,H*W,2)
+    grad = torch.gather(grad, dim = 1, index = inds)
+    Hessinv = fast_inv_2x2(Hess-torch.eye(2, device = device)[None,None])
+    step = (Hessinv @ grad[...,None])
+    return step[...,0]
+
+@torch.no_grad()
+def sample_keypoints(scoremap, num_samples = 8192, device = "cuda", use_nms = True, 
+                     sample_topk = False, return_scoremap = False, sharpen = False, upsample = False,
+                     increase_coverage = False,):
+    #scoremap = scoremap**2
+    log_scoremap = (scoremap+1e-10).log()
+    if upsample:
+        log_scoremap = F.interpolate(log_scoremap[:,None], scale_factor = 3, mode = "bicubic", align_corners = False)[:,0]#.clamp(min = 0)
+        scoremap = log_scoremap.exp()
+    B,H,W = scoremap.shape
+    if increase_coverage:
+        weights = (-torch.linspace(-2, 2, steps = 51, device = device)**2).exp()[None,None]
+        # 10000 is just some number for maybe numerical stability, who knows. :), result is invariant anyway
+        local_density_x = F.conv2d((scoremap[:,None]+1e-6)*10000,weights[...,None,:], padding = (0,51//2))
+        local_density = F.conv2d(local_density_x, weights[...,None], padding = (51//2,0))[:,0]
+        scoremap = scoremap * (local_density+1e-8)**(-1/2)
+    grid = get_grid(B,H,W, device=device).reshape(B,H*W,2)
+    if sharpen:
+        laplace_operator = torch.tensor([[[[0,1,0],[1,-4,1],[0,1,0]]]], device = device)/4
+        scoremap = scoremap[:,None] - 0.5 * F.conv2d(scoremap[:,None], weight = laplace_operator, padding = 1)
+        scoremap = scoremap[:,0].clamp(min = 0)
+    if use_nms:
+        scoremap = scoremap * (scoremap == F.max_pool2d(scoremap, (3, 3), stride = 1, padding = 1))
+    if sample_topk:
+        inds = torch.topk(scoremap.reshape(B,H*W), k = num_samples).indices
+    else:
+        inds = torch.multinomial(scoremap.reshape(B,H*W), num_samples = num_samples, replacement=False)
+    kps = torch.gather(grid, dim = 1, index = inds[...,None].expand(B,num_samples,2))
+    if return_scoremap:
+        return kps, torch.gather(scoremap.reshape(B,H*W), dim = 1, index = inds)
+    return kps
+
+@torch.no_grad()
+def jacobi_determinant(warp, certainty, R = 3, device = "cuda", dtype = torch.float32):
+    t = perf_counter()
+    *dims, _ = warp.shape
+    warp = warp.to(dtype)
+    certainty = certainty.to(dtype)
+    
+    dtype = warp.dtype
+    match_regions = torch.zeros((*dims, 4, R, R), device = device).to(dtype)
+    match_regions[:,1:-1, 1:-1] = warp.unfold(1,R,1).unfold(2,R,1)
+    match_regions = rearrange(match_regions,"B H W D R1 R2 -> B H W (R1 R2) D") - warp[...,None,:]
+    
+    match_regions_cert = torch.zeros((*dims, R, R), device = device).to(dtype)
+    match_regions_cert[:,1:-1, 1:-1] = certainty.unfold(1,R,1).unfold(2,R,1)
+    match_regions_cert = rearrange(match_regions_cert,"B H W R1 R2 -> B H W (R1 R2)")[..., None]
+
+    #print("Time for unfold", perf_counter()-t)
+    #t = perf_counter()
+    *dims, N, D = match_regions.shape
+    # standardize:
+    mu, sigma = match_regions.mean(dim=(-2,-1), keepdim = True), match_regions.std(dim=(-2,-1),keepdim=True)
+    match_regions = (match_regions-mu)/(sigma+1e-6)
+    x_a, x_b = match_regions.chunk(2,-1)
+    
+
+    A = torch.zeros((*dims,2*x_a.shape[-2],4), device = device).to(dtype)
+    A[...,::2,:2] = x_a * match_regions_cert
+    A[...,1::2,2:] = x_a * match_regions_cert
+
+    a_block = A[...,::2,:2]
+    ata = a_block.mT @ a_block
+    #print("Time for ata", perf_counter()-t)
+    #t = perf_counter()
+
+    #atainv = torch.linalg.inv(ata+1e-5*torch.eye(2,device=device).to(dtype))
+    atainv = fast_inv_2x2(ata)
+    ATA_inv = torch.zeros((*dims, 4, 4), device = device, dtype = dtype)
+    ATA_inv[...,:2,:2] = atainv
+    ATA_inv[...,2:,2:] = atainv
+    atb = A.mT @ (match_regions_cert*x_b).reshape(*dims,N*2,1)
+    theta =  ATA_inv @ atb
+    #print("Time for theta", perf_counter()-t)
+    #t = perf_counter()
+
+    J = theta.reshape(*dims, 2, 2)
+    abs_J_det = torch.linalg.det(J+1e-8*torch.eye(2,2,device = device).expand(*dims,2,2)).abs() # Note: This should always be positive for correct warps, but still taking abs here
+    abs_J_logdet = (abs_J_det+1e-12).log()
+    B = certainty.shape[0]
+    # Handle outliers
+    robust_abs_J_logdet = abs_J_logdet.clamp(-3, 3) # Shouldn't be more that exp(3) \approx 8 times zoom
+    #print("Time for logdet", perf_counter()-t)
+    #t = perf_counter()
+
+    return robust_abs_J_logdet
+
+def get_gt_warp(depth1, depth2, T_1to2, K1, K2, depth_interpolation_mode = 'bilinear', relative_depth_error_threshold = 0.05, H = None, W = None):
+    
+    if H is None:
+        B,H,W = depth1.shape
+    else:
+        B = depth1.shape[0]
+    with torch.no_grad():
+        x1_n = torch.meshgrid(
+            *[
+                torch.linspace(
+                    -1 + 1 / n, 1 - 1 / n, n, device=depth1.device
+                )
+                for n in (B, H, W)
+            ]
+        )
+        x1_n = torch.stack((x1_n[2], x1_n[1]), dim=-1).reshape(B, H * W, 2)
+        mask, x2 = warp_kpts(
+            x1_n.double(),
+            depth1.double(),
+            depth2.double(),
+            T_1to2.double(),
+            K1.double(),
+            K2.double(),
+            depth_interpolation_mode = depth_interpolation_mode,
+            relative_depth_error_threshold = relative_depth_error_threshold,
+        )
+        prob = mask.float().reshape(B, H, W)
+        x2 = x2.reshape(B, H, W, 2)
+        return torch.cat((x1_n.reshape(B,H,W,2),x2),dim=-1), prob
+
+def recover_pose(E, kpts0, kpts1, K0, K1, mask):
+    best_num_inliers = 0
+    K0inv = np.linalg.inv(K0[:2,:2])
+    K1inv = np.linalg.inv(K1[:2,:2])
+
+    kpts0_n = (K0inv @ (kpts0-K0[None,:2,2]).T).T 
+    kpts1_n = (K1inv @ (kpts1-K1[None,:2,2]).T).T
+
+    for _E in np.split(E, len(E) / 3):
+        n, R, t, _ = cv2.recoverPose(_E, kpts0_n, kpts1_n, np.eye(3), 1e9, mask=mask)
+        if n > best_num_inliers:
+            best_num_inliers = n
+            ret = (R, t, mask.ravel() > 0)
+    return ret
+
+
+
+# Code taken from https://github.com/PruneTruong/DenseMatching/blob/40c29a6b5c35e86b9509e65ab0cd12553d998e5f/validation/utils_pose_estimation.py
+# --- GEOMETRY ---
+def estimate_pose(kpts0, kpts1, K0, K1, norm_thresh, conf=0.99999, ):
+    if len(kpts0) < 5:
+        return None
+    K0inv = np.linalg.inv(K0[:2,:2])
+    K1inv = np.linalg.inv(K1[:2,:2])
+
+    kpts0 = (K0inv @ (kpts0-K0[None,:2,2]).T).T 
+    kpts1 = (K1inv @ (kpts1-K1[None,:2,2]).T).T
+    method = cv2.USAC_ACCURATE
+    E, mask = cv2.findEssentialMat(
+        kpts0, kpts1, np.eye(3), threshold=norm_thresh, prob=conf, method=method
+    )
+
+    ret = None
+    if E is not None:
+        best_num_inliers = 0
+
+        for _E in np.split(E, len(E) / 3):
+            n, R, t, _ = cv2.recoverPose(_E, kpts0, kpts1, np.eye(3), 1e9, mask=mask)
+            if n > best_num_inliers:
+                best_num_inliers = n
+                ret = (R, t, mask.ravel() > 0)
+    return ret
+
+def estimate_pose_uncalibrated(kpts0, kpts1, K0, K1, norm_thresh, conf=0.99999):
+    if len(kpts0) < 5:
+        return None
+    method = cv2.USAC_ACCURATE
+    F, mask = cv2.findFundamentalMat(
+        kpts0, kpts1, ransacReprojThreshold=norm_thresh, confidence=conf, method=method, maxIters=10000
+    )
+    E = K1.T@F@K0
+    ret = None
+    if E is not None:
+        best_num_inliers = 0
+        K0inv = np.linalg.inv(K0[:2,:2])
+        K1inv = np.linalg.inv(K1[:2,:2])
+
+        kpts0_n = (K0inv @ (kpts0-K0[None,:2,2]).T).T 
+        kpts1_n = (K1inv @ (kpts1-K1[None,:2,2]).T).T
+ 
+        for _E in np.split(E, len(E) / 3):
+            n, R, t, _ = cv2.recoverPose(_E, kpts0_n, kpts1_n, np.eye(3), 1e9, mask=mask)
+            if n > best_num_inliers:
+                best_num_inliers = n
+                ret = (R, t, mask.ravel() > 0)
+    return ret
+
+def unnormalize_coords(x_n,h,w):
+    x = torch.stack(
+        (w * (x_n[..., 0] + 1) / 2, h * (x_n[..., 1] + 1) / 2), dim=-1
+    )  # [-1+1/h, 1-1/h] -> [0.5, h-0.5]
+    return x
+
+
+def rotate_intrinsic(K, n):
+    base_rot = np.array([[0, 1, 0], [-1, 0, 0], [0, 0, 1]])
+    rot = np.linalg.matrix_power(base_rot, n)
+    return rot @ K
+
+
+def rotate_pose_inplane(i_T_w, rot):
+    rotation_matrices = [
+        np.array(
+            [
+                [np.cos(r), -np.sin(r), 0.0, 0.0],
+                [np.sin(r), np.cos(r), 0.0, 0.0],
+                [0.0, 0.0, 1.0, 0.0],
+                [0.0, 0.0, 0.0, 1.0],
+            ],
+            dtype=np.float32,
+        )
+        for r in [np.deg2rad(d) for d in (0, 270, 180, 90)]
+    ]
+    return np.dot(rotation_matrices[rot], i_T_w)
+
+
+def scale_intrinsics(K, scales):
+    scales = np.diag([1.0 / scales[0], 1.0 / scales[1], 1.0])
+    return np.dot(scales, K)
+
+def angle_error_mat(R1, R2):
+    cos = (np.trace(np.dot(R1.T, R2)) - 1) / 2
+    cos = np.clip(cos, -1.0, 1.0)  # numercial errors can make it out of bounds
+    return np.rad2deg(np.abs(np.arccos(cos)))
+
+
+def angle_error_vec(v1, v2):
+    n = np.linalg.norm(v1) * np.linalg.norm(v2)
+    return np.rad2deg(np.arccos(np.clip(np.dot(v1, v2) / n, -1.0, 1.0)))
+
+
+def compute_pose_error(T_0to1, R, t):
+    R_gt = T_0to1[:3, :3]
+    t_gt = T_0to1[:3, 3]
+    error_t = angle_error_vec(t.squeeze(), t_gt)
+    error_t = np.minimum(error_t, 180 - error_t)  # ambiguity of E estimation
+    error_R = angle_error_mat(R, R_gt)
+    return error_t, error_R
+
+
+def pose_auc(errors, thresholds):
+    sort_idx = np.argsort(errors)
+    errors = np.array(errors.copy())[sort_idx]
+    recall = (np.arange(len(errors)) + 1) / len(errors)
+    errors = np.r_[0.0, errors]
+    recall = np.r_[0.0, recall]
+    aucs = []
+    for t in thresholds:
+        last_index = np.searchsorted(errors, t)
+        r = np.r_[recall[:last_index], recall[last_index - 1]]
+        e = np.r_[errors[:last_index], t]
+        aucs.append(np.trapz(r, x=e) / t)
+    return aucs
+
+
+# From Patch2Pix https://github.com/GrumpyZhou/patch2pix
+def get_depth_tuple_transform_ops(resize=None, normalize=True, unscale=False):
+    ops = []
+    if resize:
+        ops.append(TupleResize(resize, mode=InterpolationMode.BILINEAR, antialias = False))
+    return TupleCompose(ops)
+
+
+def get_tuple_transform_ops(resize=None, normalize=True, unscale=False, clahe = False):
+    ops = []
+    if resize:
+        ops.append(TupleResize(resize, antialias = True))
+    if clahe:
+        ops.append(TupleClahe())
+    if normalize:
+        ops.append(TupleToTensorScaled())
+        ops.append(
+            TupleNormalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
+        )  # Imagenet mean/std
+    else:
+        if unscale:
+            ops.append(TupleToTensorUnscaled())
+        else:
+            ops.append(TupleToTensorScaled())
+    return TupleCompose(ops)
+
+class Clahe:
+    def __init__(self, cliplimit = 2, blocksize = 8) -> None:
+        self.clahe = cv2.createCLAHE(cliplimit,(blocksize,blocksize))
+    def __call__(self, im):
+        im_hsv = cv2.cvtColor(np.array(im),cv2.COLOR_RGB2HSV)
+        im_v = self.clahe.apply(im_hsv[:,:,2])
+        im_hsv[...,2] = im_v
+        im_clahe = cv2.cvtColor(im_hsv,cv2.COLOR_HSV2RGB)
+        return Image.fromarray(im_clahe)
+
+class TupleClahe:
+    def __init__(self, cliplimit = 8, blocksize = 8) -> None:
+        self.clahe = Clahe(cliplimit,blocksize)
+    def __call__(self, ims):
+        return [self.clahe(im) for im in ims]
+
+class ToTensorScaled(object):
+    """Convert a RGB PIL Image to a CHW ordered Tensor, scale the range to [0, 1]"""
+
+    def __call__(self, im):
+        if not isinstance(im, torch.Tensor):
+            im = np.array(im, dtype=np.float32).transpose((2, 0, 1))
+            im /= 255.0
+            return torch.from_numpy(im)
+        else:
+            return im
+
+    def __repr__(self):
+        return "ToTensorScaled(./255)"
+
+
+class TupleToTensorScaled(object):
+    def __init__(self):
+        self.to_tensor = ToTensorScaled()
+
+    def __call__(self, im_tuple):
+        return [self.to_tensor(im) for im in im_tuple]
+
+    def __repr__(self):
+        return "TupleToTensorScaled(./255)"
+
+
+class ToTensorUnscaled(object):
+    """Convert a RGB PIL Image to a CHW ordered Tensor"""
+
+    def __call__(self, im):
+        return torch.from_numpy(np.array(im, dtype=np.float32).transpose((2, 0, 1)))
+
+    def __repr__(self):
+        return "ToTensorUnscaled()"
+
+
+class TupleToTensorUnscaled(object):
+    """Convert a RGB PIL Image to a CHW ordered Tensor"""
+
+    def __init__(self):
+        self.to_tensor = ToTensorUnscaled()
+
+    def __call__(self, im_tuple):
+        return [self.to_tensor(im) for im in im_tuple]
+
+    def __repr__(self):
+        return "TupleToTensorUnscaled()"
+
+
+class TupleResize(object):
+    def __init__(self, size, mode=InterpolationMode.BICUBIC, antialias = None):
+        self.size = size
+        self.resize = transforms.Resize(size, mode, antialias = antialias)
+
+    def __call__(self, im_tuple):
+        return [self.resize(im) for im in im_tuple]
+
+    def __repr__(self):
+        return "TupleResize(size={})".format(self.size)
+
+class Normalize:
+    def __call__(self,im):
+        mean = im.mean(dim=(1,2), keepdims=True)
+        std = im.std(dim=(1,2), keepdims=True)
+        return (im-mean)/std
+
+
+class TupleNormalize(object):
+    def __init__(self, mean, std):
+        self.mean = mean
+        self.std = std
+        self.normalize = transforms.Normalize(mean=mean, std=std)
+
+    def __call__(self, im_tuple):
+        c,h,w = im_tuple[0].shape
+        if c > 3:
+            warnings.warn(f"Number of channels {c=} > 3, assuming first 3 are rgb")
+        return [self.normalize(im[:3]) for im in im_tuple]
+
+    def __repr__(self):
+        return "TupleNormalize(mean={}, std={})".format(self.mean, self.std)
+
+
+class TupleCompose(object):
+    def __init__(self, transforms):
+        self.transforms = transforms
+
+    def __call__(self, im_tuple):
+        for t in self.transforms:
+            im_tuple = t(im_tuple)
+        return im_tuple
+
+    def __repr__(self):
+        format_string = self.__class__.__name__ + "("
+        for t in self.transforms:
+            format_string += "\n"
+            format_string += "    {0}".format(t)
+        format_string += "\n)"
+        return format_string
+
+
+@torch.no_grad()
+def warp_kpts(kpts0, depth0, depth1, T_0to1, K0, K1, smooth_mask = False, return_relative_depth_error = False, depth_interpolation_mode = "bilinear", relative_depth_error_threshold = 0.05):
+    """Warp kpts0 from I0 to I1 with depth, K and Rt
+    Also check covisibility and depth consistency.
+    Depth is consistent if relative error < 0.2 (hard-coded).
+    # https://github.com/zju3dv/LoFTR/blob/94e98b695be18acb43d5d3250f52226a8e36f839/src/loftr/utils/geometry.py adapted from here
+    Args:
+        kpts0 (torch.Tensor): [N, L, 2] - <x, y>, should be normalized in (-1,1)
+        depth0 (torch.Tensor): [N, H, W],
+        depth1 (torch.Tensor): [N, H, W],
+        T_0to1 (torch.Tensor): [N, 3, 4],
+        K0 (torch.Tensor): [N, 3, 3],
+        K1 (torch.Tensor): [N, 3, 3],
+    Returns:
+        calculable_mask (torch.Tensor): [N, L]
+        warped_keypoints0 (torch.Tensor): [N, L, 2] <x0_hat, y1_hat>
+    """
+    (
+        n,
+        h,
+        w,
+    ) = depth0.shape
+    if depth_interpolation_mode == "combined":
+        # Inspired by approach in inloc, try to fill holes from bilinear interpolation by nearest neighbour interpolation
+        if smooth_mask:
+            raise NotImplementedError("Combined bilinear and NN warp not implemented")
+        valid_bilinear, warp_bilinear = warp_kpts(kpts0, depth0, depth1, T_0to1, K0, K1, 
+                  smooth_mask = smooth_mask, 
+                  return_relative_depth_error = return_relative_depth_error, 
+                  depth_interpolation_mode = "bilinear",
+                  relative_depth_error_threshold = relative_depth_error_threshold)
+        valid_nearest, warp_nearest = warp_kpts(kpts0, depth0, depth1, T_0to1, K0, K1, 
+                  smooth_mask = smooth_mask, 
+                  return_relative_depth_error = return_relative_depth_error, 
+                  depth_interpolation_mode = "nearest-exact",
+                  relative_depth_error_threshold = relative_depth_error_threshold)
+        nearest_valid_bilinear_invalid = (~valid_bilinear).logical_and(valid_nearest) 
+        warp = warp_bilinear.clone()
+        warp[nearest_valid_bilinear_invalid] = warp_nearest[nearest_valid_bilinear_invalid]
+        valid = valid_bilinear | valid_nearest
+        return valid, warp
+        
+        
+    kpts0_depth = F.grid_sample(depth0[:, None], kpts0[:, :, None], mode = depth_interpolation_mode, align_corners=False)[
+        :, 0, :, 0
+    ]
+    kpts0 = torch.stack(
+        (w * (kpts0[..., 0] + 1) / 2, h * (kpts0[..., 1] + 1) / 2), dim=-1
+    )  # [-1+1/h, 1-1/h] -> [0.5, h-0.5]
+    # Sample depth, get calculable_mask on depth != 0
+    nonzero_mask = kpts0_depth != 0
+
+    # Unproject
+    kpts0_h = (
+        torch.cat([kpts0, torch.ones_like(kpts0[:, :, [0]])], dim=-1)
+        * kpts0_depth[..., None]
+    )  # (N, L, 3)
+    kpts0_n = K0.inverse() @ kpts0_h.transpose(2, 1)  # (N, 3, L)
+    kpts0_cam = kpts0_n
+
+    # Rigid Transform
+    w_kpts0_cam = T_0to1[:, :3, :3] @ kpts0_cam + T_0to1[:, :3, [3]]  # (N, 3, L)
+    w_kpts0_depth_computed = w_kpts0_cam[:, 2, :]
+
+    # Project
+    w_kpts0_h = (K1 @ w_kpts0_cam).transpose(2, 1)  # (N, L, 3)
+    w_kpts0 = w_kpts0_h[:, :, :2] / (
+        w_kpts0_h[:, :, [2]] + 1e-4
+    )  # (N, L, 2), +1e-4 to avoid zero depth
+
+    # Covisible Check
+    h, w = depth1.shape[1:3]
+    covisible_mask = (
+        (w_kpts0[:, :, 0] > 0)
+        * (w_kpts0[:, :, 0] < w - 1)
+        * (w_kpts0[:, :, 1] > 0)
+        * (w_kpts0[:, :, 1] < h - 1)
+    )
+    w_kpts0 = torch.stack(
+        (2 * w_kpts0[..., 0] / w - 1, 2 * w_kpts0[..., 1] / h - 1), dim=-1
+    )  # from [0.5,h-0.5] -> [-1+1/h, 1-1/h]
+    # w_kpts0[~covisible_mask, :] = -5 # xd
+
+    w_kpts0_depth = F.grid_sample(
+        depth1[:, None], w_kpts0[:, :, None], mode=depth_interpolation_mode, align_corners=False
+    )[:, 0, :, 0]
+    
+    relative_depth_error = (
+        (w_kpts0_depth - w_kpts0_depth_computed) / w_kpts0_depth
+    ).abs()
+    if not smooth_mask:
+        consistent_mask = relative_depth_error < relative_depth_error_threshold
+    else:
+        consistent_mask = (-relative_depth_error/smooth_mask).exp()
+    valid_mask = nonzero_mask * covisible_mask * consistent_mask
+    if return_relative_depth_error:
+        return relative_depth_error, w_kpts0
+    else:
+        return valid_mask, w_kpts0
+
+imagenet_mean = torch.tensor([0.485, 0.456, 0.406])
+imagenet_std = torch.tensor([0.229, 0.224, 0.225])
+
+
+def numpy_to_pil(x: np.ndarray):
+    """
+    Args:
+        x: Assumed to be of shape (h,w,c)
+    """
+    if isinstance(x, torch.Tensor):
+        x = x.detach().cpu().numpy()
+    if x.max() <= 1.01:
+        x *= 255
+    x = x.astype(np.uint8)
+    return Image.fromarray(x)
+
+
+def tensor_to_pil(x, unnormalize=False, autoscale = False):
+    if unnormalize:
+        x = x * (imagenet_std[:, None, None].to(x.device)) + (imagenet_mean[:, None, None].to(x.device))
+    if autoscale:
+        if x.max() == x.min():
+            warnings.warn("x max == x min, cant autoscale")
+        else:
+            x = (x-x.min())/(x.max()-x.min())
+        
+    x = x.detach().permute(1, 2, 0).cpu().numpy()
+    x = np.clip(x, 0.0, 1.0)
+    return numpy_to_pil(x)
+
+
+def to_cuda(batch):
+    for key, value in batch.items():
+        if isinstance(value, torch.Tensor):
+            batch[key] = value.cuda()
+    return batch
+
+
+def to_cpu(batch):
+    for key, value in batch.items():
+        if isinstance(value, torch.Tensor):
+            batch[key] = value.cpu()
+    return batch
+
+
+def get_pose(calib):
+    w, h = np.array(calib["imsize"])[0]
+    return np.array(calib["K"]), np.array(calib["R"]), np.array(calib["T"]).T, h, w
+
+
+def compute_relative_pose(R1, t1, R2, t2):
+    rots = R2 @ (R1.T)
+    trans = -rots @ t1 + t2
+    return rots, trans
+
+def to_pixel_coords(flow, h1, w1):
+    flow = (
+        torch.stack(
+            (
+                w1 * (flow[..., 0] + 1) / 2,
+                h1 * (flow[..., 1] + 1) / 2,
+            ),
+            axis=-1,
+        )
+    )
+    return flow
+
+def to_normalized_coords(flow, h1, w1):
+    flow = (
+        torch.stack(
+            (
+                2 * (flow[..., 0]) / w1 - 1,
+                2 * (flow[..., 1]) / h1 - 1,
+            ),
+            axis=-1,
+        )
+    )
+    return flow
+
+
+def warp_to_pixel_coords(warp, h1, w1, h2, w2):
+    warp1 = warp[..., :2]
+    warp1 = (
+        torch.stack(
+            (
+                w1 * (warp1[..., 0] + 1) / 2,
+                h1 * (warp1[..., 1] + 1) / 2,
+            ),
+            axis=-1,
+        )
+    )
+    warp2 = warp[..., 2:]
+    warp2 = (
+        torch.stack(
+            (
+                w2 * (warp2[..., 0] + 1) / 2,
+                h2 * (warp2[..., 1] + 1) / 2,
+            ),
+            axis=-1,
+        )
+    )
+    return torch.cat((warp1,warp2), dim=-1)
+
+
+def to_homogeneous(x):
+    ones = torch.ones_like(x[...,-1:])
+    return torch.cat((x, ones), dim = -1)
+
+def from_homogeneous(xh, eps = 1e-12):
+    return xh[...,:-1] / (xh[...,-1:]+eps)
+
+def homog_transform(Homog, x):
+    xh = to_homogeneous(x)
+    yh = (Homog @ xh.mT).mT
+    y = from_homogeneous(yh)
+    return y
+
+def get_homog_warp(Homog, H, W, device = "cuda"):
+    grid = torch.meshgrid(torch.linspace(-1+1/H,1-1/H,H, device = device), torch.linspace(-1+1/W,1-1/W,W, device = device))
+    
+    x_A = torch.stack((grid[1], grid[0]), dim = -1)[None]
+    x_A_to_B = homog_transform(Homog, x_A)
+    mask = ((x_A_to_B > -1) * (x_A_to_B < 1)).prod(dim=-1).float()
+    return torch.cat((x_A.expand(*x_A_to_B.shape), x_A_to_B),dim=-1), mask
+
+def dual_log_softmax_matcher(desc_A: tuple['B','N','C'], desc_B: tuple['B','M','C'], inv_temperature = 1, normalize = False):
+    B, N, C = desc_A.shape
+    if normalize:
+        desc_A = desc_A/desc_A.norm(dim=-1,keepdim=True)
+        desc_B = desc_B/desc_B.norm(dim=-1,keepdim=True)
+        corr = torch.einsum("b n c, b m c -> b n m", desc_A, desc_B) * inv_temperature
+    else:
+        corr = torch.einsum("b n c, b m c -> b n m", desc_A, desc_B) * inv_temperature
+    logP = corr.log_softmax(dim = -2) + corr.log_softmax(dim= -1)
+    return logP
+
+def dual_softmax_matcher(desc_A: tuple['B','N','C'], desc_B: tuple['B','M','C'], inv_temperature = 1, normalize = False):
+    if len(desc_A.shape) < 3:
+        desc_A, desc_B = desc_A[None], desc_B[None]
+    B, N, C = desc_A.shape
+    if normalize:
+        desc_A = desc_A/desc_A.norm(dim=-1,keepdim=True)
+        desc_B = desc_B/desc_B.norm(dim=-1,keepdim=True)
+        corr = torch.einsum("b n c, b m c -> b n m", desc_A, desc_B) * inv_temperature
+    else:
+        corr = torch.einsum("b n c, b m c -> b n m", desc_A, desc_B) * inv_temperature
+    P = corr.softmax(dim = -2) * corr.softmax(dim= -1)
+    return P
+
+def conditional_softmax_matcher(desc_A: tuple['B','N','C'], desc_B: tuple['B','M','C'], inv_temperature = 1, normalize = False):
+    if len(desc_A.shape) < 3:
+        desc_A, desc_B = desc_A[None], desc_B[None]
+    B, N, C = desc_A.shape
+    if normalize:
+        desc_A = desc_A/desc_A.norm(dim=-1,keepdim=True)
+        desc_B = desc_B/desc_B.norm(dim=-1,keepdim=True)
+        corr = torch.einsum("b n c, b m c -> b n m", desc_A, desc_B) * inv_temperature
+    else:
+        corr = torch.einsum("b n c, b m c -> b n m", desc_A, desc_B) * inv_temperature
+    P_B_cond_A = corr.softmax(dim = -1)
+    P_A_cond_B = corr.softmax(dim = -2)
+    
+    return P_A_cond_B, P_B_cond_A 

third_party/DeDoDe/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2023 Johan Edstedt
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

third_party/DeDoDe/README.md

@@ -0,0 +1,74 @@
+<p align="center">
+  <h1 align="center"><ins>DeDoDe</ins> 🎶<br>Detect, Don't Describe, Describe, Don't Detect, <br> for Local Feature Matching</h1>
+  <p align="center">
+    <a href="https://scholar.google.com/citations?user=Ul-vMR0AAAAJ">Johan Edstedt</a>
+    ·
+    <a href="https://scholar.google.com/citations?user=FUE3Wd0AAAAJ">Georg Bökman</a>
+    ·
+    <a href="https://scholar.google.com/citations?user=6WRQpCQAAAAJ">Mårten Wadenbäck</a>
+    ·
+    <a href="https://scholar.google.com/citations?user=lkWfR08AAAAJ">Michael Felsberg</a>
+    ·
+  </p>
+  <h2 align="center"><p>
+    <a href="TODO" align="center">Paper (TODO)</a> | 
+    <a href="TODO" align="center">Project Page (TODO)</a>
+  </p></h2>
+  <div align="center"></div>
+</p>
+<p align="center">
+    <img src="assets/matches.jpg" alt="example" width=80%>
+    <br>
+    <em>The DeDoDe detector learns to detect 3D consistent repeatable keypoints, which the DeDoDe descriptor learns to match. The result is a powerful decoupled local feature matcher.</em>
+    <br>
+        <img src="assets/teaser.png" alt="example" width=40%>
+        <img src="assets/dedode_roma.png" alt="example" width=40%>
+    <br>
+    <em>
+    We experimentally find that DeDoDe significantly closes the performance gap between detector + descriptor models and fully-fledged matchers. The potential of DeDoDe is not limited to local feature matching, in fact we find that we can improve state-of-the-art matchers by incorporating DeDoDe keypoints.
+    </em>
+</p>
+
+## How to Use DeDoDe?
+Below we show how DeDoDe can be run, you can also check out the [demos](demo)
+```python
+from DeDoDe import dedode_detector_L, dedode_descriptor_B
+from DeDoDe.matchers.dual_softmax_matcher import DualSoftMaxMatcher
+
+detector = dedode_detector_L(weights = torch.load("dedode_detector_L.pth"))
+descriptor = dedode_descriptor_B(weights = torch.load("dedode_descriptor_B.pth"))
+matcher = DualSoftMaxMatcher()
+
+im_A_path = "assets/im_A.jpg"
+im_B_path = "assets/im_B.jpg"
+im_A = Image.open(im_A_path)
+im_B = Image.open(im_B_path)
+W_A, H_A = im_A.size
+W_B, H_B = im_B.size
+
+
+detections_A = detector.detect_from_path(im_A_path, num_keypoints = 10_000)
+keypoints_A, P_A = detections_A["keypoints"], detections_A["confidence"]
+
+detections_B = detector.detect_from_path(im_B_path, num_keypoints = 10_000)
+keypoints_B, P_B = detections_B["keypoints"], detections_B["confidence"]
+
+description_A = descriptor.describe_keypoints_from_path(im_A_path, keypoints_A)["descriptions"]
+description_B = descriptor.describe_keypoints_from_path(im_B_path, keypoints_B)["descriptions"]
+
+matches_A, matches_B, batch_ids = matcher.match(keypoints_A, description_A,
+    keypoints_B, description_B,
+    P_A = P_A, P_B = P_B,
+    normalize = True, inv_temp=20, threshold = 0.1)#Increasing threshold -> fewer matches, fewer outliers
+
+matches_A, matches_B = matcher.to_pixel_coords(matches_A, matches_B, H_A, W_A, H_B, W_B)
+
+```
+## Pretrained Models
+
+Right now you can find them here: https://github.com/Parskatt/DeDoDe/releases/tag/dedode_pretrained_models
+Probably we'll add some autoloading in the near future.
+
+## BibTeX
+
+Coming Soon ;)

third_party/DeDoDe/data_prep/prep_keypoints.py

@@ -0,0 +1,100 @@
+import argparse
+
+import imagesize
+
+import numpy as np
+
+import os
+
+
+base_path = "data/megadepth"
+# Remove the trailing / if need be.
+if base_path[-1] in ['/', '\\']:
+    base_path = base_path[: - 1]
+
+
+base_depth_path = os.path.join(
+    base_path, 'phoenix/S6/zl548/MegaDepth_v1'
+)
+base_undistorted_sfm_path = os.path.join(
+    base_path, 'Undistorted_SfM'
+)
+
+scene_ids = os.listdir(base_undistorted_sfm_path)
+for scene_id in scene_ids:
+    if os.path.exists(f"{base_path}/prep_scene_info/detections/detections_{scene_id}.npy"):
+        print(f"skipping {scene_id} as it exists")
+        continue
+    undistorted_sparse_path = os.path.join(
+        base_undistorted_sfm_path, scene_id, 'sparse-txt'
+    )
+    if not os.path.exists(undistorted_sparse_path):
+        print("sparse path doesnt exist")
+        continue
+
+    depths_path = os.path.join(
+        base_depth_path, scene_id, 'dense0', 'depths'
+    )
+    if not os.path.exists(depths_path):
+        print("depths doesnt exist")
+        
+        continue
+
+    images_path = os.path.join(
+        base_undistorted_sfm_path, scene_id, 'images'
+    )
+    if not os.path.exists(images_path):
+        print("images path doesnt exist")
+        continue
+
+    # Process cameras.txt
+    if not os.path.exists(os.path.join(undistorted_sparse_path, 'cameras.txt')):
+        print("no cameras")
+        continue
+    with open(os.path.join(undistorted_sparse_path, 'cameras.txt'), 'r') as f:
+        raw = f.readlines()[3 :]  # skip the header
+
+    camera_intrinsics = {}
+    for camera in raw:
+        camera = camera.split(' ')
+        camera_intrinsics[int(camera[0])] = [float(elem) for elem in camera[2 :]]
+
+    # Process points3D.txt
+    with open(os.path.join(undistorted_sparse_path, 'points3D.txt'), 'r') as f:
+        raw = f.readlines()[3 :]  # skip the header
+
+    points3D = {}
+    for point3D in raw:
+        point3D = point3D.split(' ')
+        points3D[int(point3D[0])] = np.array([
+            float(point3D[1]), float(point3D[2]), float(point3D[3])
+        ])
+        
+    # Process images.txt
+    with open(os.path.join(undistorted_sparse_path, 'images.txt'), 'r') as f:
+        raw = f.readlines()[4 :]  # skip the header
+
+    image_id_to_idx = {}
+    image_names = []
+    raw_pose = []
+    camera = []
+    points3D_id_to_2D = []
+    n_points3D = []
+    id_to_detections = {}
+    for idx, (image, points) in enumerate(zip(raw[:: 2], raw[1 :: 2])):
+        image = image.split(' ')
+        points = points.split(' ')
+
+        image_id_to_idx[int(image[0])] = idx
+
+        image_name = image[-1].strip('\n')
+        image_names.append(image_name)
+
+        raw_pose.append([float(elem) for elem in image[1 : -2]])
+        camera.append(int(image[-2]))
+        points_np = np.array(points).astype(np.float32).reshape(len(points)//3, 3)
+        visible_points = points_np[points_np[:,2] != -1]
+        id_to_detections[idx] = visible_points
+    np.save(f"{base_path}/prep_scene_info/detections/detections_{scene_id}.npy",
+            id_to_detections)
+    print(f"{scene_id} done")

third_party/DeDoDe/demo/demo_kpts.py

@@ -0,0 +1,20 @@
+import torch
+import cv2
+import numpy as np
+from PIL import Image
+from DeDoDe import dedode_detector_L
+
+def draw_kpts(im, kpts):    
+    kpts = [cv2.KeyPoint(x,y,1.) for x,y in kpts.cpu().numpy()]
+    im = np.array(im)
+    ret = cv2.drawKeypoints(im, kpts, None)
+    return ret
+
+detector = dedode_detector_L(weights = torch.load("dedode_detector_l.pth"))
+im_path = "assets/im_A.jpg"
+im = Image.open(im_path)
+out = detector.detect_from_path(im_path, num_keypoints = 10_000)
+W,H = im.size
+kps = out["keypoints"]
+kps = detector.to_pixel_coords(kps, H, W)
+Image.fromarray(draw_kpts(im, kps[0])).save("demo/keypoints.png")

third_party/DeDoDe/demo/demo_match.py

@@ -0,0 +1,45 @@
+import torch
+from DeDoDe import dedode_detector_L, dedode_descriptor_B
+from DeDoDe.matchers.dual_softmax_matcher import DualSoftMaxMatcher
+from DeDoDe.utils import *
+from PIL import Image
+import cv2
+
+def draw_matches(im_A, kpts_A, im_B, kpts_B):    
+    kpts_A = [cv2.KeyPoint(x,y,1.) for x,y in kpts_A.cpu().numpy()]
+    kpts_B = [cv2.KeyPoint(x,y,1.) for x,y in kpts_B.cpu().numpy()]
+    matches_A_to_B = [cv2.DMatch(idx, idx, 0.) for idx in range(len(kpts_A))]
+    im_A, im_B = np.array(im_A), np.array(im_B)
+    ret = cv2.drawMatches(im_A, kpts_A, im_B, kpts_B, 
+                    matches_A_to_B, None)
+    return ret
+
+detector = dedode_detector_L(weights = torch.load("dedode_detector_L.pth"))
+descriptor = dedode_descriptor_B(weights = torch.load("dedode_descriptor_B.pth"))
+matcher = DualSoftMaxMatcher()
+
+im_A_path = "assets/im_A.jpg"
+im_B_path = "assets/im_B.jpg"
+im_A = Image.open(im_A_path)
+im_B = Image.open(im_B_path)
+W_A, H_A = im_A.size
+W_B, H_B = im_B.size
+
+
+detections_A = detector.detect_from_path(im_A_path, num_keypoints = 10_000)
+keypoints_A, P_A = detections_A["keypoints"], detections_A["confidence"]
+detections_B = detector.detect_from_path(im_B_path, num_keypoints = 10_000)
+keypoints_B, P_B = detections_B["keypoints"], detections_B["confidence"]
+description_A = descriptor.describe_keypoints_from_path(im_A_path, keypoints_A)["descriptions"]
+description_B = descriptor.describe_keypoints_from_path(im_B_path, keypoints_B)["descriptions"]
+matches_A, matches_B, batch_ids = matcher.match(keypoints_A, description_A,
+    keypoints_B, description_B,
+    P_A = P_A, P_B = P_B,
+    normalize = True, inv_temp=20, threshold = 0.1)#Increasing threshold -> fewer matches, fewer outliers
+
+matches_A, matches_B = matcher.to_pixel_coords(matches_A, matches_B, H_A, W_A, H_B, W_B)
+
+import cv2
+import numpy as np
+
+Image.fromarray(draw_matches(im_A, matches_A[::5], im_B, matches_B[::5])).save("demo/matches.png")

third_party/DeDoDe/demo/demo_scoremap.py

@@ -0,0 +1,20 @@
+import torch
+from PIL import Image
+import numpy as np
+
+from DeDoDe import dedode_detector_L
+from DeDoDe.utils import tensor_to_pil
+
+detector = dedode_detector_L(weights = torch.load("dedode_detector_l.pth"))
+H, W = 768, 768
+im_path = "assets/im_A.jpg"
+
+out = detector.detect_from_path(im_path, dense = True, H = H, W = W)
+
+logit_map = out["dense_keypoint_logits"].clone()
+min = logit_map.max() - 3
+logit_map[logit_map < min] = min
+logit_map = (logit_map-min)/(logit_map.max()-min)
+logit_map = logit_map.cpu()[0].expand(3,H,W)
+im_A = torch.tensor(np.array(Image.open(im_path).resize((W,H)))/255.).permute(2,0,1)
+tensor_to_pil(logit_map * logit_map  +  0.15 * (1-logit_map) * im_A).save("demo/dense_logits.png")

third_party/DeDoDe/pretrained/dedode_descriptor_B.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:8eccfc270ec990ced60cd54a434411a47c4c504de13586f596d042e005b3022b
+size 54257185

third_party/DeDoDe/pretrained/dedode_detector_L.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:663c972a7215538ef3170ccc3183bb1019610ffde4bc7c9da6c13b143388dd64
+size 58488277

third_party/DeDoDe/requirements.txt

@@ -0,0 +1,9 @@
+numpy
+matplotlib
+torch
+torchvision
+h5py
+tqdm
+pillow
+einops
+opencv-python

third_party/DeDoDe/setup.py

@@ -0,0 +1,10 @@
+from setuptools import setup, find_packages
+
+
+setup(
+    name="DeDoDe",
+    packages=find_packages(include= ["DeDoDe*"]),
+    install_requires=open("requirements.txt", "r").read().split("\n"),
+    version="0.0.1",
+    author="Johan Edstedt",
+)

third_party/SuperGluePretrainedNetwork/.gitignore

@@ -0,0 +1,3 @@
+*.pyc
+*.DS_Store
+*.swp

third_party/SuperGluePretrainedNetwork/LICENSE

@@ -0,0 +1,48 @@
+SUPERGLUE: LEARNING FEATURE MATCHING WITH GRAPH NEURAL NETWORKS
+SOFTWARE LICENSE AGREEMENT
+ACADEMIC OR NON-PROFIT ORGANIZATION NONCOMMERCIAL RESEARCH USE ONLY
+
+BY USING OR DOWNLOADING THE SOFTWARE, YOU ARE AGREEING TO THE TERMS OF THIS LICENSE AGREEMENT.  IF YOU DO NOT AGREE WITH THESE TERMS, YOU MAY NOT USE OR DOWNLOAD THE SOFTWARE.
+
+This is a license agreement ("Agreement") between your academic institution or non-profit organization or self (called "Licensee" or "You" in this Agreement) and Magic Leap, Inc. (called "Licensor" in this Agreement).  All rights not specifically granted to you in this Agreement are reserved for Licensor. 
+
+RESERVATION OF OWNERSHIP AND GRANT OF LICENSE: 
+Licensor retains exclusive ownership of any copy of the Software (as defined below) licensed under this Agreement and hereby grants to Licensee a personal, non-exclusive, non-transferable license to use the Software for noncommercial research purposes, without the right to sublicense, pursuant to the terms and conditions of this Agreement.  As used in this Agreement, the term "Software" means (i) the actual copy of all or any portion of code for program routines made accessible to Licensee by Licensor pursuant to this Agreement, inclusive of backups, updates, and/or merged copies permitted hereunder or subsequently supplied by Licensor,  including all or any file structures, programming instructions, user interfaces and screen formats and sequences as well as any and all documentation and instructions related to it, and (ii) all or any derivatives and/or modifications created or made by You to any of the items specified in (i).
+
+CONFIDENTIALITY: Licensee acknowledges that the Software is proprietary to Licensor, and as such, Licensee agrees to receive all such materials in confidence and use the Software only in accordance with the terms of this Agreement.  Licensee agrees to use reasonable effort to protect the Software from unauthorized use, reproduction, distribution, or publication.
+
+COPYRIGHT: The Software is owned by Licensor and is protected by United  States copyright laws and applicable international treaties and/or conventions.
+
+PERMITTED USES:  The Software may be used for your own noncommercial internal research purposes. You understand and agree that Licensor is not obligated to implement any suggestions and/or feedback you might provide regarding the Software, but to the extent Licensor does so, you are not entitled to any compensation related thereto.
+
+DERIVATIVES: You may create derivatives of or make modifications to the Software, however, You agree that all and any such derivatives and modifications will be owned by Licensor and become a part of the Software licensed to You under this Agreement.  You may only use such derivatives and modifications for your own noncommercial internal research purposes, and you may not otherwise use, distribute or copy such derivatives and modifications in violation of this Agreement.
+
+BACKUPS:  If Licensee is an organization, it may make that number of copies of the Software necessary for internal noncommercial use at a single site within its organization provided that all information appearing in or on the original labels, including the copyright and trademark notices are copied onto the labels of the copies.
+
+USES NOT PERMITTED:  You may not distribute, copy or use the Software except as explicitly permitted herein. Licensee has not been granted any trademark license as part of this Agreement and may not use the name or mark "Magic Leap" or any renditions thereof without the prior written permission of Licensor.
+
+You may not sell, rent, lease, sublicense, lend, time-share or transfer, in whole or in part, or provide third parties access to prior or present versions (or any parts thereof) of the Software.
+
+ASSIGNMENT: You may not assign this Agreement or your rights hereunder without the prior written consent of Licensor. Any attempted assignment without such consent shall be null and void.
+
+TERM: The term of the license granted by this Agreement is from Licensee's acceptance of this Agreement by downloading the Software or by using the Software until terminated as provided below.
+
+The Agreement automatically terminates without notice if you fail to comply with any provision of this Agreement.  Licensee may terminate this Agreement by ceasing using the Software.  Upon any termination of this Agreement, Licensee will delete any and all copies of the Software. You agree that all provisions which operate to protect the proprietary rights of Licensor shall remain in force should breach occur and that the obligation of confidentiality described in this Agreement is binding in perpetuity and, as such, survives the term of the Agreement.
+
+FEE: Provided Licensee abides completely by the terms and conditions of this Agreement, there is no fee due to Licensor for Licensee's use of the Software in accordance with this Agreement.
+
+DISCLAIMER OF WARRANTIES:  THE SOFTWARE IS PROVIDED "AS-IS" WITHOUT WARRANTY OF ANY KIND INCLUDING ANY WARRANTIES OF PERFORMANCE OR MERCHANTABILITY OR FITNESS FOR A PARTICULAR USE OR PURPOSE OR OF NON-INFRINGEMENT.  LICENSEE BEARS ALL RISK RELATING TO QUALITY AND PERFORMANCE OF THE SOFTWARE AND RELATED MATERIALS.
+
+SUPPORT AND MAINTENANCE: No Software support or training by the Licensor is provided as part of this Agreement.  
+
+EXCLUSIVE REMEDY AND LIMITATION OF LIABILITY: To the maximum extent permitted under applicable law, Licensor shall not be liable for direct, indirect, special, incidental, or consequential damages or lost profits related to Licensee's use of and/or inability to use the Software, even if Licensor is advised of the possibility of such damage.
+
+EXPORT REGULATION: Licensee agrees to comply with any and all applicable U.S. export control laws, regulations, and/or other laws related to embargoes and sanction programs administered by the Office of Foreign Assets Control.
+
+SEVERABILITY: If any provision(s) of this Agreement shall be held to be invalid, illegal, or unenforceable by a court or other tribunal of competent jurisdiction, the validity, legality and enforceability of the remaining provisions shall not in any way be affected or impaired thereby.
+
+NO IMPLIED WAIVERS: No failure or delay by Licensor in enforcing any right or remedy under this Agreement shall be construed as a waiver of any future or other exercise of such right or remedy by Licensor.
+
+GOVERNING LAW: This Agreement shall be construed and enforced in accordance with the laws of the State of Florida without reference to conflict of laws principles.  You consent to the personal jurisdiction of the courts of this County and waive their rights to venue outside of Broward County, Florida.
+
+ENTIRE AGREEMENT AND AMENDMENTS: This Agreement constitutes the sole and entire agreement between Licensee and Licensor as to the matter set forth herein and supersedes any previous agreements, understandings, and arrangements between the parties relating hereto.

third_party/SuperGluePretrainedNetwork/README.md

@@ -0,0 +1,388 @@
+<img src="assets/magicleap.png" width="240">
+
+### Research @ Magic Leap (CVPR 2020, Oral)
+
+# SuperGlue Inference and Evaluation Demo Script
+
+## Introduction
+SuperGlue is a CVPR 2020 research project done at Magic Leap. The SuperGlue network is a Graph Neural Network combined with an Optimal Matching layer that is trained to perform matching on two sets of sparse image features. This repo includes PyTorch code and pretrained weights for running the SuperGlue matching network on top of [SuperPoint](https://arxiv.org/abs/1712.07629) keypoints and descriptors. Given a pair of images, you can use this repo to extract matching features across the image pair.
+
+<p align="center">
+  <img src="assets/teaser.png" width="500">
+</p>
+
+SuperGlue operates as a "middle-end," performing context aggregation, matching, and filtering in a single end-to-end architecture. For more details, please see:
+
+* Full paper PDF: [SuperGlue: Learning Feature Matching with Graph Neural Networks](https://arxiv.org/abs/1911.11763).
+
+* Authors: *Paul-Edouard Sarlin, Daniel DeTone, Tomasz Malisiewicz, Andrew Rabinovich*
+
+* Website: [psarlin.com/superglue](https://psarlin.com/superglue) for videos, slides, recent updates, and more visualizations.
+
+* `hloc`: a new toolbox for visual localization and SfM with SuperGlue, available at [cvg/Hierarchical-Localization](https://github.com/cvg/Hierarchical-Localization/). Winner of 3 CVPR 2020 competitions on localization and image matching!
+
+We provide two pre-trained weights files: an indoor model trained on ScanNet data, and an outdoor model trained on MegaDepth data. Both models are inside the [weights directory](./models/weights). By default, the demo will run the **indoor** model.
+
+## Dependencies
+* Python 3 >= 3.5
+* PyTorch >= 1.1
+* OpenCV >= 3.4 (4.1.2.30 recommended for best GUI keyboard interaction, see this [note](#additional-notes))
+* Matplotlib >= 3.1
+* NumPy >= 1.18
+
+Simply run the following command: `pip3 install numpy opencv-python torch matplotlib`
+
+## Contents
+There are two main top-level scripts in this repo:
+
+1. `demo_superglue.py` : runs a live demo on a webcam, IP camera, image directory or movie file
+2. `match_pairs.py`: reads image pairs from files and dumps matches to disk (also runs evaluation if ground truth relative poses are provided)
+
+## Live Matching Demo Script (`demo_superglue.py`)
+This demo runs SuperPoint + SuperGlue feature matching on an anchor image and live image. You can update the anchor image by pressing the `n` key. The demo can read image streams from a USB or IP camera, a directory containing images, or a video file. You can pass all of these inputs using the `--input` flag.
+
+### Run the demo on a live webcam
+
+Run the demo on the default USB webcam (ID #0), running on a CUDA GPU if one is found:
+
+```sh
+./demo_superglue.py
+```
+
+Keyboard control:
+
+* `n`: select the current frame as the anchor
+* `e`/`r`: increase/decrease the keypoint confidence threshold
+* `d`/`f`: increase/decrease the match filtering threshold
+* `k`: toggle the visualization of keypoints
+* `q`: quit
+
+Run the demo on 320x240 images running on the CPU:
+
+```sh
+./demo_superglue.py --resize 320 240 --force_cpu
+```
+
+The `--resize` flag can be used to resize the input image in three ways:
+
+1. `--resize` `width` `height` : will resize to exact `width` x `height` dimensions
+2. `--resize` `max_dimension` : will resize largest input image dimension to `max_dimension`
+3. `--resize` `-1` : will not resize (i.e. use original image dimensions)
+
+The default will resize images to `640x480`.
+
+### Run the demo on a directory of images
+
+The `--input` flag also accepts a path to a directory. We provide a directory of sample images from a sequence. To run the demo on the directory of images in `freiburg_sequence/` on a headless server (will not display to the screen) and write the output visualization images to `dump_demo_sequence/`:
+
+```sh
+./demo_superglue.py --input assets/freiburg_sequence/ --output_dir dump_demo_sequence --resize 320 240 --no_display
+```
+
+You should see this output on the sample Freiburg-TUM RGBD sequence:
+
+<img src="assets/freiburg_matches.gif" width="560">
+
+The matches are colored by their predicted confidence in a jet colormap (Red: more confident, Blue: less confident).
+
+### Additional useful command line parameters
+* Use `--image_glob` to change the image file extension (default: `*.png`, `*.jpg`, `*.jpeg`).
+* Use `--skip` to skip intermediate frames (default: `1`).
+* Use `--max_length` to cap the total number of frames processed (default: `1000000`).
+* Use `--show_keypoints` to visualize the detected keypoints (default: `False`).
+
+## Run Matching+Evaluation (`match_pairs.py`)
+
+This repo also contains a script `match_pairs.py` that runs the matching from a list of image pairs. With this script, you can:
+
+* Run the matcher on a set of image pairs (no ground truth needed)
+* Visualize the keypoints and matches, based on their confidence
+* Evaluate and visualize the match correctness, if the ground truth relative poses and intrinsics are provided
+* Save the keypoints, matches, and evaluation results for further processing
+* Collate evaluation results over many pairs and generate result tables
+
+### Matches only mode
+
+The simplest usage of this script will process the image pairs listed in a given text file and dump the keypoints and matches to compressed numpy `npz` files. We provide the challenging ScanNet pairs from the main paper in `assets/example_indoor_pairs/`. Running the following will run SuperPoint + SuperGlue on each image pair, and dump the results to `dump_match_pairs/`:
+
+```sh
+./match_pairs.py
+```
+
+The resulting `.npz` files can be read from Python as follows:
+
+```python
+>>> import numpy as np
+>>> path = 'dump_match_pairs/scene0711_00_frame-001680_scene0711_00_frame-001995_matches.npz'
+>>> npz = np.load(path)
+>>> npz.files
+['keypoints0', 'keypoints1', 'matches', 'match_confidence']
+>>> npz['keypoints0'].shape
+(382, 2)
+>>> npz['keypoints1'].shape
+(391, 2)
+>>> npz['matches'].shape
+(382,)
+>>> np.sum(npz['matches']>-1)
+115
+>>> npz['match_confidence'].shape
+(382,)
+```
+
+For each keypoint in `keypoints0`, the `matches` array indicates the index of the matching keypoint in `keypoints1`, or `-1` if the keypoint is unmatched.
+
+### Visualization mode
+
+You can add the flag `--viz` to dump image outputs which visualize the matches:
+
+```sh
+./match_pairs.py --viz
+```
+
+You should see images like this inside of `dump_match_pairs/` (or something very close to it, see this [note](#a-note-on-reproducibility)):
+
+<img src="assets/indoor_matches.png" width="560">
+
+The matches are colored by their predicted confidence in a jet colormap (Red: more confident, Blue: less confident).
+
+### Evaluation mode
+
+You can also estimate the pose using RANSAC + Essential Matrix decomposition and evaluate it if the ground truth relative poses and intrinsics are provided in the input `.txt` files. Each `.txt` file contains three key ground truth matrices: a 3x3 intrinsics matrix of image0: `K0`, a 3x3 intrinsics matrix of image1: `K1` , and a 4x4 matrix of the relative pose extrinsics `T_0to1`.
+
+To run the evaluation on the sample set of images (by default reading `assets/scannet_sample_pairs_with_gt.txt`), you can run:
+
+```sh
+./match_pairs.py --eval
+```
+
+
+Since you enabled `--eval`, you should see collated results printed to the terminal. For the example images provided, you should get the following numbers (or something very close to it, see this [note](#a-note-on-reproducibility)):
+
+```txt
+Evaluation Results (mean over 15 pairs):
+AUC@5    AUC@10  AUC@20  Prec    MScore
+26.99    48.40   64.47   73.52   19.60
+```
+
+The resulting `.npz` files in `dump_match_pairs/` will now contain scalar values related to the evaluation, computed on the sample images provided. Here is what you should find in one of the generated evaluation files:
+
+```python
+>>> import numpy as np
+>>> path = 'dump_match_pairs/scene0711_00_frame-001680_scene0711_00_frame-001995_evaluation.npz'
+>>> npz = np.load(path)
+>>> print(npz.files)
+['error_t', 'error_R', 'precision', 'matching_score', 'num_correct', 'epipolar_errors']
+```
+
+You can also visualize the evaluation metrics by running the following command:
+
+```sh
+./match_pairs.py --eval --viz
+```
+
+You should also now see additional images in `dump_match_pairs/` which visualize the evaluation numbers (or something very close to it, see this [note](#a-note-on-reproducibility)):
+
+<img src="assets/indoor_evaluation.png" width="560">
+
+The top left corner of the image shows the pose error and number of inliers, while the lines are colored by their epipolar error computed with the ground truth relative pose (red: higher error, green: lower error).
+
+### Running on sample outdoor pairs
+
+<details>
+  <summary>[Click to expand]</summary>
+
+In this repo, we also provide a few challenging Phototourism pairs, so that you can re-create some of the figures from the paper. Run this script to run matching and visualization (no ground truth is provided, see this [note](#reproducing-outdoor-evaluation-final-table)) on the provided pairs:
+
+```sh
+./match_pairs.py --resize 1600 --superglue outdoor --max_keypoints 2048 --nms_radius 3  --resize_float --input_dir assets/phototourism_sample_images/ --input_pairs assets/phototourism_sample_pairs.txt --output_dir dump_match_pairs_outdoor --viz
+```
+
+You should now image pairs such as these in `dump_match_pairs_outdoor/` (or something very close to it, see this [note](#a-note-on-reproducibility)):
+
+<img src="assets/outdoor_matches.png" width="560">
+
+</details>
+
+### Recommended settings for indoor / outdoor
+
+<details>
+  <summary>[Click to expand]</summary>
+
+For **indoor** images, we recommend the following settings (these are the defaults):
+
+```sh
+./match_pairs.py --resize 640 --superglue indoor --max_keypoints 1024 --nms_radius 4
+```
+
+For **outdoor** images, we recommend the following settings:
+
+```sh
+./match_pairs.py --resize 1600 --superglue outdoor --max_keypoints 2048 --nms_radius 3 --resize_float
+```
+
+You can provide your own list of pairs `--input_pairs` for images contained in `--input_dir`. Images can be resized before network inference with `--resize`. If you are re-running the same evaluation many times, you can use the `--cache` flag to reuse old computation.
+</details>
+
+### Test set pair file format explained
+
+<details>
+  <summary>[Click to expand]</summary>
+
+We provide the list of ScanNet test pairs in `assets/scannet_test_pairs_with_gt.txt` (with ground truth) and Phototourism test pairs `assets/phototourism_test_pairs.txt` (without ground truth) used to evaluate the matching from the paper. Each line corresponds to one pair and is structured as follows:
+
+```
+path_image_A path_image_B exif_rotationA exif_rotationB [KA_0 ... KA_8] [KB_0 ... KB_8] [T_AB_0 ... T_AB_15]
+```
+
+The `path_image_A` and `path_image_B` entries are paths to image A and B, respectively. The `exif_rotation` is an integer in the range [0, 3] that comes from the original EXIF metadata associated with the image, where, 0: no rotation, 1: 90 degree clockwise, 2: 180 degree clockwise, 3: 270 degree clockwise. If the EXIF data is not known, you can just provide a zero here and no rotation will be performed. `KA` and `KB` are the flattened `3x3` matrices of image A and image B intrinsics. `T_AB` is a flattened `4x4` matrix of the extrinsics between the pair.
+</details>
+
+### Reproducing the indoor evaluation on ScanNet
+
+<details>
+  <summary>[Click to expand]</summary>
+
+We provide the groundtruth for ScanNet in our format in the file `assets/scannet_test_pairs_with_gt.txt` for convenience. In order to reproduce similar tables to what was in the paper, you will need to download the dataset (we do not provide the raw test images). To download the ScanNet dataset, do the following:
+
+1. Head to the [ScanNet](https://github.com/ScanNet/ScanNet) github repo to download the ScanNet test set (100 scenes).
+2. You will need to extract the raw sensor data from the 100 `.sens` files in each scene in the test set using the [SensReader](https://github.com/ScanNet/ScanNet/tree/master/SensReader) tool.
+
+Once the ScanNet dataset is downloaded in `~/data/scannet`, you can run the following:
+
+```sh
+./match_pairs.py --input_dir ~/data/scannet --input_pairs assets/scannet_test_pairs_with_gt.txt --output_dir dump_scannet_test_results --eval
+```
+
+You should get the following table for ScanNet (or something very close to it, see this [note](#a-note-on-reproducibility)):
+
+```txt
+Evaluation Results (mean over 1500 pairs):
+AUC@5    AUC@10  AUC@20  Prec    MScore
+16.12    33.76   51.79   84.37   31.14
+```
+
+</details>
+
+### Reproducing the outdoor evaluation on YFCC
+
+<details>
+  <summary>[Click to expand]</summary>
+
+We provide the groundtruth for YFCC in our format in the file `assets/yfcc_test_pairs_with_gt.txt` for convenience. In order to reproduce similar tables to what was in the paper, you will need to download the dataset (we do not provide the raw test images). To download the YFCC dataset, you can use the [OANet](https://github.com/zjhthu/OANet) repo:
+
+```sh
+git clone https://github.com/zjhthu/OANet
+cd OANet
+bash download_data.sh raw_data raw_data_yfcc.tar.gz 0 8
+tar -xvf raw_data_yfcc.tar.gz
+mv raw_data/yfcc100m ~/data
+```
+
+Once the YFCC dataset is downloaded in `~/data/yfcc100m`, you can run the following:
+
+```sh
+./match_pairs.py --input_dir ~/data/yfcc100m --input_pairs assets/yfcc_test_pairs_with_gt.txt --output_dir dump_yfcc_test_results --eval --resize 1600 --superglue outdoor --max_keypoints 2048 --nms_radius 3 --resize_float
+```
+
+You should get the following table for YFCC (or something very close to it, see this [note](#a-note-on-reproducibility)):
+
+```txt
+Evaluation Results (mean over 4000 pairs):
+AUC@5    AUC@10  AUC@20  Prec    MScore
+39.02    59.51   75.72   98.72   23.61  
+```
+
+</details>
+
+### Reproducing outdoor evaluation on Phototourism
+
+<details>
+  <summary>[Click to expand]</summary>
+
+The Phototourism results shown in the paper were produced using similar data as the test set from the [Image Matching Challenge 2020](https://vision.uvic.ca/image-matching-challenge/), which holds the ground truth data private for the test set. We list the pairs we used in `assets/phototourism_test_pairs.txt`. To reproduce similar numbers on this test set, please submit to the challenge benchmark. While the challenge is still live, we cannot share the test set publically since we want to help maintain the integrity of the challenge. 
+
+</details>
+
+### Correcting EXIF rotation data in YFCC and Phototourism
+
+<details>
+  <summary>[Click to expand]</summary>
+
+In this repo, we provide manually corrected the EXIF rotation data for the outdoor evaluations on YFCC and Phototourism. For the YFCC dataset we found 7 images with incorrect EXIF rotation flags, resulting in 148 pairs out of 4000 being corrected. For Phototourism, we found 36 images with incorrect EXIF rotation flags, resulting in 212 out of 2200 pairs being corrected.
+
+The SuperGlue paper reports the results of SuperGlue **without** the corrected rotations, while the numbers in this README are reported **with** the corrected rotations. We found that our final conclusions from the evaluation still hold with or without the corrected rotations. For backwards compatability, we included the original, uncorrected EXIF rotation data in `assets/phototourism_test_pairs_original.txt` and `assets/yfcc_test_pairs_with_gt_original.txt` respectively.
+
+</details>
+
+### Outdoor training / validation scene splits of MegaDepth
+
+<details>
+  <summary>[Click to expand]</summary>
+
+For training and validation of the outdoor model, we used scenes from the [MegaDepth dataset](http://www.cs.cornell.edu/projects/megadepth/). We provide the list of scenes used to train the outdoor model in the `assets/` directory:
+
+* Training set: `assets/megadepth_train_scenes.txt`
+* Validation set: `assets/megadepth_validation_scenes.txt`
+
+</details>
+
+### A note on reproducibility
+
+<details>
+  <summary>[Click to expand]</summary>
+
+After simplifying the model code and evaluation code and preparing it for release, we made some improvements and tweaks that result in slightly different numbers than what was reported in the paper. The numbers and figures reported in the README were done using Ubuntu 16.04, OpenCV 3.4.5, and PyTorch 1.1.0. Even with matching the library versions, we observed some slight differences across Mac and Ubuntu, which we believe are due to differences in OpenCV's image resize function implementation and randomization of RANSAC.
+</details>
+
+### Creating high-quality PDF visualizations and faster visualization with --fast_viz
+
+<details>
+  <summary>[Click to expand]</summary>
+
+When generating output images with `match_pairs.py`, the default `--viz` flag uses a Matplotlib renderer which allows for the generation of camera-ready PDF visualizations if you additionally use `--viz_extension pdf` instead of the default png extension.
+
+```
+./match_pairs.py --viz --viz_extension pdf
+```
+
+Alternatively, you might want to save visualization images but have the generation be much faster.  You can use the `--fast_viz` flag to use an OpenCV-based image renderer as follows:
+
+```
+./match_pairs.py --viz --fast_viz
+```
+
+If you would also like an OpenCV display window to preview the results (you must use non-pdf output and use fast_fiz), simply run:
+
+```
+./match_pairs.py --viz --fast_viz --opencv_display
+```
+
+</details>
+
+
+## BibTeX Citation
+If you use any ideas from the paper or code from this repo, please consider citing:
+
+```txt
+@inproceedings{sarlin20superglue,
+  author    = {Paul-Edouard Sarlin and
+               Daniel DeTone and
+               Tomasz Malisiewicz and
+               Andrew Rabinovich},
+  title     = {{SuperGlue}: Learning Feature Matching with Graph Neural Networks},
+  booktitle = {CVPR},
+  year      = {2020},
+  url       = {https://arxiv.org/abs/1911.11763}
+}
+```
+
+## Additional Notes
+* For the demo, we found that the keyboard interaction works well with OpenCV 4.1.2.30, older versions were less responsive and the newest version had a [OpenCV bug on Mac](https://stackoverflow.com/questions/60032540/opencv-cv2-imshow-is-not-working-because-of-the-qt)
+* We generally do not recommend to run SuperPoint+SuperGlue below 160x120 resolution (QQVGA) and above 2000x1500
+* We do not intend to release the SuperGlue training code.
+* We do not intend to release the SIFT-based or homography SuperGlue models.
+
+## Legal Disclaimer
+Magic Leap is proud to provide its latest samples, toolkits, and research projects on Github to foster development and gather feedback from the spatial computing community. Use of the resources within this repo is subject to (a) the license(s) included herein, or (b) if no license is included, Magic Leap's [Developer Agreement](https://id.magicleap.com/terms/developer), which is available on our [Developer Portal](https://developer.magicleap.com/).
+If you need more, just ask on the [forums](https://forum.magicleap.com/hc/en-us/community/topics)!
+We're thrilled to be part of a well-meaning, friendly and welcoming community of millions.