vismatch 1.1.1__py3-none-any.whl

vismatch/third_party/mast3r/dust3r/dust3r_visloc/datasets/utils.py

@@ -0,0 +1,118 @@
+# Copyright (C) 2024-present Naver Corporation. All rights reserved.
+# Licensed under CC BY-NC-SA 4.0 (non-commercial use only).
+#
+# --------------------------------------------------------
+# dataset utilities
+# --------------------------------------------------------
+import numpy as np
+import quaternion
+import torchvision.transforms as tvf
+from dust3r.utils.geometry import geotrf
+
+
+def cam_to_world_from_kapture(kdata, timestamp, camera_id):
+    camera_to_world = kdata.trajectories[timestamp, camera_id].inverse()
+    camera_pose = np.eye(4, dtype=np.float32)
+    camera_pose[:3, :3] = quaternion.as_rotation_matrix(camera_to_world.r)
+    camera_pose[:3, 3] = camera_to_world.t_raw
+    return camera_pose
+
+
+ratios_resolutions = {
+    224: {1.0: [224, 224]},
+    512: {4 / 3: [512, 384], 32 / 21: [512, 336], 16 / 9: [512, 288], 2 / 1: [512, 256], 16 / 5: [512, 160]}
+}
+
+
+def get_HW_resolution(H, W, maxdim, patchsize=16):
+    assert maxdim in ratios_resolutions, "Error, maxdim can only be 224 or 512 for now. Other maxdims not implemented yet."
+    ratios_resolutions_maxdim = ratios_resolutions[maxdim]
+    mindims = set([min(res) for res in ratios_resolutions_maxdim.values()])
+    ratio = W / H
+    ref_ratios = np.array([*(ratios_resolutions_maxdim.keys())])
+    islandscape = (W >= H)
+    if islandscape:
+        diff = np.abs(ratio - ref_ratios)
+    else:
+        diff = np.abs(ratio - (1 / ref_ratios))
+    selkey = ref_ratios[np.argmin(diff)]
+    res = ratios_resolutions_maxdim[selkey]
+    # check patchsize and make sure output resolution is a multiple of patchsize
+    if isinstance(patchsize, tuple):
+        assert len(patchsize) == 2 and isinstance(patchsize[0], int) and isinstance(
+            patchsize[1], int), "What is your patchsize format? Expected a single int or a tuple of two ints."
+        assert patchsize[0] == patchsize[1], "Error, non square patches not managed"
+        patchsize = patchsize[0]
+    assert max(res) == maxdim
+    assert min(res) in mindims
+    return res[::-1] if islandscape else res  # return HW
+
+
+def get_resize_function(maxdim, patch_size, H, W, is_mask=False):
+    if [max(H, W), min(H, W)] in ratios_resolutions[maxdim].values():
+        return lambda x: x, np.eye(3), np.eye(3)
+    else:
+        target_HW = get_HW_resolution(H, W, maxdim=maxdim, patchsize=patch_size)
+
+        ratio = W / H
+        target_ratio = target_HW[1] / target_HW[0]
+        to_orig_crop = np.eye(3)
+        to_rescaled_crop = np.eye(3)
+        if abs(ratio - target_ratio) < np.finfo(np.float32).eps:
+            crop_W = W
+            crop_H = H
+        elif ratio - target_ratio < 0:
+            crop_W = W
+            crop_H = int(W / target_ratio)
+            to_orig_crop[1, 2] = (H - crop_H) / 2.0
+            to_rescaled_crop[1, 2] = -(H - crop_H) / 2.0
+        else:
+            crop_W = int(H * target_ratio)
+            crop_H = H
+            to_orig_crop[0, 2] = (W - crop_W) / 2.0
+            to_rescaled_crop[0, 2] = - (W - crop_W) / 2.0
+
+        crop_op = tvf.CenterCrop([crop_H, crop_W])
+
+        if is_mask:
+            resize_op = tvf.Resize(size=target_HW, interpolation=tvf.InterpolationMode.NEAREST_EXACT)
+        else:
+            resize_op = tvf.Resize(size=target_HW)
+        to_orig_resize = np.array([[crop_W / target_HW[1], 0, 0],
+                                   [0, crop_H / target_HW[0], 0],
+                                   [0, 0, 1]])
+        to_rescaled_resize = np.array([[target_HW[1] / crop_W, 0, 0],
+                                       [0, target_HW[0] / crop_H, 0],
+                                       [0, 0, 1]])
+
+        op = tvf.Compose([crop_op, resize_op])
+
+        return op, to_rescaled_resize @ to_rescaled_crop, to_orig_crop @ to_orig_resize
+
+
+def rescale_points3d(pts2d, pts3d, to_resize, HR, WR):
+    # rescale pts2d as floats
+    # to colmap, so that the image is in [0, D] -> [0, NewD]
+    pts2d = pts2d.copy()
+    pts2d[:, 0] += 0.5
+    pts2d[:, 1] += 0.5
+
+    pts2d_rescaled = geotrf(to_resize, pts2d, norm=True)
+
+    pts2d_rescaled_int = pts2d_rescaled.copy()
+    # convert back to cv2 before round [-0.5, 0.5] -> pixel 0
+    pts2d_rescaled_int[:, 0] -= 0.5
+    pts2d_rescaled_int[:, 1] -= 0.5
+    pts2d_rescaled_int = pts2d_rescaled_int.round().astype(np.int64)
+
+    # update valid (remove cropped regions)
+    valid_rescaled = (pts2d_rescaled_int[:, 0] >= 0) & (pts2d_rescaled_int[:, 0] < WR) & (
+        pts2d_rescaled_int[:, 1] >= 0) & (pts2d_rescaled_int[:, 1] < HR)
+
+    pts2d_rescaled_int = pts2d_rescaled_int[valid_rescaled]
+
+    # rebuild pts3d from rescaled ps2d poses
+    pts3d_rescaled = np.full((HR, WR, 3), np.nan, dtype=np.float32)  # pts3d in 512 x something
+    pts3d_rescaled[pts2d_rescaled_int[:, 1], pts2d_rescaled_int[:, 0]] = pts3d[valid_rescaled]
+
+    return pts2d_rescaled, pts2d_rescaled_int, pts3d_rescaled, np.isfinite(pts3d_rescaled.sum(axis=-1))

vismatch/third_party/mast3r/dust3r/dust3r_visloc/evaluation.py

@@ -0,0 +1,65 @@
+# Copyright (C) 2024-present Naver Corporation. All rights reserved.
+# Licensed under CC BY-NC-SA 4.0 (non-commercial use only).
+#
+# --------------------------------------------------------
+# evaluation utilities
+# --------------------------------------------------------
+import numpy as np
+import quaternion
+import torch
+import roma
+import collections
+import os
+
+
+def aggregate_stats(info_str, pose_errors, angular_errors):
+    stats = collections.Counter()
+    median_pos_error = np.median(pose_errors)
+    median_angular_error = np.median(angular_errors)
+    out_str = f'{info_str}: {len(pose_errors)} images - {median_pos_error=}, {median_angular_error=}'
+
+    for trl_thr, ang_thr in [(0.1, 1), (0.25, 2), (0.5, 5), (5, 10)]:
+        for pose_error, angular_error in zip(pose_errors, angular_errors):
+            correct_for_this_threshold = (pose_error < trl_thr) and (angular_error < ang_thr)
+            stats[trl_thr, ang_thr] += correct_for_this_threshold
+    stats = {f'acc@{key[0]:g}m,{key[1]}deg': 100 * val / len(pose_errors) for key, val in stats.items()}
+    for metric, perf in stats.items():
+        out_str += f'  - {metric:12s}={float(perf):.3f}'
+    return out_str
+
+
+def get_pose_error(pr_camtoworld, gt_cam_to_world):
+    abs_transl_error = torch.linalg.norm(torch.tensor(pr_camtoworld[:3, 3]) - torch.tensor(gt_cam_to_world[:3, 3]))
+    abs_angular_error = roma.rotmat_geodesic_distance(torch.tensor(pr_camtoworld[:3, :3]),
+                                                      torch.tensor(gt_cam_to_world[:3, :3])) * 180 / np.pi
+    return abs_transl_error, abs_angular_error
+
+
+def export_results(output_dir, xp_label, query_names, poses_pred):
+    if output_dir is not None:
+        os.makedirs(output_dir, exist_ok=True)
+
+        lines = ""
+        lines_ltvl = ""
+        for query_name, pr_querycam_to_world in zip(query_names, poses_pred):
+            if pr_querycam_to_world is None:
+                pr_world_to_querycam = np.eye(4)
+            else:
+                pr_world_to_querycam = np.linalg.inv(pr_querycam_to_world)
+            query_shortname = os.path.basename(query_name)
+            pr_world_to_querycam_q = quaternion.from_rotation_matrix(pr_world_to_querycam[:3, :3])
+            pr_world_to_querycam_t = pr_world_to_querycam[:3, 3]
+
+            line_pose = quaternion.as_float_array(pr_world_to_querycam_q).tolist() + \
+                pr_world_to_querycam_t.flatten().tolist()
+
+            line_content = [query_name] + line_pose
+            lines += ' '.join(str(v) for v in line_content) + '\n'
+
+            line_content_ltvl = [query_shortname] + line_pose
+            lines_ltvl += ' '.join(str(v) for v in line_content_ltvl) + '\n'
+
+        with open(os.path.join(output_dir, xp_label + '_results.txt'), 'wt') as f:
+            f.write(lines)
+        with open(os.path.join(output_dir, xp_label + '_ltvl.txt'), 'wt') as f:
+            f.write(lines_ltvl)

vismatch/third_party/mast3r/dust3r/dust3r_visloc/localization.py

@@ -0,0 +1,140 @@
+# Copyright (C) 2024-present Naver Corporation. All rights reserved.
+# Licensed under CC BY-NC-SA 4.0 (non-commercial use only).
+#
+# --------------------------------------------------------
+# main pnp code
+# --------------------------------------------------------
+import numpy as np
+import quaternion
+import cv2
+from packaging import version
+
+from dust3r.utils.geometry import opencv_to_colmap_intrinsics
+
+try:
+    import poselib  # noqa
+    HAS_POSELIB = True
+except Exception as e:
+    HAS_POSELIB = False
+
+try:
+    import pycolmap  # noqa
+    version_number = pycolmap.__version__
+    if version.parse(version_number) < version.parse("0.5.0"):
+        HAS_PYCOLMAP = False
+    else:
+        HAS_PYCOLMAP = True
+except Exception as e:
+    HAS_PYCOLMAP = False
+    
+def run_pnp(pts2D, pts3D, K, distortion = None, mode='cv2', reprojectionError=5, img_size = None):
+    """
+    use OPENCV model for distortion (4 values)
+    """
+    assert mode in ['cv2', 'poselib', 'pycolmap']
+    try:
+        if len(pts2D) > 4 and mode == "cv2":
+            confidence = 0.9999
+            iterationsCount = 10_000
+            if distortion is not None:
+                cv2_pts2ds = np.copy(pts2D)
+                cv2_pts2ds = cv2.undistortPoints(cv2_pts2ds, K, np.array(distortion), R=None, P=K)
+                pts2D = cv2_pts2ds.reshape((-1, 2))
+
+            success, r_pose, t_pose, _ = cv2.solvePnPRansac(pts3D, pts2D, K, None, flags=cv2.SOLVEPNP_SQPNP,
+                                                            iterationsCount=iterationsCount,
+                                                            reprojectionError=reprojectionError,
+                                                            confidence=confidence)
+            if not success:
+                return False, None
+            r_pose = cv2.Rodrigues(r_pose)[0]  # world2cam == world2cam2
+            RT = np.r_[np.c_[r_pose, t_pose], [(0,0,0,1)]] # world2cam2
+            return True, np.linalg.inv(RT)  # cam2toworld
+        elif len(pts2D) > 4 and mode == "poselib":
+            assert HAS_POSELIB
+            confidence = 0.9999
+            iterationsCount = 10_000
+            # NOTE: `Camera` struct currently contains `width`/`height` fields,
+            # however these are not used anywhere in the code-base and are provided simply to be consistent with COLMAP.
+            # so we put garbage in there
+            colmap_intrinsics = opencv_to_colmap_intrinsics(K)
+            fx = colmap_intrinsics[0, 0]
+            fy = colmap_intrinsics[1, 1]
+            cx = colmap_intrinsics[0, 2]
+            cy = colmap_intrinsics[1, 2]
+            width = img_size[0] if img_size is not None else int(cx*2)
+            height = img_size[1] if img_size is not None else int(cy*2)
+
+            if distortion is None:
+                camera = {'model': 'PINHOLE', 'width': width, 'height': height, 'params': [fx, fy, cx, cy]}
+            else:
+                camera = {'model': 'OPENCV', 'width': width, 'height': height,
+                          'params': [fx, fy, cx, cy] + distortion}
+            
+            pts2D = np.copy(pts2D)
+            pts2D[:, 0] += 0.5
+            pts2D[:, 1] += 0.5
+            pose, _ = poselib.estimate_absolute_pose(pts2D, pts3D, camera,
+                                                        {'max_reproj_error': reprojectionError,
+                                                        'max_iterations': iterationsCount,
+                                                        'success_prob': confidence}, {})
+            if pose is None:
+                return False, None
+            RT = pose.Rt  # (3x4)
+            RT = np.r_[RT, [(0,0,0,1)]]  # world2cam
+            return True, np.linalg.inv(RT)  # cam2toworld
+        elif len(pts2D) > 4 and mode == "pycolmap":
+            assert HAS_PYCOLMAP
+            assert img_size is not None
+            
+            pts2D = np.copy(pts2D)
+            pts2D[:, 0] += 0.5
+            pts2D[:, 1] += 0.5
+            colmap_intrinsics = opencv_to_colmap_intrinsics(K)
+            fx = colmap_intrinsics[0, 0]
+            fy = colmap_intrinsics[1, 1]
+            cx = colmap_intrinsics[0, 2]
+            cy = colmap_intrinsics[1, 2]
+            width = img_size[0]
+            height = img_size[1]
+            if distortion is None:
+                camera_dict = {'model': 'PINHOLE', 'width': width, 'height': height, 'params': [fx, fy, cx, cy]}
+            else:
+                camera_dict = {'model': 'OPENCV', 'width': width, 'height': height,
+                               'params': [fx, fy, cx, cy] + distortion}
+
+            pycolmap_camera = pycolmap.Camera(
+            model=camera_dict['model'], width=camera_dict['width'], height=camera_dict['height'],
+            params=camera_dict['params'])
+
+            pycolmap_estimation_options = dict(ransac=dict(max_error=reprojectionError, min_inlier_ratio=0.01,
+                                               min_num_trials=1000, max_num_trials=100000,
+                                            confidence=0.9999))
+            pycolmap_refinement_options=dict(refine_focal_length=False, refine_extra_params=False)
+            ret = pycolmap.absolute_pose_estimation(pts2D, pts3D, pycolmap_camera,
+                                                    estimation_options=pycolmap_estimation_options,
+                                                    refinement_options=pycolmap_refinement_options)
+            if ret is None:
+                ret = {'success': False}
+            else:
+                ret['success'] = True
+                if callable(ret['cam_from_world'].matrix):
+                    retmat = ret['cam_from_world'].matrix()
+                else:
+                    retmat = ret['cam_from_world'].matrix
+                ret['qvec'] = quaternion.from_rotation_matrix(retmat[:3, :3])
+                ret['tvec'] = retmat[:3, 3]
+                
+            if not (ret['success'] and ret['num_inliers'] > 0):
+                success = False
+                pose = None
+            else:
+                success = True
+                pr_world_to_querycam = np.r_[ret['cam_from_world'].matrix(), [(0,0,0,1)]]
+                pose = np.linalg.inv(pr_world_to_querycam)
+            return success, pose
+        else:
+            return False, None
+    except Exception as e:
+        print(f'error during pnp: {e}')
+        return False, None

vismatch/third_party/mast3r/dust3r/train.py

@@ -0,0 +1,13 @@
+#!/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).
+#
+# --------------------------------------------------------
+# training executable for DUSt3R
+# --------------------------------------------------------
+from dust3r.training import get_args_parser, train
+
+if __name__ == '__main__':
+    args = get_args_parser()
+    args = args.parse_args()
+    train(args)

vismatch/third_party/mast3r/dust3r/visloc.py

@@ -0,0 +1,193 @@
+#!/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).
+#
+# --------------------------------------------------------
+# Simple visloc script
+# --------------------------------------------------------
+import numpy as np
+import random
+import argparse
+from tqdm import tqdm
+import math
+
+from dust3r.inference import inference
+from dust3r.model import AsymmetricCroCo3DStereo
+from dust3r.utils.geometry import find_reciprocal_matches, xy_grid, geotrf
+
+from dust3r_visloc.datasets import *
+from dust3r_visloc.localization import run_pnp
+from dust3r_visloc.evaluation import get_pose_error, aggregate_stats, export_results
+
+
+def get_args_parser():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--dataset", type=str, required=True, help="visloc dataset to eval")
+    parser_weights = parser.add_mutually_exclusive_group(required=True)
+    parser_weights.add_argument("--weights", type=str, help="path to the model weights", default=None)
+    parser_weights.add_argument("--model_name", type=str, help="name of the model weights",
+                                choices=["DUSt3R_ViTLarge_BaseDecoder_512_dpt",
+                                         "DUSt3R_ViTLarge_BaseDecoder_512_linear",
+                                         "DUSt3R_ViTLarge_BaseDecoder_224_linear"])
+    parser.add_argument("--confidence_threshold", type=float, default=3.0,
+                        help="confidence values higher than threshold are invalid")
+    parser.add_argument("--device", type=str, default='cuda', help="pytorch device")
+    parser.add_argument("--pnp_mode", type=str, default="cv2", choices=['cv2', 'poselib', 'pycolmap'],
+                        help="pnp lib to use")
+    parser_reproj = parser.add_mutually_exclusive_group()
+    parser_reproj.add_argument("--reprojection_error", type=float, default=5.0, help="pnp reprojection error")
+    parser_reproj.add_argument("--reprojection_error_diag_ratio", type=float, default=None,
+                               help="pnp reprojection error as a ratio of the diagonal of the image")
+
+    parser.add_argument("--pnp_max_points", type=int, default=100_000, help="pnp maximum number of points kept")
+    parser.add_argument("--viz_matches", type=int, default=0, help="debug matches")
+
+    parser.add_argument("--output_dir", type=str, default=None, help="output path")
+    parser.add_argument("--output_label", type=str, default='', help="prefix for results files")
+    return parser
+
+
+if __name__ == '__main__':
+    parser = get_args_parser()
+    args = parser.parse_args()
+    conf_thr = args.confidence_threshold
+    device = args.device
+    pnp_mode = args.pnp_mode
+    reprojection_error = args.reprojection_error
+    reprojection_error_diag_ratio = args.reprojection_error_diag_ratio
+    pnp_max_points = args.pnp_max_points
+    viz_matches = args.viz_matches
+
+    if args.weights is not None:
+        weights_path = args.weights
+    else:
+        weights_path = "naver/" + args.model_name
+    model = AsymmetricCroCo3DStereo.from_pretrained(weights_path).to(args.device)
+
+    dataset = eval(args.dataset)
+    dataset.set_resolution(model)
+
+    query_names = []
+    poses_pred = []
+    pose_errors = []
+    angular_errors = []
+    for idx in tqdm(range(len(dataset))):
+        views = dataset[(idx)]  # 0 is the query
+        query_view = views[0]
+        map_views = views[1:]
+        query_names.append(query_view['image_name'])
+
+        query_pts2d = []
+        query_pts3d = []
+        for map_view in map_views:
+            # prepare batch
+            imgs = []
+            for idx, img in enumerate([query_view['rgb_rescaled'], map_view['rgb_rescaled']]):
+                imgs.append(dict(img=img.unsqueeze(0), true_shape=np.int32([img.shape[1:]]),
+                                 idx=idx, instance=str(idx)))
+            output = inference([tuple(imgs)], model, device, batch_size=1, verbose=False)
+            pred1, pred2 = output['pred1'], output['pred2']
+            confidence_masks = [pred1['conf'].squeeze(0) >= conf_thr,
+                                (pred2['conf'].squeeze(0) >= conf_thr) & map_view['valid_rescaled']]
+            pts3d = [pred1['pts3d'].squeeze(0), pred2['pts3d_in_other_view'].squeeze(0)]
+
+            # find 2D-2D matches between the two images
+            pts2d_list, pts3d_list = [], []
+            for i in range(2):
+                conf_i = confidence_masks[i].cpu().numpy()
+                true_shape_i = imgs[i]['true_shape'][0]
+                pts2d_list.append(xy_grid(true_shape_i[1], true_shape_i[0])[conf_i])
+                pts3d_list.append(pts3d[i].detach().cpu().numpy()[conf_i])
+
+            PQ, PM = pts3d_list[0], pts3d_list[1]
+            if len(PQ) == 0 or len(PM) == 0:
+                continue
+            reciprocal_in_PM, nnM_in_PQ, num_matches = find_reciprocal_matches(PQ, PM)
+            if viz_matches > 0:
+                print(f'found {num_matches} matches')
+            matches_im1 = pts2d_list[1][reciprocal_in_PM]
+            matches_im0 = pts2d_list[0][nnM_in_PQ][reciprocal_in_PM]
+            valid_pts3d = map_view['pts3d_rescaled'][matches_im1[:, 1], matches_im1[:, 0]]
+
+            # from cv2 to colmap
+            matches_im0 = matches_im0.astype(np.float64)
+            matches_im1 = matches_im1.astype(np.float64)
+            matches_im0[:, 0] += 0.5
+            matches_im0[:, 1] += 0.5
+            matches_im1[:, 0] += 0.5
+            matches_im1[:, 1] += 0.5
+            # rescale coordinates
+            matches_im0 = geotrf(query_view['to_orig'], matches_im0, norm=True)
+            matches_im1 = geotrf(query_view['to_orig'], matches_im1, norm=True)
+            # from colmap back to cv2
+            matches_im0[:, 0] -= 0.5
+            matches_im0[:, 1] -= 0.5
+            matches_im1[:, 0] -= 0.5
+            matches_im1[:, 1] -= 0.5
+
+            # visualize a few matches
+            if viz_matches > 0:
+                viz_imgs = [np.array(query_view['rgb']), np.array(map_view['rgb'])]
+                from matplotlib import pyplot as pl
+                n_viz = viz_matches
+                match_idx_to_viz = np.round(np.linspace(0, num_matches - 1, n_viz)).astype(int)
+                viz_matches_im0, viz_matches_im1 = matches_im0[match_idx_to_viz], matches_im1[match_idx_to_viz]
+
+                H0, W0, H1, W1 = *viz_imgs[0].shape[:2], *viz_imgs[1].shape[:2]
+                img0 = np.pad(viz_imgs[0], ((0, max(H1 - H0, 0)), (0, 0), (0, 0)), 'constant', constant_values=0)
+                img1 = np.pad(viz_imgs[1], ((0, max(H0 - H1, 0)), (0, 0), (0, 0)), 'constant', constant_values=0)
+                img = np.concatenate((img0, img1), axis=1)
+                pl.figure()
+                pl.imshow(img)
+                cmap = pl.get_cmap('jet')
+                for i in range(n_viz):
+                    (x0, y0), (x1, y1) = viz_matches_im0[i].T, viz_matches_im1[i].T
+                    pl.plot([x0, x1 + W0], [y0, y1], '-+', color=cmap(i / (n_viz - 1)), scalex=False, scaley=False)
+                pl.show(block=True)
+
+            if len(valid_pts3d) == 0:
+                pass
+            else:
+                query_pts3d.append(valid_pts3d.cpu().numpy())
+                query_pts2d.append(matches_im0)
+
+        if len(query_pts2d) == 0:
+            success = False
+            pr_querycam_to_world = None
+        else:
+            query_pts2d = np.concatenate(query_pts2d, axis=0).astype(np.float32)
+            query_pts3d = np.concatenate(query_pts3d, axis=0)
+            if len(query_pts2d) > pnp_max_points:
+                idxs = random.sample(range(len(query_pts2d)), pnp_max_points)
+                query_pts3d = query_pts3d[idxs]
+                query_pts2d = query_pts2d[idxs]
+
+            W, H = query_view['rgb'].size
+            if reprojection_error_diag_ratio is not None:
+                reprojection_error_img = reprojection_error_diag_ratio * math.sqrt(W**2 + H**2)
+            else:
+                reprojection_error_img = reprojection_error
+            success, pr_querycam_to_world = run_pnp(query_pts2d, query_pts3d,
+                                                    query_view['intrinsics'], query_view['distortion'],
+                                                    pnp_mode, reprojection_error_img, img_size=[W, H])
+
+        if not success:
+            abs_transl_error = float('inf')
+            abs_angular_error = float('inf')
+        else:
+            abs_transl_error, abs_angular_error = get_pose_error(pr_querycam_to_world, query_view['cam_to_world'])
+
+        pose_errors.append(abs_transl_error)
+        angular_errors.append(abs_angular_error)
+        poses_pred.append(pr_querycam_to_world)
+
+    xp_label = f'tol_conf_{conf_thr}'
+    if args.output_label:
+        xp_label = args.output_label + '_' + xp_label
+    if reprojection_error_diag_ratio is not None:
+        xp_label = xp_label + f'_reproj_diag_{reprojection_error_diag_ratio}'
+    else:
+        xp_label = xp_label + f'_reproj_err_{reprojection_error}'
+    export_results(args.output_dir, xp_label, query_names, poses_pred)
+    out_string = aggregate_stats(f'{args.dataset}', pose_errors, angular_errors)
+    print(out_string)

vismatch/third_party/mast3r/kapture_mast3r_mapping.py

@@ -0,0 +1,127 @@
+#!/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).
+#
+# --------------------------------------------------------
+# mast3r exec for running standard sfm
+# --------------------------------------------------------
+import pycolmap
+import os
+import os.path as path
+import argparse
+
+from mast3r.model import AsymmetricMASt3R
+from mast3r.colmap.mapping import (kapture_import_image_folder_or_list, run_mast3r_matching, pycolmap_run_triangulator,
+                                   pycolmap_run_mapper, glomap_run_mapper)
+from kapture.io.csv import kapture_from_dir
+
+from kapture.converter.colmap.database_extra import kapture_to_colmap, generate_priors_for_reconstruction
+from kapture_localization.utils.pairsfile import get_pairs_from_file
+from kapture.io.records import get_image_fullpath
+from kapture.converter.colmap.database import COLMAPDatabase
+
+
+def get_argparser():
+    parser = argparse.ArgumentParser(description='point triangulator with mast3r from kapture data')
+    parser_weights = parser.add_mutually_exclusive_group(required=True)
+    parser_weights.add_argument("--weights", type=str, help="path to the model weights", default=None)
+    parser_weights.add_argument("--model_name", type=str, help="name of the model weights",
+                                choices=["MASt3R_ViTLarge_BaseDecoder_512_catmlpdpt_metric"])
+
+    parser_input = parser.add_mutually_exclusive_group(required=True)
+    parser_input.add_argument('-i', '--input', default=None, help='kdata')
+    parser_input.add_argument('--dir', default=None, help='image dir (individual intrinsics)')
+    parser_input.add_argument('--dir_same_camera', default=None, help='image dir (shared intrinsics)')
+
+    parser.add_argument('-o', '--output', required=True, help='output path to reconstruction')
+    parser.add_argument('--pairsfile_path', required=True, help='pairsfile')
+
+    parser.add_argument('--glomap_bin', default='glomap', type=str, help='glomap bin')
+
+    parser_mapper = parser.add_mutually_exclusive_group()
+    parser_mapper.add_argument('--ignore_pose', action='store_true', default=False)
+    parser_mapper.add_argument('--use_glomap_mapper', action='store_true', default=False)
+
+    parser_matching = parser.add_mutually_exclusive_group()
+    parser_matching.add_argument('--dense_matching', action='store_true', default=False)
+    parser_matching.add_argument('--pixel_tol', default=0, type=int)
+    parser.add_argument('--device', default='cuda')
+
+    parser.add_argument('--conf_thr', default=1.001, type=float)
+    parser.add_argument('--skip_geometric_verification', action='store_true', default=False)
+    parser.add_argument('--min_len_track', default=5, type=int)
+
+    return parser
+
+
+if __name__ == '__main__':
+    parser = get_argparser()
+    args = parser.parse_args()
+    if args.weights is not None:
+        weights_path = args.weights
+    else:
+        weights_path = "naver/" + args.model_name
+    model = AsymmetricMASt3R.from_pretrained(weights_path).to(args.device)
+    maxdim = max(model.patch_embed.img_size)
+    patch_size = model.patch_embed.patch_size
+
+    if args.input is not None:
+        kdata = kapture_from_dir(args.input)
+        records_data_path = get_image_fullpath(args.input)
+    else:
+        if args.dir_same_camera is not None:
+            use_single_camera = True
+            records_data_path = args.dir_same_camera
+        elif args.dir is not None:
+            use_single_camera = False
+            records_data_path = args.dir
+        else:
+            raise ValueError('all inputs choices are None')
+        kdata = kapture_import_image_folder_or_list(records_data_path, use_single_camera)
+    has_pose = kdata.trajectories is not None
+    image_pairs = get_pairs_from_file(args.pairsfile_path, kdata.records_camera, kdata.records_camera)
+
+    colmap_db_path = path.join(args.output, 'colmap.db')
+    reconstruction_path = path.join(args.output, "reconstruction")
+    priors_txt_path = path.join(args.output, "priors_for_reconstruction")
+    for path_i in [reconstruction_path, priors_txt_path]:
+        os.makedirs(path_i, exist_ok=True)
+    assert not os.path.isfile(colmap_db_path)
+
+    colmap_db = COLMAPDatabase.connect(colmap_db_path)
+    try:
+        kapture_to_colmap(kdata, args.input, tar_handler=None, database=colmap_db,
+                          keypoints_type=None, descriptors_type=None, export_two_view_geometry=False)
+        if has_pose:
+            generate_priors_for_reconstruction(kdata, colmap_db, priors_txt_path)
+
+        colmap_image_pairs = run_mast3r_matching(model, maxdim, patch_size, args.device,
+                                                 kdata, records_data_path, image_pairs, colmap_db,
+                                                 args.dense_matching, args.pixel_tol, args.conf_thr,
+                                                 args.skip_geometric_verification, args.min_len_track)
+        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 = {}
+    if not args.skip_geometric_verification:
+        print("verify_matches")
+        f = open(args.output + '/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, args.output + '/pairs.txt')
+
+    print("running mapping")
+    if has_pose and not args.ignore_pose and not args.use_glomap_mapper:
+        pycolmap_run_triangulator(colmap_db_path, priors_txt_path, reconstruction_path, records_data_path)
+    elif not args.use_glomap_mapper:
+        pycolmap_run_mapper(colmap_db_path, reconstruction_path, records_data_path)
+    else:
+        glomap_run_mapper(args.glomap_bin, colmap_db_path, reconstruction_path, records_data_path)