third_party/DarkFeat/nets/loss.py

import torch
import torch.nn.functional as F

from .geom import rnd_sample, interpolate, get_dist_mat


def make_detector_loss(pos0, pos1, dense_feat_map0, dense_feat_map1,
               score_map0, score_map1, batch_size, num_corr, loss_type, config):
    joint_loss = 0.
    accuracy = 0.
    all_valid_pos0 = []
    all_valid_pos1 = []
    all_valid_match = []
    for i in range(batch_size):
        # random sample
        valid_pos0, valid_pos1 = rnd_sample([pos0[i], pos1[i]], num_corr)
        valid_num = valid_pos0.shape[0]

        valid_feat0 = interpolate(valid_pos0 / 4, dense_feat_map0[i])
        valid_feat1 = interpolate(valid_pos1 / 4, dense_feat_map1[i])

        valid_feat0 = F.normalize(valid_feat0, p=2, dim=-1)
        valid_feat1 = F.normalize(valid_feat1, p=2, dim=-1)

        valid_score0 = interpolate(valid_pos0, torch.squeeze(score_map0[i], dim=-1), nd=False)
        valid_score1 = interpolate(valid_pos1, torch.squeeze(score_map1[i], dim=-1), nd=False)
            
        if config['network']['det']['corr_weight']:
            corr_weight = valid_score0 * valid_score1
        else:
            corr_weight = None

        safe_radius = config['network']['det']['safe_radius']
        if safe_radius > 0:
            radius_mask_row = get_dist_mat(
                valid_pos1, valid_pos1, "euclidean_dist_no_norm")
            radius_mask_row = torch.le(radius_mask_row, safe_radius)
            radius_mask_col = get_dist_mat(
                valid_pos0, valid_pos0, "euclidean_dist_no_norm")
            radius_mask_col = torch.le(radius_mask_col, safe_radius)
            radius_mask_row = radius_mask_row.float() - torch.eye(valid_num, device=radius_mask_row.device)
            radius_mask_col = radius_mask_col.float() - torch.eye(valid_num, device=radius_mask_col.device)
        else:
            radius_mask_row = None
            radius_mask_col = None

        if valid_num < 32:
            si_loss, si_accuracy, matched_mask = 0., 1., torch.zeros((1, valid_num)).bool()
        else:
            si_loss, si_accuracy, matched_mask = make_structured_loss(
                torch.unsqueeze(valid_feat0, 0), torch.unsqueeze(valid_feat1, 0),
                loss_type=loss_type,
                radius_mask_row=radius_mask_row, radius_mask_col=radius_mask_col,
                corr_weight=torch.unsqueeze(corr_weight, 0) if corr_weight is not None else None
            )

        joint_loss += si_loss / batch_size
        accuracy += si_accuracy / batch_size
        all_valid_match.append(torch.squeeze(matched_mask, dim=0))
        all_valid_pos0.append(valid_pos0)
        all_valid_pos1.append(valid_pos1)

    return joint_loss, accuracy


def make_structured_loss(feat_anc, feat_pos,
                         loss_type='RATIO', inlier_mask=None,
                         radius_mask_row=None, radius_mask_col=None,
                         corr_weight=None, dist_mat=None):
    """
    Structured loss construction.
    Args:
        feat_anc, feat_pos: Feature matrix.
        loss_type: Loss type.
        inlier_mask:
    Returns:

    """
    batch_size = feat_anc.shape[0]
    num_corr = feat_anc.shape[1]
    if inlier_mask is None:
        inlier_mask = torch.ones((batch_size, num_corr), device=feat_anc.device).bool()
    inlier_num = torch.count_nonzero(inlier_mask.float(), dim=-1)

    if loss_type == 'L2NET' or loss_type == 'CIRCLE':
        dist_type = 'cosine_dist'
    elif loss_type.find('HARD') >= 0:
        dist_type = 'euclidean_dist'
    else:
        raise NotImplementedError()

    if dist_mat is None:
        dist_mat = get_dist_mat(feat_anc.squeeze(0), feat_pos.squeeze(0), dist_type).unsqueeze(0)
    pos_vec = dist_mat[0].diag().unsqueeze(0)

    if loss_type.find('HARD') >= 0:
        neg_margin = 1
        dist_mat_without_min_on_diag = dist_mat + \
            10 * torch.unsqueeze(torch.eye(num_corr, device=dist_mat.device), dim=0)
        mask = torch.le(dist_mat_without_min_on_diag, 0.008).float()
        dist_mat_without_min_on_diag += mask*10

        if radius_mask_row is not None:
            hard_neg_dist_row = dist_mat_without_min_on_diag + 10 * radius_mask_row
        else:
            hard_neg_dist_row = dist_mat_without_min_on_diag
        if radius_mask_col is not None:
            hard_neg_dist_col = dist_mat_without_min_on_diag + 10 * radius_mask_col
        else:
            hard_neg_dist_col = dist_mat_without_min_on_diag

        hard_neg_dist_row = torch.min(hard_neg_dist_row, dim=-1)[0]
        hard_neg_dist_col = torch.min(hard_neg_dist_col, dim=-2)[0]

        if loss_type == 'HARD_TRIPLET':
            loss_row = torch.clamp(neg_margin + pos_vec - hard_neg_dist_row, min=0)
            loss_col = torch.clamp(neg_margin + pos_vec - hard_neg_dist_col, min=0)
        elif loss_type == 'HARD_CONTRASTIVE':
            pos_margin = 0.2
            pos_loss = torch.clamp(pos_vec - pos_margin, min=0)
            loss_row = pos_loss + torch.clamp(neg_margin - hard_neg_dist_row, min=0)
            loss_col = pos_loss + torch.clamp(neg_margin - hard_neg_dist_col, min=0)
        else:
            raise NotImplementedError()
    
    elif loss_type == 'CIRCLE':
        log_scale = 512
        m = 0.1
        neg_mask_row = torch.unsqueeze(torch.eye(num_corr, device=feat_anc.device), 0)
        if radius_mask_row is not None:
            neg_mask_row += radius_mask_row
        neg_mask_col = torch.unsqueeze(torch.eye(num_corr, device=feat_anc.device), 0)
        if radius_mask_col is not None:
            neg_mask_col += radius_mask_col

        pos_margin = 1 - m
        neg_margin = m
        pos_optimal = 1 + m
        neg_optimal = -m

        neg_mat_row = dist_mat - 128 * neg_mask_row
        neg_mat_col = dist_mat - 128 * neg_mask_col

        lse_positive = torch.logsumexp(-log_scale * (pos_vec[..., None] - pos_margin) * \
                    torch.clamp(pos_optimal - pos_vec[..., None], min=0).detach(), dim=-1)
        
        lse_negative_row = torch.logsumexp(log_scale * (neg_mat_row - neg_margin) * \
                    torch.clamp(neg_mat_row - neg_optimal, min=0).detach(), dim=-1)

        lse_negative_col = torch.logsumexp(log_scale * (neg_mat_col - neg_margin) * \
                    torch.clamp(neg_mat_col - neg_optimal, min=0).detach(), dim=-2)

        loss_row = F.softplus(lse_positive + lse_negative_row) / log_scale
        loss_col = F.softplus(lse_positive + lse_negative_col) / log_scale

    else:
        raise NotImplementedError()

    if dist_type == 'cosine_dist':
        err_row = dist_mat - torch.unsqueeze(pos_vec, -1)
        err_col = dist_mat - torch.unsqueeze(pos_vec, -2)
    elif dist_type == 'euclidean_dist' or dist_type == 'euclidean_dist_no_norm':
        err_row = torch.unsqueeze(pos_vec, -1) - dist_mat
        err_col = torch.unsqueeze(pos_vec, -2) - dist_mat
    else:
        raise NotImplementedError()
    if radius_mask_row is not None:
        err_row = err_row - 10 * radius_mask_row
    if radius_mask_col is not None:
        err_col = err_col - 10 * radius_mask_col
    err_row = torch.sum(torch.clamp(err_row, min=0), dim=-1)
    err_col = torch.sum(torch.clamp(err_col, min=0), dim=-2)

    loss = 0
    accuracy = 0

    tot_loss = (loss_row + loss_col) / 2
    if corr_weight is not None:
        tot_loss = tot_loss * corr_weight

    for i in range(batch_size):
        if corr_weight is not None:
            loss += torch.sum(tot_loss[i][inlier_mask[i]]) / \
                (torch.sum(corr_weight[i][inlier_mask[i]]) + 1e-6)
        else:
            loss += torch.mean(tot_loss[i][inlier_mask[i]])
        cnt_err_row = torch.count_nonzero(err_row[i][inlier_mask[i]]).float()
        cnt_err_col = torch.count_nonzero(err_col[i][inlier_mask[i]]).float()
        tot_err = cnt_err_row + cnt_err_col
        if inlier_num[i] != 0:
            accuracy += 1. - tot_err / inlier_num[i] / batch_size / 2.
        else:
            accuracy += 1.

    matched_mask = torch.logical_and(torch.eq(err_row, 0), torch.eq(err_col, 0))
    matched_mask = torch.logical_and(matched_mask, inlier_mask)

    loss /= batch_size
    accuracy /= batch_size

    return loss, accuracy, matched_mask


# for the neighborhood areas of keypoints extracted from normal image, the score from noise_score_map should be close
# for the rest, the noise image's score should less than normal image
# input: score_map [batch_size, H, W, 1]; indices [2, k, 2]
# output: loss [scalar]
def make_noise_score_map_loss(score_map, noise_score_map, indices, batch_size, thld=0.):
    H, W = score_map.shape[1:3]
    loss = 0
    for i in range(batch_size):
        kpts_coords = indices[i].T # (2, num_kpts)
        mask = torch.zeros([H, W], device=score_map.device)
        mask[kpts_coords.cpu().numpy()] = 1

        # using 3x3 kernel to put kpts' neightborhood area into the mask
        kernel = torch.ones([1, 1, 3, 3], device=score_map.device)
        mask = F.conv2d(mask.unsqueeze(0).unsqueeze(0), kernel, padding=1)[0, 0] > 0

        loss1 = torch.sum(torch.abs(score_map[i] - noise_score_map[i]).squeeze() * mask) / torch.sum(mask)
        loss2 = torch.sum(torch.clamp(noise_score_map[i] - score_map[i] - thld, min=0).squeeze() * torch.logical_not(mask)) / (H * W - torch.sum(mask))

        loss += loss1
        loss += loss2

        if i == 0:
            first_mask = mask

    return loss, first_mask


def make_noise_score_map_loss_labelmap(score_map, noise_score_map, labelmap, batch_size, thld=0.):
    H, W = score_map.shape[1:3]
    loss = 0
    for i in range(batch_size):
        # using 3x3 kernel to put kpts' neightborhood area into the mask
        kernel = torch.ones([1, 1, 3, 3], device=score_map.device)
        mask = F.conv2d(labelmap[i].unsqueeze(0).to(score_map.device).float(), kernel, padding=1)[0, 0] > 0

        loss1 = torch.sum(torch.abs(score_map[i] - noise_score_map[i]).squeeze() * mask) / torch.sum(mask)
        loss2 = torch.sum(torch.clamp(noise_score_map[i] - score_map[i] - thld, min=0).squeeze() * torch.logical_not(mask)) / (H * W - torch.sum(mask))

        loss += loss1
        loss += loss2

        if i == 0:
            first_mask = mask

    return loss, first_mask


def make_score_map_peakiness_loss(score_map, scores, batch_size):
    H, W = score_map.shape[1:3]
    loss = 0

    for i in range(batch_size):
        loss += torch.mean(scores[i]) - torch.mean(score_map[i])

    loss /= batch_size
    return 1 - loss