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

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