vismatch 1.1.1__py3-none-any.whl

vismatch/third_party/affine-steerers/affine_steerers/descriptors/descriptor_loss.py

@@ -0,0 +1,358 @@
+import torch
+import torch.nn as nn
+import math
+import torch.nn.functional as F
+
+from affine_steerers.utils import *
+from affine_steerers.matchers.dual_softmax_matcher import DualSoftMaxMatcher
+import affine_steerers
+
+class DescriptorLoss(nn.Module):
+    def __init__(self,
+                 detector,
+                 num_keypoints = 5000,
+                 steerer = None,
+                 kptwise_affine_loss = False,
+                 filter_kptwise_affine_loss = False,
+                 affine_loss_factor = 1.,
+                 nll_loss_factor = 1.,
+                 steer_kptwise = False,
+                 steer_with_estimated_affine = False,
+                 ordinary_steer_and_estimated = False,
+                 sing_value_cutoff = 5.,
+                 matcher = DualSoftMaxMatcher(),
+                 filter_inliers_with_mask = False,
+                 artificial_incorrect_scale = False,
+                 steer_base_affine = True,
+                 nbr_random_prototypes = 2,
+                 steer_closest_prototype = False,
+                 use_equiv_loss = False,
+                 equiv_loss_factor = 1.,
+                 device = get_best_device()) -> None:
+        super().__init__()
+        if filter_kptwise_affine_loss and not steer_with_estimated_affine:
+            raise ValueError()
+        if ordinary_steer_and_estimated and not steer_with_estimated_affine:
+            raise ValueError()
+        self.detector = detector
+        self.tracked_metrics = {}
+        self.num_keypoints = num_keypoints
+        self.steerer = steerer
+        self.kptwise_affine_loss = kptwise_affine_loss
+        self.filter_kptwise_affine_loss = filter_kptwise_affine_loss
+        self.affine_loss_factor = affine_loss_factor
+        self.nll_loss_factor = nll_loss_factor
+        self.steer_kptwise = steer_kptwise
+        self.steer_with_estimated_affine = steer_with_estimated_affine
+        self.ordinary_steer_and_estimated = ordinary_steer_and_estimated
+        self.sing_value_cutoff = sing_value_cutoff
+        self.matcher = matcher
+        self.filter_inliers_with_mask = filter_inliers_with_mask
+        self.artificial_incorrect_scale = artificial_incorrect_scale
+        self.steer_base_affine = steer_base_affine
+        self.nbr_random_prototypes = nbr_random_prototypes
+        self.steer_closest_prototype = steer_closest_prototype
+        self.use_equiv_loss = use_equiv_loss
+        self.equiv_loss_factor = equiv_loss_factor
+    
+    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 affine_from_depth(self, batch, kpts_A, kpts_B):
+        affine_A_to_B, _, kpts_A_to_B, mask_A_to_B = local_affine(
+            kpts_A,
+            batch["im_A_depth"],
+            batch["im_B_depth"],
+            batch["T_1to2"],
+            batch["K1"],
+            batch["K2"],
+            default=batch["affine_AtoB"],
+            sing_value_cutoff=self.sing_value_cutoff,
+        )
+        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, affine_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 affine_from_homog(self, batch, kpts_A, kpts_B):
+        affine_A_to_B, _, kpts_A_to_B, mask_A_to_B = local_affine_from_homography(
+            kpts_A,
+            batch["Homog_A_to_B"],
+        )
+        kpts_B_to_A = homog_transform(batch["Homog_A_to_B"].inverse(), kpts_B)
+        return (mask_A_to_B, kpts_A_to_B, affine_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:
+            if self.steer_kptwise or self.kptwise_affine_loss or self.steer_closest_prototype:
+                (mask_A_to_B, kpts_A_to_B, affine_A_to_B), (mask_B_to_A, kpts_B_to_A) = (
+                    self.affine_from_depth(batch, kpts_A, kpts_B)
+                )
+            else:
+                (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:
+            if self.steer_kptwise or self.kptwise_affine_loss or self.steer_closest_prototype:
+                (mask_A_to_B, kpts_A_to_B, affine_A_to_B), (mask_B_to_A, kpts_B_to_A) = (
+                    self.affine_from_homog(batch, kpts_A, kpts_B)
+                )
+            else:
+                (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)
+            mask = (
+                (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)
+            )
+
+            if self.filter_inliers_with_mask:
+                mask *= mask_A_to_B[:, :, None]
+
+            if self.artificial_incorrect_scale:
+                p = torch.ones_like(kpts_A[0, :, 0])
+                incorrect_ids = p.multinomial(num_samples=self.num_keypoints//2)
+                angle1 = math.pi * torch.rand(self.num_keypoints//2, device=kpts_A.device)
+                rot = torch.stack([
+                    torch.stack([torch.cos(angle1), torch.sin(angle1)], dim=-1),
+                    torch.stack([-torch.sin(angle1), torch.cos(angle1)], dim=-1),
+                ], dim=-1)
+                diag = 0.5 + torch.rand([self.num_keypoints//2, 2], device=kpts_A.device)
+                affine_A_to_B[:, incorrect_ids] = (
+                    rot.mT @ torch.diag_embed(diag) @ rot @ affine_A_to_B[:, incorrect_ids]
+                )
+                mask[:, incorrect_ids, :] = 0
+
+            inds = torch.nonzero(mask)
+
+        affine_loss = 0.
+        if self.kptwise_affine_loss and not self.filter_kptwise_affine_loss:
+            est_affine_A_to_ref = self.steerer.estimate_affine(
+                desc_A[None, inds[:,0], inds[:,1]], to_ref=True,
+            )[0]
+            affine_loss = torch.norm(
+                # est_affine_A_to_B - affine_A_to_B[inds[:,0], inds[:,1]],
+                est_affine_A_to_ref @ torch.linalg.inv(
+                    self.steerer.estimate_affine(
+                        desc_B[None, inds[:, 0], inds[:, 2]], to_ref=True,
+                    )[0].clone().detach()
+                    @ affine_A_to_B[inds[:,0], inds[:,1]]
+                )
+                - torch.eye(2, device=desc_A.device, dtype=desc_A.dtype),
+                dim=(-2, -1),
+            ).mean()
+            self.tracked_metrics["affine_loss"] = (
+                0.99 * self.tracked_metrics.get(
+                    "affine_loss", affine_loss.detach().item()
+                ) + 0.01 * affine_loss.detach().item()
+            )
+
+        if self.steer_with_estimated_affine:
+            est_affine_A_to_ref = self.steerer.estimate_affine(
+                desc_A, to_ref=True,
+            )
+            est_affine_B_to_ref = self.steerer.estimate_affine(
+                desc_B, to_ref=True,
+            )
+
+            if self.ordinary_steer_and_estimated:
+                desc_A_ref = self.steerer(
+                    desc_A,
+                    est_affine_A_to_ref,
+                )
+                desc_B_ref = self.steerer(
+                    desc_B,
+                    est_affine_B_to_ref,
+                )
+            else:
+                desc_A = self.steerer(
+                    desc_A,
+                    est_affine_A_to_ref,
+                )
+                desc_B = self.steerer(
+                    desc_B,
+                    est_affine_B_to_ref,
+                )
+
+            with torch.no_grad():
+                est_A_inl = est_affine_A_to_ref[inds[:, 0], inds[:, 1]]
+                est_B_inl = est_affine_B_to_ref[inds[:, 0], inds[:, 2]]
+                sv_A = torch.linalg.svdvals(est_A_inl)
+                sv_B = torch.linalg.svdvals(est_B_inl)
+                good_idx = torch.nonzero(
+                    torch.isfinite(sv_A.sum(dim=-1))
+                    * torch.isfinite(sv_B.sum(dim=-1))
+                    * (sv_A.max(dim=-1).values < self.sing_value_cutoff)
+                    * (sv_B.max(dim=-1).values < self.sing_value_cutoff)
+                    * (sv_A.min(dim=-1).values > 1. / self.sing_value_cutoff)
+                    * (sv_B.min(dim=-1).values > 1. / self.sing_value_cutoff)
+                    * (  # Same sign determinant as GT
+                        torch.linalg.det(est_A_inl)
+                        * torch.linalg.det(est_B_inl)
+                        * torch.linalg.det(batch["affine_AtoB"][inds[:, 0]])
+                        > 0
+                    ) 
+                )
+                inds = inds[good_idx[:, 0]]
+                if self.kptwise_affine_loss and self.filter_kptwise_affine_loss:
+                    affine_loss = torch.norm(
+                        (
+                            fast_inv_2x2(
+                                est_affine_A_to_ref[inds[:, 0], inds[:, 1]]
+                            )
+                            @ est_affine_B_to_ref[inds[:, 0], inds[:, 2]]
+                            @ affine_A_to_B[inds[:, 0], inds[:, 1]]
+                        )
+                        - torch.eye(2, device=desc_A.device, dtype=desc_A.dtype),
+                        dim=(-2, -1),
+                    ).mean()
+                    self.tracked_metrics["affine_loss"] = (
+                        0.99 * self.tracked_metrics.get(
+                            "affine_loss", affine_loss.detach().item()
+                        ) + 0.01 * affine_loss.detach().item()
+                    )
+        elif self.steer_kptwise:
+            desc_A = self.steerer(desc_A, affine_A_to_B)
+        elif self.steerer is not None and "affine_AtoB" in batch and self.steer_base_affine:
+            # same steerer for all kpts
+            affine_AtoB = batch["affine_AtoB"][:, None]
+            desc_A = self.steerer(desc_A, affine_AtoB)
+
+        equiv_loss = 0.
+        if self.use_equiv_loss:
+            equiv_loss = torch.linalg.norm(
+                desc_A[inds[:, 0], inds[:, 1]] - desc_B[inds[:, 0], inds[:, 2]],
+                dim=-1,
+            ).mean()
+            self.tracked_metrics["equiv_loss"] = (
+                0.99 * self.tracked_metrics.get(
+                    "equiv_loss", equiv_loss.detach().item()
+                ) + 0.01 * equiv_loss.detach().item()
+            )
+
+        if self.steerer is not None and self.steerer.use_prototype_affines:
+            corr = self.matcher.similarity(desc_A, desc_B)
+            if self.steerer.prototype_affines.shape[0] > self.nbr_random_prototypes:
+                for _ in range(self.nbr_random_prototypes):
+                    corr = torch.maximum(
+                        corr,
+                        self.matcher.similarity(
+                            self.steerer.steer_with_random_prototypes(desc_A),
+                            desc_B,
+                        )
+                    )
+            else:
+                for i in range(self.steerer.prototype_affines.shape[0]):
+                    corr = torch.maximum(
+                        corr,
+                        self.matcher.similarity(
+                            self.steerer(
+                                desc_A,
+                                self.steerer.prototype_affines[i, None, None],
+                            ),
+                            desc_B,
+                        )
+                    )
+            if self.steer_closest_prototype:
+                with torch.no_grad():
+                    prototype_inds = torch.cdist(
+                        affine_A_to_B[inds[:, 0], inds[:, 1]].flatten(start_dim=-2),
+                        self.steerer.prototype_affines.flatten(start_dim=-2),
+                    ).min(dim=-1).indices
+                corr[inds[:, 0], inds[:, 1], inds[:, 2]] = torch.maximum(
+                    -torch.linalg.norm(  # Fallback option of no steer
+                        desc_A[inds[:, 0], inds[:, 1]]
+                        - desc_B[inds[:, 0], inds[:, 2]],
+                        dim=-1,
+                    ),
+                    -torch.linalg.norm(  # Optimal steer acc to annotations
+                        self.steerer(
+                            desc_A[None, inds[:, 0], inds[:, 1]],
+                            self.steerer.prototype_affines[None, prototype_inds],
+                        )[0] - desc_B[inds[:, 0], inds[:, 2]],
+                        dim=-1,
+                    )
+                )
+            corr = self.matcher.inv_temp * corr
+            logP_A_B = corr.log_softmax(dim=-2) + corr.log_softmax(dim=-1)
+        elif self.ordinary_steer_and_estimated:
+            logP_A_B = 0.5 * (
+                self.matcher.log_softmax(
+                    self.steerer(desc_A, affine_A_to_B),
+                    desc_B,
+                )
+                + self.matcher.log_softmax(desc_A_ref, desc_B_ref)
+            )
+        else:
+            logP_A_B = self.matcher.log_softmax(desc_A, desc_B)
+        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(f'nll: {self.tracked_metrics["neg_log_likelihood"]}')
+            if self.steerer is not None and self.steerer.use_prototype_affines:
+                print(f'prototypes: {self.steerer.prototype_affines[:8].clone().detach().cpu().numpy()}')
+            if self.use_equiv_loss:
+                print(f'equiv_loss: {self.tracked_metrics["equiv_loss"]}')
+            if self.kptwise_affine_loss:
+                print(f'aff_loss: {self.tracked_metrics["affine_loss"]}')
+            if self.steerer is not None and self.steerer.learnable_reference_direction:
+                print(f'steerer_ref_dir: {self.steerer.reference_direction.clone().detach().cpu().numpy()}')
+            if self.steerer is not None and self.steerer.learnable_determinant_scaling:
+                print(f'steerer_det_scalings: {[(n, x.data.max().item(), x.data.median().item(), x.data.min().item()) for n, x in enumerate(self.steerer.determinant_scalings)]}')
+
+        loss = (
+            self.nll_loss_factor * neg_log_likelihood 
+            + self.affine_loss_factor * affine_loss
+            + self.equiv_loss_factor * equiv_loss
+        )
+
+        return loss
+
+    def forward(self,
+                outputs,
+                batch):
+        losses = self.supervised_loss(outputs, batch)
+        return losses

vismatch/third_party/affine-steerers/affine_steerers/detectors/dedode_detector.py

@@ -0,0 +1,75 @@
+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 affine_steerers.utils import sample_keypoints, to_pixel_coords, to_normalized_coords, get_best_device
+
+
+class affine_steerersDetector(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)).convert('RGB')
+        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/affine-steerers/affine_steerers/detectors/keypoint_loss.py

@@ -0,0 +1,215 @@
+import torch
+import torch.nn as nn
+import torchvision.transforms.functional as TTF
+import math
+
+from affine_steerers.utils import *
+import affine_steerers
+
+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, rot_A=0, rot_B=0, continuous_rot=False):
+        keypoint_logits_A, keypoint_logits_B = outputs["keypoint_logits"].chunk(2)
+        # rotate feature maps back for standard loss computation using upright GT
+        if rot_A > 0:
+            if continuous_rot:
+                keypoint_logits_A = TTF.rotate(
+                    keypoint_logits_A,
+                    360. - np.rad2deg(rot_A),
+                    interpolation=TTF.InterpolationMode.BILINEAR,
+                )
+            else:
+                keypoint_logits_A = keypoint_logits_A.rot90(k=-rot_A, dims=[-2, -1])
+        if rot_B > 0:
+            if continuous_rot:
+                keypoint_logits_B = TTF.rotate(
+                    keypoint_logits_B,
+                    360. - np.rad2deg(rot_B),
+                    interpolation=TTF.InterpolationMode.BILINEAR,
+                )
+            else:
+                keypoint_logits_B = keypoint_logits_B.rot90(k=-rot_B, dims=[-2, -1])
+
+        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, rot_A=0, rot_B=0, steerer=None, projector=None, continuous_rot=False):
+        if steerer is not None:
+            raise ValueError()
+        if projector is not None:
+            raise ValueError()
+        
+        if not isinstance(outputs, list):
+            outputs = [outputs]
+        losses = 0
+        for output in outputs:
+            losses = losses + self.supervised_loss(output,
+                                                   batch,
+                                                   rot_A=rot_A,
+                                                   rot_B=rot_B,
+                                                   continuous_rot=continuous_rot,
+                                                  )
+        return losses