third_party/lanet/loss_function.py

import torch

def build_descriptor_loss(source_des, target_des, tar_points_un, top_kk=None, relax_field=4, eval_only=False):
    """
    Desc Head Loss, per-pixel level triplet loss from https://arxiv.org/pdf/1902.11046.pdf.

    Parameters
    ----------
    source_des: torch.Tensor (B,256,H/8,W/8)
        Source image descriptors.
    target_des: torch.Tensor (B,256,H/8,W/8)
        Target image descriptors.
    source_points: torch.Tensor (B,H/8,W/8,2) 
        Source image keypoints 
    tar_points: torch.Tensor (B,H/8,W/8,2)
        Target image keypoints 
    tar_points_un: torch.Tensor (B,2,H/8,W/8)
        Target image keypoints unnormalized 
    eval_only: bool
        Computes only recall without the loss.
    Returns
    -------
    loss: torch.Tensor
        Descriptor loss.
    recall: torch.Tensor
        Descriptor match recall.
    """
    device = source_des.device
    loss = 0
    batch_size = source_des.size(0)
    recall = 0.

    relax_field_size = [relax_field]
    margins          = [1.0]
    weights          = [1.0]

    isource_dense = top_kk is None

    for b_id in range(batch_size):

        if isource_dense:
            ref_desc = source_des[b_id].squeeze().view(256, -1)
            tar_desc = target_des[b_id].squeeze().view(256, -1)
            tar_points_raw = tar_points_un[b_id].view(2, -1)
        else:
            top_k = top_kk[b_id].squeeze()

            n_feat = top_k.sum().item()
            if n_feat < 20:
                continue

            ref_desc = source_des[b_id].squeeze()[:, top_k]
            tar_desc = target_des[b_id].squeeze()[:, top_k]         
            tar_points_raw = tar_points_un[b_id][:, top_k]

        # Compute dense descriptor distance matrix and find nearest neighbor
        ref_desc = ref_desc.div(torch.norm(ref_desc, p=2, dim=0))
        tar_desc = tar_desc.div(torch.norm(tar_desc, p=2, dim=0))
        dmat = torch.mm(ref_desc.t(), tar_desc)

        dmat = torch.sqrt(2 - 2 * torch.clamp(dmat, min=-1, max=1))
        _, idx = torch.sort(dmat, dim=1)


        # Compute triplet loss and recall
        for pyramid in range(len(relax_field_size)):

            candidates = idx.t()

            match_k_x = tar_points_raw[0, candidates]
            match_k_y = tar_points_raw[1, candidates]

            tru_x = tar_points_raw[0]
            tru_y = tar_points_raw[1]

            if pyramid == 0:
                correct2 = (abs(match_k_x[0]-tru_x) == 0) & (abs(match_k_y[0]-tru_y) == 0)
                correct2_cnt = correct2.float().sum()
                recall += float(1.0 / batch_size) * (float(correct2_cnt) / float( ref_desc.size(1)))

            if eval_only:
                continue
            correct_k = (abs(match_k_x - tru_x) <= relax_field_size[pyramid]) & (abs(match_k_y - tru_y) <= relax_field_size[pyramid])

            incorrect_index = torch.arange(start=correct_k.shape[0]-1, end=-1, step=-1).unsqueeze(1).repeat(1,correct_k.shape[1]).to(device)
            incorrect_first = torch.argmax(incorrect_index * (1 - correct_k.long()), dim=0)

            incorrect_first_index = candidates.gather(0, incorrect_first.unsqueeze(0)).squeeze()

            anchor_var = ref_desc
            posource_var = tar_desc
            neg_var = tar_desc[:, incorrect_first_index]

            loss += float(1.0 / batch_size) * torch.nn.functional.triplet_margin_loss(anchor_var.t(), posource_var.t(), neg_var.t(), margin=margins[pyramid]).mul(weights[pyramid])

    return loss, recall


class KeypointLoss(object):
    """
    Loss function class encapsulating the location loss, the descriptor loss, and the score loss.
    """
    def __init__(self, config):
        self.score_weight = config.score_weight
        self.loc_weight = config.loc_weight
        self.desc_weight = config.desc_weight
        self.corres_weight = config.corres_weight
        self.corres_threshold = config.corres_threshold
        
    def __call__(self, data):
        B, _, hc, wc = data['source_score'].shape
        
        loc_mat_abs = torch.abs(data['target_coord_warped'].view(B, 2, -1).unsqueeze(3) - data['target_coord'].view(B, 2, -1).unsqueeze(2))
        l2_dist_loc_mat = torch.norm(loc_mat_abs, p=2, dim=1)
        l2_dist_loc_min, l2_dist_loc_min_index = l2_dist_loc_mat.min(dim=2)

        # construct pseudo ground truth matching matrix
        loc_min_mat = torch.repeat_interleave(l2_dist_loc_min.unsqueeze(dim=-1), repeats=l2_dist_loc_mat.shape[-1], dim=-1)
        pos_mask = l2_dist_loc_mat.eq(loc_min_mat) & l2_dist_loc_mat.le(1.)
        neg_mask = l2_dist_loc_mat.ge(4.)

        pos_corres = - torch.log(data['confidence_matrix'][pos_mask])
        neg_corres = - torch.log(1.0 - data['confidence_matrix'][neg_mask])
        corres_loss = pos_corres.mean() + 5e5 * neg_corres.mean()

        # corresponding distance threshold is 4
        dist_norm_valid_mask = l2_dist_loc_min.lt(self.corres_threshold) & data['border_mask'].view(B, hc * wc)
        
        # location loss
        loc_loss = l2_dist_loc_min[dist_norm_valid_mask].mean()
        
        # desc Head Loss, per-pixel level triplet loss from https://arxiv.org/pdf/1902.11046.pdf.
        desc_loss, _ = build_descriptor_loss(data['source_desc'], data['target_desc_warped'], data['target_coord_warped'].detach(), top_kk=data['border_mask'], relax_field=8)
        
        # score loss
        target_score_associated = data['target_score'].view(B, hc * wc).gather(1, l2_dist_loc_min_index).view(B, hc, wc).unsqueeze(1)
        dist_norm_valid_mask = dist_norm_valid_mask.view(B, hc, wc).unsqueeze(1) & data['border_mask'].unsqueeze(1) 
        l2_dist_loc_min = l2_dist_loc_min.view(B, hc, wc).unsqueeze(1)
        loc_err = l2_dist_loc_min[dist_norm_valid_mask]
        
        # repeatable_constrain in score loss
        repeatable_constrain = ((target_score_associated[dist_norm_valid_mask] + data['source_score'][dist_norm_valid_mask]) * (loc_err - loc_err.mean())).mean()

        # consistent_constrain in score_loss
        consistent_constrain = torch.nn.functional.mse_loss(data['target_score_warped'][data['border_mask'].unsqueeze(1)], data['source_score'][data['border_mask'].unsqueeze(1)]).mean() * 2
        aware_consistent_loss = torch.nn.functional.mse_loss(data['target_aware_warped'][data['border_mask'].unsqueeze(1).repeat(1, 2, 1, 1)], data['source_aware'][data['border_mask'].unsqueeze(1).repeat(1, 2, 1, 1)]).mean() * 2
        
        score_loss = repeatable_constrain + consistent_constrain + aware_consistent_loss
        
        loss = self.loc_weight * loc_loss + self.desc_weight * desc_loss + self.score_weight * score_loss + self.corres_weight * corres_loss
        
        return loss, self.loc_weight * loc_loss, self.desc_weight * desc_loss, self.score_weight * score_loss, self.corres_weight * corres_loss