vismatch 1.1.1__py3-none-any.whl

vismatch/third_party/DeDoDe/DeDoDe/descriptors/dedode_descriptor.py

@@ -0,0 +1,50 @@
+import torch
+from PIL import Image
+import torch.nn as nn
+import torchvision.models as tvm
+import torch.nn.functional as F
+import numpy as np
+from DeDoDe.utils import get_best_device
+
+class DeDoDeDescriptor(nn.Module):
+    def __init__(self, encoder, decoder, *args, **kwargs) -> None:
+        super().__init__(*args, **kwargs)
+        self.encoder = encoder
+        self.decoder = decoder
+        import torchvision.transforms as transforms
+        self.normalizer = transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
+        
+    def forward(
+        self,
+        batch,
+    ):
+        if "im_A" in batch:
+            images = torch.cat((batch["im_A"], batch["im_B"]))
+        else:
+            images = batch["image"]
+        features, sizes = self.encoder(images)
+        descriptor = 0
+        context = None
+        scales = self.decoder.scales
+        for idx, (feature_map, scale) in enumerate(zip(reversed(features), scales)):
+            delta_descriptor, context = self.decoder(feature_map, scale = scale, context = context)
+            descriptor = descriptor + delta_descriptor
+            if idx < len(scales) - 1:
+                size = sizes[-(idx+2)]
+                descriptor = F.interpolate(descriptor, size = size, mode = "bilinear", align_corners = False)
+                context = F.interpolate(context, size = size, mode = "bilinear", align_corners = False)
+        return {"description_grid" : descriptor}
+    
+    @torch.inference_mode()
+    def describe_keypoints(self, batch, keypoints):
+        self.train(False)
+        description_grid = self.forward(batch)["description_grid"]
+        described_keypoints = F.grid_sample(description_grid.float(), keypoints[:,None], mode = "bilinear", align_corners = False)[:,:,0].mT
+        return {"descriptions": described_keypoints}
+    
+    def read_image(self, im_path, H = 784, W = 784, device=get_best_device()):
+        return self.normalizer(torch.from_numpy(np.array(Image.open(im_path).resize((W,H)))/255.).permute(2,0,1)).float().to(device)[None]
+
+    def describe_keypoints_from_path(self, im_path, keypoints, H = 784, W = 784):
+        batch = {"image": self.read_image(im_path, H = H, W = W)}
+        return self.describe_keypoints(batch, keypoints)

vismatch/third_party/DeDoDe/DeDoDe/descriptors/descriptor_loss.py

@@ -0,0 +1,68 @@
+import torch
+import torch.nn as nn
+import math
+import torch.nn.functional as F
+
+from DeDoDe.utils import *
+import DeDoDe
+
+class DescriptorLoss(nn.Module):
+    
+    def __init__(self, detector, num_keypoints = 5000, normalize_descriptions = False, inv_temp = 1, device = get_best_device()) -> None:
+        super().__init__()
+        self.detector = detector
+        self.tracked_metrics = {}
+        self.num_keypoints = num_keypoints
+        self.normalize_descriptions = normalize_descriptions
+        self.inv_temp = inv_temp
+    
+    def warp_from_depth(self, batch, kpts_A, kpts_B):
+        mask_A_to_B, kpts_A_to_B = warp_kpts(kpts_A, 
+                    batch["im_A_depth"],
+                    batch["im_B_depth"],
+                    batch["T_1to2"],
+                    batch["K1"],
+                    batch["K2"],)
+        mask_B_to_A, kpts_B_to_A = warp_kpts(kpts_B, 
+                    batch["im_B_depth"],
+                    batch["im_A_depth"],
+                    batch["T_1to2"].inverse(),
+                    batch["K2"],
+                    batch["K1"],)
+        return (mask_A_to_B, kpts_A_to_B), (mask_B_to_A, kpts_B_to_A)
+    
+    def warp_from_homog(self, batch, kpts_A, kpts_B):
+        kpts_A_to_B = homog_transform(batch["Homog_A_to_B"], kpts_A)
+        kpts_B_to_A = homog_transform(batch["Homog_A_to_B"].inverse(), kpts_B)
+        return (None, kpts_A_to_B), (None, kpts_B_to_A)
+
+    def supervised_loss(self, outputs, batch):
+        kpts_A, kpts_B = self.detector.detect(batch, num_keypoints = self.num_keypoints)['keypoints'].clone().chunk(2)
+        desc_grid_A, desc_grid_B = outputs["description_grid"].chunk(2)
+        desc_A = F.grid_sample(desc_grid_A.float(), kpts_A[:,None], mode = "bilinear", align_corners = False)[:,:,0].mT
+        desc_B = F.grid_sample(desc_grid_B.float(), kpts_B[:,None], mode = "bilinear", align_corners = False)[:,:,0].mT
+        if "im_A_depth" in batch:
+            (mask_A_to_B, kpts_A_to_B), (mask_B_to_A, kpts_B_to_A) = self.warp_from_depth(batch, kpts_A, kpts_B)
+        elif "Homog_A_to_B" in batch:
+            (mask_A_to_B, kpts_A_to_B), (mask_B_to_A, kpts_B_to_A) = self.warp_from_homog(batch, kpts_A, kpts_B)
+            
+        with torch.no_grad():
+            D_B = torch.cdist(kpts_A_to_B, kpts_B)
+            D_A = torch.cdist(kpts_A, kpts_B_to_A)
+            inds = torch.nonzero((D_B == D_B.min(dim=-1, keepdim = True).values) 
+                                 * (D_A == D_A.min(dim=-2, keepdim = True).values)
+                                 * (D_B < 0.01)
+                                 * (D_A < 0.01))
+            
+        logP_A_B = dual_log_softmax_matcher(desc_A, desc_B, 
+                                            normalize = self.normalize_descriptions,
+                                            inv_temperature = self.inv_temp)
+        neg_log_likelihood = -logP_A_B[inds[:,0], inds[:,1], inds[:,2]].mean()
+        self.tracked_metrics["neg_log_likelihood"] = (0.99 * self.tracked_metrics.get("neg_log_likelihood", neg_log_likelihood.detach().item()) + 0.01 * neg_log_likelihood.detach().item())
+        if np.random.rand() > 0.99:
+            print(self.tracked_metrics["neg_log_likelihood"])
+        return neg_log_likelihood
+    
+    def forward(self, outputs, batch):
+        losses = self.supervised_loss(outputs, batch)
+        return losses

vismatch/third_party/DeDoDe/DeDoDe/detectors/dedode_detector.py

@@ -0,0 +1,76 @@
+import torch
+from PIL import Image
+import torch.nn as nn
+import torchvision.models as tvm
+import torch.nn.functional as F
+import numpy as np
+
+from DeDoDe.utils import sample_keypoints, to_pixel_coords, to_normalized_coords, get_best_device
+
+
+
+class DeDoDeDetector(nn.Module):
+    def __init__(self, encoder, decoder, *args, remove_borders = False, **kwargs) -> None:
+        super().__init__(*args, **kwargs)
+        self.encoder = encoder
+        self.decoder = decoder
+        import torchvision.transforms as transforms
+        self.normalizer = transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
+        self.remove_borders = remove_borders
+        
+    def forward(
+        self,
+        batch,
+    ):
+        if "im_A" in batch:
+            images = torch.cat((batch["im_A"], batch["im_B"]))
+        else:
+            images = batch["image"]
+        features, sizes = self.encoder(images)
+        logits = 0
+        context = None
+        scales = ["8", "4", "2", "1"]
+        for idx, (feature_map, scale) in enumerate(zip(reversed(features), scales)):
+            delta_logits, context = self.decoder(feature_map, context = context, scale = scale)
+            logits = logits + delta_logits.float() # ensure float (need bf16 doesnt have f.interpolate)
+            if idx < len(scales) - 1:
+                size = sizes[-(idx+2)]
+                logits = F.interpolate(logits, size = size, mode = "bicubic", align_corners = False)
+                context = F.interpolate(context.float(), size = size, mode = "bilinear", align_corners = False)
+        return {"keypoint_logits" : logits.float()}
+    
+    @torch.inference_mode()
+    def detect(self, batch, num_keypoints = 10_000):
+        self.train(False)
+        keypoint_logits = self.forward(batch)["keypoint_logits"]
+        B,K,H,W = keypoint_logits.shape
+        keypoint_p = keypoint_logits.reshape(B, K*H*W).softmax(dim=-1).reshape(B, K, H*W).sum(dim=1)
+        keypoints, confidence = sample_keypoints(keypoint_p.reshape(B,H,W), 
+                                  use_nms = False, sample_topk = True, num_samples = num_keypoints, 
+                                  return_scoremap=True, sharpen = False, upsample = False,
+                                  increase_coverage=True, remove_borders = self.remove_borders)
+        return {"keypoints": keypoints, "confidence": confidence}
+
+    @torch.inference_mode()
+    def detect_dense(self, batch):
+        self.train(False)
+        keypoint_logits = self.forward(batch)["keypoint_logits"]
+        return {"dense_keypoint_logits": keypoint_logits}
+
+    def read_image(self, im_path, H = 784, W = 784, device=get_best_device()):
+        pil_im = Image.open(im_path).resize((W, H))
+        standard_im = np.array(pil_im)/255.
+        return self.normalizer(torch.from_numpy(standard_im).permute(2,0,1)).float().to(device)[None]
+
+    def detect_from_path(self, im_path, num_keypoints = 30_000, H = 784, W = 784, dense = False):
+        batch = {"image": self.read_image(im_path, H = H, W = W)}
+        if dense:
+            return self.detect_dense(batch)
+        else:
+            return self.detect(batch, num_keypoints = num_keypoints)
+
+    def to_pixel_coords(self, x, H, W):
+        return to_pixel_coords(x, H, W)
+    
+    def to_normalized_coords(self, x, H, W):
+        return to_normalized_coords(x, H, W)

vismatch/third_party/DeDoDe/DeDoDe/detectors/keypoint_loss.py

@@ -0,0 +1,185 @@
+import torch
+import torch.nn as nn
+import math
+
+from DeDoDe.utils import *
+import DeDoDe
+
+class KeyPointLoss(nn.Module):
+    
+    def __init__(self, smoothing_size = 1, use_max_logit = False, entropy_target = 80, 
+                 num_matches = 1024, jacobian_density_adjustment = False,
+                 matchability_weight = 1, device = "cuda") -> None:
+        super().__init__()
+        X = torch.linspace(-1,1,smoothing_size, device = device)
+        G = (-X**2 / (2 *1/2**2)).exp()
+        G = G/G.sum()
+        self.use_max_logit = use_max_logit
+        self.entropy_target = entropy_target
+        self.smoothing_kernel = G[None, None, None,:]
+        self.smoothing_size = smoothing_size
+        self.tracked_metrics = {}
+        self.center = None
+        self.num_matches = num_matches
+        self.jacobian_density_adjustment = jacobian_density_adjustment
+        self.matchability_weight = matchability_weight
+        
+    def compute_consistency(self, logits_A, logits_B_to_A, mask = None):
+        
+        masked_logits_A = torch.full_like(logits_A, -torch.inf)
+        masked_logits_A[mask] = logits_A[mask]
+
+        masked_logits_B_to_A = torch.full_like(logits_B_to_A, -torch.inf)
+        masked_logits_B_to_A[mask] = logits_B_to_A[mask]
+
+        log_p_A = masked_logits_A.log_softmax(dim=-1)[mask]
+        log_p_B_to_A = masked_logits_B_to_A.log_softmax(dim=-1)[mask]
+
+        return self.compute_jensen_shannon_div(log_p_A, log_p_B_to_A)
+    
+    def compute_joint_neg_log_likelihood(self, logits_A, logits_B_to_A, detections_A = None, detections_B_to_A = None, mask = None, device = "cuda", dtype = torch.float32, num_matches = None):
+        B, K, HW = logits_A.shape
+        logits_A, logits_B_to_A = logits_A.to(dtype), logits_B_to_A.to(dtype)
+        mask = mask[:,None].expand(B, K, HW).reshape(B, K*HW)
+        log_p_B_to_A = self.masked_log_softmax(logits_B_to_A.reshape(B,K*HW), mask = mask)
+        log_p_A = self.masked_log_softmax(logits_A.reshape(B,K*HW), mask = mask)
+        log_p = log_p_A + log_p_B_to_A
+        if detections_A is None:
+            detections_A = torch.zeros_like(log_p_A)
+        if detections_B_to_A is None:
+            detections_B_to_A = torch.zeros_like(log_p_B_to_A)
+        detections_A = detections_A.reshape(B, HW)
+        detections_A[~mask] = 0         
+        detections_B_to_A = detections_B_to_A.reshape(B, HW)
+        detections_B_to_A[~mask] = 0
+        log_p_target = log_p.detach() + 50*detections_A + 50*detections_B_to_A
+        num_matches = self.num_matches if num_matches is None else num_matches
+        best_k = -(-log_p_target).flatten().kthvalue(k = B * num_matches, dim=-1).values
+        p_target = (log_p_target > best_k[..., None]).float().reshape(B,K*HW)/num_matches
+        return self.compute_cross_entropy(log_p_A[mask], p_target[mask]) + self.compute_cross_entropy(log_p_B_to_A[mask], p_target[mask])
+                
+    def compute_jensen_shannon_div(self, log_p, log_q):
+        return 1/2 * (self.compute_kl_div(log_p, log_q) + self.compute_kl_div(log_q, log_p))
+    
+    def compute_kl_div(self, log_p, log_q):
+        return (log_p.exp()*(log_p-log_q)).sum(dim=-1)
+    
+    def masked_log_softmax(self, logits, mask):
+        masked_logits = torch.full_like(logits, -torch.inf)
+        masked_logits[mask] = logits[mask]
+        log_p = masked_logits.log_softmax(dim=-1)
+        return log_p
+    
+    def masked_softmax(self, logits, mask):
+        masked_logits = torch.full_like(logits, -torch.inf)
+        masked_logits[mask] = logits[mask]
+        log_p = masked_logits.softmax(dim=-1)
+        return log_p
+    
+    def compute_detection_img(self, detections, mask, B, H, W, device = "cuda"):
+        kernel_size = 5
+        X = torch.linspace(-2,2,kernel_size, device = device)
+        G = (-X**2 / (2 * (1/2)**2)).exp() # half pixel std
+        G = G/G.sum()
+        det_smoothing_kernel = G[None, None, None,:]
+        det_img = torch.zeros((B,1,H,W), device = device)
+        for b in range(B):
+            valid_detections = (detections[b][mask[b]]).int()
+            det_img[b,0][valid_detections[:,1], valid_detections[:,0]] = 1
+        det_img = F.conv2d(det_img, weight = det_smoothing_kernel, padding = (kernel_size//2, 0))
+        det_img = F.conv2d(det_img, weight = det_smoothing_kernel.mT, padding = (0, kernel_size//2))
+        return det_img
+
+    def compute_cross_entropy(self, log_p_hat, p):
+        return -(log_p_hat * p).sum(dim=-1)
+
+    def compute_matchability(self, keypoint_p, has_depth, B, K, H, W, device = "cuda"):
+        smooth_keypoint_p = F.conv2d(keypoint_p.reshape(B,1,H,W), weight = self.smoothing_kernel, padding = (self.smoothing_size//2,0))
+        smooth_keypoint_p = F.conv2d(smooth_keypoint_p, weight = self.smoothing_kernel.mT, padding = (0,self.smoothing_size//2))
+        log_p_hat = (smooth_keypoint_p+1e-8).log().reshape(B,H*W).log_softmax(dim=-1)
+        smooth_has_depth = F.conv2d(has_depth.reshape(B,1,H,W), weight = self.smoothing_kernel, padding = (0,self.smoothing_size//2))
+        smooth_has_depth = F.conv2d(smooth_has_depth, weight = self.smoothing_kernel.mT, padding = (self.smoothing_size//2,0)).reshape(B,H*W)
+        p = smooth_has_depth/smooth_has_depth.sum(dim=-1,keepdim=True)
+        return self.compute_cross_entropy(log_p_hat, p) - self.compute_cross_entropy((p+1e-12).log(), p)
+
+    def supervised_loss(self, outputs, batch):
+        keypoint_logits_A, keypoint_logits_B = outputs["keypoint_logits"].chunk(2)
+        B, K, H, W = keypoint_logits_A.shape
+
+        detections_A, detections_B = batch["detections_A"], batch["detections_B"]
+        
+        gt_warp_A_to_B, valid_mask_A_to_B = get_gt_warp(                
+                    batch["im_A_depth"],
+                    batch["im_B_depth"],
+                    batch["T_1to2"],
+                    batch["K1"],
+                    batch["K2"],
+                    H=H,
+                    W=W,
+                )
+        gt_warp_B_to_A, valid_mask_B_to_A = get_gt_warp(                
+            batch["im_B_depth"],
+            batch["im_A_depth"],
+            batch["T_1to2"].inverse(),
+            batch["K2"],
+            batch["K1"],
+            H=H,
+            W=W,
+        )
+        keypoint_logits_A = keypoint_logits_A.reshape(B, K, H*W)
+        keypoint_logits_B = keypoint_logits_B.reshape(B, K, H*W)
+        keypoint_logits = torch.cat((keypoint_logits_A, keypoint_logits_B))
+
+        B = 2*B
+        gt_warp = torch.cat((gt_warp_A_to_B, gt_warp_B_to_A))
+        valid_mask = torch.cat((valid_mask_A_to_B, valid_mask_B_to_A))
+        valid_mask = valid_mask.reshape(B,H*W)
+        binary_mask = valid_mask == 1
+        detections = torch.cat((detections_A, detections_B))
+        legit_detections = ((detections > 0).prod(dim = -1) * (detections[...,0] < W) * (detections[...,1] < H)).bool()
+        det_imgs_A, det_imgs_B = self.compute_detection_img(detections, legit_detections, B, H, W).chunk(2)
+        det_imgs = torch.cat((det_imgs_A, det_imgs_B))
+        det_imgs_backwarped = F.grid_sample(torch.cat((det_imgs_B, det_imgs_A)).reshape(B,1,H,W), 
+            gt_warp[...,-2:].reshape(B,H,W,2).float(), align_corners = False, mode = "bicubic")
+
+        keypoint_logits_backwarped = F.grid_sample(torch.cat((keypoint_logits_B, keypoint_logits_A)).reshape(B,K,H,W), 
+            gt_warp[...,-2:].reshape(B,H,W,2).float(), align_corners = False, mode = "bicubic")
+        
+        keypoint_logits_backwarped = (keypoint_logits_backwarped).reshape(B,K,H*W)
+
+
+        depth = F.interpolate(torch.cat((batch["im_A_depth"][:,None],batch["im_B_depth"][:,None]),dim=0), size = (H,W), mode = "bilinear", align_corners=False)
+        has_depth = (depth > 0).float().reshape(B,H*W)
+        
+        joint_log_likelihood_loss = self.compute_joint_neg_log_likelihood(keypoint_logits, keypoint_logits_backwarped, 
+                                                                          mask = binary_mask, detections_A = det_imgs, 
+                                                                          detections_B_to_A = det_imgs_backwarped).mean()
+        keypoint_p = keypoint_logits.reshape(B, K*H*W).softmax(dim=-1).reshape(B, K, H*W).sum(dim=1)
+        matchability_loss = self.compute_matchability(keypoint_p, has_depth, B, K, H, W).mean()
+        B = B//2
+        kpts_A = sample_keypoints(keypoint_p[:B].reshape(B,H,W), 
+                                use_nms = False, sample_topk = True, num_samples = 4*2048)
+        kpts_B = sample_keypoints(keypoint_p[B:].reshape(B,H,W),
+                                use_nms = False, sample_topk = True, num_samples = 4*2048)
+        kpts_A_to_B = F.grid_sample(gt_warp_A_to_B[...,2:].float().permute(0,3,1,2), kpts_A[...,None,:], 
+                                    align_corners=False, mode = 'bilinear')[...,0].mT
+        legit_A_to_B = F.grid_sample(valid_mask_A_to_B.reshape(B,1,H,W), kpts_A[...,None,:], 
+                                    align_corners=False, mode = 'bilinear')[...,0,:,0]
+        percent_inliers = (torch.cdist(kpts_A_to_B, kpts_B).min(dim=-1).values[legit_A_to_B > 0] < 0.01).float().mean()
+        self.tracked_metrics["mega_percent_inliers"] = (0.9 * self.tracked_metrics.get("mega_percent_inliers", percent_inliers) + 0.1 * percent_inliers)
+
+        tot_loss = joint_log_likelihood_loss + self.matchability_weight * matchability_loss# 
+        if torch.rand(1) > 1:
+            print(f"Precent Inlier: {self.tracked_metrics.get('mega_percent_inliers', 0)}")
+            print(f"{joint_log_likelihood_loss=} {matchability_loss=}")
+            print(f"Total Loss: {tot_loss.item()}")
+        return  tot_loss
+    
+    def forward(self, outputs, batch):
+        
+        if not isinstance(outputs, list):
+            outputs = [outputs]
+        losses = 0
+        for output in outputs:
+            losses = losses + self.supervised_loss(output, batch)
+        return losses

vismatch/third_party/DeDoDe/DeDoDe/encoder.py

@@ -0,0 +1,87 @@
+import torch
+import torch.nn as nn
+import torchvision.models as tvm
+
+
+class VGG19(nn.Module):
+    def __init__(self, pretrained=False, amp = False, amp_dtype = torch.float16) -> None:
+        super().__init__()
+        self.layers = nn.ModuleList(tvm.vgg19_bn(pretrained=pretrained).features[:40])
+        # Maxpool layers: 6, 13, 26, 39
+        self.amp = amp
+        self.amp_dtype = amp_dtype
+
+    def forward(self, x, **kwargs):
+        with torch.autocast("cuda", enabled=self.amp, dtype = self.amp_dtype):
+            feats = []
+            sizes = []
+            for layer in self.layers:
+                if isinstance(layer, nn.MaxPool2d):
+                    feats.append(x)
+                    sizes.append(x.shape[-2:])
+                x = layer(x)
+            return feats, sizes
+
+class VGG(nn.Module):
+    def __init__(self, size = "19", pretrained=False, amp = False, amp_dtype = torch.float16) -> None:
+        super().__init__()
+        if size == "11":
+            self.layers = nn.ModuleList(tvm.vgg11_bn(pretrained=pretrained).features[:22])
+        elif size == "13": 
+            self.layers = nn.ModuleList(tvm.vgg13_bn(pretrained=pretrained).features[:28])
+        elif size == "19": 
+            self.layers = nn.ModuleList(tvm.vgg19_bn(pretrained=pretrained).features[:40])
+        # Maxpool layers: 6, 13, 26, 39
+        self.amp = amp
+        self.amp_dtype = amp_dtype
+
+    def forward(self, x, **kwargs):
+        with torch.autocast("cuda", enabled=self.amp, dtype = self.amp_dtype):
+            feats = []
+            sizes = []
+            for layer in self.layers:
+                if isinstance(layer, nn.MaxPool2d):
+                    feats.append(x)
+                    sizes.append(x.shape[-2:])
+                x = layer(x)
+            return feats, sizes
+
+class FrozenDINOv2(nn.Module):
+    def __init__(self, amp = True, amp_dtype = torch.float16, dinov2_weights = None):
+        super().__init__()
+        if dinov2_weights is None:
+            dinov2_weights = torch.hub.load_state_dict_from_url("https://dl.fbaipublicfiles.com/dinov2/dinov2_vitl14/dinov2_vitl14_pretrain.pth", map_location="cpu")
+        from .transformer import vit_large
+        vit_kwargs = dict(img_size= 518,
+            patch_size= 14,
+            init_values = 1.0,
+            ffn_layer = "mlp",
+            block_chunks = 0,
+        )
+        dinov2_vitl14 = vit_large(**vit_kwargs).eval()
+        dinov2_vitl14.load_state_dict(dinov2_weights)
+        self.amp = amp
+        self.amp_dtype = amp_dtype
+        if self.amp:
+            dinov2_vitl14 = dinov2_vitl14.to(self.amp_dtype)
+        self.dinov2_vitl14 = [dinov2_vitl14] # ugly hack to not show parameters to DDP
+    def forward(self, x):
+        B, C, H, W = x.shape
+        if self.dinov2_vitl14[0].device != x.device:
+            self.dinov2_vitl14[0] = self.dinov2_vitl14[0].to(x.device).to(self.amp_dtype)
+        with torch.inference_mode():
+            dinov2_features_16 = self.dinov2_vitl14[0].forward_features(x.to(self.amp_dtype))
+            features_16 = dinov2_features_16['x_norm_patchtokens'].permute(0,2,1).reshape(B,1024,H//14, W//14)
+        return [features_16.clone()], [(H//14, W//14)] # clone from inference mode to use in autograd
+
+class VGG_DINOv2(nn.Module):
+    def __init__(self, vgg_kwargs = None, dinov2_kwargs = None):
+        assert vgg_kwargs is not None and dinov2_kwargs is not None, "Input kwargs pls"
+        super().__init__()        
+        self.vgg = VGG(**vgg_kwargs)
+        self.frozen_dinov2 = FrozenDINOv2(**dinov2_kwargs)
+        
+    def forward(self, x):
+        feats_vgg, sizes_vgg = self.vgg(x)
+        feat_dinov2, size_dinov2 = self.frozen_dinov2(x)
+        return feats_vgg + feat_dinov2, sizes_vgg + size_dinov2

vismatch/third_party/DeDoDe/DeDoDe/matchers/dual_softmax_matcher.py

@@ -0,0 +1,38 @@
+import torch
+from PIL import Image
+import torch.nn as nn
+import torchvision.models as tvm
+import torch.nn.functional as F
+import numpy as np
+from DeDoDe.utils import dual_softmax_matcher, to_pixel_coords, to_normalized_coords
+
+class DualSoftMaxMatcher(nn.Module):        
+    @torch.inference_mode()
+    def match(self, keypoints_A, descriptions_A, 
+              keypoints_B, descriptions_B, P_A = None, P_B = None, 
+              normalize = False, inv_temp = 1, threshold = 0.0):
+        if isinstance(descriptions_A, list):
+            matches = [self.match(k_A[None], d_A[None], k_B[None], d_B[None], normalize = normalize,
+                               inv_temp = inv_temp, threshold = threshold) 
+                    for k_A,d_A,k_B,d_B in
+                    zip(keypoints_A, descriptions_A, keypoints_B, descriptions_B)]
+            matches_A = torch.cat([m[0] for m in matches])
+            matches_B = torch.cat([m[1] for m in matches])
+            inds = torch.cat([m[2] + b for b, m in enumerate(matches)])
+            return matches_A, matches_B, inds
+        
+        P = dual_softmax_matcher(descriptions_A, descriptions_B, 
+                                 normalize = normalize, inv_temperature=inv_temp,
+                                 )
+        inds = torch.nonzero((P == P.max(dim=-1, keepdim = True).values) 
+                        * (P == P.max(dim=-2, keepdim = True).values) * (P > threshold))
+        batch_inds = inds[:,0]
+        matches_A = keypoints_A[batch_inds, inds[:,1]]
+        matches_B = keypoints_B[batch_inds, inds[:,2]]
+        return matches_A, matches_B, batch_inds
+
+    def to_pixel_coords(self, x_A, x_B, H_A, W_A, H_B, W_B):
+        return to_pixel_coords(x_A, H_A, W_A), to_pixel_coords(x_B, H_B, W_B)
+    
+    def to_normalized_coords(self, x_A, x_B, H_A, W_A, H_B, W_B):
+        return to_normalized_coords(x_A, H_A, W_A), to_normalized_coords(x_B, H_B, W_B)

vismatch/third_party/DeDoDe/DeDoDe/model_zoo/__init__.py

@@ -0,0 +1,3 @@
+from .dedode_models import dedode_detector_B, dedode_detector_L, dedode_descriptor_B, dedode_descriptor_G
+    
+