# Copyright (C) 2024-present Naver Corporation. All rights reserved.
# Licensed under CC BY-NC-SA 4.0 (non-commercial use only).
#
# --------------------------------------------------------
# Implementation of MASt3R training losses
# --------------------------------------------------------
import torch
import torch.nn as nn
import numpy as np
from sklearn.metrics import average_precision_score
import mast3r.utils.path_to_dust3r # noqa
from dust3r.losses import BaseCriterion, Criterion, MultiLoss, Sum, ConfLoss
from dust3r.losses import Regr3D as Regr3D_dust3r
from dust3r.utils.geometry import (geotrf, inv, normalize_pointcloud)
from dust3r.inference import get_pred_pts3d
from dust3r.utils.geometry import get_joint_pointcloud_depth, get_joint_pointcloud_center_scale
def apply_log_to_norm(xyz):
d = xyz.norm(dim=-1, keepdim=True)
xyz = xyz / d.clip(min=1e-8)
xyz = xyz * torch.log1p(d)
return xyz
class Regr3D (Regr3D_dust3r):
def __init__(self, criterion, norm_mode='avg_dis', gt_scale=False, opt_fit_gt=False,
sky_loss_value=2, max_metric_scale=False, loss_in_log=False):
self.loss_in_log = loss_in_log
if norm_mode.startswith('?'):
# do no norm pts from metric scale datasets
self.norm_all = False
self.norm_mode = norm_mode[1:]
else:
self.norm_all = True
self.norm_mode = norm_mode
super().__init__(criterion, self.norm_mode, gt_scale)
self.sky_loss_value = sky_loss_value
self.max_metric_scale = max_metric_scale
def get_all_pts3d(self, gt1, gt2, pred1, pred2, dist_clip=None):
# everything is normalized w.r.t. camera of view1
in_camera1 = inv(gt1['camera_pose'])
gt_pts1 = geotrf(in_camera1, gt1['pts3d']) # B,H,W,3
gt_pts2 = geotrf(in_camera1, gt2['pts3d']) # B,H,W,3
valid1 = gt1['valid_mask'].clone()
valid2 = gt2['valid_mask'].clone()
if dist_clip is not None:
# points that are too far-away == invalid
dis1 = gt_pts1.norm(dim=-1) # (B, H, W)
dis2 = gt_pts2.norm(dim=-1) # (B, H, W)
valid1 = valid1 & (dis1 <= dist_clip)
valid2 = valid2 & (dis2 <= dist_clip)
if self.loss_in_log == 'before':
# this only make sense when depth_mode == 'linear'
gt_pts1 = apply_log_to_norm(gt_pts1)
gt_pts2 = apply_log_to_norm(gt_pts2)
pr_pts1 = get_pred_pts3d(gt1, pred1, use_pose=False).clone()
pr_pts2 = get_pred_pts3d(gt2, pred2, use_pose=True).clone()
if not self.norm_all:
if self.max_metric_scale:
B = valid1.shape[0]
# valid1: B, H, W
# torch.linalg.norm(gt_pts1, dim=-1) -> B, H, W
# dist1_to_cam1 -> reshape to B, H*W
dist1_to_cam1 = torch.where(valid1, torch.linalg.norm(gt_pts1, dim=-1), 0).view(B, -1)
dist2_to_cam1 = torch.where(valid2, torch.linalg.norm(gt_pts2, dim=-1), 0).view(B, -1)
# is_metric_scale: B
# dist1_to_cam1.max(dim=-1).values -> B
gt1['is_metric_scale'] = gt1['is_metric_scale'] \
& (dist1_to_cam1.max(dim=-1).values < self.max_metric_scale) \
& (dist2_to_cam1.max(dim=-1).values < self.max_metric_scale)
gt2['is_metric_scale'] = gt1['is_metric_scale']
mask = ~gt1['is_metric_scale']
else:
mask = torch.ones_like(gt1['is_metric_scale'])
# normalize 3d points
if self.norm_mode and mask.any():
pr_pts1[mask], pr_pts2[mask] = normalize_pointcloud(pr_pts1[mask], pr_pts2[mask], self.norm_mode,
valid1[mask], valid2[mask])
if self.norm_mode and not self.gt_scale:
gt_pts1, gt_pts2, norm_factor = normalize_pointcloud(gt_pts1, gt_pts2, self.norm_mode,
valid1, valid2, ret_factor=True)
# apply the same normalization to prediction
pr_pts1[~mask] = pr_pts1[~mask] / norm_factor[~mask]
pr_pts2[~mask] = pr_pts2[~mask] / norm_factor[~mask]
# return sky segmentation, making sure they don't include any labelled 3d points
sky1 = gt1['sky_mask'] & (~valid1)
sky2 = gt2['sky_mask'] & (~valid2)
return gt_pts1, gt_pts2, pr_pts1, pr_pts2, valid1, valid2, sky1, sky2, {}
def compute_loss(self, gt1, gt2, pred1, pred2, **kw):
gt_pts1, gt_pts2, pred_pts1, pred_pts2, mask1, mask2, sky1, sky2, monitoring = \
self.get_all_pts3d(gt1, gt2, pred1, pred2, **kw)
if self.sky_loss_value > 0:
assert self.criterion.reduction == 'none', 'sky_loss_value should be 0 if no conf loss'
# add the sky pixel as "valid" pixels...
mask1 = mask1 | sky1
mask2 = mask2 | sky2
# loss on img1 side
pred_pts1 = pred_pts1[mask1]
gt_pts1 = gt_pts1[mask1]
if self.loss_in_log and self.loss_in_log != 'before':
# this only make sense when depth_mode == 'exp'
pred_pts1 = apply_log_to_norm(pred_pts1)
gt_pts1 = apply_log_to_norm(gt_pts1)
l1 = self.criterion(pred_pts1, gt_pts1)
# loss on gt2 side
pred_pts2 = pred_pts2[mask2]
gt_pts2 = gt_pts2[mask2]
if self.loss_in_log and self.loss_in_log != 'before':
pred_pts2 = apply_log_to_norm(pred_pts2)
gt_pts2 = apply_log_to_norm(gt_pts2)
l2 = self.criterion(pred_pts2, gt_pts2)
if self.sky_loss_value > 0:
assert self.criterion.reduction == 'none', 'sky_loss_value should be 0 if no conf loss'
# ... but force the loss to be high there
l1 = torch.where(sky1[mask1], self.sky_loss_value, l1)
l2 = torch.where(sky2[mask2], self.sky_loss_value, l2)
self_name = type(self).__name__
details = {self_name + '_pts3d_1': float(l1.mean()), self_name + '_pts3d_2': float(l2.mean())}
return Sum((l1, mask1), (l2, mask2)), (details | monitoring)
class Regr3D_ShiftInv (Regr3D):
""" Same than Regr3D but invariant to depth shift.
"""
def get_all_pts3d(self, gt1, gt2, pred1, pred2):
# compute unnormalized points
gt_pts1, gt_pts2, pred_pts1, pred_pts2, mask1, mask2, sky1, sky2, monitoring = \
super().get_all_pts3d(gt1, gt2, pred1, pred2)
# compute median depth
gt_z1, gt_z2 = gt_pts1[..., 2], gt_pts2[..., 2]
pred_z1, pred_z2 = pred_pts1[..., 2], pred_pts2[..., 2]
gt_shift_z = get_joint_pointcloud_depth(gt_z1, gt_z2, mask1, mask2)[:, None, None]
pred_shift_z = get_joint_pointcloud_depth(pred_z1, pred_z2, mask1, mask2)[:, None, None]
# subtract the median depth
gt_z1 -= gt_shift_z
gt_z2 -= gt_shift_z
pred_z1 -= pred_shift_z
pred_z2 -= pred_shift_z
# monitoring = dict(monitoring, gt_shift_z=gt_shift_z.mean().detach(), pred_shift_z=pred_shift_z.mean().detach())
return gt_pts1, gt_pts2, pred_pts1, pred_pts2, mask1, mask2, sky1, sky2, monitoring
class Regr3D_ScaleInv (Regr3D):
""" Same than Regr3D but invariant to depth scale.
if gt_scale == True: enforce the prediction to take the same scale than GT
"""
def get_all_pts3d(self, gt1, gt2, pred1, pred2):
# compute depth-normalized points
gt_pts1, gt_pts2, pred_pts1, pred_pts2, mask1, mask2, sky1, sky2, monitoring = \
super().get_all_pts3d(gt1, gt2, pred1, pred2)
# measure scene scale
_, gt_scale = get_joint_pointcloud_center_scale(gt_pts1, gt_pts2, mask1, mask2)
_, pred_scale = get_joint_pointcloud_center_scale(pred_pts1, pred_pts2, mask1, mask2)
# prevent predictions to be in a ridiculous range
pred_scale = pred_scale.clip(min=1e-3, max=1e3)
# subtract the median depth
if self.gt_scale:
pred_pts1 *= gt_scale / pred_scale
pred_pts2 *= gt_scale / pred_scale
# monitoring = dict(monitoring, pred_scale=(pred_scale/gt_scale).mean())
else:
gt_pts1 /= gt_scale
gt_pts2 /= gt_scale
pred_pts1 /= pred_scale
pred_pts2 /= pred_scale
# monitoring = dict(monitoring, gt_scale=gt_scale.mean(), pred_scale=pred_scale.mean().detach())
return gt_pts1, gt_pts2, pred_pts1, pred_pts2, mask1, mask2, sky1, sky2, monitoring
class Regr3D_ScaleShiftInv (Regr3D_ScaleInv, Regr3D_ShiftInv):
# calls Regr3D_ShiftInv first, then Regr3D_ScaleInv
pass
def get_similarities(desc1, desc2, euc=False):
if euc: # euclidean distance in same range than similarities
dists = (desc1[:, :, None] - desc2[:, None]).norm(dim=-1)
sim = 1 / (1 + dists)
else:
# Compute similarities
sim = desc1 @ desc2.transpose(-2, -1)
return sim
class MatchingCriterion(BaseCriterion):
def __init__(self, reduction='mean', fp=torch.float32):
super().__init__(reduction)
self.fp = fp
def forward(self, a, b, valid_matches=None, euc=False):
assert a.ndim >= 2 and 1 <= a.shape[-1], f'Bad shape = {a.shape}'
dist = self.loss(a.to(self.fp), b.to(self.fp), valid_matches, euc=euc)
# one dimension less or reduction to single value
assert (valid_matches is None and dist.ndim == a.ndim -
1) or self.reduction in ['mean', 'sum', '1-mean', 'none']
if self.reduction == 'none':
return dist
if self.reduction == 'sum':
return dist.sum()
if self.reduction == 'mean':
return dist.mean() if dist.numel() > 0 else dist.new_zeros(())
if self.reduction == '1-mean':
return 1. - dist.mean() if dist.numel() > 0 else dist.new_ones(())
raise ValueError(f'bad {self.reduction=} mode')
def loss(self, a, b, valid_matches=None):
raise NotImplementedError
class InfoNCE(MatchingCriterion):
def __init__(self, temperature=0.07, eps=1e-8, mode='all', **kwargs):
super().__init__(**kwargs)
self.temperature = temperature
self.eps = eps
assert mode in ['all', 'proper', 'dual']
self.mode = mode
def loss(self, desc1, desc2, valid_matches=None, euc=False):
# valid positives are along diagonals
B, N, D = desc1.shape
B2, N2, D2 = desc2.shape
assert B == B2 and D == D2
if valid_matches is None:
valid_matches = torch.ones([B, N], dtype=bool)
# torch.all(valid_matches.sum(dim=-1) > 0) some pairs have no matches????
assert valid_matches.shape == torch.Size([B, N]) and valid_matches.sum() > 0
# Tempered similarities
sim = get_similarities(desc1, desc2, euc) / self.temperature
sim[sim.isnan()] = -torch.inf # ignore nans
# Softmax of positives with temperature
sim = sim.exp_() # save peak memory
positives = sim.diagonal(dim1=-2, dim2=-1)
# Loss
if self.mode == 'all': # Previous InfoNCE
loss = -torch.log((positives / sim.sum(dim=-1).sum(dim=-1, keepdim=True)).clip(self.eps))
elif self.mode == 'proper': # Proper InfoNCE
loss = -(torch.log((positives / sim.sum(dim=-2)).clip(self.eps)) +
torch.log((positives / sim.sum(dim=-1)).clip(self.eps)))
elif self.mode == 'dual': # Dual Softmax
loss = -(torch.log((positives**2 / sim.sum(dim=-1) / sim.sum(dim=-2)).clip(self.eps)))
else:
raise ValueError("This should not happen...")
return loss[valid_matches]
class APLoss (MatchingCriterion):
""" AP loss.
"""
def __init__(self, nq='torch', min=0, max=1, euc=False, **kw):
super().__init__(**kw)
# Exact/True AP loss (not differentiable)
if nq == 0:
nq = 'sklearn' # special case
try:
self.compute_AP = eval('self.compute_true_AP_' + nq)
except:
raise ValueError("Unknown mode %s for AP loss" % nq)
@staticmethod
def compute_true_AP_sklearn(scores, labels):
def compute_AP(label, score):
return average_precision_score(label, score)
aps = scores.new_zeros((scores.shape[0], scores.shape[1]))
label_np = labels.cpu().numpy().astype(bool)
scores_np = scores.cpu().numpy()
for bi in range(scores_np.shape[0]):
for i in range(scores_np.shape[1]):
labels = label_np[bi, i, :]
if labels.sum() < 1:
continue
aps[bi, i] = compute_AP(labels, scores_np[bi, i, :])
return aps
@staticmethod
def compute_true_AP_torch(scores, labels):
assert scores.shape == labels.shape
B, N, M = labels.shape
dev = labels.device
with torch.no_grad():
# sort scores
_, order = scores.sort(dim=-1, descending=True)
# sort labels accordingly
labels = labels[torch.arange(B, device=dev)[:, None, None].expand(order.shape),
torch.arange(N, device=dev)[None, :, None].expand(order.shape),
order]
# compute number of positives per query
npos = labels.sum(dim=-1)
assert torch.all(torch.isclose(npos, npos[0, 0])
), "only implemented for constant number of positives per query"
npos = int(npos[0, 0])
# compute precision at each recall point
posrank = labels.nonzero()[:, -1].view(B, N, npos)
recall = torch.arange(1, 1 + npos, dtype=torch.float32, device=dev)[None, None, :].expand(B, N, npos)
precision = recall / (1 + posrank).float()
# average precision values at all recall points
aps = precision.mean(dim=-1)
return aps
def loss(self, desc1, desc2, valid_matches=None, euc=False): # if matches is None, positives are the diagonal
B, N1, D = desc1.shape
B2, N2, D2 = desc2.shape
assert B == B2 and D == D2
scores = get_similarities(desc1, desc2, euc)
labels = torch.zeros([B, N1, N2], dtype=scores.dtype, device=scores.device)
# allow all diagonal positives and only mask afterwards
labels.diagonal(dim1=-2, dim2=-1)[...] = 1.
apscore = self.compute_AP(scores, labels)
if valid_matches is not None:
apscore = apscore[valid_matches]
return apscore
class MatchingLoss (Criterion, MultiLoss):
"""
Matching loss per image
only compare pixels inside an image but not in the whole batch as what would be done usually
"""
def __init__(self, criterion, withconf=False, use_pts3d=False, negatives_padding=0, blocksize=4096):
super().__init__(criterion)
self.negatives_padding = negatives_padding
self.use_pts3d = use_pts3d
self.blocksize = blocksize
self.withconf = withconf
def add_negatives(self, outdesc2, desc2, batchid, x2, y2):
if self.negatives_padding:
B, H, W, D = desc2.shape
negatives = torch.ones([B, H, W], device=desc2.device, dtype=bool)
negatives[batchid, y2, x2] = False
sel = negatives & (negatives.view([B, -1]).cumsum(dim=-1).view(B, H, W)
<= self.negatives_padding) # take the N-first negatives
outdesc2 = torch.cat([outdesc2, desc2[sel].view([B, -1, D])], dim=1)
return outdesc2
def get_confs(self, pred1, pred2, sel1, sel2):
if self.withconf:
if self.use_pts3d:
outconfs1 = pred1['conf'][sel1]
outconfs2 = pred2['conf'][sel2]
else:
outconfs1 = pred1['desc_conf'][sel1]
outconfs2 = pred2['desc_conf'][sel2]
else:
outconfs1 = outconfs2 = None
return outconfs1, outconfs2
def get_descs(self, pred1, pred2):
if self.use_pts3d:
desc1, desc2 = pred1['pts3d'], pred2['pts3d_in_other_view']
else:
desc1, desc2 = pred1['desc'], pred2['desc']
return desc1, desc2
def get_matching_descs(self, gt1, gt2, pred1, pred2, **kw):
outdesc1 = outdesc2 = outconfs1 = outconfs2 = None
# Recover descs, GT corres and valid mask
desc1, desc2 = self.get_descs(pred1, pred2)
(x1, y1), (x2, y2) = gt1['corres'].unbind(-1), gt2['corres'].unbind(-1)
valid_matches = gt1['valid_corres']
# Select descs that have GT matches
B, N = x1.shape
batchid = torch.arange(B)[:, None].repeat(1, N) # B, N
outdesc1, outdesc2 = desc1[batchid, y1, x1], desc2[batchid, y2, x2] # B, N, D
# Padd with unused negatives
outdesc2 = self.add_negatives(outdesc2, desc2, batchid, x2, y2)
# Gather confs if needed
sel1 = batchid, y1, x1
sel2 = batchid, y2, x2
outconfs1, outconfs2 = self.get_confs(pred1, pred2, sel1, sel2)
return outdesc1, outdesc2, outconfs1, outconfs2, valid_matches, {'use_euclidean_dist': self.use_pts3d}
def blockwise_criterion(self, descs1, descs2, confs1, confs2, valid_matches, euc, rng=np.random, shuffle=True):
loss = None
details = {}
B, N, D = descs1.shape
if N <= self.blocksize: # Blocks are larger than provided descs, compute regular loss
loss = self.criterion(descs1, descs2, valid_matches, euc=euc)
else: # Compute criterion on the blockdiagonal only, after shuffling
# Shuffle if necessary
matches_perm = slice(None)
if shuffle:
matches_perm = np.stack([rng.choice(range(N), size=N, replace=False) for _ in range(B)])
batchid = torch.tile(torch.arange(B), (N, 1)).T
matches_perm = batchid, matches_perm
descs1 = descs1[matches_perm]
descs2 = descs2[matches_perm]
valid_matches = valid_matches[matches_perm]
assert N % self.blocksize == 0, "Error, can't chunk block-diagonal, please check blocksize"
n_chunks = N // self.blocksize
descs1 = descs1.reshape([B * n_chunks, self.blocksize, D]) # [B*(N//blocksize), blocksize, D]
descs2 = descs2.reshape([B * n_chunks, self.blocksize, D]) # [B*(N//blocksize), blocksize, D]
valid_matches = valid_matches.view([B * n_chunks, self.blocksize])
loss = self.criterion(descs1, descs2, valid_matches, euc=euc)
if self.withconf:
confs1, confs2 = map(lambda x: x[matches_perm], (confs1, confs2)) # apply perm to confidences if needed
if self.withconf:
# split confidences between positives/negatives for loss computation
details['conf_pos'] = map(lambda x: x[valid_matches.view(B, -1)], (confs1, confs2))
details['conf_neg'] = map(lambda x: x[~valid_matches.view(B, -1)], (confs1, confs2))
details['Conf1_std'] = confs1.std()
details['Conf2_std'] = confs2.std()
return loss, details
def compute_loss(self, gt1, gt2, pred1, pred2, **kw):
# Gather preds and GT
descs1, descs2, confs1, confs2, valid_matches, monitoring = self.get_matching_descs(
gt1, gt2, pred1, pred2, **kw)
# loss on matches
loss, details = self.blockwise_criterion(descs1, descs2, confs1, confs2,
valid_matches, euc=monitoring.pop('use_euclidean_dist', False))
details[type(self).__name__] = float(loss.mean())
return loss, (details | monitoring)
class ConfMatchingLoss(ConfLoss):
""" Weight matching by learned confidence. Same as ConfLoss but for a matching criterion
Assuming the input matching_loss is a match-level loss.
"""
def __init__(self, pixel_loss, alpha=1., confmode='prod', neg_conf_loss_quantile=False):
super().__init__(pixel_loss, alpha)
self.pixel_loss.withconf = True
self.confmode = confmode
self.neg_conf_loss_quantile = neg_conf_loss_quantile
def aggregate_confs(self, confs1, confs2): # get the confidences resulting from the two view predictions
if self.confmode == 'prod':
confs = confs1 * confs2 if confs1 is not None and confs2 is not None else 1.
elif self.confmode == 'mean':
confs = .5 * (confs1 + confs2) if confs1 is not None and confs2 is not None else 1.
else:
raise ValueError(f"Unknown conf mode {self.confmode}")
return confs
def compute_loss(self, gt1, gt2, pred1, pred2, **kw):
# compute per-pixel loss
loss, details = self.pixel_loss(gt1, gt2, pred1, pred2, **kw)
# Recover confidences for positive and negative samples
conf1_pos, conf2_pos = details.pop('conf_pos')
conf1_neg, conf2_neg = details.pop('conf_neg')
conf_pos = self.aggregate_confs(conf1_pos, conf2_pos)
# weight Matching loss by confidence on positives
conf_pos, log_conf_pos = self.get_conf_log(conf_pos)
conf_loss = loss * conf_pos - self.alpha * log_conf_pos
# average + nan protection (in case of no valid pixels at all)
conf_loss = conf_loss.mean() if conf_loss.numel() > 0 else 0
# Add negative confs loss to give some supervision signal to confidences for pixels that are not matched in GT
if self.neg_conf_loss_quantile:
conf_neg = torch.cat([conf1_neg, conf2_neg])
conf_neg, log_conf_neg = self.get_conf_log(conf_neg)
# recover quantile that will be used for negatives loss value assignment
neg_loss_value = torch.quantile(loss, self.neg_conf_loss_quantile).detach()
neg_loss = neg_loss_value * conf_neg - self.alpha * log_conf_neg
neg_loss = neg_loss.mean() if neg_loss.numel() > 0 else 0
conf_loss = conf_loss + neg_loss
return conf_loss, dict(matching_conf_loss=float(conf_loss), **details)