vismatch 1.1.1__py3-none-any.whl

vismatch/im_models/silk.py

@@ -0,0 +1,405 @@
+"""
+Standalone SiLK (Self-supervised Interest point Learning with Keypoints) implementation.
+Reference: https://github.com/facebookresearch/silk
+Paper: https://arxiv.org/abs/2304.06194
+"""
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from kornia.color import rgb_to_grayscale
+from safetensors.torch import load_file
+
+from huggingface_hub import snapshot_download
+from vismatch import BaseMatcher
+
+
+def simple_nms(scores, radius=4):
+    """Non-maximum suppression on a 2D score map."""
+    kernel_size = 2 * radius + 1
+    max_scores = F.max_pool2d(scores, kernel_size=kernel_size, stride=1, padding=radius)
+    mask = scores == max_scores
+    return scores * mask.float()
+
+
+def extract_keypoints(logits, detection_threshold=0.0, nms_radius=4, top_k=10000, border_dist=4):
+    """Extract keypoints from detector logits.
+
+    Args:
+        logits (torch.Tensor): (B, 1, H, W) detector output
+        detection_threshold (float, optional): minimum score threshold. Defaults to 0.0.
+        nms_radius (int, optional): radius for non-maximum suppression. Defaults to 4.
+        top_k (int, optional): maximum number of keypoints to return. Defaults to 10000.
+        border_dist (int, optional): distance from border to exclude keypoints. Defaults to 4.
+
+    Returns:
+        list[torch.Tensor]: list of (N, 3) tensors with (row, col, score) for each batch item
+    """
+    B, _, H, W = logits.shape
+    scores = torch.sigmoid(logits)
+
+    # Apply NMS
+    scores = simple_nms(scores, radius=nms_radius)
+
+    # Remove border
+    if border_dist > 0:
+        scores[:, :, :border_dist, :] = 0
+        scores[:, :, -border_dist:, :] = 0
+        scores[:, :, :, :border_dist] = 0
+        scores[:, :, :, -border_dist:] = 0
+
+    keypoints_list = []
+    for b in range(B):
+        score_map = scores[b, 0]  # (H, W)
+
+        # Threshold
+        mask = score_map > detection_threshold
+
+        # Get coordinates and scores
+        coords = torch.nonzero(mask, as_tuple=False)  # (N, 2) - (row, col)
+        kpt_scores = score_map[mask]  # (N,)
+
+        # Sort by score and take top-k
+        if len(kpt_scores) > top_k:
+            _, indices = torch.topk(kpt_scores, top_k)
+            coords = coords[indices]
+            kpt_scores = kpt_scores[indices]
+
+        # Stack: (row, col, score)
+        if len(coords) > 0:
+            keypoints = torch.cat([coords.float(), kpt_scores.unsqueeze(1)], dim=1)
+        else:
+            keypoints = torch.zeros((0, 3), device=logits.device)
+
+        keypoints_list.append(keypoints)
+
+    return keypoints_list
+
+
+def sample_descriptors(descriptors, keypoints, mode="bilinear"):
+    """Sample descriptors at keypoint locations.
+
+    Args:
+        descriptors (torch.Tensor): (B, C, H, W) descriptor map
+        keypoints (list[torch.Tensor]): list of (N, 3) tensors with (row, col, score)
+        mode (str, optional): interpolation mode. Defaults to "bilinear".
+
+    Returns:
+        list[torch.Tensor]: list of (N, C) descriptor tensors
+    """
+    B, C, H, W = descriptors.shape
+    desc_list = []
+
+    for b in range(B):
+        kpts = keypoints[b][:, :2]  # (N, 2) - (row, col)
+
+        if len(kpts) == 0:
+            desc_list.append(torch.zeros((0, C), device=descriptors.device))
+            continue
+
+        # Normalize to [-1, 1] for grid_sample
+        # grid_sample expects (x, y) in [-1, 1] where x is width, y is height
+        grid_x = 2.0 * kpts[:, 1] / (W - 1) - 1.0  # col -> x
+        grid_y = 2.0 * kpts[:, 0] / (H - 1) - 1.0  # row -> y
+        grid = torch.stack([grid_x, grid_y], dim=1).view(1, -1, 1, 2)  # (1, N, 1, 2)
+
+        # Sample
+        sampled = F.grid_sample(
+            descriptors[b : b + 1], grid, mode=mode, align_corners=True, padding_mode="border"
+        )  # (1, C, N, 1)
+        sampled = sampled.squeeze(0).squeeze(-1).T  # (N, C)
+
+        # Normalize
+        sampled = F.normalize(sampled, p=2, dim=1)
+        desc_list.append(sampled)
+
+    return desc_list
+
+
+def match_descriptors_mnn(desc0, desc1, threshold=0.8, mode="ratio-test"):
+    """Match descriptors using mutual nearest neighbors with optional ratio test.
+
+    Args:
+        desc0 (torch.Tensor): (N, C) descriptors from image 0
+        desc1 (torch.Tensor): (M, C) descriptors from image 1
+        threshold (float, optional): ratio threshold (for ratio-test) or distance threshold. Defaults to 0.8.
+        mode (str, optional): "ratio-test", "mnn", or "double-softmax". Defaults to "ratio-test".
+
+    Returns:
+        torch.Tensor: (K, 2) tensor of match indices
+    """
+    if len(desc0) == 0 or len(desc1) == 0:
+        return torch.zeros((0, 2), dtype=torch.long, device=desc0.device)
+
+    # Compute cosine similarity (descriptors should be normalized)
+    sim = desc0 @ desc1.T  # (N, M)
+
+    if mode == "ratio-test":
+        # Forward matches with ratio test
+        sorted_sim0, indices0 = sim.topk(min(2, sim.shape[1]), dim=1)
+        nn0 = indices0[:, 0]
+        if sim.shape[1] >= 2:
+            ratio0 = sorted_sim0[:, 0] / (sorted_sim0[:, 1] + 1e-8)
+            valid0 = ratio0 > (1.0 / threshold)  # Higher ratio = better match
+        else:
+            valid0 = torch.ones(len(nn0), dtype=torch.bool, device=desc0.device)
+
+        # Backward matches with ratio test
+        sorted_sim1, indices1 = sim.topk(min(2, sim.shape[0]), dim=0)
+        nn1 = indices1[0]
+        if sim.shape[0] >= 2:
+            ratio1 = sorted_sim1[0] / (sorted_sim1[1] + 1e-8)
+            valid1 = ratio1 > (1.0 / threshold)
+        else:
+            valid1 = torch.ones(len(nn1), dtype=torch.bool, device=desc1.device)
+
+        # Mutual check
+        ids0 = torch.arange(len(nn0), device=desc0.device)
+        mutual = (nn1[nn0] == ids0) & valid0 & valid1[nn0]
+
+        matches = torch.stack([ids0[mutual], nn0[mutual]], dim=1)
+
+    elif mode == "mnn":
+        # Simple mutual nearest neighbor
+        nn0 = sim.argmax(dim=1)  # (N,)
+        nn1 = sim.argmax(dim=0)  # (M,)
+
+        ids0 = torch.arange(len(nn0), device=desc0.device)
+        mutual = nn1[nn0] == ids0
+
+        matches = torch.stack([ids0[mutual], nn0[mutual]], dim=1)
+
+    elif mode == "double-softmax":
+        # Double softmax matching
+        temperature = 0.1
+        prob0 = F.softmax(sim / temperature, dim=1)
+        prob1 = F.softmax(sim / temperature, dim=0)
+        prob = prob0 * prob1
+
+        # Get best matches
+        max_prob, nn0 = prob.max(dim=1)
+        valid = max_prob > threshold
+
+        ids0 = torch.arange(len(nn0), device=desc0.device)
+        matches = torch.stack([ids0[valid], nn0[valid]], dim=1)
+    else:
+        raise ValueError(f"Unknown matching mode: {mode}")
+
+    return matches
+
+
+class VGGBlock(nn.Module):
+    """A single VGG-style conv block: Conv -> ReLU -> BatchNorm."""
+
+    def __init__(self, in_channels, out_channels, kernel_size=3):
+        super().__init__()
+        padding = kernel_size // 2
+        self.conv = nn.Conv2d(in_channels, out_channels, kernel_size, padding=padding, bias=True)
+        self.bn = nn.BatchNorm2d(out_channels)
+
+    def forward(self, x):
+        return self.bn(F.relu(self.conv(x)))
+
+
+class VGGBackbone(nn.Module):
+    """VGG-style backbone for SiLK."""
+
+    def __init__(self):
+        super().__init__()
+        # Layer 0: 1 -> 64
+        self.layer0 = nn.Sequential(
+            VGGBlock(1, 64),
+            VGGBlock(64, 64),
+        )
+        # Layer 1: 64 -> 64
+        self.layer1 = nn.Sequential(
+            VGGBlock(64, 64),
+            VGGBlock(64, 64),
+        )
+        # Layer 2: 64 -> 128
+        self.layer2 = nn.Sequential(
+            VGGBlock(64, 128),
+            VGGBlock(128, 128),
+        )
+        # Layer 3: 128 -> 128
+        self.layer3 = nn.Sequential(
+            VGGBlock(128, 128),
+            VGGBlock(128, 128),
+        )
+
+    def forward(self, x):
+        x = self.layer0(x)
+        x = self.layer1(x)
+        x = self.layer2(x)
+        x = self.layer3(x)
+        return x
+
+
+class DetectorHead(nn.Module):
+    """Detector head: outputs logits for keypoint detection."""
+
+    def __init__(self, in_channels=128):
+        super().__init__()
+        self.conv1 = nn.Conv2d(in_channels, 128, 3, padding=1, bias=True)
+        self.bn1 = nn.BatchNorm2d(128)
+        self.conv2 = nn.Conv2d(128, 1, 1, bias=True)
+        self.bn2 = nn.BatchNorm2d(1)
+
+    def forward(self, x):
+        x = self.bn1(F.relu(self.conv1(x)))
+        x = self.bn2(self.conv2(x))
+        return x
+
+
+class DescriptorHead(nn.Module):
+    """Descriptor head: outputs dense descriptors."""
+
+    def __init__(self, in_channels=128, out_channels=128):
+        super().__init__()
+        self.conv1 = nn.Conv2d(in_channels, 128, 3, padding=1, bias=True)
+        self.bn1 = nn.BatchNorm2d(128)
+        self.conv2 = nn.Conv2d(128, out_channels, 1, bias=True)
+        self.bn2 = nn.BatchNorm2d(out_channels)
+
+    def forward(self, x):
+        x = self.bn1(F.relu(self.conv1(x)))
+        x = self.bn2(self.conv2(x))
+        return x
+
+
+class SiLKModel(nn.Module):
+    """
+    Self-supervised Interest point Learning with Keypoints (SiLK).
+    """
+
+    def __init__(self, descriptor_dim=128):
+        super().__init__()
+        self.backbone = VGGBackbone()
+        self.detector_head = DetectorHead(128)
+        self.descriptor_head = DescriptorHead(128, descriptor_dim)
+        self.descriptor_scale_factor = nn.Parameter(torch.tensor(1.0))
+
+    def forward(self, x):
+        """Run SiLK model forward pass.
+
+        Args:
+            x (torch.Tensor): (B, 1, H, W) grayscale image
+
+        Returns:
+            tuple: (logits, descriptors) where:
+                - logits (torch.Tensor): (B, 1, H, W) detector logits
+                - descriptors (torch.Tensor): (B, C, H, W) descriptor map
+        """
+        features = self.backbone(x)
+        logits = self.detector_head(features)
+        descriptors = self.descriptor_head(features)
+        descriptors = descriptors * self.descriptor_scale_factor
+        return logits, descriptors
+
+
+def load_silk_model(weights_path, device="cpu"):
+    """Load SiLK model from safetensors."""
+
+    model = SiLKModel(descriptor_dim=128)
+    state_dict = load_file(weights_path)
+    model.load_state_dict(state_dict)
+    model = model.to(device).eval()
+    return model
+
+
+class SilkMatcher(BaseMatcher):
+    """
+    SiLK matcher using standalone implementation.
+    Reference: https://github.com/facebookresearch/silk
+    """
+
+    MATCHER_POSTPROCESS_OPTIONS = ["ratio-test", "mnn", "double-softmax"]
+
+    def __init__(
+        self,
+        device="cpu",
+        max_num_keypoints=10000,
+        detection_threshold=0.0,
+        nms_radius=4,
+        border_dist=4,
+        matcher_post_processing="ratio-test",
+        matcher_thresh=0.8,
+        **kwargs,
+    ):
+        super().__init__(device, **kwargs)
+        assert self.device != "mps", (
+            f"Device must be 'cpu' or 'cuda', 'mps' not yet supported for {self.name}. Device = {self.device}"
+        )
+
+        self.max_num_keypoints = max_num_keypoints
+        self.detection_threshold = detection_threshold
+        self.nms_radius = nms_radius
+        self.border_dist = border_dist
+
+        assert matcher_post_processing in SilkMatcher.MATCHER_POSTPROCESS_OPTIONS, (
+            f"Matcher postprocessing must be one of {SilkMatcher.MATCHER_POSTPROCESS_OPTIONS}"
+        )
+        self.matcher_mode = matcher_post_processing
+        self.matcher_thresh = matcher_thresh
+
+        # Load model
+        weights_path = f"{snapshot_download('vismatch/silk')}/silk.safetensors"
+        self.model = load_silk_model(weights_path, device)
+
+    def preprocess(self, img):
+        """Convert to grayscale and add batch dimension if needed."""
+        if img.ndim == 3:
+            img = img.unsqueeze(0)
+        return rgb_to_grayscale(img)
+
+    def _forward(self, img0, img1):
+        img0 = self.preprocess(img0)
+        img1 = self.preprocess(img1)
+
+        # Extract features
+        with torch.no_grad():
+            logits0, descriptors0 = self.model(img0)
+            logits1, descriptors1 = self.model(img1)
+
+        # Extract keypoints
+        kpts0_list = extract_keypoints(
+            logits0,
+            detection_threshold=self.detection_threshold,
+            nms_radius=self.nms_radius,
+            top_k=self.max_num_keypoints,
+            border_dist=self.border_dist,
+        )
+        kpts1_list = extract_keypoints(
+            logits1,
+            detection_threshold=self.detection_threshold,
+            nms_radius=self.nms_radius,
+            top_k=self.max_num_keypoints,
+            border_dist=self.border_dist,
+        )
+
+        # Sample descriptors at keypoint locations
+        desc0_list = sample_descriptors(descriptors0, kpts0_list)
+        desc1_list = sample_descriptors(descriptors1, kpts1_list)
+
+        # Get keypoints and descriptors for first image in batch
+        kpts0 = kpts0_list[0]  # (N, 3) - row, col, score
+        kpts1 = kpts1_list[0]  # (M, 3)
+        desc0 = desc0_list[0]  # (N, C)
+        desc1 = desc1_list[0]  # (M, C)
+
+        # Match descriptors
+        matches = match_descriptors_mnn(desc0, desc1, threshold=self.matcher_thresh, mode=self.matcher_mode)
+
+        # Get matched keypoints - convert from (row, col) to (x, y) = (col, row)
+        if len(matches) > 0:
+            mkpts0 = kpts0[matches[:, 0], :2][:, [1, 0]]  # (row, col) -> (col, row) = (x, y)
+            mkpts1 = kpts1[matches[:, 1], :2][:, [1, 0]]
+        else:
+            mkpts0 = torch.zeros((0, 2), device=self.device)
+            mkpts1 = torch.zeros((0, 2), device=self.device)
+
+        # Convert all keypoints to (x, y) format
+        all_kpts0 = kpts0[:, :2][:, [1, 0]] if len(kpts0) > 0 else torch.zeros((0, 2), device=self.device)
+        all_kpts1 = kpts1[:, :2][:, [1, 0]] if len(kpts1) > 0 else torch.zeros((0, 2), device=self.device)
+
+        return mkpts0, mkpts1, all_kpts0, all_kpts1, desc0, desc1

vismatch/im_models/sphereglue.py

@@ -0,0 +1,97 @@
+import torch
+from safetensors.torch import load_file
+
+from vismatch.im_models.lightglue import SIFT, SuperPoint
+from vismatch.utils import add_to_path
+from huggingface_hub import snapshot_download
+from vismatch import THIRD_PARTY_DIR, BaseMatcher
+
+add_to_path(THIRD_PARTY_DIR.joinpath("SphereGlue"))
+
+from model.sphereglue import SphereGlue
+from utils.Utils import sphericalToCartesian
+
+
+def unit_cartesian(points):
+    phi, theta = torch.split(torch.as_tensor(points), 1, dim=1)
+    unitCartesian = sphericalToCartesian(phi, theta, 1).squeeze(dim=2)
+    return unitCartesian
+
+
+class SphereGlueBase(BaseMatcher):
+    """
+    This class is the parent for all methods that use LightGlue as a matcher,
+    with different local features. It implements the forward which is the same
+    regardless of the feature extractor of choice.
+    Therefore this class should *NOT* be instatiated, as it needs its children to define
+    the extractor and the matcher.
+    """
+
+    def __init__(self, device="cpu", **kwargs):
+        super().__init__(device, **kwargs)
+        self.sphereglue_cfg = {
+            "K": kwargs.get("K", 2),
+            "GNN_layers": kwargs.get("GNN_layers", ["cross"]),
+            "match_threshold": kwargs.get("match_threshold", 0.2),
+            "sinkhorn_iterations": kwargs.get("sinkhorn_iterations", 20),
+            "aggr": kwargs.get("aggr", "add"),
+            "knn": kwargs.get("knn", 20),
+        }
+
+        self.skip_ransac = True
+
+    def _forward(self, img0, img1):
+        """
+        "extractor" and "matcher" are instantiated by the subclasses.
+        """
+        feats0 = self.extractor.extract(img0)
+        feats1 = self.extractor.extract(img1)
+
+        unit_cartesian1 = unit_cartesian(feats0["keypoints"][0]).unsqueeze(dim=0).to(self.device)
+        unit_cartesian2 = unit_cartesian(feats1["keypoints"][0]).unsqueeze(dim=0).to(self.device)
+
+        inputs = {
+            "h1": feats0["descriptors"],
+            "h2": feats1["descriptors"],
+            "scores1": feats0["keypoint_scores"],
+            "scores2": feats1["keypoint_scores"],
+            "unitCartesian1": unit_cartesian1,
+            "unitCartesian2": unit_cartesian2,
+        }
+        outputs = self.matcher(inputs)
+
+        kpts0, kpts1, matches = (
+            feats0["keypoints"].squeeze(dim=0),
+            feats1["keypoints"].squeeze(dim=0),
+            outputs["matches0"].squeeze(dim=0),
+        )
+        desc0 = feats0["descriptors"].squeeze(dim=0)
+        desc1 = feats1["descriptors"].squeeze(dim=0)
+
+        mask = matches.ge(0)
+        kpts0_idx = torch.masked_select(torch.arange(matches.shape[0]).to(mask.device), mask)
+        kpts1_idx = torch.masked_select(matches, mask)
+        mkpts0 = kpts0[kpts0_idx]
+        mkpts1 = kpts1[kpts1_idx]
+
+        return mkpts0, mkpts1, kpts0, kpts1, desc0, desc1
+
+
+class SiftSphereGlue(SphereGlueBase):
+    def __init__(self, device="cpu", max_num_keypoints=2048, *args, **kwargs):
+        super().__init__(device, **kwargs)
+        self.sphereglue_cfg.update({"descriptor_dim": 128, "output_dim": 128 * 2, "max_kpts": max_num_keypoints})
+        self.extractor = SIFT(max_num_keypoints=max_num_keypoints).eval().to(self.device)
+        self.matcher = SphereGlue(config=self.sphereglue_cfg).to(self.device)
+        weights_path = f"{snapshot_download('vismatch/sift-sphereglue')}/model.safetensors"
+        self.matcher.load_state_dict(load_file(weights_path))
+
+
+class SuperpointSphereGlue(SphereGlueBase):
+    def __init__(self, device="cpu", max_num_keypoints=2048, *args, **kwargs):
+        super().__init__(device, **kwargs)
+        self.sphereglue_cfg.update({"descriptor_dim": 256, "output_dim": 256 * 2, "max_kpts": max_num_keypoints})
+        self.extractor = SuperPoint(max_num_keypoints=max_num_keypoints).eval().to(self.device)
+        self.matcher = SphereGlue(config=self.sphereglue_cfg).to(self.device)
+        weights_path = f"{snapshot_download('vismatch/superpoint-sphereglue')}/model.safetensors"
+        self.matcher.load_state_dict(load_file(weights_path))

vismatch/im_models/steerers.py

@@ -0,0 +1,140 @@
+import torch
+import torchvision.transforms as tfm
+from huggingface_hub import snapshot_download
+from vismatch import BaseMatcher, THIRD_PARTY_DIR
+from vismatch.utils import resize_to_divisible, add_to_path
+
+
+add_to_path(THIRD_PARTY_DIR.joinpath("DeDoDe"))
+from DeDoDe import (
+    dedode_detector_L,
+    dedode_descriptor_B,
+)
+
+add_to_path(THIRD_PARTY_DIR.joinpath("Steerers"))
+from rotation_steerers.steerers import DiscreteSteerer, ContinuousSteerer
+from rotation_steerers.matchers.max_similarity import (
+    MaxSimilarityMatcher,
+    ContinuousMaxSimilarityMatcher,
+)
+
+
+class SteererMatcher(BaseMatcher):
+    dino_patch_size = 14
+
+    def __init__(
+        self,
+        device="cpu",
+        max_num_keypoints=2048,
+        dedode_thresh=0.05,
+        steerer_type="C8",
+        *args,
+        **kwargs,
+    ):
+        super().__init__(device, **kwargs)
+        assert "cuda" in self.device, f"Device must be 'cuda' for {self.name}. Device='{self.device}' not supported"
+
+        # Download weights from HuggingFace Hub
+        repo = snapshot_download("vismatch/steerers")
+        self.detector_path_L = f"{repo}/dedode_detector_L.pth"
+        self.descriptor_path_G = f"{repo}/dedode_descriptor_G.pth"
+        self.descriptor_path_B_C4 = f"{repo}/B_C4_Perm_descriptor_setting_C.pth"
+        self.descriptor_path_B_SO2 = f"{repo}/B_SO2_Spread_descriptor_setting_B.pth"
+        self.steerer_path_C = f"{repo}/B_C4_Perm_steerer_setting_C.pth"
+        self.steerer_path_B = f"{repo}/B_SO2_Spread_steerer_setting_B.pth"
+
+        self.max_keypoints = max_num_keypoints
+        self.threshold = dedode_thresh
+
+        self.normalize = tfm.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
+
+        self.detector, self.descriptor, self.steerer, self.matcher = self.build_matcher(steerer_type, device=device)
+
+    def build_matcher(self, steerer_type="C8", device="cpu"):
+        if steerer_type == "C4":
+            detector = dedode_detector_L(
+                weights=torch.load(self.detector_path_L, map_location=device, weights_only=True)
+            )
+            descriptor = dedode_descriptor_B(
+                weights=torch.load(self.descriptor_path_B_C4, map_location=device, weights_only=True)
+            )
+            steerer = DiscreteSteerer(generator=torch.load(self.steerer_path_C, map_location=device, weights_only=True))
+            steerer_order = 4
+        elif steerer_type == "C8":
+            detector = dedode_detector_L(
+                weights=torch.load(self.detector_path_L, map_location=device, weights_only=True)
+            )
+            descriptor = dedode_descriptor_B(
+                weights=torch.load(self.descriptor_path_B_SO2, map_location=device, weights_only=True)
+            )
+            steerer_order = 8
+            steerer = DiscreteSteerer(
+                generator=torch.matrix_exp(
+                    (2 * 3.14159 / steerer_order)
+                    * torch.load(self.steerer_path_B, map_location=device, weights_only=True)
+                )
+            )
+
+        elif steerer_type == "S02":
+            descriptor = dedode_descriptor_B(
+                weights=torch.load(self.descriptor_path_B_SO2, map_location=device, weights_only=True)
+            )
+            steerer = ContinuousSteerer(
+                generator=torch.load(self.steerer_path_B, map_location=device, weights_only=True)
+            )
+
+        else:
+            print(f"Steerer type {steerer_type} not yet implemented")
+
+        if steerer_type == "SO2":
+            matcher = ContinuousMaxSimilarityMatcher(steerer=steerer, angles=[0.2, 1.2879, 3.14])
+        else:
+            matcher = MaxSimilarityMatcher(steerer=steerer, steerer_order=steerer_order)
+
+        return detector, descriptor, steerer, matcher
+
+    def preprocess(self, img):
+        # ensure that the img has the proper w/h to be compatible with patch sizes
+        _, h, w = img.shape
+        orig_shape = h, w
+        img = resize_to_divisible(img, self.dino_patch_size)
+
+        img = self.normalize(img).unsqueeze(0).to(self.device)
+        return img, orig_shape
+
+    def _forward(self, img0, img1):
+        img0, img0_orig_shape = self.preprocess(img0)
+        img1, img1_orig_shape = self.preprocess(img1)
+
+        batch_0 = {"image": img0}
+        detections_0 = self.detector.detect(batch_0, num_keypoints=self.max_keypoints)
+        keypoints_0, P_0 = detections_0["keypoints"], detections_0["confidence"]
+
+        batch_1 = {"image": img1}
+        detections_1 = self.detector.detect(batch_1, num_keypoints=self.max_keypoints)
+        keypoints_1, P_1 = detections_1["keypoints"], detections_1["confidence"]
+
+        description_0 = self.descriptor.describe_keypoints(batch_0, keypoints_0)["descriptions"]
+        description_1 = self.descriptor.describe_keypoints(batch_1, keypoints_1)["descriptions"]
+
+        matches_0, matches_1, _ = self.matcher.match(
+            keypoints_0,
+            description_0,
+            keypoints_1,
+            description_1,
+            P_A=P_0,
+            P_B=P_1,
+            normalize=True,
+            inv_temp=20,
+            threshold=self.threshold,  # Increasing threshold -> fewer matches, fewer outliers
+        )
+
+        H0, W0, H1, W1 = *img0.shape[-2:], *img1.shape[-2:]
+        mkpts0, mkpts1 = self.matcher.to_pixel_coords(matches_0, matches_1, H0, W0, H1, W1)
+
+        # dedode sometimes requires reshaping an image to fit vit patch size evenly, so we need to
+        # rescale kpts to the original img
+        mkpts0 = self.rescale_coords(mkpts0, *img0_orig_shape, H0, W0)
+        mkpts1 = self.rescale_coords(mkpts1, *img1_orig_shape, H1, W1)
+
+        return mkpts0, mkpts1, keypoints_0[0], keypoints_1[0], description_0[0], description_1[0]