Vincentqyw
fix: roma
8b973ee
raw
history blame
11 kB
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