vismatch 1.1.1__py3-none-any.whl

vismatch/third_party/rdd/training/losses/detector_loss.py

@@ -0,0 +1,499 @@
+import torch
+import torch.nn.functional as F
+
+from torch import nn
+import cv2
+import numpy as np
+from copy import deepcopy
+from RDD.dataset.megadepth.utils import warp
+
+def plot_keypoints(image, kpts, radius=2, color=(255, 0, 0)):
+    image = image.cpu().detach().numpy() if isinstance(image, torch.Tensor) else image
+    kpts = kpts.cpu().detach().numpy() if isinstance(kpts, torch.Tensor) else kpts
+
+    if image.dtype is not np.dtype('uint8'):
+        image = image * 255
+        image = image.astype(np.uint8)
+
+    if len(image.shape) == 2 or image.shape[2] == 1:
+        image = cv2.cvtColor(image, cv2.COLOR_GRAY2RGB)
+
+    out = np.ascontiguousarray(deepcopy(image))
+    kpts = np.round(kpts).astype(int)
+
+    for kpt in kpts:
+        y0, x0 = kpt
+        cv2.drawMarker(out, (x0, y0), color, cv2.MARKER_CROSS, radius)
+
+        # cv2.circle(out, (x0, y0), radius, color, -1, lineType=cv2.LINE_4)
+    return out
+
+class DetectorLoss(nn.Module):
+    def __init__(self, temperature = 0.1, scores_th = 0.1, peaky_weight = 0.5, reprojection_weight = 1, scoremap_weight = 0.5):
+        super().__init__()
+        self.temperature = temperature
+        self.scores_th = scores_th
+        self.peaky_weight = peaky_weight
+        self.reprojection_weight = reprojection_weight
+        self.scoremap_weight = scoremap_weight
+        
+        self.PeakyLoss = PeakyLoss(scores_th = scores_th)
+        self.ReprojectionLocLoss = ReprojectionLocLoss(scores_th = scores_th)
+        self.ScoreMapRepLoss = ScoreMapRepLoss(temperature = temperature)
+        
+    def forward(self, correspondences, pred0_with_rand, pred1_with_rand):
+        loss_peaky0 = self.PeakyLoss(pred0_with_rand)
+        loss_peaky1 = self.PeakyLoss(pred1_with_rand)
+        loss_peaky = (loss_peaky0 + loss_peaky1) / 2.
+        
+        loss_reprojection = self.ReprojectionLocLoss(pred0_with_rand, pred1_with_rand, correspondences)
+        
+        loss_score_map_rp = self.ScoreMapRepLoss(pred0_with_rand, pred1_with_rand, correspondences)
+        
+        loss_kp = loss_peaky * self.peaky_weight + loss_reprojection * self.reprojection_weight + loss_score_map_rp * self.scoremap_weight
+        
+        return loss_kp
+
+class PeakyLoss(object):
+    """ PeakyLoss to avoid an uniform score map """
+
+    def __init__(self, scores_th: float = 0.1):
+        super().__init__()
+        self.scores_th = scores_th
+
+    def __call__(self, pred):
+        b, c, h, w = pred['scores_map'].shape
+        loss_mean = 0
+        CNT = 0
+        
+        for idx in range(b):
+            n_original = len(pred['score_dispersity'][idx])
+            scores_kpts = pred['scores'][idx][:n_original]
+            valid = scores_kpts > self.scores_th
+            loss_peaky = pred['score_dispersity'][idx][valid]
+
+            loss_mean = loss_mean + loss_peaky.sum()
+            CNT = CNT + len(loss_peaky)
+
+        loss_mean = loss_mean / CNT if CNT != 0 else pred['scores_map'].new_tensor(0)
+        assert not torch.isnan(loss_mean)
+        return loss_mean
+
+
+class ReprojectionLocLoss(object):
+    """
+    Reprojection location errors of keypoints to train repeatable detector.
+    """
+
+    def __init__(self, norm: int = 1, scores_th: float = 0.1):
+        super().__init__()
+        self.norm = norm
+        self.scores_th = scores_th
+
+    def __call__(self, pred0, pred1, correspondences):
+        b, c, h, w = pred0['scores_map'].shape
+        loss_mean = 0
+        CNT = 0
+        for idx in range(b):
+            if correspondences[idx]['correspondence0'] is None:
+                continue
+
+            if self.norm == 2:
+                dist = correspondences[idx]['dist']
+            elif self.norm == 1:
+                dist = correspondences[idx]['dist_l1']
+            else:
+                raise TypeError('No such norm in correspondence.')
+
+            ids0_d = correspondences[idx]['ids0_d']
+            ids1_d = correspondences[idx]['ids1_d']
+
+            scores0 = correspondences[idx]['scores0'].detach()[ids0_d]
+            scores1 = correspondences[idx]['scores1'].detach()[ids1_d]
+            valid = (scores0 > self.scores_th) * (scores1 > self.scores_th)
+            reprojection_errors = dist[ids0_d, ids1_d][valid]
+
+            loss_mean = loss_mean + reprojection_errors.sum()
+            CNT = CNT + len(reprojection_errors)
+
+        loss_mean = loss_mean / CNT if CNT != 0 else correspondences[0]['dist'].new_tensor(0)
+
+        assert not torch.isnan(loss_mean)
+        return loss_mean
+
+
+def local_similarity(descriptor_map, descriptors, kpts_wh, radius):
+    """
+    :param descriptor_map: CxHxW
+    :param descriptors: NxC
+    :param kpts_wh: Nx2 (W,H)
+    :return:
+    """
+    _, h, w = descriptor_map.shape
+    ksize = 2 * radius + 1
+
+    descriptor_map_unflod = torch.nn.functional.unfold(descriptor_map.unsqueeze(0),
+                                                       kernel_size=(ksize, ksize),
+                                                       padding=(radius, radius))
+    descriptor_map_unflod = descriptor_map_unflod[0].t().reshape(h * w, -1, ksize * ksize)
+    # find the correspondence patch
+    kpts_wh_long = kpts_wh.detach().long()
+    patch_ids = kpts_wh_long[:, 0] + kpts_wh_long[:, 1] * h
+    desc_patches = descriptor_map_unflod[patch_ids].permute(0, 2, 1).detach()  # N_kpts x s*s x 128
+
+    local_sim = torch.einsum('nsd,nd->ns', desc_patches, descriptors)
+    local_sim_sort = torch.sort(local_sim, dim=1, descending=True).values
+    local_sim_sort_mean = local_sim_sort[:, 4:].mean(dim=1)  # 4 is safe radius for bilinear interplation
+
+    return local_sim_sort_mean
+
+
+class ScoreMapRepLoss(object):
+    """ Scoremap repetability"""
+
+    def __init__(self, temperature: float = 0.1):
+        super().__init__()
+        self.temperature = temperature
+        self.radius = 2
+
+    def __call__(self, pred0, pred1, correspondences):
+        b, c, h, w = pred0['scores_map'].shape
+        wh = pred0['keypoints'][0].new_tensor([[w - 1, h - 1]])
+        loss_mean = 0
+        CNT = 0
+
+        for idx in range(b):
+            if correspondences[idx]['correspondence0'] is None:
+                continue
+
+            scores_map0 = pred0['scores_map'][idx]
+            scores_map1 = pred1['scores_map'][idx]
+            kpts01 = correspondences[idx]['kpts01']
+            kpts10 = correspondences[idx]['kpts10']  # valid warped keypoints
+
+            # =====================
+            scores_kpts10 = torch.nn.functional.grid_sample(scores_map0.unsqueeze(0), kpts10.view(1, 1, -1, 2),
+                                                            mode='bilinear', align_corners=True)[0, 0, 0, :]
+            scores_kpts01 = torch.nn.functional.grid_sample(scores_map1.unsqueeze(0), kpts01.view(1, 1, -1, 2),
+                                                            mode='bilinear', align_corners=True)[0, 0, 0, :]
+
+            s0 = scores_kpts01 * correspondences[idx]['scores0']  # repeatability
+            s1 = scores_kpts10 * correspondences[idx]['scores1']  # repeatability
+
+            # ===================== repetability
+            similarity_map_01 = correspondences[idx]['similarity_map_01_valid']
+            similarity_map_10 = correspondences[idx]['similarity_map_10_valid']
+
+            pmf01 = ((similarity_map_01.detach() - 1) / self.temperature).exp()
+            pmf10 = ((similarity_map_10.detach() - 1) / self.temperature).exp()
+            kpts01 = kpts01.detach()
+            kpts10 = kpts10.detach()
+
+            pmf01_kpts = torch.nn.functional.grid_sample(pmf01.unsqueeze(0), kpts01.view(1, 1, -1, 2),
+                                                         mode='bilinear', align_corners=True)[0, :, 0, :]
+            pmf10_kpts = torch.nn.functional.grid_sample(pmf10.unsqueeze(0), kpts10.view(1, 1, -1, 2),
+                                                         mode='bilinear', align_corners=True)[0, :, 0, :]
+            repetability01 = torch.diag(pmf01_kpts)
+            repetability10 = torch.diag(pmf10_kpts)
+            
+            # ===================== reliability
+            # ids0, ids1 = correspondences[idx]['ids0'], correspondences[idx]['ids1']
+            # descriptor_map0 = pred0['descriptor_map'][idx].detach()
+            # descriptor_map1 = pred1['descriptor_map'][idx].detach()
+            # descriptors0 = pred0['descriptors'][idx][ids0].detach()
+            # descriptors1 = pred1['descriptors'][idx][ids1].detach()
+            # kpts0 = pred0['keypoints'][idx][ids0].detach()
+            # kpts1 = pred1['keypoints'][idx][ids1].detach()
+            # kpts0_wh = (kpts0 / 2 + 0.5) * wh
+            # kpts1_wh = (kpts1 / 2 + 0.5) * wh
+            # ls0 = local_similarity(descriptor_map0, descriptors0, kpts0_wh, self.radius)
+            # ls1 = local_similarity(descriptor_map1, descriptors1, kpts1_wh, self.radius)
+            # reliability0 = 1 - ((ls0 - 1) / self.temperature).exp()
+            # reliability1 = 1 - ((ls1 - 1) / self.temperature).exp()
+
+            fs0 = repetability01  # * reliability0
+            fs1 = repetability10  # * reliability1
+
+            if s0.sum() != 0:
+                loss01 = (1 - fs0) * s0 * len(s0) / s0.sum()
+                loss_mean = loss_mean + loss01.sum()
+                CNT = CNT + len(loss01)
+            if s1.sum() != 0:
+                loss10 = (1 - fs1) * s1 * len(s1) / s1.sum()
+                loss_mean = loss_mean + loss10.sum()
+                CNT = CNT + len(loss10)
+        
+        loss_mean = loss_mean / CNT if CNT != 0 else pred0['scores_map'].new_tensor(0)
+        assert not torch.isnan(loss_mean)
+        return loss_mean
+    
+    
+
+#+++++++++++++++++++++++++++++++++++++++++++++++++++Taken from ALIKE+++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
+
+def compute_keypoints_distance(kpts0, kpts1, p=2, debug=False):
+    """
+    Args:
+        kpts0: torch.tensor [M,2]
+        kpts1: torch.tensor [N,2]
+        p: (int, float, inf, -inf, 'fro', 'nuc', optional): the order of norm
+
+    Returns:
+        dist, torch.tensor [N,M]
+    """
+    dist = kpts0[:, None, :] - kpts1[None, :, :]  # [M,N,2]
+    dist = torch.norm(dist, p=p, dim=2)  # [M,N]
+    return dist
+
+def mutual_argmax(value, mask=None, as_tuple=True):
+    """
+    Args:
+        value: MxN
+        mask:  MxN
+
+    Returns:
+
+    """
+    value = value - value.min()  # convert to non-negative tensor
+    if mask is not None:
+        value = value * mask
+
+    max0 = value.max(dim=1, keepdim=True)  # the col index the max value in each row
+    max1 = value.max(dim=0, keepdim=True)
+
+    valid_max0 = value == max0[0]
+    valid_max1 = value == max1[0]
+
+    mutual = valid_max0 * valid_max1
+    if mask is not None:
+        mutual = mutual * mask
+
+    return mutual.nonzero(as_tuple=as_tuple)
+
+
+def mutual_argmin(value, mask=None):
+    return mutual_argmax(-value, mask)
+
+def compute_correspondence(model, pred0, pred1, batch, radius = 2, rand=True, train_gt_th = 5, debug=False):
+    b, c, h, w = pred0['scores_map'].shape
+    wh = pred0['scores_map'][0].new_tensor([[w - 1, h - 1]])
+    
+    pred0_with_rand = pred0
+    pred1_with_rand = pred1
+    pred0_with_rand['scores'] = []
+    pred1_with_rand['scores'] = []
+    pred0_with_rand['descriptors'] = []
+    pred1_with_rand['descriptors'] = []
+    pred0_with_rand['num_det'] = []
+    pred1_with_rand['num_det'] = []
+
+    kps, score_dispersity, scores = model.softdetect.detect_keypoints(pred0['scores_map'])
+    
+    pred0_with_rand['keypoints'] = kps
+    pred0_with_rand['score_dispersity'] = score_dispersity
+    
+    
+    kps, score_dispersity, scores = model.softdetect.detect_keypoints(pred1['scores_map'])
+    pred1_with_rand['keypoints'] = kps
+    pred1_with_rand['score_dispersity'] = score_dispersity
+
+    correspondences = []
+    for idx in range(b):
+        # =========================== prepare keypoints
+        kpts0, kpts1 = pred0['keypoints'][idx], pred1['keypoints'][idx]  # (x,y), shape: Nx2
+        
+        # additional random keypoints
+        if rand:
+            rand0 = torch.rand(len(kpts0), 2, device=kpts0.device) * 2 - 1  # -1~1
+            rand1 = torch.rand(len(kpts1), 2, device=kpts1.device) * 2 - 1  # -1~1
+            kpts0 = torch.cat([kpts0, rand0])
+            kpts1 = torch.cat([kpts1, rand1])
+
+            pred0_with_rand['keypoints'][idx] = kpts0
+            pred1_with_rand['keypoints'][idx] = kpts1
+        
+        scores_map0 = pred0['scores_map'][idx]
+        scores_map1 = pred1['scores_map'][idx]
+        scores_kpts0 = torch.nn.functional.grid_sample(scores_map0.unsqueeze(0), kpts0.view(1, 1, -1, 2),
+                                                        mode='bilinear', align_corners=True).squeeze()
+        scores_kpts1 = torch.nn.functional.grid_sample(scores_map1.unsqueeze(0), kpts1.view(1, 1, -1, 2),
+                                                        mode='bilinear', align_corners=True).squeeze()
+
+        kpts0_wh_ = (kpts0 / 2 + 0.5) * wh  # N0x2, (w,h)
+        kpts1_wh_ = (kpts1 / 2 + 0.5) * wh  # N1x2, (w,h)
+
+        # ========================= nms
+        dist = compute_keypoints_distance(kpts0_wh_.detach(), kpts0_wh_.detach())
+        local_mask = dist < radius
+        valid_cnt = torch.sum(local_mask, dim=1)
+        indices_need_nms = torch.where(valid_cnt > 1)[0]
+        
+        for i in indices_need_nms:
+            if valid_cnt[i] > 0:
+                kpt_indices = torch.where(local_mask[i])[0]
+                scs_max_idx = scores_kpts0[kpt_indices].argmax()
+
+                tmp_mask = kpt_indices.new_ones(len(kpt_indices)).bool()
+                tmp_mask[scs_max_idx] = False
+                suppressed_indices = kpt_indices[tmp_mask]
+
+                valid_cnt[suppressed_indices] = 0
+        
+        valid_mask = valid_cnt > 0
+        kpts0_wh = kpts0_wh_[valid_mask]
+        kpts0 = kpts0[valid_mask]
+        scores_kpts0 = scores_kpts0[valid_mask]
+        pred0_with_rand['keypoints'][idx] = kpts0
+
+        valid_mask = valid_mask[:len(pred0_with_rand['score_dispersity'][idx])]
+        pred0_with_rand['score_dispersity'][idx] = pred0_with_rand['score_dispersity'][idx][valid_mask]
+        pred0_with_rand['num_det'].append(valid_mask.sum())
+
+        dist = compute_keypoints_distance(kpts1_wh_.detach(), kpts1_wh_.detach())
+        local_mask = dist < radius
+        valid_cnt = torch.sum(local_mask, dim=1)
+        indices_need_nms = torch.where(valid_cnt > 1)[0]
+        for i in indices_need_nms:
+            if valid_cnt[i] > 0:
+                kpt_indices = torch.where(local_mask[i])[0]
+                scs_max_idx = scores_kpts1[kpt_indices].argmax()
+
+                tmp_mask = kpt_indices.new_ones(len(kpt_indices)).bool()
+                tmp_mask[scs_max_idx] = False
+                suppressed_indices = kpt_indices[tmp_mask]
+
+                valid_cnt[suppressed_indices] = 0
+                
+        valid_mask = valid_cnt > 0
+        kpts1_wh = kpts1_wh_[valid_mask]
+        kpts1 = kpts1[valid_mask]
+        scores_kpts1 = scores_kpts1[valid_mask]
+        pred1_with_rand['keypoints'][idx] = kpts1
+
+        valid_mask = valid_mask[:len(pred1_with_rand['score_dispersity'][idx])]
+        pred1_with_rand['score_dispersity'][idx] = pred1_with_rand['score_dispersity'][idx][valid_mask]
+        pred1_with_rand['num_det'].append(valid_mask.sum())
+
+        # del dist, local_mask, valid_cnt, indices_need_nms, scs_max_idx, tmp_mask, suppressed_indices, valid_mask
+        # torch.cuda.empty_cache()
+        # ========================= nms
+
+        pred0_with_rand['scores'].append(scores_kpts0)
+        pred1_with_rand['scores'].append(scores_kpts1)
+        
+        descriptor_map0, descriptor_map1 = pred0['descriptor_map'][idx], pred1['descriptor_map'][idx]
+        descriptor_map0 = F.normalize(descriptor_map0, dim=0)
+        descriptor_map1 = F.normalize(descriptor_map1, dim=0)
+
+        desc0 = torch.nn.functional.grid_sample(descriptor_map0.unsqueeze(0), kpts0.view(1, 1, -1, 2),
+                                                mode='bilinear', align_corners=True)[0, :, 0, :].t()
+        desc1 = torch.nn.functional.grid_sample(descriptor_map1.unsqueeze(0), kpts1.view(1, 1, -1, 2),
+                                                mode='bilinear', align_corners=True)[0, :, 0, :].t()
+        desc0 = F.normalize(desc0, dim=-1)
+        desc1 = F.normalize(desc1, dim=-1)
+
+        pred0_with_rand['descriptors'].append(desc0)
+        pred1_with_rand['descriptors'].append(desc1)
+        
+        # =========================== prepare warp parameters
+        warp01_params = {}
+        for k, v in batch['warp01_params'].items():
+            warp01_params[k] = v[idx]
+        warp10_params = {}
+        for k, v in batch['warp10_params'].items():
+            warp10_params[k] = v[idx]
+
+        # =========================== warp keypoints across images
+        try:
+            kpts0_wh, kpts01_wh, ids0, ids0_out = warp(kpts0_wh, warp01_params)
+            kpts1_wh, kpts10_wh, ids1, ids1_out = warp(kpts1_wh, warp10_params)
+        except:
+            correspondences.append({'correspondence0': None, 'correspondence1': None,
+                                    'dist': kpts0_wh.new_tensor(0),
+                                    })
+            continue
+
+        if debug:
+            from training.utils import save_image_in_actual_size
+            
+            image0 = batch['image0'][idx].cpu().detach().numpy().transpose(1, 2, 0)
+            image1 = batch['image1'][idx].cpu().detach().numpy().transpose(1, 2, 0)
+            
+            p0 = kpts0_wh[:, [1, 0]].cpu().detach().numpy()
+            img_kpts0 = plot_keypoints(image0, p0, radius=5, color=(255, 0, 0))
+            # display_image_in_actual_size(img_kpts0)
+
+            p1 = kpts1_wh[:, [1, 0]].cpu().detach().numpy()
+            img_kpts1 = plot_keypoints(image1, p1, radius=5, color=(255, 0, 0))
+            # display_image_in_actual_size(img_kpts1)
+
+            p01 = kpts01_wh[:, [1, 0]].cpu().detach().numpy()
+            img_kpts01 = plot_keypoints(img_kpts1, p01, radius=5, color=(0, 255, 0))
+            save_image_in_actual_size(img_kpts01, name='kpts01.png')
+            
+            p10 = kpts10_wh[:, [1, 0]].cpu().detach().numpy()
+            img_kpts10 = plot_keypoints(img_kpts0, p10, radius=5, color=(0, 255, 0))
+            save_image_in_actual_size(img_kpts10, name='kpts10.png')
+
+        # ============================= compute reprojection error
+        dist01 = compute_keypoints_distance(kpts0_wh, kpts10_wh)
+        dist10 = compute_keypoints_distance(kpts1_wh, kpts01_wh)
+
+        dist_l2 = (dist01 + dist10.t()) / 2.
+        # find mutual correspondences by calculating the distance
+        # between keypoints (I1) and warpped keypoints (I2->I1)
+        mutual_min_indices = mutual_argmin(dist_l2)
+
+        dist_mutual_min = dist_l2[mutual_min_indices]
+        valid_dist_mutual_min = dist_mutual_min.detach() < train_gt_th
+
+        ids0_d = mutual_min_indices[0][valid_dist_mutual_min]
+        ids1_d = mutual_min_indices[1][valid_dist_mutual_min]
+
+        correspondence0 = ids0[ids0_d]
+        correspondence1 = ids1[ids1_d]
+
+        # L1 distance
+        dist01_l1 = compute_keypoints_distance(kpts0_wh, kpts10_wh, p=1)
+        dist10_l1 = compute_keypoints_distance(kpts1_wh, kpts01_wh, p=1)
+
+        dist_l1 = (dist01_l1 + dist10_l1.t()) / 2.
+
+        # =========================== compute cross image descriptor similarity_map
+        similarity_map_01 = (desc0 @ descriptor_map1.reshape(h*w, 256).t()) * 20
+        similarity_map_01 = similarity_map_01.softmax(dim = -2) * similarity_map_01.softmax(dim= -1)
+        similarity_map_01 = similarity_map_01.reshape(desc0.shape[0], h, w)
+        similarity_map_01 = torch.clamp(similarity_map_01, 1e-6, 1-1e-6)
+        
+        similarity_map_10 = (desc1 @ descriptor_map0.reshape(h*w, 256).t()) * 20
+        similarity_map_10 = similarity_map_10.softmax(dim = -2) * similarity_map_10.softmax(dim= -1)
+        similarity_map_10 = similarity_map_10.reshape(desc1.shape[0], h, w)
+        similarity_map_10 = torch.clamp(similarity_map_10, 1e-6, 1-1e-6)
+
+        similarity_map_01_valid = similarity_map_01[ids0]  # valid descriptors
+        similarity_map_10_valid = similarity_map_10[ids1]
+
+        kpts01 = 2 * kpts01_wh.detach() / wh - 1  # N0x2, (x,y), [-1,1]
+        kpts10 = 2 * kpts10_wh.detach() / wh - 1  # N0x2, (x,y), [-1,1]
+
+        correspondences.append({'correspondence0': correspondence0,  # indices of matched kpts0 in all kpts
+                                'correspondence1': correspondence1,  # indices of matched kpts1 in all kpts
+                                'scores0': scores_kpts0[ids0],
+                                'scores1': scores_kpts1[ids1],
+                                'kpts01': kpts01, 'kpts10': kpts10,  # warped valid kpts
+                                'ids0': ids0, 'ids1': ids1,  # valid indices of kpts0 and kpts1
+                                'ids0_out': ids0_out, 'ids1_out': ids1_out,
+                                'ids0_d': ids0_d, 'ids1_d': ids1_d,  # match indices of valid kpts0 and kpts1
+                                'dist_l1': dist_l1,  # cross distance matrix of valid kpts using L1 norm
+                                'dist': dist_l2,  # cross distance matrix of valid kpts using L2 norm
+                                'similarity_map_01': similarity_map_01,  # all
+                                'similarity_map_10': similarity_map_10,  # all
+                                'similarity_map_01_valid': similarity_map_01_valid,  # valid
+                                'similarity_map_10_valid': similarity_map_10_valid,  # valid
+                                })
+
+    return correspondences, pred0_with_rand, pred1_with_rand
+
+
+class EmptyTensorError(Exception):
+    pass