vismatch 1.1.1__py3-none-any.whl

vismatch/third_party/mast3r/mast3r/demo_glomap.py

@@ -0,0 +1,343 @@
+#!/usr/bin/env python3
+# Copyright (C) 2024-present Naver Corporation. All rights reserved.
+# Licensed under CC BY-NC-SA 4.0 (non-commercial use only).
+#
+# --------------------------------------------------------
+# gradio demo functions
+# --------------------------------------------------------
+import pycolmap
+import gradio
+import os
+import numpy as np
+import functools
+import trimesh
+import copy
+from scipy.spatial.transform import Rotation
+import tempfile
+import shutil
+import PIL.Image
+import torch
+
+from kapture.converter.colmap.database_extra import kapture_to_colmap
+from kapture.converter.colmap.database import COLMAPDatabase
+
+from mast3r.colmap.mapping import kapture_import_image_folder_or_list, run_mast3r_matching, glomap_run_mapper
+from mast3r.demo import set_scenegraph_options
+try:
+    from mast3r.retrieval.processor import Retriever
+    has_retrieval = True
+except Exception as e:
+    has_retrieval = False
+from mast3r.image_pairs import make_pairs
+
+import mast3r.utils.path_to_dust3r  # noqa
+from dust3r.utils.image import load_images
+from dust3r.viz import add_scene_cam, CAM_COLORS, OPENGL
+from dust3r.demo import get_args_parser as dust3r_get_args_parser
+
+import matplotlib.pyplot as pl
+
+
+class GlomapRecon:
+    def __init__(self, world_to_cam, intrinsics, points3d, imgs):
+        self.world_to_cam = world_to_cam
+        self.intrinsics = intrinsics
+        self.points3d = points3d
+        self.imgs = imgs
+
+
+class GlomapReconState:
+    def __init__(self, glomap_recon, should_delete=False, cache_dir=None, outfile_name=None):
+        self.glomap_recon = glomap_recon
+        self.cache_dir = cache_dir
+        self.outfile_name = outfile_name
+        self.should_delete = should_delete
+
+    def __del__(self):
+        if not self.should_delete:
+            return
+        if self.cache_dir is not None and os.path.isdir(self.cache_dir):
+            shutil.rmtree(self.cache_dir)
+        self.cache_dir = None
+        if self.outfile_name is not None and os.path.isfile(self.outfile_name):
+            os.remove(self.outfile_name)
+        self.outfile_name = None
+
+
+def get_args_parser():
+    parser = dust3r_get_args_parser()
+    parser.add_argument('--share', action='store_true')
+    parser.add_argument('--gradio_delete_cache', default=None, type=int,
+                        help='age/frequency at which gradio removes the file. If >0, matching cache is purged')
+    parser.add_argument('--glomap_bin', default='glomap', type=str, help='glomap bin')
+    parser.add_argument('--retrieval_model', default=None, type=str, help="retrieval_model to be loaded")
+
+    actions = parser._actions
+    for action in actions:
+        if action.dest == 'model_name':
+            action.choices = ["MASt3R_ViTLarge_BaseDecoder_512_catmlpdpt_metric"]
+    # change defaults
+    parser.prog = 'mast3r demo'
+    return parser
+
+
+def get_reconstructed_scene(glomap_bin, outdir, gradio_delete_cache, model, retrieval_model, device, silent, image_size,
+                            current_scene_state, filelist, transparent_cams, cam_size, scenegraph_type, winsize,
+                            win_cyclic, refid, shared_intrinsics, **kw):
+    """
+    from a list of images, run mast3r inference, sparse global aligner.
+    then run get_3D_model_from_scene
+    """
+    imgs = load_images(filelist, size=image_size, verbose=not silent)
+    if len(imgs) == 1:
+        imgs = [imgs[0], copy.deepcopy(imgs[0])]
+        imgs[1]['idx'] = 1
+        filelist = [filelist[0], filelist[0]]
+
+    scene_graph_params = [scenegraph_type]
+    if scenegraph_type in ["swin", "logwin"]:
+        scene_graph_params.append(str(winsize))
+    elif scenegraph_type == "oneref":
+        scene_graph_params.append(str(refid))
+    elif scenegraph_type == "retrieval":
+        scene_graph_params.append(str(winsize))  # Na
+        scene_graph_params.append(str(refid))  # k
+
+    if scenegraph_type in ["swin", "logwin"] and not win_cyclic:
+        scene_graph_params.append('noncyclic')
+    scene_graph = '-'.join(scene_graph_params)
+
+    sim_matrix = None
+    if 'retrieval' in scenegraph_type:
+        assert has_retrieval
+        assert retrieval_model is not None
+        retriever = Retriever(retrieval_model, backbone=model, device=device)
+        with torch.no_grad():
+            sim_matrix = retriever(filelist)
+
+        # Cleanup
+        del retriever
+        torch.cuda.empty_cache()
+
+    pairs = make_pairs(imgs, scene_graph=scene_graph, prefilter=None, symmetrize=True, sim_mat=sim_matrix)
+
+    if current_scene_state is not None and \
+        not current_scene_state.should_delete and \
+            current_scene_state.cache_dir is not None:
+        cache_dir = current_scene_state.cache_dir
+    elif gradio_delete_cache:
+        cache_dir = tempfile.mkdtemp(suffix='_cache', dir=outdir)
+    else:
+        cache_dir = os.path.join(outdir, 'cache')
+
+    root_path = os.path.commonpath(filelist)
+    filelist_relpath = [
+        os.path.relpath(filename, root_path).replace('\\', '/')
+        for filename in filelist
+    ]
+    kdata = kapture_import_image_folder_or_list((root_path, filelist_relpath), shared_intrinsics)
+    image_pairs = [
+        (filelist_relpath[img1['idx']], filelist_relpath[img2['idx']])
+        for img1, img2 in pairs
+    ]
+
+    colmap_db_path = os.path.join(cache_dir, 'colmap.db')
+    if os.path.isfile(colmap_db_path):
+        os.remove(colmap_db_path)
+
+    os.makedirs(os.path.dirname(colmap_db_path), exist_ok=True)
+    colmap_db = COLMAPDatabase.connect(colmap_db_path)
+    try:
+        kapture_to_colmap(kdata, root_path, tar_handler=None, database=colmap_db,
+                          keypoints_type=None, descriptors_type=None, export_two_view_geometry=False)
+        colmap_image_pairs = run_mast3r_matching(model, image_size, 16, device,
+                                                 kdata, root_path, image_pairs, colmap_db,
+                                                 False, 5, 1.001,
+                                                 False, 3)
+        colmap_db.close()
+    except Exception as e:
+        print(f'Error {e}')
+        colmap_db.close()
+        exit(1)
+
+    if len(colmap_image_pairs) == 0:
+        raise Exception("no matches were kept")
+
+    # colmap db is now full, run colmap
+    colmap_world_to_cam = {}
+    print("verify_matches")
+    f = open(cache_dir + '/pairs.txt', "w")
+    for image_path1, image_path2 in colmap_image_pairs:
+        f.write("{} {}\n".format(image_path1, image_path2))
+    f.close()
+    pycolmap.verify_matches(colmap_db_path, cache_dir + '/pairs.txt')
+
+    reconstruction_path = os.path.join(cache_dir, "reconstruction")
+    if os.path.isdir(reconstruction_path):
+        shutil.rmtree(reconstruction_path)
+    os.makedirs(reconstruction_path, exist_ok=True)
+    glomap_run_mapper(glomap_bin, colmap_db_path, reconstruction_path, root_path)
+
+    if current_scene_state is not None and \
+        not current_scene_state.should_delete and \
+            current_scene_state.outfile_name is not None:
+        outfile_name = current_scene_state.outfile_name
+    else:
+        outfile_name = tempfile.mktemp(suffix='_scene.glb', dir=outdir)
+
+    ouput_recon = pycolmap.Reconstruction(os.path.join(reconstruction_path, '0'))
+    print(ouput_recon.summary())
+
+    colmap_world_to_cam = {}
+    colmap_intrinsics = {}
+    colmap_image_id_to_name = {}
+    images = {}
+    num_reg_images = ouput_recon.num_reg_images()
+    for idx, (colmap_imgid, colmap_image) in enumerate(ouput_recon.images.items()):
+        colmap_image_id_to_name[colmap_imgid] = colmap_image.name
+        if callable(colmap_image.cam_from_world.matrix):
+            colmap_world_to_cam[colmap_imgid] = colmap_image.cam_from_world.matrix(
+            )
+        else:
+            colmap_world_to_cam[colmap_imgid] = colmap_image.cam_from_world.matrix
+        camera = ouput_recon.cameras[colmap_image.camera_id]
+        K = np.eye(3)
+        K[0, 0] = camera.focal_length_x
+        K[1, 1] = camera.focal_length_y
+        K[0, 2] = camera.principal_point_x
+        K[1, 2] = camera.principal_point_y
+        colmap_intrinsics[colmap_imgid] = K
+
+        with PIL.Image.open(os.path.join(root_path, colmap_image.name)) as im:
+            images[colmap_imgid] = np.asarray(im)
+
+        if idx + 1 == num_reg_images:
+            break  # bug with the iterable ?
+    points3D = []
+    num_points3D = ouput_recon.num_points3D()
+    for idx, (pt3d_id, pts3d) in enumerate(ouput_recon.points3D.items()):
+        points3D.append((pts3d.xyz, pts3d.color))
+        if idx + 1 == num_points3D:
+            break  # bug with the iterable ?
+    scene = GlomapRecon(colmap_world_to_cam, colmap_intrinsics, points3D, images)
+    scene_state = GlomapReconState(scene, gradio_delete_cache, cache_dir, outfile_name)
+    outfile = get_3D_model_from_scene(silent, scene_state, transparent_cams, cam_size)
+    return scene_state, outfile
+
+
+def get_3D_model_from_scene(silent, scene_state, transparent_cams=False, cam_size=0.05):
+    """
+    extract 3D_model (glb file) from a reconstructed scene
+    """
+    if scene_state is None:
+        return None
+    outfile = scene_state.outfile_name
+    if outfile is None:
+        return None
+
+    recon = scene_state.glomap_recon
+
+    scene = trimesh.Scene()
+    pts = np.stack([p[0] for p in recon.points3d], axis=0)
+    col = np.stack([p[1] for p in recon.points3d], axis=0)
+    pct = trimesh.PointCloud(pts, colors=col)
+    scene.add_geometry(pct)
+
+    # add each camera
+    cams2world = []
+    for i, (id, pose_w2c_3x4) in enumerate(recon.world_to_cam.items()):
+        intrinsics = recon.intrinsics[id]
+        focal = (intrinsics[0, 0] + intrinsics[1, 1]) / 2.0
+        camera_edge_color = CAM_COLORS[i % len(CAM_COLORS)]
+        pose_w2c = np.eye(4)
+        pose_w2c[:3, :] = pose_w2c_3x4
+        pose_c2w = np.linalg.inv(pose_w2c)
+        cams2world.append(pose_c2w)
+        add_scene_cam(scene, pose_c2w, camera_edge_color,
+                      None if transparent_cams else recon.imgs[id], focal,
+                      imsize=recon.imgs[id].shape[1::-1], screen_width=cam_size)
+
+    rot = np.eye(4)
+    rot[:3, :3] = Rotation.from_euler('y', np.deg2rad(180)).as_matrix()
+    scene.apply_transform(np.linalg.inv(cams2world[0] @ OPENGL @ rot))
+    if not silent:
+        print('(exporting 3D scene to', outfile, ')')
+    scene.export(file_obj=outfile)
+
+    return outfile
+
+
+def main_demo(glomap_bin, tmpdirname, model, retrieval_model, device, image_size, server_name, server_port,
+              silent=False, share=False, gradio_delete_cache=False):
+    if not silent:
+        print('Outputing stuff in', tmpdirname)
+
+    recon_fun = functools.partial(get_reconstructed_scene, glomap_bin, tmpdirname, gradio_delete_cache, model,
+                                  retrieval_model, device, silent, image_size)
+    model_from_scene_fun = functools.partial(get_3D_model_from_scene, silent)
+
+    available_scenegraph_type = [("complete: all possible image pairs", "complete"),
+                                 ("swin: sliding window", "swin"),
+                                 ("logwin: sliding window with long range", "logwin"),
+                                 ("oneref: match one image with all", "oneref")]
+    if retrieval_model is not None:
+        available_scenegraph_type.insert(1, ("retrieval: connect views based on similarity", "retrieval"))
+
+    def get_context(delete_cache):
+        css = """.gradio-container {margin: 0 !important; min-width: 100%};"""
+        title = "MASt3R Demo"
+        if delete_cache:
+            return gradio.Blocks(css=css, title=title, delete_cache=(delete_cache, delete_cache))
+        else:
+            return gradio.Blocks(css=css, title="MASt3R Demo")  # for compatibility with older versions
+
+    with get_context(gradio_delete_cache) as demo:
+        # scene state is save so that you can change conf_thr, cam_size... without rerunning the inference
+        scene = gradio.State(None)
+        gradio.HTML('<h2 style="text-align: center;">MASt3R Demo</h2>')
+        with gradio.Column():
+            inputfiles = gradio.File(file_count="multiple")
+            with gradio.Row():
+                shared_intrinsics = gradio.Checkbox(value=False, label="Shared intrinsics",
+                                                    info="Only optimize one set of intrinsics for all views")
+                scenegraph_type = gradio.Dropdown(available_scenegraph_type,
+                                                  value='complete', label="Scenegraph",
+                                                  info="Define how to make pairs",
+                                                  interactive=True)
+                with gradio.Column(visible=False) as win_col:
+                    winsize = gradio.Slider(label="Scene Graph: Window Size", value=1,
+                                            minimum=1, maximum=1, step=1)
+                    win_cyclic = gradio.Checkbox(value=False, label="Cyclic sequence")
+                refid = gradio.Slider(label="Scene Graph: Id", value=0,
+                                      minimum=0, maximum=0, step=1, visible=False)
+            run_btn = gradio.Button("Run")
+
+            with gradio.Row():
+                # adjust the camera size in the output pointcloud
+                cam_size = gradio.Slider(label="cam_size", value=0.01, minimum=0.001, maximum=1.0, step=0.001)
+            with gradio.Row():
+                transparent_cams = gradio.Checkbox(value=False, label="Transparent cameras")
+
+            outmodel = gradio.Model3D()
+
+            # events
+            scenegraph_type.change(set_scenegraph_options,
+                                   inputs=[inputfiles, win_cyclic, refid, scenegraph_type],
+                                   outputs=[win_col, winsize, win_cyclic, refid])
+            inputfiles.change(set_scenegraph_options,
+                              inputs=[inputfiles, win_cyclic, refid, scenegraph_type],
+                              outputs=[win_col, winsize, win_cyclic, refid])
+            win_cyclic.change(set_scenegraph_options,
+                              inputs=[inputfiles, win_cyclic, refid, scenegraph_type],
+                              outputs=[win_col, winsize, win_cyclic, refid])
+            run_btn.click(fn=recon_fun,
+                          inputs=[scene, inputfiles, transparent_cams, cam_size,
+                                  scenegraph_type, winsize, win_cyclic, refid, shared_intrinsics],
+                          outputs=[scene, outmodel])
+            cam_size.change(fn=model_from_scene_fun,
+                            inputs=[scene, transparent_cams, cam_size],
+                            outputs=outmodel)
+            transparent_cams.change(model_from_scene_fun,
+                                    inputs=[scene, transparent_cams, cam_size],
+                                    outputs=outmodel)
+    demo.launch(share=share, server_name=server_name, server_port=server_port)

vismatch/third_party/mast3r/mast3r/fast_nn.py

@@ -0,0 +1,223 @@
+# Copyright (C) 2024-present Naver Corporation. All rights reserved.
+# Licensed under CC BY-NC-SA 4.0 (non-commercial use only).
+#
+# --------------------------------------------------------
+# MASt3R Fast Nearest Neighbor
+# --------------------------------------------------------
+import torch
+import numpy as np
+import math
+from scipy.spatial import KDTree
+
+import mast3r.utils.path_to_dust3r  # noqa
+from dust3r.utils.device import to_numpy, todevice  # noqa
+
+
+@torch.no_grad()
+def bruteforce_reciprocal_nns(A, B, device='cuda', block_size=None, dist='l2'):
+    if isinstance(A, np.ndarray):
+        A = torch.from_numpy(A).to(device)
+    if isinstance(B, np.ndarray):
+        B = torch.from_numpy(B).to(device)
+
+    A = A.to(device)
+    B = B.to(device)
+
+    if dist == 'l2':
+        dist_func = torch.cdist
+        argmin = torch.min
+    elif dist == 'dot':
+        def dist_func(A, B):
+            return A @ B.T
+
+        def argmin(X, dim):
+            sim, nn = torch.max(X, dim=dim)
+            return sim.neg_(), nn
+    else:
+        raise ValueError(f'Unknown {dist=}')
+
+    if block_size is None or len(A) * len(B) <= block_size**2:
+        dists = dist_func(A, B)
+        _, nn_A = argmin(dists, dim=1)
+        _, nn_B = argmin(dists, dim=0)
+    else:
+        dis_A = torch.full((A.shape[0],), float('inf'), device=device, dtype=A.dtype)
+        dis_B = torch.full((B.shape[0],), float('inf'), device=device, dtype=B.dtype)
+        nn_A = torch.full((A.shape[0],), -1, device=device, dtype=torch.int64)
+        nn_B = torch.full((B.shape[0],), -1, device=device, dtype=torch.int64)
+        number_of_iteration_A = math.ceil(A.shape[0] / block_size)
+        number_of_iteration_B = math.ceil(B.shape[0] / block_size)
+
+        for i in range(number_of_iteration_A):
+            A_i = A[i * block_size:(i + 1) * block_size]
+            for j in range(number_of_iteration_B):
+                B_j = B[j * block_size:(j + 1) * block_size]
+                dists_blk = dist_func(A_i, B_j)  # A, B, 1
+                # dists_blk = dists[i * block_size:(i+1)*block_size, j * block_size:(j+1)*block_size]
+                min_A_i, argmin_A_i = argmin(dists_blk, dim=1)
+                min_B_j, argmin_B_j = argmin(dists_blk, dim=0)
+
+                col_mask = min_A_i < dis_A[i * block_size:(i + 1) * block_size]
+                line_mask = min_B_j < dis_B[j * block_size:(j + 1) * block_size]
+
+                dis_A[i * block_size:(i + 1) * block_size][col_mask] = min_A_i[col_mask]
+                dis_B[j * block_size:(j + 1) * block_size][line_mask] = min_B_j[line_mask]
+
+                nn_A[i * block_size:(i + 1) * block_size][col_mask] = argmin_A_i[col_mask] + (j * block_size)
+                nn_B[j * block_size:(j + 1) * block_size][line_mask] = argmin_B_j[line_mask] + (i * block_size)
+    nn_A = nn_A.cpu().numpy()
+    nn_B = nn_B.cpu().numpy()
+    return nn_A, nn_B
+
+
+class cdistMatcher:
+    def __init__(self, db_pts, device='cuda'):
+        self.db_pts = db_pts.to(device)
+        self.device = device
+
+    def query(self, queries, k=1, **kw):
+        assert k == 1
+        if queries.numel() == 0:
+            return None, []
+        nnA, nnB = bruteforce_reciprocal_nns(queries, self.db_pts, device=self.device, **kw)
+        dis = None
+        return dis, nnA
+
+
+def merge_corres(idx1, idx2, shape1=None, shape2=None, ret_xy=True, ret_index=False):
+    assert idx1.dtype == idx2.dtype == np.int32
+
+    # unique and sort along idx1
+    corres = np.unique(np.c_[idx2, idx1].view(np.int64), return_index=ret_index)
+    if ret_index:
+        corres, indices = corres
+    xy2, xy1 = corres[:, None].view(np.int32).T
+
+    if ret_xy:
+        assert shape1 and shape2
+        xy1 = np.unravel_index(xy1, shape1)
+        xy2 = np.unravel_index(xy2, shape2)
+        if ret_xy != 'y_x':
+            xy1 = xy1[0].base[:, ::-1]
+            xy2 = xy2[0].base[:, ::-1]
+
+    if ret_index:
+        return xy1, xy2, indices
+    return xy1, xy2
+
+
+def fast_reciprocal_NNs(pts1, pts2, subsample_or_initxy1=8, ret_xy=True, pixel_tol=0, ret_basin=False,
+                        device='cuda', **matcher_kw):
+    H1, W1, DIM1 = pts1.shape
+    H2, W2, DIM2 = pts2.shape
+    assert DIM1 == DIM2
+
+    pts1 = pts1.reshape(-1, DIM1)
+    pts2 = pts2.reshape(-1, DIM2)
+
+    if isinstance(subsample_or_initxy1, int) and pixel_tol == 0:
+        S = subsample_or_initxy1
+        y1, x1 = np.mgrid[S // 2:H1:S, S // 2:W1:S].reshape(2, -1)
+        max_iter = 10
+    else:
+        x1, y1 = subsample_or_initxy1
+        if isinstance(x1, torch.Tensor):
+            x1 = x1.cpu().numpy()
+        if isinstance(y1, torch.Tensor):
+            y1 = y1.cpu().numpy()
+        max_iter = 1
+
+    xy1 = np.int32(np.unique(x1 + W1 * y1))  # make sure there's no doublons
+    xy2 = np.full_like(xy1, -1)
+    old_xy1 = xy1.copy()
+    old_xy2 = xy2.copy()
+
+    if 'dist' in matcher_kw or 'block_size' in matcher_kw \
+            or (isinstance(device, str) and device.startswith('cuda')) \
+            or (isinstance(device, torch.device) and device.type.startswith('cuda')):
+        pts1 = pts1.to(device)
+        pts2 = pts2.to(device)
+        tree1 = cdistMatcher(pts1, device=device)
+        tree2 = cdistMatcher(pts2, device=device)
+    else:
+        pts1, pts2 = to_numpy((pts1, pts2))
+        tree1 = KDTree(pts1)
+        tree2 = KDTree(pts2)
+
+    notyet = np.ones(len(xy1), dtype=bool)
+    basin = np.full((H1 * W1 + 1,), -1, dtype=np.int32) if ret_basin else None
+
+    niter = 0
+    # n_notyet = [len(notyet)]
+    while notyet.any():
+        _, xy2[notyet] = to_numpy(tree2.query(pts1[xy1[notyet]], **matcher_kw))
+        if not ret_basin:
+            notyet &= (old_xy2 != xy2)  # remove points that have converged
+
+        _, xy1[notyet] = to_numpy(tree1.query(pts2[xy2[notyet]], **matcher_kw))
+        if ret_basin:
+            basin[old_xy1[notyet]] = xy1[notyet]
+        notyet &= (old_xy1 != xy1)  # remove points that have converged
+
+        # n_notyet.append(notyet.sum())
+        niter += 1
+        if niter >= max_iter:
+            break
+
+        old_xy2[:] = xy2
+        old_xy1[:] = xy1
+
+    # print('notyet_stats:', ' '.join(map(str, (n_notyet+[0]*10)[:max_iter])))
+
+    if pixel_tol > 0:
+        # in case we only want to match some specific points
+        # and still have some way of checking reciprocity
+        old_yx1 = np.unravel_index(old_xy1, (H1, W1))[0].base
+        new_yx1 = np.unravel_index(xy1, (H1, W1))[0].base
+        dis = np.linalg.norm(old_yx1 - new_yx1, axis=-1)
+        converged = dis < pixel_tol
+        if not isinstance(subsample_or_initxy1, int):
+            xy1 = old_xy1  # replace new points by old ones
+    else:
+        converged = ~notyet  # converged correspondences
+
+    # keep only unique correspondences, and sort on xy1
+    xy1, xy2 = merge_corres(xy1[converged], xy2[converged], (H1, W1), (H2, W2), ret_xy=ret_xy)
+    if ret_basin:
+        return xy1, xy2, basin
+    return xy1, xy2
+
+
+def extract_correspondences_nonsym(A, B, confA, confB, subsample=8, device=None, ptmap_key='pred_desc', pixel_tol=0):
+    if '3d' in ptmap_key:
+        opt = dict(device='cpu', workers=32)
+    else:
+        opt = dict(device=device, dist='dot', block_size=2**13)
+
+    # matching the two pairs
+    idx1 = []
+    idx2 = []
+    # merge corres from opposite pairs
+    HA, WA = A.shape[:2]
+    HB, WB = B.shape[:2]
+    if pixel_tol == 0:
+        nn1to2 = fast_reciprocal_NNs(A, B, subsample_or_initxy1=subsample, ret_xy=False, **opt)
+        nn2to1 = fast_reciprocal_NNs(B, A, subsample_or_initxy1=subsample, ret_xy=False, **opt)
+    else:
+        S = subsample
+        yA, xA = np.mgrid[S // 2:HA:S, S // 2:WA:S].reshape(2, -1)
+        yB, xB = np.mgrid[S // 2:HB:S, S // 2:WB:S].reshape(2, -1)
+
+        nn1to2 = fast_reciprocal_NNs(A, B, subsample_or_initxy1=(xA, yA), ret_xy=False, pixel_tol=pixel_tol, **opt)
+        nn2to1 = fast_reciprocal_NNs(B, A, subsample_or_initxy1=(xB, yB), ret_xy=False, pixel_tol=pixel_tol, **opt)
+
+    idx1 = np.r_[nn1to2[0], nn2to1[1]]
+    idx2 = np.r_[nn1to2[1], nn2to1[0]]
+
+    c1 = confA.ravel()[idx1]
+    c2 = confB.ravel()[idx2]
+
+    xy1, xy2, idx = merge_corres(idx1, idx2, (HA, WA), (HB, WB), ret_xy=True, ret_index=True)
+    conf = np.minimum(c1[idx], c2[idx])
+    corres = (xy1.copy(), xy2.copy(), conf)
+    return todevice(corres, device)

vismatch/third_party/mast3r/mast3r/image_pairs.py

@@ -0,0 +1,115 @@
+# Copyright (C) 2024-present Naver Corporation. All rights reserved.
+# Licensed under CC BY-NC-SA 4.0 (non-commercial use only).
+#
+# --------------------------------------------------------
+# utilities needed to load image pairs
+# --------------------------------------------------------
+import numpy as np
+import torch
+from mast3r.retrieval.graph import make_pairs_fps
+
+def make_pairs(imgs, scene_graph='complete', prefilter=None, symmetrize=True, sim_mat=None):
+    pairs = []
+    if scene_graph == 'complete':  # complete graph
+        for i in range(len(imgs)):
+            for j in range(i):
+                pairs.append((imgs[i], imgs[j]))
+    elif scene_graph.startswith('swin'):
+        iscyclic = not scene_graph.endswith('noncyclic')
+        try:
+            winsize = int(scene_graph.split('-')[1])
+        except Exception as e:
+            winsize = 3
+        pairsid = set()
+        for i in range(len(imgs)):
+            for j in range(1, winsize + 1):
+                idx = (i + j)
+                if iscyclic:
+                    idx = idx % len(imgs)  # explicit loop closure
+                if idx >= len(imgs):
+                    continue
+                pairsid.add((i, idx) if i < idx else (idx, i))
+        for i, j in pairsid:
+            pairs.append((imgs[i], imgs[j]))
+    elif scene_graph.startswith('logwin'):
+        iscyclic = not scene_graph.endswith('noncyclic')
+        try:
+            winsize = int(scene_graph.split('-')[1])
+        except Exception as e:
+            winsize = 3
+        offsets = [2**i for i in range(winsize)]
+        pairsid = set()
+        for i in range(len(imgs)):
+            ixs_l = [i - off for off in offsets]
+            ixs_r = [i + off for off in offsets]
+            for j in ixs_l + ixs_r:
+                if iscyclic:
+                    j = j % len(imgs)  # Explicit loop closure
+                if j < 0 or j >= len(imgs) or j == i:
+                    continue
+                pairsid.add((i, j) if i < j else (j, i))
+        for i, j in pairsid:
+            pairs.append((imgs[i], imgs[j]))
+    elif scene_graph.startswith('oneref'):
+        refid = int(scene_graph.split('-')[1]) if '-' in scene_graph else 0
+        for j in range(len(imgs)):
+            if j != refid:
+                pairs.append((imgs[refid], imgs[j]))
+    elif scene_graph.startswith('retrieval'):
+        mode, Na, k = scene_graph.split('-')
+        assert sim_mat is not None, "sim_mat is required for retrieval mode"
+
+        fps_pairs, anchor_idxs = make_pairs_fps(sim_mat, Na=int(Na), tokK=int(k), dist_thresh=None)
+
+        for i, j in fps_pairs:
+            pairs.append((imgs[i], imgs[j]))
+    else:
+        raise ValueError(f'unrecognized value for {scene_graph=}')
+
+    if symmetrize:
+        pairs += [(img2, img1) for img1, img2 in pairs]
+
+    # now, remove edges
+    if isinstance(prefilter, str) and prefilter.startswith('seq'):
+        pairs = filter_pairs_seq(pairs, int(prefilter[3:]))
+
+    if isinstance(prefilter, str) and prefilter.startswith('cyc'):
+        pairs = filter_pairs_seq(pairs, int(prefilter[3:]), cyclic=True)
+
+    return pairs
+
+
+def sel(x, kept):
+    if isinstance(x, dict):
+        return {k: sel(v, kept) for k, v in x.items()}
+    if isinstance(x, (torch.Tensor, np.ndarray)):
+        return x[kept]
+    if isinstance(x, (tuple, list)):
+        return type(x)([x[k] for k in kept])
+
+
+def _filter_edges_seq(edges, seq_dis_thr, cyclic=False):
+    # number of images
+    n = max(max(e) for e in edges) + 1
+
+    kept = []
+    for e, (i, j) in enumerate(edges):
+        dis = abs(i - j)
+        if cyclic:
+            dis = min(dis, abs(i + n - j), abs(i - n - j))
+        if dis <= seq_dis_thr:
+            kept.append(e)
+    return kept
+
+
+def filter_pairs_seq(pairs, seq_dis_thr, cyclic=False):
+    edges = [(img1['idx'], img2['idx']) for img1, img2 in pairs]
+    kept = _filter_edges_seq(edges, seq_dis_thr, cyclic=cyclic)
+    return [pairs[i] for i in kept]
+
+
+def filter_edges_seq(view1, view2, pred1, pred2, seq_dis_thr, cyclic=False):
+    edges = [(int(i), int(j)) for i, j in zip(view1['idx'], view2['idx'])]
+    kept = _filter_edges_seq(edges, seq_dis_thr, cyclic=cyclic)
+    print(f'>> Filtering edges more than {seq_dis_thr} frames apart: kept {len(kept)}/{len(edges)} edges')
+    return sel(view1, kept), sel(view2, kept), sel(pred1, kept), sel(pred2, kept)