vismatch 1.1.1__py3-none-any.whl

vismatch/third_party/MINIMA/third_party/XoFTR/src/utils/plotting.py

@@ -0,0 +1,227 @@
+import bisect
+import numpy as np
+import matplotlib.pyplot as plt
+import matplotlib
+plt.switch_backend('agg')
+from einops.einops import rearrange
+import torch.nn.functional as F
+
+
+def _compute_conf_thresh(data):
+    dataset_name = data['dataset_name'][0].lower()
+    if dataset_name == 'scannet':
+        thr = 5e-4
+    elif dataset_name == 'megadepth':
+        thr = 1e-4
+    elif dataset_name == 'vistir':
+        thr = 5e-4
+    else:
+        raise ValueError(f'Unknown dataset: {dataset_name}')
+    return thr
+
+
+# --- VISUALIZATION --- #
+
+def make_matching_figure(
+        img0, img1, mkpts0, mkpts1, color,
+        kpts0=None, kpts1=None, text=[], dpi=75, path=None):
+    # draw image pair
+    assert mkpts0.shape[0] == mkpts1.shape[0], f'mkpts0: {mkpts0.shape[0]} v.s. mkpts1: {mkpts1.shape[0]}'
+    fig, axes = plt.subplots(1, 2, figsize=(10, 6), dpi=dpi)
+    axes[0].imshow(img0, cmap='gray')
+    axes[1].imshow(img1, cmap='gray')
+    for i in range(2):   # clear all frames
+        axes[i].get_yaxis().set_ticks([])
+        axes[i].get_xaxis().set_ticks([])
+        for spine in axes[i].spines.values():
+            spine.set_visible(False)
+    plt.tight_layout(pad=1)
+    
+    if kpts0 is not None:
+        assert kpts1 is not None
+        axes[0].scatter(kpts0[:, 0], kpts0[:, 1], c='w', s=2)
+        axes[1].scatter(kpts1[:, 0], kpts1[:, 1], c='w', s=2)
+
+    # draw matches
+    if mkpts0.shape[0] != 0 and mkpts1.shape[0] != 0:
+        fig.canvas.draw()
+        transFigure = fig.transFigure.inverted()
+        fkpts0 = transFigure.transform(axes[0].transData.transform(mkpts0))
+        fkpts1 = transFigure.transform(axes[1].transData.transform(mkpts1))
+        fig.lines = [matplotlib.lines.Line2D((fkpts0[i, 0], fkpts1[i, 0]),
+                                            (fkpts0[i, 1], fkpts1[i, 1]),
+                                            transform=fig.transFigure, c=color[i], linewidth=1)
+                                        for i in range(len(mkpts0))]
+        
+        axes[0].scatter(mkpts0[:, 0], mkpts0[:, 1], c=color, s=4)
+        axes[1].scatter(mkpts1[:, 0], mkpts1[:, 1], c=color, s=4)
+
+    # put txts
+    txt_color = 'k' if img0[:100, :200].mean() > 200 else 'w'
+    fig.text(
+        0.01, 0.99, '\n'.join(text), transform=fig.axes[0].transAxes,
+        fontsize=15, va='top', ha='left', color=txt_color)
+
+    # save or return figure
+    if path:
+        plt.savefig(str(path), bbox_inches='tight', pad_inches=0)
+        plt.close()
+    else:
+        return fig
+
+
+def _make_evaluation_figure(data, b_id, alpha='dynamic', ret_dict=None):
+    b_mask = data['m_bids'] == b_id
+    conf_thr = _compute_conf_thresh(data)
+    
+    img0 = (data['image0'][b_id][0].cpu().numpy() * 255).round().astype(np.int32)
+    img1 = (data['image1'][b_id][0].cpu().numpy() * 255).round().astype(np.int32)
+    kpts0 = data['mkpts0_f'][b_mask].cpu().numpy()
+    kpts1 = data['mkpts1_f'][b_mask].cpu().numpy()
+    
+    # for megadepth, we visualize matches on the resized image
+    if 'scale0' in data:
+        kpts0 = kpts0 / data['scale0'][b_id].cpu().numpy()[[1, 0]]
+        kpts1 = kpts1 / data['scale1'][b_id].cpu().numpy()[[1, 0]]
+
+    epi_errs = data['epi_errs'][b_mask].cpu().numpy()
+    correct_mask = epi_errs < conf_thr
+    precision = np.mean(correct_mask) if len(correct_mask) > 0 else 0
+    n_correct = np.sum(correct_mask)
+
+    # matching info
+    if alpha == 'dynamic':
+        alpha = dynamic_alpha(len(correct_mask))
+    color = error_colormap(epi_errs, conf_thr, alpha=alpha)
+    
+    text = [
+        f'#Matches {len(kpts0)}',
+        f'Precision({conf_thr:.2e}) ({100 * precision:.1f}%): {n_correct}/{len(kpts0)}']
+    if ret_dict is not None:
+        text += [f"t_err: {ret_dict['metrics']['t_errs'][b_id]:.2f}",
+                f"R_err: {ret_dict['metrics']['R_errs'][b_id]:.2f}"]
+    
+    # make the figure
+    figure = make_matching_figure(img0, img1, kpts0, kpts1,
+                                  color, text=text)
+    return figure
+
+def _make_confidence_figure(data, b_id):
+    # TODO: Implement confidence figure
+    raise NotImplementedError()
+
+
+def make_matching_figures(data, config, mode='evaluation', ret_dict=None):
+    """ Make matching figures for a batch.
+    
+    Args:
+        data (Dict): a batch updated by PL_XoFTR.
+        config (Dict): matcher config
+    Returns:
+        figures (Dict[str, List[plt.figure]]
+    """
+    assert mode in ['evaluation', 'confidence']  # 'confidence'
+    figures = {mode: []}
+    for b_id in range(data['image0'].size(0)):
+        if mode == 'evaluation':
+            fig = _make_evaluation_figure(
+                data, b_id,
+                alpha=config.TRAINER.PLOT_MATCHES_ALPHA, ret_dict=ret_dict)
+        elif mode == 'confidence':
+            fig = _make_confidence_figure(data, b_id)
+        else:
+            raise ValueError(f'Unknown plot mode: {mode}')
+        figures[mode].append(fig)
+    return figures
+
+def make_mae_figures(data):
+    """ Make mae figures for a batch.
+    
+    Args:
+        data (Dict): a batch updated by PL_XoFTR_Pretrain.
+    Returns:
+        figures (List[plt.figure])
+    """
+    
+    scale = data['hw0_i'][0] // data['hw0_f'][0]
+    W_f = data["W_f"]
+
+    pred0, pred1 = data["pred0"], data["pred1"]
+    target0, target1 = data["target0"], data["target1"]
+
+    # replace masked regions with predictions
+    target0[data['b_ids'][data["ids_image0"]], data['i_ids'][data["ids_image0"]]] = pred0[data["ids_image0"]]
+    target1[data['b_ids'][data["ids_image1"]], data['j_ids'][data["ids_image1"]]] = pred1[data["ids_image1"]]
+
+    # remove excess parts, since the 10x10 windows have overlaping regions
+    target0 = rearrange(target0, 'n l (h w) (p q c) -> n c (h p) (w q) l', h=W_f, w=W_f, p=scale, q=scale, c=1)
+    target1 = rearrange(target1, 'n l (h w) (p q c) -> n c (h p) (w q) l', h=W_f, w=W_f, p=scale, q=scale, c=1) 
+    # target0[:,:,-scale:,:] = 0.0
+    # target0[:,:,:,-scale:] = 0.0
+    # target1[:,:,-scale:,:] = 0.0
+    # target1[:,:,:,-scale:] = 0.0
+    gap = scale //2
+    target0[:,:,-gap:,:] = 0.0
+    target0[:,:,:,-gap:] = 0.0
+    target1[:,:,-gap:,:] = 0.0
+    target1[:,:,:,-gap:] = 0.0
+    target0[:,:,:gap,:] = 0.0
+    target0[:,:,:,:gap] = 0.0
+    target1[:,:,:gap,:] = 0.0
+    target1[:,:,:,:gap] = 0.0
+    target0 = rearrange(target0, 'n c (h p) (w q) l -> n (c h p w q) l', h=W_f, w=W_f, p=scale, q=scale, c=1)
+    target1 = rearrange(target1, 'n c (h p) (w q) l -> n (c h p w q) l', h=W_f, w=W_f, p=scale, q=scale, c=1)
+
+    # windows to image 
+    kernel_size = [int(W_f*scale), int(W_f*scale)]
+    padding = kernel_size[0]//2 -1 if kernel_size[0] % 2 == 0 else kernel_size[0]//2
+    stride = data['hw0_i'][0] // data['hw0_c'][0]
+    target0 = F.fold(target0, output_size=data["image0"].shape[2:], kernel_size=kernel_size, stride=stride, padding=padding)
+    target1 = F.fold(target1, output_size=data["image1"].shape[2:], kernel_size=kernel_size, stride=stride, padding=padding)
+
+    # add mean and std of original image for visualization
+    if ("image0_norm" in data) and ("image1_norm" in data):
+        target0 = target0 * data["image0_std"] + data["image0_mean"]
+        target1 = target1 * data["image1_std"] + data["image1_mean"]
+        masked_image0 = data["masked_image0"] * data["image0_std"].to("cpu") + data["image0_mean"].to("cpu")
+        masked_image1 = data["masked_image1"] * data["image1_std"].to("cpu") + data["image1_mean"].to("cpu")
+    else:
+        masked_image0 = data["masked_image0"] 
+        masked_image1 = data["masked_image1"] 
+
+    figures = []
+    # Create a list of these tensors
+    image_groups = [[data["image0"], masked_image0, target0],
+                     [data["image1"], masked_image1, target1]]
+
+    # Iterate through the batches
+    for batch_idx in range(image_groups[0][0].shape[0]):  # Assuming batch dimension is the first dimension
+        fig, axs = plt.subplots(2, 3, figsize=(9, 6))  
+        for i, image_tensors in enumerate(image_groups):
+            for j, img_tensor in enumerate(image_tensors):
+                img = img_tensor[batch_idx, 0, :, :].detach().cpu().numpy()  # Get the image data as a NumPy array
+                axs[i,j].imshow(img, cmap='gray', vmin=0, vmax=1)  # Display the image in a subplot with correct colormap
+                axs[i,j].axis('off')  # Turn off axis labels
+        fig.tight_layout()
+        figures.append(fig)
+    return figures
+
+def dynamic_alpha(n_matches,
+                  milestones=[0, 300, 1000, 2000],
+                  alphas=[1.0, 0.8, 0.4, 0.2]):
+    if n_matches == 0:
+        return 1.0
+    ranges = list(zip(alphas, alphas[1:] + [None]))
+    loc = bisect.bisect_right(milestones, n_matches) - 1
+    _range = ranges[loc]
+    if _range[1] is None:
+        return _range[0]
+    return _range[1] + (milestones[loc + 1] - n_matches) / (
+        milestones[loc + 1] - milestones[loc]) * (_range[0] - _range[1])
+
+
+def error_colormap(err, thr, alpha=1.0):
+    assert alpha <= 1.0 and alpha > 0, f"Invaid alpha value: {alpha}"
+    x = 1 - np.clip(err / (thr * 2), 0, 1)
+    return np.clip(
+        np.stack([2-x*2, x*2, np.zeros_like(x), np.ones_like(x)*alpha], -1), 0, 1)

vismatch/third_party/MINIMA/third_party/XoFTR/src/utils/pretrain_utils.py

@@ -0,0 +1,83 @@
+import torch 
+import torch.nn as nn
+from einops.einops import rearrange
+import torch.nn.functional as F
+
+def generate_random_masks(batch, patch_size, mask_ratio, generator=None, margins=[0,0,0,0]):
+    mae_mask0 = _gen_random_mask(batch['image0'], patch_size, mask_ratio, generator, margins=margins)
+    mae_mask1 = _gen_random_mask(batch['image1'], patch_size, mask_ratio, generator, margins=margins)
+    batch.update({"mae_mask0" : mae_mask0, "mae_mask1": mae_mask1})
+
+def _gen_random_mask(image, patch_size, mask_ratio, generator=None, margins=[0, 0, 0, 0]): 
+    """ Random mask generator
+    Args:
+        image (torch.Tensor): [N, C, H, W]
+        patch_size (int)
+        mask_ratio (float)
+        generator (torch.Generator): RNG to create the same random masks for validation  
+        margins [float, float, float, float]: unused part for masking (up bottom left right)
+    Returns:
+        mask (torch.Tensor): (N, L)
+    """ 
+    N = image.shape[0]
+    l = (image.shape[2] // patch_size)
+    L = l ** 2
+    len_keep = int(L *  (1 - mask_ratio * (1 - sum(margins))))
+
+    margins = [int(margin * l) for margin in margins]
+
+    noise = torch.rand(N, l, l, device=image.device, generator=generator)
+    if margins[0] > 0 : noise[:,:margins[0],:] = 0
+    if margins[1] > 0 : noise[:,-margins[1]:,:] = 0
+    if margins[2] > 0 : noise[:,:,:margins[2]] = 0
+    if margins[3] > 0 : noise[:,:,-margins[3]:] = 0
+    noise = noise.flatten(1)
+
+    # sort noise for each sample
+    ids_shuffle = torch.argsort(noise, dim=1)
+    ids_restore = torch.argsort(ids_shuffle, dim=1)
+
+    # generate the binary mask: 0 is keep 1 is remove
+    mask = torch.ones([N, L], device=image.device)
+    mask[:, :len_keep] = 0
+    # unshuffle to get the binary mask
+    mask = torch.gather(mask, dim=1, index=ids_restore)
+    return mask
+
+def patchify(data):
+    """ Split images into small overlapped patches
+    Args:
+        data (dict):{
+            'image0_norm' (torch.Tensor): [N, C, H, W] normalized image,
+            'image1_norm' (torch.Tensor): [N, C, H, W] normalized image,
+    Returns:
+        image0 (torch.Tensor): [N, K, W_f**2, -1] (K: num of windows)
+        image1 (torch.Tensor): [N, K, W_f**2, -1] (K: num of windows)
+    """
+    stride = data['hw0_i'][0] // data['hw0_c'][0]
+    scale = data['hw0_i'][0] // data['hw0_f'][0]
+    W_f = data["W_f"]
+    kernel_size = [int(W_f*scale), int(W_f*scale)]
+    padding = kernel_size[0]//2 -1 if kernel_size[0] % 2 == 0 else kernel_size[0]//2
+
+    image0 = data["image0_norm"] if "image0_norm" in data else data["image0"]
+    image1 = data["image1_norm"] if "image1_norm" in data else data["image1"]
+    
+    image0 = F.unfold(image0, kernel_size=kernel_size, stride=stride, padding=padding)
+    image0 = rearrange(image0, 'n (c h p w q) l -> n l h w p q c', h=W_f, w=W_f, p=scale, q=scale)
+    image0 = image0.flatten(4)
+    image0 = image0.reshape(*image0.shape[:2], W_f**2, -1)
+    
+    image1 = F.unfold(image1, kernel_size=kernel_size, stride=stride, padding=padding)
+    image1 = rearrange(image1, 'n (c h p w q) l -> n l h w p q c', h=W_f, w=W_f, p=scale, q=scale)
+    image1 = image1.flatten(4)
+    image1 = image1.reshape(*image1.shape[:2], W_f**2, -1)
+
+    return image0, image1
+
+def get_target(data):
+    """Create target patches for mae"""
+    target0, target1 = patchify(data)
+    data.update({"target0":target0, "target1":target1})
+
+

vismatch/third_party/MINIMA/third_party/XoFTR/src/utils/profiler.py

@@ -0,0 +1,39 @@
+import torch
+from pytorch_lightning.profiler import SimpleProfiler, PassThroughProfiler
+from contextlib import contextmanager
+from pytorch_lightning.utilities import rank_zero_only
+
+
+class InferenceProfiler(SimpleProfiler):
+    """
+    This profiler records duration of actions with cuda.synchronize()
+    Use this in test time. 
+    """
+
+    def __init__(self):
+        super().__init__()
+        self.start = rank_zero_only(self.start)
+        self.stop = rank_zero_only(self.stop)
+        self.summary = rank_zero_only(self.summary)
+
+    @contextmanager
+    def profile(self, action_name: str) -> None:
+        try:
+            torch.cuda.synchronize()
+            self.start(action_name)
+            yield action_name
+        finally:
+            torch.cuda.synchronize()
+            self.stop(action_name)
+
+
+def build_profiler(name):
+    if name == 'inference':
+        return InferenceProfiler()
+    elif name == 'pytorch':
+        from pytorch_lightning.profiler import PyTorchProfiler
+        return PyTorchProfiler(use_cuda=True, profile_memory=True, row_limit=100)
+    elif name is None:
+        return PassThroughProfiler()
+    else:
+        raise ValueError(f'Invalid profiler: {name}')

vismatch/third_party/MINIMA/third_party/XoFTR/src/xoftr/__init__.py

@@ -0,0 +1,2 @@
+from .xoftr import XoFTR
+from .xoftr_pretrain import XoFTR_Pretrain

vismatch/third_party/MINIMA/third_party/XoFTR/src/xoftr/backbone/__init__.py

@@ -0,0 +1 @@
+from .resnet import ResNet_8_2

vismatch/third_party/MINIMA/third_party/XoFTR/src/xoftr/backbone/resnet.py

@@ -0,0 +1,95 @@
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+def conv1x1(in_planes, out_planes, stride=1):
+    """1x1 convolution without padding"""
+    return nn.Conv2d(in_planes, out_planes, kernel_size=1, stride=stride, padding=0, bias=False)
+
+
+def conv3x3(in_planes, out_planes, stride=1):
+    """3x3 convolution with padding"""
+    return nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride, padding=1, bias=False)
+
+
+class BasicBlock(nn.Module):
+    def __init__(self, in_planes, planes, stride=1):
+        super().__init__()
+        self.conv1 = conv3x3(in_planes, planes, stride)
+        self.conv2 = conv3x3(planes, planes)
+        self.bn1 = nn.BatchNorm2d(planes)
+        self.bn2 = nn.BatchNorm2d(planes)
+        self.relu = nn.ReLU(inplace=True)
+
+        if stride == 1:
+            self.downsample = None
+        else:
+            self.downsample = nn.Sequential(
+                conv1x1(in_planes, planes, stride=stride),
+                nn.BatchNorm2d(planes)
+            )
+
+    def forward(self, x):
+        y = x
+        y = self.relu(self.bn1(self.conv1(y)))
+        y = self.bn2(self.conv2(y))
+
+        if self.downsample is not None:
+            x = self.downsample(x)
+
+        return self.relu(x+y)
+
+class ResNet_8_2(nn.Module):
+    """
+    ResNet, output resolution are 1/8 and 1/2.
+    Each block has 2 layers.
+    """
+
+    def __init__(self, config):
+        super().__init__()
+        # Config
+        block = BasicBlock
+        initial_dim = config['initial_dim']
+        block_dims = config['block_dims']
+
+        # Class Variable
+        self.in_planes = initial_dim
+
+        # Networks
+        self.conv1 = nn.Conv2d(1, initial_dim, kernel_size=7, stride=2, padding=3, bias=False)
+        self.bn1 = nn.BatchNorm2d(initial_dim)
+        self.relu = nn.ReLU(inplace=True)
+
+        self.layer1 = self._make_layer(block, block_dims[0], stride=1)  # 1/2
+        self.layer2 = self._make_layer(block, block_dims[1], stride=2)  # 1/4
+        self.layer3 = self._make_layer(block, block_dims[2], stride=2)  # 1/8
+
+        self.layer3_outconv = conv1x1(block_dims[2], block_dims[2])
+        
+
+        for m in self.modules():
+            if isinstance(m, nn.Conv2d):
+                nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
+            elif isinstance(m, (nn.BatchNorm2d, nn.GroupNorm)):
+                nn.init.constant_(m.weight, 1)
+                nn.init.constant_(m.bias, 0)
+
+    def _make_layer(self, block, dim, stride=1):
+        layer1 = block(self.in_planes, dim, stride=stride)
+        layer2 = block(dim, dim, stride=1)
+        layers = (layer1, layer2)
+
+        self.in_planes = dim
+        return nn.Sequential(*layers)
+
+    def forward(self, x):
+        # ResNet Backbone
+        x0 = self.relu(self.bn1(self.conv1(x)))
+        x1 = self.layer1(x0)  # 1/2
+        x2 = self.layer2(x1)  # 1/4
+        x3 = self.layer3(x2)  # 1/8
+
+        x3_out = self.layer3_outconv(x3)
+
+        return x3_out, x2, x1
+

vismatch/third_party/MINIMA/third_party/XoFTR/src/xoftr/utils/geometry.py

@@ -0,0 +1,107 @@
+import torch
+
+
+@torch.no_grad()
+def warp_kpts(kpts0, depth0, depth1, T_0to1, K0, K1):
+    """ Warp kpts0 from I0 to I1 with depth, K and Rt
+    Also check covisibility and depth consistency.
+    Depth is consistent if relative error < 0.2 (hard-coded).
+    
+    Args:
+        kpts0 (torch.Tensor): [N, L, 2] - <x, y>,
+        depth0 (torch.Tensor): [N, H, W],
+        depth1 (torch.Tensor): [N, H, W],
+        T_0to1 (torch.Tensor): [N, 3, 4],
+        K0 (torch.Tensor): [N, 3, 3],
+        K1 (torch.Tensor): [N, 3, 3],
+    Returns:
+        calculable_mask (torch.Tensor): [N, L]
+        warped_keypoints0 (torch.Tensor): [N, L, 2] <x0_hat, y1_hat>
+    """
+    kpts0_long = kpts0.round().long()
+
+    # Sample depth, get calculable_mask on depth != 0
+    kpts0_depth = torch.stack(
+        [depth0[i, kpts0_long[i, :, 1], kpts0_long[i, :, 0]] for i in range(kpts0.shape[0])], dim=0
+    )  # (N, L)
+    nonzero_mask = kpts0_depth != 0
+
+    # Unproject
+    kpts0_h = torch.cat([kpts0, torch.ones_like(kpts0[:, :, [0]])], dim=-1) * kpts0_depth[..., None]  # (N, L, 3)
+    kpts0_cam = K0.inverse() @ kpts0_h.transpose(2, 1)  # (N, 3, L)
+
+    # Rigid Transform
+    w_kpts0_cam = T_0to1[:, :3, :3] @ kpts0_cam + T_0to1[:, :3, [3]]    # (N, 3, L)
+    w_kpts0_depth_computed = w_kpts0_cam[:, 2, :]
+
+    # Project
+    w_kpts0_h = (K1 @ w_kpts0_cam).transpose(2, 1)  # (N, L, 3)
+    w_kpts0 = w_kpts0_h[:, :, :2] / (w_kpts0_h[:, :, [2]] + 1e-4)  # (N, L, 2), +1e-4 to avoid zero depth
+
+    # Covisible Check
+    h, w = depth1.shape[1:3]
+    covisible_mask = (w_kpts0[:, :, 0] > 0) * (w_kpts0[:, :, 0] < w-1) * \
+        (w_kpts0[:, :, 1] > 0) * (w_kpts0[:, :, 1] < h-1)
+    w_kpts0_long = w_kpts0.long()
+    w_kpts0_long[~covisible_mask, :] = 0
+
+    w_kpts0_depth = torch.stack(
+        [depth1[i, w_kpts0_long[i, :, 1], w_kpts0_long[i, :, 0]] for i in range(w_kpts0_long.shape[0])], dim=0
+    )  # (N, L)
+    consistent_mask = ((w_kpts0_depth - w_kpts0_depth_computed) / w_kpts0_depth).abs() < 0.2
+    valid_mask = nonzero_mask * covisible_mask * consistent_mask
+
+    return valid_mask, w_kpts0
+
+@torch.no_grad()
+def warp_kpts_fine(kpts0, depth0, depth1, T_0to1, K0, K1, b_ids):
+    """ Warp kpts0 from I0 to I1 with depth, K and Rt for give batch ids
+    Also check covisibility and depth consistency.
+    Depth is consistent if relative error < 0.2 (hard-coded).
+    
+    Args:
+        kpts0 (torch.Tensor): [N, L, 2] - <x, y>,
+        depth0 (torch.Tensor): [N, H, W],
+        depth1 (torch.Tensor): [N, H, W],
+        T_0to1 (torch.Tensor): [N, 3, 4],
+        K0 (torch.Tensor): [N, 3, 3],
+        K1 (torch.Tensor): [N, 3, 3],
+        b_ids (torch.Tensor): [M], selected batch ids for fine-level matching
+    Returns:
+        calculable_mask (torch.Tensor): [N, L]
+        warped_keypoints0 (torch.Tensor): [N, L, 2] <x0_hat, y1_hat>
+    """
+    kpts0_long = kpts0.round().long()
+
+    # Sample depth, get calculable_mask on depth != 0
+    kpts0_depth = torch.stack(
+        [depth0[b_ids[i], kpts0_long[i, :, 1], kpts0_long[i, :, 0]] for i in range(kpts0.shape[0])], dim=0
+    )  # (N, L)
+    nonzero_mask = kpts0_depth != 0
+
+    # Unproject
+    kpts0_h = torch.cat([kpts0, torch.ones_like(kpts0[:, :, [0]])], dim=-1) * kpts0_depth[..., None]  # (N, L, 3)
+    kpts0_cam = K0[b_ids].inverse() @ kpts0_h.transpose(2, 1)  # (N, 3, L)
+
+    # Rigid Transform
+    w_kpts0_cam = T_0to1[b_ids, :3, :3] @ kpts0_cam + T_0to1[b_ids, :3, [3]][...,None]    # (N, 3, L)
+    w_kpts0_depth_computed = w_kpts0_cam[:, 2, :]
+
+    # Project
+    w_kpts0_h = (K1[b_ids] @ w_kpts0_cam).transpose(2, 1)  # (N, L, 3)
+    w_kpts0 = w_kpts0_h[:, :, :2] / (w_kpts0_h[:, :, [2]] + 1e-4)  # (N, L, 2), +1e-4 to avoid zero depth
+
+    # Covisible Check
+    h, w = depth1.shape[1:3]
+    covisible_mask = (w_kpts0[:, :, 0] > 0) * (w_kpts0[:, :, 0] < w-1) * \
+        (w_kpts0[:, :, 1] > 0) * (w_kpts0[:, :, 1] < h-1)
+    w_kpts0_long = w_kpts0.long()
+    w_kpts0_long[~covisible_mask, :] = 0
+
+    w_kpts0_depth = torch.stack(
+        [depth1[b_ids[i], w_kpts0_long[i, :, 1], w_kpts0_long[i, :, 0]] for i in range(w_kpts0_long.shape[0])], dim=0
+    )  # (N, L)
+    consistent_mask = ((w_kpts0_depth - w_kpts0_depth_computed) / w_kpts0_depth).abs() < 0.2
+    valid_mask = nonzero_mask * covisible_mask * consistent_mask
+
+    return valid_mask, w_kpts0

vismatch/third_party/MINIMA/third_party/XoFTR/src/xoftr/utils/position_encoding.py

@@ -0,0 +1,36 @@
+import math
+import torch
+from torch import nn
+
+
+class PositionEncodingSine(nn.Module):
+    """
+    This is a sinusoidal position encoding that generalized to 2-dimensional images
+    """
+
+    def __init__(self, d_model, max_shape=(256, 256)):
+        """
+        Args:
+            max_shape (tuple): for 1/8 featmap, the max length of 256 corresponds to 2048 pixels
+        """
+        super().__init__()
+
+        pe = torch.zeros((d_model, *max_shape))
+        y_position = torch.ones(max_shape).cumsum(0).float().unsqueeze(0)
+        x_position = torch.ones(max_shape).cumsum(1).float().unsqueeze(0)
+        div_term = torch.exp(torch.arange(0, d_model//2, 2).float() * (-math.log(10000.0) / (d_model//2)))
+
+        div_term = div_term[:, None, None]  # [C//4, 1, 1]
+        pe[0::4, :, :] = torch.sin(x_position * div_term)
+        pe[1::4, :, :] = torch.cos(x_position * div_term)
+        pe[2::4, :, :] = torch.sin(y_position * div_term)
+        pe[3::4, :, :] = torch.cos(y_position * div_term)
+
+        self.register_buffer('pe', pe.unsqueeze(0), persistent=False)  # [1, C, H, W]
+
+    def forward(self, x):
+        """
+        Args:
+            x: [N, C, H, W]
+        """
+        return x + self.pe[:, :, :x.size(2), :x.size(3)]