vismatch 1.1.1__py3-none-any.whl

vismatch/third_party/mast3r/mast3r/retrieval/processor.py

@@ -0,0 +1,129 @@
+# Copyright (C) 2024-present Naver Corporation. All rights reserved.
+# Licensed under CC BY-NC-SA 4.0 (non-commercial use only).
+#
+# --------------------------------------------------------
+# Main Retriever class
+# --------------------------------------------------------
+import os
+import argparse
+import numpy as np
+import torch
+
+from mast3r.model import AsymmetricMASt3R
+from mast3r.retrieval.model import RetrievalModel, extract_local_features
+
+try:
+    import faiss
+    faiss.StandardGpuResources()  # when loading the checkpoint, it will try to instanciate FaissGpuL2Index
+except AttributeError as e:
+    import asmk.index
+
+    class FaissCpuL2Index(asmk.index.FaissL2Index):
+        def __init__(self, gpu_id):
+            super().__init__()
+            self.gpu_id = gpu_id
+
+        def _faiss_index_flat(self, dim):
+            """Return initialized faiss.IndexFlatL2"""
+            return faiss.IndexFlatL2(dim)
+
+    asmk.index.FaissGpuL2Index = FaissCpuL2Index
+
+from asmk import asmk_method  # noqa
+
+
+def get_args_parser():
+    parser = argparse.ArgumentParser('Retrieval scores from a set of retrieval', add_help=False, allow_abbrev=False)
+    parser.add_argument('--model', type=str, required=True,
+                        help="shortname of a retrieval model or path to the corresponding .pth")
+    parser.add_argument('--input', type=str, required=True,
+                        help="directory containing images or a file containing a list of image paths")
+    parser.add_argument('--outfile', type=str, required=True, help="numpy file where to store the matrix score")
+    return parser
+
+
+def get_impaths(imlistfile):
+    with open(imlistfile, 'r') as fid:
+        impaths = [f for f in imlistfile.read().splitlines() if not f.startswith('#')
+                   and len(f) > 0]  # ignore comments and empty lines
+    return impaths
+
+
+def get_impaths_from_imdir(imdir, extensions=['png', 'jpg', 'PNG', 'JPG']):
+    assert os.path.isdir(imdir)
+    impaths = [os.path.join(imdir, f) for f in sorted(os.listdir(imdir)) if any(f.endswith(ext) for ext in extensions)]
+    return impaths
+
+
+def get_impaths_from_imdir_or_imlistfile(input_imdir_or_imlistfile):
+    if os.path.isfile(input_imdir_or_imlistfile):
+        return get_impaths(input_imdir_or_imlistfile)
+    else:
+        return get_impaths_from_imdir(input_imdir_or_imlistfile)
+
+
+class Retriever(object):
+    def __init__(self, modelname, backbone=None, device='cuda'):
+        # load the model
+        assert os.path.isfile(modelname), modelname
+        print(f'Loading retrieval model from {modelname}')
+        ckpt = torch.load(modelname, 'cpu')  # TODO from pretrained to download it automatically
+        ckpt_args = ckpt['args']
+        if backbone is None:
+            backbone = AsymmetricMASt3R.from_pretrained(ckpt_args.pretrained)
+        self.model = RetrievalModel(
+            backbone, freeze_backbone=ckpt_args.freeze_backbone, prewhiten=ckpt_args.prewhiten,
+            hdims=list(map(int, ckpt_args.hdims.split('_'))) if len(ckpt_args.hdims) > 0 else "",
+            residual=getattr(ckpt_args, 'residual', False), postwhiten=ckpt_args.postwhiten,
+            featweights=ckpt_args.featweights, nfeat=ckpt_args.nfeat
+        ).to(device)
+        self.device = device
+        msg = self.model.load_state_dict(ckpt['model'], strict=False)
+        assert all(k.startswith('backbone') for k in msg.missing_keys)
+        assert len(msg.unexpected_keys) == 0
+        self.imsize = ckpt_args.imsize
+
+        # load the asmk codebook
+        dname, bname = os.path.split(modelname)  # TODO they should both be in the same file ?
+        bname_splits = bname.split('_')
+        cache_codebook_fname = os.path.join(dname, '_'.join(bname_splits[:-1]) + '_codebook.pkl')
+        assert os.path.isfile(cache_codebook_fname), cache_codebook_fname
+        asmk_params = {'index': {'gpu_id': 0}, 'train_codebook': {'codebook': {'size': '64k'}},
+                       'build_ivf': {'kernel': {'binary': True}, 'ivf': {'use_idf': False},
+                                     'quantize': {'multiple_assignment': 1}, 'aggregate': {}},
+                       'query_ivf': {'quantize': {'multiple_assignment': 5}, 'aggregate': {},
+                                     'search': {'topk': None},
+                                     'similarity': {'similarity_threshold': 0.0, 'alpha': 3.0}}}
+        asmk_params['train_codebook']['codebook']['size'] = ckpt_args.nclusters
+        self.asmk = asmk_method.ASMKMethod.initialize_untrained(asmk_params)
+        self.asmk = self.asmk.train_codebook(None, cache_path=cache_codebook_fname)
+
+    def __call__(self, input_imdir_or_imlistfile, outfile=None):
+        # get impaths
+        if isinstance(input_imdir_or_imlistfile, str):
+            impaths = get_impaths_from_imdir_or_imlistfile(input_imdir_or_imlistfile)
+        else:
+            impaths = input_imdir_or_imlistfile  # we're assuming a list has been passed
+        print(f'Found {len(impaths)} images')
+
+        # build the database
+        feat, ids = extract_local_features(self.model, impaths, self.imsize, tocpu=True, device=self.device)
+        feat = feat.cpu().numpy()
+        ids = ids.cpu().numpy()
+        asmk_dataset = self.asmk.build_ivf(feat, ids)
+
+        # we actually retrieve the same set of images
+        metadata, query_ids, ranks, ranked_scores = asmk_dataset.query_ivf(feat, ids)
+
+        # well ... scores are actually reordered according to ranks ...
+        # so we redo it the other way around...
+        scores = np.empty_like(ranked_scores)
+        scores[np.arange(ranked_scores.shape[0])[:, None], ranks] = ranked_scores
+
+        # save
+        if outfile is not None:
+            if os.path.isdir(os.path.dirname(outfile)):
+                os.makedirs(os.path.dirname(outfile), exist_ok=True)
+            np.save(outfile, scores)
+            print(f'Scores matrix saved in {outfile}')
+        return scores

vismatch/third_party/mast3r/mast3r/utils/__init__.py

@@ -0,0 +1,2 @@
+# Copyright (C) 2024-present Naver Corporation. All rights reserved.
+# Licensed under CC BY-NC-SA 4.0 (non-commercial use only).

vismatch/third_party/mast3r/mast3r/utils/coarse_to_fine.py

@@ -0,0 +1,214 @@
+# Copyright (C) 2024-present Naver Corporation. All rights reserved.
+# Licensed under CC BY-NC-SA 4.0 (non-commercial use only).
+#
+# --------------------------------------------------------
+# coarse to fine utilities
+# --------------------------------------------------------
+import numpy as np
+
+
+def crop_tag(cell):
+    return f'[{cell[1]}:{cell[3]},{cell[0]}:{cell[2]}]'
+
+
+def crop_slice(cell):
+    return slice(cell[1], cell[3]), slice(cell[0], cell[2])
+
+
+def _start_pos(total_size, win_size, overlap):
+    # we must have AT LEAST overlap between segments
+    # first segment starts at 0, last segment starts at total_size-win_size
+    assert 0 <= overlap < 1
+    assert total_size >= win_size
+    spacing = win_size * (1 - overlap)
+    last_pt = total_size - win_size
+    n_windows = 2 + int((last_pt - 1) // spacing)
+    return np.linspace(0, last_pt, n_windows).round().astype(int)
+
+
+def multiple_of_16(x):
+    return (x // 16) * 16
+
+
+def _make_overlapping_grid(H, W, size, overlap):
+    H_win = multiple_of_16(H * size // max(H, W))
+    W_win = multiple_of_16(W * size // max(H, W))
+    x = _start_pos(W, W_win, overlap)
+    y = _start_pos(H, H_win, overlap)
+    grid = np.stack(np.meshgrid(x, y, indexing='xy'), axis=-1)
+    grid = np.concatenate((grid, grid + (W_win, H_win)), axis=-1)
+    return grid.reshape(-1, 4)
+
+
+def _cell_size(cell2):
+    width, height = cell2[:, 2] - cell2[:, 0], cell2[:, 3] - cell2[:, 1]
+    assert width.min() >= 0
+    assert height.min() >= 0
+    return width, height
+
+
+def _norm_windows(cell2, H2, W2, forced_resolution=None):
+    # make sure the window aspect ratio is 3/4, or the output resolution is forced_resolution  if defined
+    outcell = cell2.copy()
+    width, height = _cell_size(cell2)
+    width2, height2 = width.clip(max=W2), height.clip(max=H2)
+    if forced_resolution is None:
+        width2[width < height] = (height2[width < height] * 3.01 / 4).clip(max=W2)
+        height2[width >= height] = (width2[width >= height] * 3.01 / 4).clip(max=H2)
+    else:
+        forced_H, forced_W = forced_resolution
+        width2[:] = forced_W
+        height2[:] = forced_H
+
+    half = (width2 - width) / 2
+    outcell[:, 0] -= half
+    outcell[:, 2] += half
+    half = (height2 - height) / 2
+    outcell[:, 1] -= half
+    outcell[:, 3] += half
+
+    # proj to integers
+    outcell = np.floor(outcell).astype(int)
+    # Take care of flooring errors
+    tmpw, tmph = _cell_size(outcell)
+    outcell[:, 0] += tmpw.astype(tmpw.dtype) - width2.astype(tmpw.dtype)
+    outcell[:, 1] += tmph.astype(tmpw.dtype) - height2.astype(tmpw.dtype)
+
+    # make sure 0 <= x < W2 and 0 <= y < H2
+    outcell[:, 0::2] -= outcell[:, [0]].clip(max=0)
+    outcell[:, 1::2] -= outcell[:, [1]].clip(max=0)
+    outcell[:, 0::2] -= outcell[:, [2]].clip(min=W2) - W2
+    outcell[:, 1::2] -= outcell[:, [3]].clip(min=H2) - H2
+
+    width, height = _cell_size(outcell)
+    assert np.all(width == width2.astype(width.dtype)) and np.all(
+        height == height2.astype(height.dtype)), "Error, output is not of the expected shape."
+    assert np.all(width <= W2)
+    assert np.all(height <= H2)
+    return outcell
+
+
+def _weight_pixels(cell, pix, assigned, gauss_var=2):
+    center = cell.reshape(-1, 2, 2).mean(axis=1)
+    width, height = _cell_size(cell)
+
+    # square distance between each cell center and each point
+    dist = (center[:, None] - pix[None]) / np.c_[width, height][:, None]
+    dist2 = np.square(dist).sum(axis=-1)
+
+    assert assigned.shape == dist2.shape
+    res = np.where(assigned, np.exp(-gauss_var * dist2), 0)
+    return res
+
+
+def pos2d_in_rect(p1, cell1):
+    x, y = p1.T
+    l, t, r, b = cell1
+    assigned = (l <= x) & (x < r) & (t <= y) & (y < b)
+    return assigned
+
+
+def _score_cell(cell1, H2, W2, p1, p2, min_corres=10, forced_resolution=None):
+    assert p1.shape == p2.shape
+
+    # compute keypoint assignment
+    assigned = pos2d_in_rect(p1, cell1[None].T)
+    assert assigned.shape == (len(cell1), len(p1))
+
+    # remove cells without correspondences
+    valid_cells = assigned.sum(axis=1) >= min_corres
+    cell1 = cell1[valid_cells]
+    assigned = assigned[valid_cells]
+    if not valid_cells.any():
+        return cell1, cell1, assigned
+
+    # fill-in the assigned points in both image
+    assigned_p1 = np.empty((len(cell1), len(p1), 2), dtype=np.float32)
+    assigned_p2 = np.empty((len(cell1), len(p2), 2), dtype=np.float32)
+    assigned_p1[:] = p1[None]
+    assigned_p2[:] = p2[None]
+    assigned_p1[~assigned] = np.nan
+    assigned_p2[~assigned] = np.nan
+
+    # find the median center and scale of assigned points in each cell
+    # cell_center1 = np.nanmean(assigned_p1, axis=1)
+    cell_center2 = np.nanmean(assigned_p2, axis=1)
+    im1_q25, im1_q75 = np.nanquantile(assigned_p1, (0.1, 0.9), axis=1)
+    im2_q25, im2_q75 = np.nanquantile(assigned_p2, (0.1, 0.9), axis=1)
+
+    robust_std1 = (im1_q75 - im1_q25).clip(20.)
+    robust_std2 = (im2_q75 - im2_q25).clip(20.)
+
+    cell_size1 = (cell1[:, 2:4] - cell1[:, 0:2])
+    cell_size2 = cell_size1 * robust_std2 / robust_std1
+    cell2 = np.c_[cell_center2 - cell_size2 / 2, cell_center2 + cell_size2 / 2]
+
+    # make sure cell bounds are valid
+    cell2 = _norm_windows(cell2, H2, W2, forced_resolution=forced_resolution)
+
+    # compute correspondence weights
+    corres_weights = _weight_pixels(cell1, p1, assigned) * _weight_pixels(cell2, p2, assigned)
+
+    # return a list of window pairs and assigned correspondences
+    return cell1, cell2, corres_weights
+
+
+def greedy_selection(corres_weights, target=0.9):
+    # corres_weight = (n_cell_pair, n_corres) matrix.
+    # If corres_weight[c,p]>0, means that correspondence p is visible in cell pair p
+    assert 0 < target <= 1
+    corres_weights = corres_weights.copy()
+
+    total = corres_weights.max(axis=0).sum()
+    target *= total
+
+    # init = empty
+    res = []
+    cur = np.zeros(corres_weights.shape[1])  # current selection
+
+    while cur.sum() < target:
+        # pick the nex best cell pair
+        best = corres_weights.sum(axis=1).argmax()
+        res.append(best)
+
+        # update current
+        cur += corres_weights[best]
+        # print('appending', best, 'with score', corres_weights[best].sum(), '-->', cur.sum())
+
+        # remove from all other views
+        corres_weights = (corres_weights - corres_weights[best]).clip(min=0)
+
+    return res
+
+
+def select_pairs_of_crops(img_q, img_b, pos2d_in_query, pos2d_in_ref, maxdim=512, overlap=.5, forced_resolution=None):
+    # prepare the overlapping cells
+    grid_q = _make_overlapping_grid(*img_q.shape[:2], maxdim, overlap)
+    grid_b = _make_overlapping_grid(*img_b.shape[:2], maxdim, overlap)
+
+    assert forced_resolution is None or len(forced_resolution) == 2
+    if isinstance(forced_resolution[0], int) or not len(forced_resolution[0]) == 2:
+        forced_resolution1 = forced_resolution2 = forced_resolution
+    else:
+        assert len(forced_resolution[1]) == 2
+        forced_resolution1 = forced_resolution[0]
+        forced_resolution2 = forced_resolution[1]
+
+    # Make sure crops respect constraints
+    grid_q = _norm_windows(grid_q.astype(float), *img_q.shape[:2], forced_resolution=forced_resolution1)
+    grid_b = _norm_windows(grid_b.astype(float), *img_b.shape[:2], forced_resolution=forced_resolution2)
+
+    # score cells
+    pairs_q = _score_cell(grid_q, *img_b.shape[:2], pos2d_in_query, pos2d_in_ref, forced_resolution=forced_resolution2)
+    pairs_b = _score_cell(grid_b, *img_q.shape[:2], pos2d_in_ref, pos2d_in_query, forced_resolution=forced_resolution1)
+    pairs_b = pairs_b[1], pairs_b[0], pairs_b[2]  # cellq, cellb, corres_weights
+
+    # greedy selection until all correspondences are generated
+    cell1, cell2, corres_weights = map(np.concatenate, zip(pairs_q, pairs_b))
+    if len(corres_weights) == 0:
+        return  # tolerated for empty generators
+    order = greedy_selection(corres_weights, target=0.9)
+
+    for i in order:
+        def pair_tag(qi, bi): return (str(qi) + crop_tag(cell1[i]), str(bi) + crop_tag(cell2[i]))
+        yield cell1[i], cell2[i], pair_tag

vismatch/third_party/mast3r/mast3r/utils/collate.py

@@ -0,0 +1,62 @@
+# Copyright (C) 2024-present Naver Corporation. All rights reserved.
+# Licensed under CC BY-NC-SA 4.0 (non-commercial use only).
+#
+# --------------------------------------------------------
+# Collate extensions
+# --------------------------------------------------------
+
+import torch
+import collections
+from torch.utils.data._utils.collate import default_collate_fn_map, default_collate_err_msg_format
+from typing import Callable, Dict, Optional, Tuple, Type, Union, List
+
+
+def cat_collate_tensor_fn(batch, *, collate_fn_map):
+    return torch.cat(batch, dim=0)
+
+
+def cat_collate_list_fn(batch, *, collate_fn_map: Optional[Dict[Union[Type, Tuple[Type, ...]], Callable]] = None):
+    return [item for bb in batch for item in bb]  # concatenate all lists
+
+
+cat_collate_fn_map = default_collate_fn_map.copy()
+cat_collate_fn_map[torch.Tensor] = cat_collate_tensor_fn
+cat_collate_fn_map[List] = cat_collate_list_fn
+cat_collate_fn_map[type(None)] = lambda _, **kw: None  # When some Nones, simply return a single None
+
+
+def cat_collate(batch, *, collate_fn_map: Optional[Dict[Union[Type, Tuple[Type, ...]], Callable]] = None):
+    r"""Custom collate function that concatenates stuff instead of stacking them, and handles NoneTypes """
+    elem = batch[0]
+    elem_type = type(elem)
+
+    if collate_fn_map is not None:
+        if elem_type in collate_fn_map:
+            return collate_fn_map[elem_type](batch, collate_fn_map=collate_fn_map)
+
+        for collate_type in collate_fn_map:
+            if isinstance(elem, collate_type):
+                return collate_fn_map[collate_type](batch, collate_fn_map=collate_fn_map)
+
+    if isinstance(elem, collections.abc.Mapping):
+        try:
+            return elem_type({key: cat_collate([d[key] for d in batch], collate_fn_map=collate_fn_map) for key in elem})
+        except TypeError:
+            # The mapping type may not support `__init__(iterable)`.
+            return {key: cat_collate([d[key] for d in batch], collate_fn_map=collate_fn_map) for key in elem}
+    elif isinstance(elem, tuple) and hasattr(elem, '_fields'):  # namedtuple
+        return elem_type(*(cat_collate(samples, collate_fn_map=collate_fn_map) for samples in zip(*batch)))
+    elif isinstance(elem, collections.abc.Sequence):
+        transposed = list(zip(*batch))  # It may be accessed twice, so we use a list.
+
+        if isinstance(elem, tuple):
+            # Backwards compatibility.
+            return [cat_collate(samples, collate_fn_map=collate_fn_map) for samples in transposed]
+        else:
+            try:
+                return elem_type([cat_collate(samples, collate_fn_map=collate_fn_map) for samples in transposed])
+            except TypeError:
+                # The sequence type may not support `__init__(iterable)` (e.g., `range`).
+                return [cat_collate(samples, collate_fn_map=collate_fn_map) for samples in transposed]
+
+    raise TypeError(default_collate_err_msg_format.format(elem_type))

vismatch/third_party/mast3r/mast3r/utils/misc.py

@@ -0,0 +1,17 @@
+# Copyright (C) 2024-present Naver Corporation. All rights reserved.
+# Licensed under CC BY-NC-SA 4.0 (non-commercial use only).
+#
+# --------------------------------------------------------
+# utilitary functions for MASt3R
+# --------------------------------------------------------
+import os
+import hashlib
+
+
+def mkdir_for(f):
+    os.makedirs(os.path.dirname(f), exist_ok=True)
+    return f
+
+
+def hash_md5(s):
+    return hashlib.md5(s.encode('utf-8')).hexdigest()

vismatch/third_party/mast3r/mast3r/utils/path_to_dust3r.py

@@ -0,0 +1,19 @@
+# Copyright (C) 2024-present Naver Corporation. All rights reserved.
+# Licensed under CC BY-NC-SA 4.0 (non-commercial use only).
+#
+# --------------------------------------------------------
+# dust3r submodule import
+# --------------------------------------------------------
+
+import sys
+import os.path as path
+HERE_PATH = path.normpath(path.dirname(__file__))
+DUSt3R_REPO_PATH = path.normpath(path.join(HERE_PATH, '../../dust3r'))
+DUSt3R_LIB_PATH = path.join(DUSt3R_REPO_PATH, 'dust3r')
+# check the presence of models directory in repo to be sure its cloned
+if path.isdir(DUSt3R_LIB_PATH):
+    # workaround for sibling import
+    sys.path.insert(0, DUSt3R_REPO_PATH)
+else:
+    raise ImportError(f"dust3r is not initialized, could not find: {DUSt3R_LIB_PATH}.\n "
+                      "Did you forget to run 'git submodule update --init --recursive' ?")

vismatch/third_party/mast3r/train.py

@@ -0,0 +1,48 @@
+#!/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 MASt3R
+# --------------------------------------------------------
+from mast3r.model import AsymmetricMASt3R
+from mast3r.losses import ConfMatchingLoss, MatchingLoss, APLoss, Regr3D, InfoNCE, Regr3D_ScaleShiftInv
+from mast3r.datasets import ARKitScenes, BlendedMVS, Co3d, MegaDepth, ScanNetpp, StaticThings3D, Waymo, WildRGBD
+
+import mast3r.utils.path_to_dust3r  # noqa
+# add mast3r classes to dust3r imports
+import dust3r.training
+dust3r.training.AsymmetricMASt3R = AsymmetricMASt3R
+dust3r.training.Regr3D = Regr3D
+dust3r.training.Regr3D_ScaleShiftInv = Regr3D_ScaleShiftInv
+dust3r.training.MatchingLoss = MatchingLoss
+dust3r.training.ConfMatchingLoss = ConfMatchingLoss
+dust3r.training.InfoNCE = InfoNCE
+dust3r.training.APLoss = APLoss
+
+import dust3r.datasets
+dust3r.datasets.ARKitScenes = ARKitScenes
+dust3r.datasets.BlendedMVS = BlendedMVS
+dust3r.datasets.Co3d = Co3d
+dust3r.datasets.MegaDepth = MegaDepth
+dust3r.datasets.ScanNetpp = ScanNetpp
+dust3r.datasets.StaticThings3D = StaticThings3D
+dust3r.datasets.Waymo = Waymo
+dust3r.datasets.WildRGBD = WildRGBD
+
+from dust3r.training import get_args_parser as dust3r_get_args_parser  # noqa
+from dust3r.training import train  # noqa
+
+
+def get_args_parser():
+    parser = dust3r_get_args_parser()
+    # change defaults
+    parser.prog = 'MASt3R training'
+    parser.set_defaults(model="AsymmetricMASt3R(patch_embed_cls='ManyAR_PatchEmbed')")
+    return parser
+
+
+if __name__ == '__main__':
+    args = get_args_parser()
+    args = args.parse_args()
+    train(args)