vismatch 1.1.1__py3-none-any.whl

vismatch/third_party/TopicFM/src/models/utils/coarse_matching.py

@@ -0,0 +1,213 @@
+import math
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from einops.einops import rearrange
+
+
+INF = 1e9
+
+
+def mask_border(m, b: int, v):
+    """ Mask borders with value
+    Args:
+        m (torch.Tensor): [N, H0, W0, H1, W1]
+        b (int)
+        v (m.dtype)
+    """
+    if b <= 0:
+        return
+
+    m[:, :b] = v
+    m[:, :, :b] = v
+    m[:, :, :, :b] = v
+    m[:, :, :, :, :b] = v
+    m[:, -b:] = v
+    m[:, :, -b:] = v
+    m[:, :, :, -b:] = v
+    m[:, :, :, :, -b:] = v
+
+
+def mask_border_with_padding(m, bd, v, p_m0, p_m1):
+    if bd <= 0:
+        return
+
+    m[:, :bd] = v
+    m[:, :, :bd] = v
+    m[:, :, :, :bd] = v
+    m[:, :, :, :, :bd] = v
+
+    h0s, w0s = p_m0.sum(1).max(-1)[0].int(), p_m0.sum(-1).max(-1)[0].int()
+    h1s, w1s = p_m1.sum(1).max(-1)[0].int(), p_m1.sum(-1).max(-1)[0].int()
+    for b_idx, (h0, w0, h1, w1) in enumerate(zip(h0s, w0s, h1s, w1s)):
+        m[b_idx, h0 - bd:] = v
+        m[b_idx, :, w0 - bd:] = v
+        m[b_idx, :, :, h1 - bd:] = v
+        m[b_idx, :, :, :, w1 - bd:] = v
+
+
+def compute_max_candidates(p_m0, p_m1):
+    """Compute the max candidates of all pairs within a batch
+    
+    Args:
+        p_m0, p_m1 (torch.Tensor): padded masks
+    """
+    h0s, w0s = p_m0.sum(1).max(-1)[0], p_m0.sum(-1).max(-1)[0]
+    h1s, w1s = p_m1.sum(1).max(-1)[0], p_m1.sum(-1).max(-1)[0]
+    max_cand = torch.sum(
+        torch.min(torch.stack([h0s * w0s, h1s * w1s], -1), -1)[0])
+    return max_cand
+
+
+class CoarseMatching(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.config = config
+        # general config
+        self.thr = config['thr']
+        self.border_rm = config['border_rm']
+        # -- # for trainig fine-level LoFTR
+        self.train_coarse_percent = config['train_coarse_percent']
+        self.train_pad_num_gt_min = config['train_pad_num_gt_min']
+
+    def forward(self, feat_c0, feat_c1, mask_c0, mask_c1, data):
+        """
+        Args:
+            data (dict)
+        Update:
+            data (dict): {
+                'b_ids' (torch.Tensor): [M'],
+                'i_ids' (torch.Tensor): [M'],
+                'j_ids' (torch.Tensor): [M'],
+                'gt_mask' (torch.Tensor): [M'],
+                'mkpts0_c' (torch.Tensor): [M, 2],
+                'mkpts1_c' (torch.Tensor): [M, 2],
+                'mconf' (torch.Tensor): [M]}
+            NOTE: M' != M during training.
+        """
+        conf_matrix = torch.einsum("nlc,nsc->nls", feat_c0, feat_c1) / feat_c0.shape[-1] ** .5  # (C * temperature)
+        if mask_c0 is not None:
+            outlier_mask = (mask_c0[..., None] * mask_c1[:, None])
+            conf_matrix.masked_fill_(~outlier_mask.bool(), -1e9)
+        conf_matrix = F.softmax(conf_matrix, 1) * F.softmax(conf_matrix, 2)
+        # conf_matrix = data['conf_matrix']
+        axes_lengths = {
+            'h0c': data['hw0_c'][0],
+            'w0c': data['hw0_c'][1],
+            'h1c': data['hw1_c'][0],
+            'w1c': data['hw1_c'][1]
+        }
+        _device = conf_matrix.device
+        # 1. confidence thresholding
+        conf_mask = conf_matrix > self.thr
+        conf_mask = rearrange(conf_mask, 'b (h0c w0c) (h1c w1c) -> b h0c w0c h1c w1c', **axes_lengths)
+        if 'mask0' not in data:
+            mask_border(conf_mask, self.border_rm, False)
+        else:
+            mask_border_with_padding(conf_mask, self.border_rm, False,
+                                     data['mask0'], data['mask1'])
+        conf_mask = rearrange(conf_mask, 'b h0c w0c h1c w1c -> b (h0c w0c) (h1c w1c)', **axes_lengths)
+
+        # 2. mutual nearest
+
+        # conf_mask = conf_mask \
+        #     * torch.zeros_like(conf_matrix).scatter_(1, conf_matrix.topk(k=1, dim=1)[1], 1.0) \
+        #     * torch.zeros_like(conf_matrix).scatter_(2, conf_matrix.topk(k=1, dim=2)[1], 1.0)
+        # b_ids, i_ids, j_ids = torch.nonzero(mask, as_tuple=True)
+        conf_mask = conf_mask \
+            * (conf_matrix == conf_matrix.max(dim=2, keepdim=True)[0]) \
+            * (conf_matrix == conf_matrix.max(dim=1, keepdim=True)[0])
+
+        # predict coarse matches from conf_matrix
+        data["conf_matrix"] = conf_matrix
+        data.update(**self.get_coarse_match(conf_mask, data))
+
+    @torch.no_grad()
+    def get_coarse_match(self, conf_mask, data):
+        """
+        Args:
+            conf_matrix (torch.Tensor): [N, L, S]
+            data (dict): with keys ['hw0_i', 'hw1_i', 'hw0_c', 'hw1_c']
+        Returns:
+            coarse_matches (dict): {
+                'b_ids' (torch.Tensor): [M'],
+                'i_ids' (torch.Tensor): [M'],
+                'j_ids' (torch.Tensor): [M'],
+                'gt_mask' (torch.Tensor): [M'],
+                'm_bids' (torch.Tensor): [M],
+                'mkpts0_c' (torch.Tensor): [M, 2],
+                'mkpts1_c' (torch.Tensor): [M, 2],
+                'mconf' (torch.Tensor): [M]}
+        """
+        _device = conf_mask.device
+
+        b_ids, i_ids, j_ids = torch.nonzero(conf_mask, as_tuple=True)
+        mconf = data["conf_matrix"][b_ids, i_ids, j_ids]
+        mconf, sorted_ids = torch.sort(mconf, dim=0, descending=True)
+        b_ids, i_ids, j_ids = b_ids[sorted_ids], i_ids[sorted_ids], j_ids[sorted_ids]
+
+        # 4. Random sampling of training samples for fine-level LoFTR
+        # (optional) pad samples with gt coarse-level matches
+        if self.training:
+            # NOTE:
+            # The sampling is performed across all pairs in a batch without manually balancing
+            # #samples for fine-level increases w.r.t. batch_size
+            if 'mask0' not in data:
+                num_candidates_max = conf_mask.size(0) * max(conf_mask.size(1), conf_mask.size(2))
+            else:
+                num_candidates_max = compute_max_candidates(
+                    data['mask0'], data['mask1'])
+            num_matches_train = int(num_candidates_max *
+                                    self.train_coarse_percent)
+            num_matches_pred = len(b_ids)
+            assert self.train_pad_num_gt_min < num_matches_train, "min-num-gt-pad should be less than num-train-matches"
+
+            # pred_indices is to select from prediction
+            if num_matches_pred <= num_matches_train - self.train_pad_num_gt_min:
+                pred_indices = torch.arange(num_matches_pred, device=_device)
+            else:
+                pred_indices = torch.randint(
+                    num_matches_pred,
+                    (num_matches_train - self.train_pad_num_gt_min, ),
+                    device=_device)
+
+            # gt_pad_indices is to select from gt padding. e.g. max(3787-4800, 200)
+            gt_pad_indices = torch.randint(
+                    len(data['spv_b_ids']),
+                    (max(num_matches_train - num_matches_pred,
+                        self.train_pad_num_gt_min), ),
+                    device=_device)
+            mconf_gt = torch.zeros(len(data['spv_b_ids']), device=_device)  # set conf of gt paddings to all zero
+
+            b_ids, i_ids, j_ids, mconf = map(
+                lambda x, y: torch.cat([x[pred_indices], y[gt_pad_indices]],
+                                       dim=0),
+                *zip([b_ids, data['spv_b_ids']], [i_ids, data['spv_i_ids']],
+                     [j_ids, data['spv_j_ids']], [mconf, mconf_gt]))
+
+        # These matches select patches that feed into fine-level network
+        coarse_matches = {'b_ids': b_ids, 'i_ids': i_ids, 'j_ids': j_ids}
+
+        # 4. Update with matches in original image resolution
+        scale = data['hw0_i'][0] / data['hw0_c'][0]
+        scale0 = scale * data['scale0'][b_ids] if 'scale0' in data else scale
+        scale1 = scale * data['scale1'][b_ids] if 'scale1' in data else scale
+        mkpts0_c = torch.stack(
+            [i_ids % data['hw0_c'][1], torch.div(i_ids, data['hw0_c'][1], rounding_mode="floor")],
+            dim=1) * scale0
+        mkpts1_c = torch.stack(
+            [j_ids % data['hw1_c'][1], torch.div(j_ids, data['hw1_c'][1], rounding_mode="floor")],
+            dim=1) * scale1
+
+        # These matches is the current prediction (for visualization)
+        coarse_matches.update({
+            'gt_mask': mconf == 0,
+            'm_bids': b_ids[mconf != 0],  # mconf == 0 => gt matches
+            'mkpts0_c': mkpts0_c[mconf != 0],
+            'mkpts1_c': mkpts1_c[mconf != 0],
+            'all_mkpts0_c': mkpts0_c,
+            'all_mkpts1_c': mkpts1_c,
+            'mconf': mconf[mconf != 0]
+        })
+
+        return coarse_matches

vismatch/third_party/TopicFM/src/models/utils/fine_matching.py

@@ -0,0 +1,172 @@
+import math
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from kornia.geometry.subpix import dsnt
+from kornia.utils.grid import create_meshgrid
+
+
+class FineMatching(nn.Module):
+    """FineMatching with s2d paradigm"""
+
+    def __init__(self):
+        super().__init__()
+
+    def forward(self, feat_f0, feat_f1, data):
+        """
+        Args:
+            feat0 (torch.Tensor): [M, WW, C]
+            feat1 (torch.Tensor): [M, WW, C]
+            data (dict)
+        Update:
+            data (dict):{
+                'expec_f' (torch.Tensor): [M, 3],
+                'mkpts0_f' (torch.Tensor): [M, 2],
+                'mkpts1_f' (torch.Tensor): [M, 2]}
+        """
+        M, WW, C = feat_f0.shape
+        W = int(math.sqrt(WW))
+        scale = data['hw0_i'][0] / data['hw0_f'][0]
+        self.M, self.W, self.WW, self.C, self.scale = M, W, WW, C, scale
+
+        # corner case: if no coarse matches found
+        if M == 0:
+            assert self.training == False, "M is always >0, when training, see coarse_matching.py"
+            # logger.warning('No matches found in coarse-level.')
+            data.update({
+                'expec_f': torch.empty(0, 3, device=feat_f0.device),
+                'mkpts0_f': data['mkpts0_c'],
+                'mkpts1_f': data['mkpts1_c'],
+            })
+            return
+
+        feat_f0_picked = feat_f0[:, WW//2, :]
+
+        sim_matrix = torch.einsum('mc,mrc->mr', feat_f0_picked, feat_f1)
+        softmax_temp = 1. / C**.5
+        heatmap = torch.softmax(softmax_temp * sim_matrix, dim=1)
+        feat_f1_picked = (feat_f1 * heatmap.unsqueeze(-1)).sum(dim=1) # [M, C]
+        heatmap = heatmap.view(-1, W, W)
+
+        # compute coordinates from heatmap
+        coords1_normalized = dsnt.spatial_expectation2d(heatmap[None], True)[0]  # [M, 2]
+        grid_normalized = create_meshgrid(W, W, True, heatmap.device).reshape(1, -1, 2)  # [1, WW, 2]
+
+        # compute std over <x, y>
+        var = torch.sum(grid_normalized**2 * heatmap.view(-1, WW, 1), dim=1) - coords1_normalized**2  # [M, 2]
+        std = torch.sum(torch.sqrt(torch.clamp(var, min=1e-10)), -1)  # [M]  clamp needed for numerical stability
+        
+        # for fine-level supervision
+        data.update({'expec_f': torch.cat([coords1_normalized, std.unsqueeze(1)], -1),
+                     'descriptors0': feat_f0_picked.detach(), 'descriptors1': feat_f1_picked.detach()})
+
+        # compute absolute kpt coords
+        self.get_fine_match(coords1_normalized, data)
+
+    @torch.no_grad()
+    def get_fine_match(self, coords1_normed, data):
+        W, WW, C, scale = self.W, self.WW, self.C, self.scale
+
+        # mkpts0_f and mkpts1_f
+        # scale0 = scale * data['scale0'][data['b_ids']] if 'scale0' in data else scale
+        mkpts0_f = data['mkpts0_c'] # + (coords0_normed * (W // 2) * scale0 )[:len(data['mconf'])]
+        scale1 = scale * data['scale1'][data['b_ids']] if 'scale1' in data else scale
+        mkpts1_f = data['mkpts1_c'] + (coords1_normed * (W // 2) * scale1)[:len(data['mconf'])]
+
+        data.update({
+            "mkpts0_f": mkpts0_f,
+            "mkpts1_f": mkpts1_f
+        })
+
+
+class DynamicFineMatching(nn.Module):
+    """FineMatching with s2d paradigm"""
+
+    def __init__(self):
+        super().__init__()
+
+    def forward(self, feat_f0, feat_f1, data):
+        """
+        Args:
+            feat0 (torch.Tensor): [M, WW, C]
+            feat1 (torch.Tensor): [M, WW, C]
+            data (dict)
+        Update:
+            data (dict):{
+                'expec_f' (torch.Tensor): [M, 3],
+                'mkpts0_f' (torch.Tensor): [M, 2],
+                'mkpts1_f' (torch.Tensor): [M, 2]}
+        """
+        M, WW, C = feat_f0.shape
+        W = int(math.sqrt(WW))
+        scale = data['hw0_i'][0] / data['hw0_f'][0]
+        self.M, self.W, self.WW, self.C= M, W, WW, C
+        self.scale0 = scale * data['scale0'][data['b_ids']] if 'scale0' in data else scale
+        self.scale1 = scale * data['scale1'][data['b_ids']] if 'scale1' in data else scale
+        # grid_WW = create_meshgrid(W, W, normalized_coordinates=True, device=feat_f0.device).reshape(1, -1, 2).repeat(M, 1, 1)
+
+        # corner case: if no coarse matches found
+        if M == 0:
+            assert self.training == False, "M is always >0, when training, see coarse_matching.py"
+            # logger.warning('No matches found in coarse-level.')
+            data.update({
+                'f_b_ids': data['b_ids'],
+                'all_mkpts0_f': data['all_mkpts0_c'],
+                'all_mkpts1_f': data['all_mkpts1_c'],
+                'mkpts0_f': data['mkpts0_c'],
+                'mkpts1_f': data['mkpts1_c'],
+            })
+            return
+
+        temperature = 0.1
+        score_map0 = data["score_map0"]
+        if self.training:
+            temperature = math.pow(2, -0.5*data["epoch_idx"]) if data["epoch_idx"] < 8 else 0.1
+        heatmap0 = F.softmax(score_map0 / temperature, dim=-1)
+        feat_f0_picked = (feat_f0 * heatmap0.unsqueeze(-1)).sum(dim=1)
+        coords0_normed = dsnt.spatial_expectation2d(heatmap0.view(-1, W, W)[None], True)[0]
+        out = self.get_fine_match(data, coords0_normed, feat_f0_picked, feat_f1)
+        out["mconf"] *= data["mconf"]
+        data.update(out)
+
+    def get_fine_match(self, data, coords0_normed, feat_f0_picked, feat_f1, mask_b_ids=None):
+        if mask_b_ids is None:
+            mask_b_ids = torch.ones_like(data["b_ids"]).bool()
+        coords0_normed = coords0_normed[mask_b_ids]
+        feat_f0_picked = feat_f0_picked[mask_b_ids]
+        feat_f1 = feat_f1[mask_b_ids]
+
+        sim_matrix = torch.einsum('mc,mrc->mr', feat_f0_picked, feat_f1) / self.C ** .5
+        heatmap1 = torch.softmax(sim_matrix, dim=1)
+        mconf, _ = torch.max(heatmap1, dim=1)
+        # data.update({"heatmap1": heatmap1})
+        # feat_f1_picked = (feat_f1 * heatmap1.unsqueeze(-1)).sum(dim=1)  # [M, C]
+        heatmap1 = heatmap1.view(-1, self.W, self.W)
+
+        # compute coordinates from heatmap
+        coords1_normed = dsnt.spatial_expectation2d(heatmap1[None], True)[0]
+        scale0 = self.scale0[mask_b_ids] if isinstance(self.scale0, torch.Tensor) else self.scale0
+        all_mkpts0_f = data['all_mkpts0_c'][mask_b_ids] + (coords0_normed * (self.W // 2) * scale0)
+        scale1 = self.scale1[mask_b_ids] if isinstance(self.scale1, torch.Tensor) else self.scale1
+        all_mkpts1_f = data['all_mkpts1_c'][mask_b_ids] + (coords1_normed * (self.W // 2) * scale1)
+
+        # mkpts0_f and mkpts1_f
+        true_matches = (~data["gt_mask"])[mask_b_ids]
+        mkpts0_f = all_mkpts0_f.detach()[true_matches]
+        mkpts1_f = all_mkpts1_f.detach()[true_matches]
+        scale0_f = scale0[true_matches] if isinstance(scale0, torch.Tensor) else scale0
+        scale1_f = scale1[true_matches] if isinstance(scale1, torch.Tensor) else scale1
+        mconf = mconf.detach()[true_matches]
+        f_b_ids = data["b_ids"][mask_b_ids]
+
+        return {
+            "f_b_ids": f_b_ids,
+            "m_bids": f_b_ids[true_matches],
+            "mconf": mconf,
+            "all_mkpts0_f": all_mkpts0_f,
+            "all_mkpts1_f": all_mkpts1_f,
+            "mkpts0_f": mkpts0_f,
+            "mkpts1_f": mkpts1_f,
+            "scale0_f": scale0_f, "scale1_f": scale1_f
+        }

vismatch/third_party/TopicFM/src/models/utils/geometry.py

@@ -0,0 +1,54 @@
+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

vismatch/third_party/TopicFM/src/models/utils/supervision.py

@@ -0,0 +1,167 @@
+from math import log
+from loguru import logger
+
+import torch
+import torch.nn.functional as F
+from einops import repeat
+from kornia.utils import create_meshgrid
+from kornia.geometry.epipolar import fundamental_from_projections
+
+from .geometry import warp_kpts
+
+##############  ↓  Coarse-Level supervision  ↓  ##############
+
+
+@torch.no_grad()
+def mask_pts_at_padded_regions(grid_pt, mask):
+    """For megadepth dataset, zero-padding exists in images"""
+    mask = repeat(mask, 'n h w -> n (h w) c', c=2)
+    grid_pt[~mask.bool()] = 0
+    return grid_pt
+
+
+@torch.no_grad()
+def spvs_coarse(data, config):
+    """
+    Update:
+        data (dict): {
+            "conf_matrix_gt": [N, hw0, hw1],
+            'spv_b_ids': [M]
+            'spv_i_ids': [M]
+            'spv_j_ids': [M]
+            'spv_w_pt0_i': [N, hw0, 2], in original image resolution
+            'spv_pt1_i': [N, hw1, 2], in original image resolution
+        }
+        
+    NOTE:
+        - for scannet dataset, there're 3 kinds of resolution {i, c, f}
+        - for megadepth dataset, there're 4 kinds of resolution {i, i_resize, c, f}
+    """
+    # 1. misc
+    device = data['image0'].device
+    N, _, H0, W0 = data['image0'].shape
+    _, _, H1, W1 = data['image1'].shape
+    scale = config['MODEL']['RESOLUTION'][0]
+    scale0 = scale * data['scale0'][:, None] if 'scale0' in data else scale
+    scale1 = scale * data['scale1'][:, None] if 'scale0' in data else scale
+    h0, w0, h1, w1 = map(lambda x: x // scale, [H0, W0, H1, W1])
+
+    # 2. warp grids
+    # create kpts in meshgrid and resize them to image resolution
+    grid_pt0_c = create_meshgrid(h0, w0, False, device).reshape(1, h0*w0, 2).repeat(N, 1, 1)    # [N, hw, 2]
+    grid_pt0_i = scale0 * grid_pt0_c
+    grid_pt1_c = create_meshgrid(h1, w1, False, device).reshape(1, h1*w1, 2).repeat(N, 1, 1)
+    grid_pt1_i = scale1 * grid_pt1_c
+
+    # mask padded region to (0, 0), so no need to manually mask conf_matrix_gt
+    if 'mask0' in data:
+        grid_pt0_i = mask_pts_at_padded_regions(grid_pt0_i, data['mask0'])
+        grid_pt1_i = mask_pts_at_padded_regions(grid_pt1_i, data['mask1'])
+
+    # warp kpts bi-directionally and resize them to coarse-level resolution
+    # (no depth consistency check, since it leads to worse results experimentally)
+    # (unhandled edge case: points with 0-depth will be warped to the left-up corner)
+    _, w_pt0_i = warp_kpts(grid_pt0_i, data['depth0'], data['depth1'], data['T_0to1'], data['K0'], data['K1'])
+    _, w_pt1_i = warp_kpts(grid_pt1_i, data['depth1'], data['depth0'], data['T_1to0'], data['K1'], data['K0'])
+    w_pt0_c = w_pt0_i / scale1
+    w_pt1_c = w_pt1_i / scale0
+
+    # 3. check if mutual nearest neighbor
+    w_pt0_c_round = w_pt0_c[:, :, :].round().long()
+    nearest_index1 = w_pt0_c_round[..., 0] + w_pt0_c_round[..., 1] * w1
+    w_pt1_c_round = w_pt1_c[:, :, :].round().long()
+    nearest_index0 = w_pt1_c_round[..., 0] + w_pt1_c_round[..., 1] * w0
+
+    # corner case: out of boundary
+    def out_bound_mask(pt, w, h):
+        return (pt[..., 0] < 0) + (pt[..., 0] >= w) + (pt[..., 1] < 0) + (pt[..., 1] >= h)
+    nearest_index1[out_bound_mask(w_pt0_c_round, w1, h1)] = 0
+    nearest_index0[out_bound_mask(w_pt1_c_round, w0, h0)] = 0
+
+    loop_back = torch.stack([nearest_index0[_b][_i] for _b, _i in enumerate(nearest_index1)], dim=0)
+    correct_0to1 = loop_back == torch.arange(h0*w0, device=device)[None].repeat(N, 1)
+    correct_0to1[:, 0] = False  # ignore the top-left corner
+
+    # estimate coarse-level gt matches using F.grid_sample
+    """w_grid0_c = w_pt0_c.reshape(N, h0, w0, 2)
+    w_grid0_c[..., 0] = (w_grid0_c[..., 0] / (w1 - 1) - 0.5) * 2
+    w_grid0_c[..., 1] = (w_grid0_c[..., 1] / (h1 - 1) - 0.5) * 2
+    w_grid1_c = w_pt1_c.reshape(N, h1, w1, 2).permute(0, 3, 1, 2)
+    reproj_grid0_c = F.grid_sample(w_grid1_c, w_grid0_c, mode='nearest', align_corners=True) # [N, 2, h0, w0]
+    reproj_grid0_c = reproj_grid0_c.reshape(N, 2, -1).permute(0, 2, 1)
+    correct_0to1 = ((reproj_grid0_c - grid_pt0_c)**2).sum(dim=-1).sqrt() < 1
+    correct_0to1[:, 0] = False"""
+
+    # 4. construct a gt conf_matrix
+    conf_matrix_gt = torch.zeros(N, h0*w0, h1*w1, device=device)
+    b_ids, i_ids = torch.where(correct_0to1 != 0)
+    j_ids = nearest_index1[b_ids, i_ids]
+
+    conf_matrix_gt[b_ids, i_ids, j_ids] = 1
+    data.update({'conf_matrix_gt': conf_matrix_gt})
+
+    # 5. save coarse matches(gt) for training fine level
+    if len(b_ids) == 0:
+        logger.warning(f"No groundtruth coarse match found for: {data['pair_names']}")
+        # this won't affect fine-level loss calculation
+        b_ids = torch.tensor([0], device=device)
+        i_ids = torch.tensor([0], device=device)
+        j_ids = torch.tensor([0], device=device)
+
+    data.update({
+        'spv_b_ids': b_ids,
+        'spv_i_ids': i_ids,
+        'spv_j_ids': j_ids
+    })
+
+    # 6. save intermediate results (for fast fine-level computation)
+    data.update({
+        'spv_w_pt0_i': w_pt0_i,
+        'spv_pt1_i': grid_pt1_i
+    })
+
+
+def compute_supervision_coarse(data, config):
+    assert len(set(data['dataset_name'])) == 1, "Do not support mixed datasets training!"
+    data_source = data['dataset_name'][0]
+    if data_source.lower() in ['scannet', 'megadepth']:
+        spvs_coarse(data, config)
+    else:
+        raise ValueError(f'Unknown data source: {data_source}')
+
+
+##############  ↓  Fine-Level supervision  ↓  ##############
+
+@torch.no_grad()
+def spvs_fine(data, config):
+    """
+    Update:
+        data (dict):{
+            "expec_f_gt": [M, 2]}
+    """
+    # 1. misc
+    # w_pt0_i, pt1_i = data.pop('spv_w_pt0_i'), data.pop('spv_pt1_i')
+    if config["MODEL"]["LOSS"]["FINE_TYPE"] in ['l2_with_std', 'l2']:
+        w_pt0_i, pt1_i = data['spv_w_pt0_i'], data['spv_pt1_i']
+        scale = config['MODEL']['RESOLUTION'][1]
+        radius = config['MODEL']['FINE_WINDOW_SIZE'] // 2
+
+        # 2. get coarse prediction
+        b_ids, i_ids, j_ids = data['b_ids'], data['i_ids'], data['j_ids']
+
+        # 3. compute gt
+        scale = scale * data['scale1'][b_ids] if 'scale0' in data else scale
+        # `expec_f_gt` might exceed the window, i.e. abs(*) > 1, which would be filtered later
+        expec_f_gt = (w_pt0_i[b_ids, i_ids] - pt1_i[b_ids, j_ids]) / scale / radius  # [M, 2]
+        data.update({"expec_f_gt": expec_f_gt})
+    else:
+        FMat = fundamental_from_projections(data["proj_mat0"], data["proj_mat1"])
+        data.update({"FMat_f": FMat[data["f_b_ids"]]})
+
+
+def compute_supervision_fine(data, config):
+    data_source = data['dataset_name'][0]
+    if data_source.lower() in ['scannet', 'megadepth']:
+        spvs_fine(data, config)
+    else:
+        raise NotImplementedError

vismatch/third_party/TopicFM/src/optimizers/__init__.py

@@ -0,0 +1,42 @@
+import torch
+from torch.optim.lr_scheduler import MultiStepLR, CosineAnnealingLR, ExponentialLR
+
+
+def build_optimizer(model, config):
+    name = config.TRAINER.OPTIMIZER
+    lr = config.TRAINER.TRUE_LR
+
+    if name == "adam":
+        return torch.optim.Adam(model.parameters(), lr=lr, weight_decay=config.TRAINER.ADAM_DECAY)
+    elif name == "adamw":
+        return torch.optim.AdamW(model.parameters(), lr=lr, weight_decay=config.TRAINER.ADAMW_DECAY)
+    else:
+        raise ValueError(f"TRAINER.OPTIMIZER = {name} is not a valid optimizer!")
+
+
+def build_scheduler(config, optimizer):
+    """
+    Returns:
+        scheduler (dict):{
+            'scheduler': lr_scheduler,
+            'interval': 'step',  # or 'epoch'
+            'monitor': 'val_f1', (optional)
+            'frequency': x, (optional)
+        }
+    """
+    scheduler = {'interval': config.TRAINER.SCHEDULER_INTERVAL}
+    name = config.TRAINER.SCHEDULER
+
+    if name == 'MultiStepLR':
+        scheduler.update(
+            {'scheduler': MultiStepLR(optimizer, config.TRAINER.MSLR_MILESTONES, gamma=config.TRAINER.MSLR_GAMMA)})
+    elif name == 'CosineAnnealing':
+        scheduler.update(
+            {'scheduler': CosineAnnealingLR(optimizer, config.TRAINER.COSA_TMAX)})
+    elif name == 'ExponentialLR':
+        scheduler.update(
+            {'scheduler': ExponentialLR(optimizer, config.TRAINER.ELR_GAMMA)})
+    else:
+        raise NotImplementedError()
+
+    return scheduler