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.0
    accuracy = 0.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.0,
                1.0,
                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.0 - tot_err / inlier_num[i] / batch_size / 2.0
        else:
            accuracy += 1.0

    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.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.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