vismatch 1.1.1__py3-none-any.whl

vismatch/third_party/MatchAnything/imcui/third_party/MatchAnything/third_party/ROMA/roma/benchmarks/megadepth_pose_estimation_benchmark.py

@@ -0,0 +1,140 @@
+import numpy as np
+import torch
+from roma.utils import *
+from PIL import Image
+from tqdm import tqdm
+import torch.nn.functional as F
+import roma
+import kornia.geometry.epipolar as kepi
+
+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, model, model_name = None, resolution = None, scale_intrinsics = True, calibrated = True):
+        H,W = 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]}"
+                    dense_matches, dense_certainty = model.match(
+                        im_A_path, im_B_path, K1.copy(), K2.copy(), T1_to_2.copy()
+                    )
+                    sparse_matches,_ = model.sample(
+                        dense_matches, dense_certainty, 5000
+                    )
+                    
+                    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 = sparse_matches[:, :2]
+                    kpts1 = (
+                        np.stack(
+                            (
+                                w1 * (kpts1[:, 0] + 1) / 2,
+                                h1 * (kpts1[:, 1] + 1) / 2,
+                            ),
+                            axis=-1,
+                        )
+                    )
+                    kpts2 = sparse_matches[:, 2:]
+                    kpts2 = (
+                        np.stack(
+                            (
+                                w2 * (kpts2[:, 0] + 1) / 2,
+                                h2 * (kpts2[:, 1] + 1) / 2,
+                            ),
+                            axis=-1,
+                        )
+                    )
+
+                    for _ in range(5):
+                        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,
+                                    kpts2,
+                                    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,
+            }

vismatch/third_party/MatchAnything/imcui/third_party/MatchAnything/third_party/ROMA/roma/benchmarks/scannet_benchmark.py

@@ -0,0 +1,143 @@
+import os.path as osp
+import numpy as np
+import torch
+from roma.utils import *
+from PIL import Image
+from tqdm import tqdm
+
+
+class ScanNetBenchmark:
+    def __init__(self, data_root="data/scannet") -> None:
+        self.data_root = data_root
+
+    def benchmark(self, model, model_name = None):
+        model.train(False)
+        with torch.no_grad():
+            data_root = self.data_root
+            tmp = np.load(osp.join(data_root, "test.npz"))
+            pairs, rel_pose = tmp["name"], tmp["rel_pose"]
+            tot_e_t, tot_e_R, tot_e_pose = [], [], []
+            pair_inds = np.random.choice(
+                range(len(pairs)), size=len(pairs), replace=False
+            )
+            for pairind in tqdm(pair_inds, smoothing=0.9):
+                scene = pairs[pairind]
+                scene_name = f"scene0{scene[0]}_00"
+                im_A_path = osp.join(
+                        self.data_root,
+                        "scans_test",
+                        scene_name,
+                        "color",
+                        f"{scene[2]}.jpg",
+                    )
+                im_A = Image.open(im_A_path)
+                im_B_path = osp.join(
+                        self.data_root,
+                        "scans_test",
+                        scene_name,
+                        "color",
+                        f"{scene[3]}.jpg",
+                    )
+                im_B = Image.open(im_B_path)
+                T_gt = rel_pose[pairind].reshape(3, 4)
+                R, t = T_gt[:3, :3], T_gt[:3, 3]
+                K = np.stack(
+                    [
+                        np.array([float(i) for i in r.split()])
+                        for r in open(
+                            osp.join(
+                                self.data_root,
+                                "scans_test",
+                                scene_name,
+                                "intrinsic",
+                                "intrinsic_color.txt",
+                            ),
+                            "r",
+                        )
+                        .read()
+                        .split("\n")
+                        if r
+                    ]
+                )
+                w1, h1 = im_A.size
+                w2, h2 = im_B.size
+                K1 = K.copy()
+                K2 = K.copy()
+                dense_matches, dense_certainty = model.match(im_A_path, im_B_path)
+                sparse_matches, sparse_certainty = model.sample(
+                    dense_matches, dense_certainty, 5000
+                )
+                scale1 = 480 / min(w1, h1)
+                scale2 = 480 / min(w2, h2)
+                w1, h1 = scale1 * w1, scale1 * h1
+                w2, h2 = scale2 * w2, scale2 * h2
+                K1 = K1 * scale1
+                K2 = K2 * scale2
+
+                offset = 0.5
+                kpts1 = sparse_matches[:, :2]
+                kpts1 = (
+                    np.stack(
+                        (
+                            w1 * (kpts1[:, 0] + 1) / 2 - offset,
+                            h1 * (kpts1[:, 1] + 1) / 2 - offset,
+                        ),
+                        axis=-1,
+                    )
+                )
+                kpts2 = sparse_matches[:, 2:]
+                kpts2 = (
+                    np.stack(
+                        (
+                            w2 * (kpts2[:, 0] + 1) / 2 - offset,
+                            h2 * (kpts2[:, 1] + 1) / 2 - offset,
+                        ),
+                        axis=-1,
+                    )
+                )
+                for _ in range(5):
+                    shuffling = np.random.permutation(np.arange(len(kpts1)))
+                    kpts1 = kpts1[shuffling]
+                    kpts2 = kpts2[shuffling]
+                    try:
+                        norm_threshold = 0.5 / (
+                        np.mean(np.abs(K1[:2, :2])) + np.mean(np.abs(K2[:2, :2])))
+                        R_est, t_est, mask = estimate_pose(
+                            kpts1,
+                            kpts2,
+                            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_t.append(e_t)
+                tot_e_R.append(e_R)
+                tot_e_pose.append(e_pose)
+            tot_e_pose = np.array(tot_e_pose)
+            thresholds = [5, 10, 20]
+            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])
+            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,
+            }

vismatch/third_party/MatchAnything/imcui/third_party/MatchAnything/third_party/ROMA/roma/checkpointing/__init__.py

@@ -0,0 +1 @@
+from .checkpoint import CheckPoint

vismatch/third_party/MatchAnything/imcui/third_party/MatchAnything/third_party/ROMA/roma/checkpointing/checkpoint.py

@@ -0,0 +1,60 @@
+import os
+import torch
+from torch.nn.parallel.data_parallel import DataParallel
+from torch.nn.parallel.distributed import DistributedDataParallel
+from loguru import logger
+import gc
+
+import roma
+
+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 roma.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")
+            logger.info(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 roma.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

vismatch/third_party/MatchAnything/imcui/third_party/MatchAnything/third_party/ROMA/roma/datasets/__init__.py

@@ -0,0 +1,2 @@
+from .megadepth import MegadepthBuilder
+from .scannet import ScanNetBuilder

vismatch/third_party/MatchAnything/imcui/third_party/MatchAnything/third_party/ROMA/roma/datasets/megadepth.py

@@ -0,0 +1,230 @@
+import os
+from PIL import Image
+import h5py
+import numpy as np
+import torch
+import torchvision.transforms.functional as tvf
+import kornia.augmentation as K
+from roma.utils import get_depth_tuple_transform_ops, get_tuple_transform_ops
+import roma
+from roma.utils import *
+import math
+
+class MegadepthScene:
+    def __init__(
+        self,
+        data_root,
+        scene_info,
+        ht=384,
+        wt=512,
+        min_overlap=0.0,
+        max_overlap=1.0,
+        shake_t=0,
+        rot_prob=0.0,
+        normalize=True,
+        max_num_pairs = 100_000,
+        scene_name = None,
+        use_horizontal_flip_aug = False,
+        use_single_horizontal_flip_aug = False,
+        colorjiggle_params = None,
+        random_eraser = None,
+        use_randaug = False,
+        randaug_params = None,
+        randomize_size = 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]
+        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]
+        if randomize_size:
+            area = ht * wt
+            s = int(16 * (math.sqrt(area)//16))
+            sizes = ((ht,wt), (s,s), (wt,ht))
+            choice = roma.RANK % 3
+            ht, wt = sizes[choice] 
+        # counts, bins = np.histogram(self.overlaps,20)
+        # print(counts)
+        self.im_transform_ops = get_tuple_transform_ops(
+            resize=(ht, wt), normalize=normalize, colorjiggle_params = colorjiggle_params,
+        )
+        self.depth_transform_ops = get_depth_tuple_transform_ops(
+                resize=(ht, wt)
+            )
+        self.wt, self.ht = wt, ht
+        self.shake_t = shake_t
+        self.random_eraser = random_eraser
+        if use_horizontal_flip_aug and use_single_horizontal_flip_aug:
+            raise ValueError("Can't both flip both images and only flip one")
+        self.use_horizontal_flip_aug = use_horizontal_flip_aug
+        self.use_single_horizontal_flip_aug = use_single_horizontal_flip_aug
+        self.use_randaug = use_randaug
+
+    def load_im(self, im_path):
+        im = Image.open(im_path)
+        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 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 __getitem__(self, pair_idx):
+        # 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)
+        im_A = self.load_im(im_A_ref)
+        im_B = self.load_im(im_B_ref)
+        K1 = self.scale_intrinsic(K1, im_A.width, im_A.height)
+        K2 = self.scale_intrinsic(K2, im_B.width, im_B.height)
+
+        if self.use_randaug:
+            im_A, im_B = self.rand_augment(im_A, im_B)
+
+        depth_A = self.load_depth(depth_A_ref)
+        depth_B = self.load_depth(depth_B_ref)
+        # 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, im_B, depth_A, depth_B], t = self.rand_shake(im_A, im_B, depth_A, depth_B)
+        K1[:2, 2] += t
+        K2[:2, 2] += t
+        
+        im_A, im_B = im_A[None], im_B[None]
+        if self.random_eraser is not None:
+            im_A, depth_A = self.random_eraser(im_A, depth_A)
+            im_B, depth_B = self.random_eraser(im_B, depth_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)
+        if self.use_single_horizontal_flip_aug:
+            if np.random.rand() > 0.5:
+                im_B, depth_B, K2 = self.single_horizontal_flip(im_B, depth_B, K2)
+        
+        if roma.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")
+            
+        data_dict = {
+            "im_A": im_A[0],
+            "im_A_identifier": self.image_paths[idx1].split("/")[-1].split(".jpg")[0],
+            "im_B": im_B[0],
+            "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],
+            "K1": K1,
+            "K2": K2,
+            "T_1to2": T_1to2,
+            "im_A_path": im_A_ref,
+            "im_B_path": im_B_ref,
+            
+        }
+        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 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
+            scene_info = np.load(
+                os.path.join(self.scene_info_root, scene_name), allow_pickle=True
+            ).item()
+            scenes.append(
+                MegadepthScene(
+                    self.data_root, scene_info, 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