spann3r/loss.py

import torch 
from dust3r.losses import Criterion, MultiLoss
from dust3r.inference import get_pred_pts3d
from dust3r.utils.misc import invalid_to_zeros, invalid_to_nans
from dust3r.utils.geometry import inv, geotrf


def Sum(losses, masks, conf=None):
    loss, mask = losses[0], masks[0]
    if loss.ndim > 0:
        # we are actually returning the loss for every pixels
        if conf is not None:
            return losses, masks, conf
        return losses, masks
    else:
        # we are returning the global loss
        for loss2 in losses[1:]:
            loss = loss + loss2
        return loss


def get_norm_factor(pts, norm_mode='avg_dis', valids=None, fix_first=True):
    assert pts[0].ndim >= 3 and pts[0].shape[-1] == 3
    assert pts[1] is None or (pts[1].ndim >= 3 and pts[1].shape[-1] == 3)
    norm_mode, dis_mode = norm_mode.split('_')
    
    nan_pts = []
    nnzs = []

    if norm_mode == 'avg':
        # gather all points together (joint normalization)
        
        for i, pt in enumerate(pts):
            nan_pt, nnz = invalid_to_zeros(pt, valids[i], ndim=3)
            nan_pts.append(nan_pt)
            nnzs.append(nnz)
            
            if fix_first:
                break
        all_pts = torch.cat(nan_pts, dim=1)

        # compute distance to origin
        all_dis = all_pts.norm(dim=-1)
        if dis_mode == 'dis':
            pass  # do nothing
        elif dis_mode == 'log1p':
            all_dis = torch.log1p(all_dis)
        else:
            raise ValueError(f'bad {dis_mode=}')

        norm_factor = all_dis.sum(dim=1) / (torch.cat(nnzs).sum() + 1e-8)
    else:
        raise ValueError(f'Not implemented {norm_mode=}')
    
    norm_factor = norm_factor.clip(min=1e-8)
    while norm_factor.ndim < pts[0].ndim:
        norm_factor.unsqueeze_(-1)
    
    return norm_factor


def normalize_pointcloud_t(pts, norm_mode='avg_dis', valids=None, fix_first=True, gt=False):
    if gt:
        norm_factor = get_norm_factor(pts, norm_mode, valids, fix_first)
        res = []
    
        for i, pt in enumerate(pts):
            res.append(pt / norm_factor)
    
    else:
        pts_l, pts_r = pts
        # use pts_l and pts_r[-1] as pts to normalize
        norm_factor = get_norm_factor(pts_l + [pts_r[-1]], norm_mode, valids, fix_first)

        res_l = []
        res_r = []

        for i in range(len(pts_l)):
            res_l.append(pts_l[i] / norm_factor)
            res_r.append(pts_r[i] / norm_factor)
        
        res = [res_l, res_r]
    
    return res, norm_factor


@torch.no_grad()
def get_joint_pointcloud_depth(zs, valid_masks=None, quantile=0.5):
    # set invalid points to NaN
    _zs = []
    for i in range(len(zs)):
        valid_mask = valid_masks[i] if valid_masks is not None else None
        _z = invalid_to_nans(zs[i], valid_mask).reshape(len(zs[i]), -1)
        _zs.append(_z)
    
    _zs = torch.cat(_zs, dim=-1)

    # compute median depth overall (ignoring nans)
    if quantile == 0.5:
        shift_z = torch.nanmedian(_zs, dim=-1).values
    else:
        shift_z = torch.nanquantile(_zs, quantile, dim=-1)
    return shift_z  # (B,)


@torch.no_grad()
def get_joint_pointcloud_center_scale(pts, valid_masks=None, z_only=False, center=True):
    # set invalid points to NaN
    
    _pts = []
    for i in range(len(pts)):
        valid_mask = valid_masks[i] if valid_masks is not None else None
        _pt = invalid_to_nans(pts[i], valid_mask).reshape(len(pts[i]), -1, 3)
        _pts.append(_pt)
    
    _pts = torch.cat(_pts, dim=1)

    # compute median center
    _center = torch.nanmedian(_pts, dim=1, keepdim=True).values  # (B,1,3)
    if z_only:
        _center[..., :2] = 0  # do not center X and Y

    # compute median norm
    _norm = ((_pts - _center) if center else _pts).norm(dim=-1)
    scale = torch.nanmedian(_norm, dim=1).values
    return _center[:, None, :, :], scale[:, None, None, None]


class Regr3D_t(Criterion, MultiLoss):
    def __init__(self, criterion, norm_mode='avg_dis', 
                 gt_scale=False, fix_first=True):
        super().__init__(criterion)
        self.norm_mode = norm_mode
        self.gt_scale = gt_scale
        self.fix_first = fix_first
    
    def get_all_pts3d_t(self, gts, preds, dist_clip=None):
        # everything is normalized w.r.t. camera of view1
        in_camera1 = inv(gts[0]['camera_pose'])
        
        gt_pts = []
        valids = []
        pr_pts = []
        pr_pts_l = []
        pr_pts_r = []
        
        for i, gt in enumerate(gts):
            # in_camera1: Bs, 4, 4 gt['pts3d']: Bs, H, W, 3
            gt_pts.append(geotrf(in_camera1, gt['pts3d']))
            
            valid = gt['valid_mask'].clone()
            
            if dist_clip is not None:
                # points that are too far-away == invalid
                dis = gt['pts3d'].norm(dim=-1)
                valid = valid & (dis <= dist_clip)
            
            valids.append(valid)
            
            if i != len(gts)-1:
                pr_pts_l.append(get_pred_pts3d(gt, preds[i][0], use_pose=(i!=0)))
            
            if i != 0:
                pr_pts_r.append(get_pred_pts3d(gt, preds[i-1][1], use_pose=(i!=0)))
                
        
        pr_pts = (pr_pts_l, pr_pts_r)

        if self.norm_mode:
            pr_pts, pr_factor = normalize_pointcloud_t(pr_pts, self.norm_mode, valids, fix_first=self.fix_first, gt=False)
        else:
            pr_factor = None
        

        if self.norm_mode and not self.gt_scale:
            gt_pts, gt_factor = normalize_pointcloud_t(gt_pts, self.norm_mode, valids, fix_first=self.fix_first, gt=True)
        else:
            gt_factor = None

        return gt_pts, pr_pts, gt_factor, pr_factor, valids, {}


    
    def compute_frame_loss(self, gts, preds, **kw):
        gt_pts, pred_pts, gt_factor, pr_factor, masks, monitoring = \
            self.get_all_pts3d_t(gts, preds, **kw)
        
        pred_pts_l, pred_pts_r = pred_pts

        
        loss_all = []
        mask_all = []
        conf_all = []

        loss_left = 0
        loss_right = 0
        pred_conf_l = 0
        pred_conf_r = 0


        for i in range(len(gt_pts)):

            # Left (Reference)
             if i != len(gt_pts)-1:
                frame_loss = self.criterion(pred_pts_l[i][masks[i]], gt_pts[i][masks[i]])

                loss_all.append(frame_loss)
                mask_all.append(masks[i])
                conf_all.append(preds[i][0]['conf'])

                # To compare target/reference loss
                if i != 0:
                    loss_left += frame_loss.cpu().detach().numpy().mean()
                    pred_conf_l += preds[i][0]['conf'].cpu().detach().numpy().mean()
            
            # Right (Target)
             if i != 0:
                frame_loss = self.criterion(pred_pts_r[i-1][masks[i]], gt_pts[i][masks[i]])

                loss_all.append(frame_loss)
                mask_all.append(masks[i])
                conf_all.append(preds[i-1][1]['conf'])

                # To compare target/reference loss
                if i != len(gt_pts)-1:
                    loss_right += frame_loss.cpu().detach().numpy().mean()
                    pred_conf_r += preds[i-1][1]['conf'].cpu().detach().numpy().mean()
        
        if pr_factor is not None and gt_factor is not None:
            filter_factor = pr_factor[pr_factor > gt_factor]
        else:
            filter_factor = []
        
        if len(filter_factor) > 0:
            factor_loss = (filter_factor - gt_factor).abs().mean()
        else:
            factor_loss = 0.0
        
        self_name = type(self).__name__
        details = {self_name+'_pts3d_1': float(loss_all[0].mean()), 
                   self_name+'_pts3d_2': float(loss_all[1].mean()), 
                   self_name+'loss_left': float(loss_left), 
                   self_name+'loss_right': float(loss_right),
                   self_name+'conf_left': float(pred_conf_l),
                   self_name+'conf_right': float(pred_conf_r)}

        return Sum(loss_all, mask_all, conf_all), (details | monitoring), factor_loss


class ConfLoss_t(MultiLoss):
    """ Weighted regression by learned confidence.
        Assuming the input pixel_loss is a pixel-level regression loss.

    Principle:
        high-confidence means high conf = 0.1 ==> conf_loss = x / 10 + alpha*log(10)
        low  confidence means low  conf = 10  ==> conf_loss = x * 10 - alpha*log(10) 

        alpha: hyperparameter
    """

    def __init__(self, pixel_loss, alpha=1):
        super().__init__()
        assert alpha > 0
        self.alpha = alpha
        self.pixel_loss = pixel_loss.with_reduction('none')

    def get_name(self):
        return f'ConfLoss({self.pixel_loss})'

    def get_conf_log(self, x):
        return x, torch.log(x)

    def compute_frame_loss(self, gts, preds, **kw):
        # compute per-pixel loss
        (losses, masks, confs), details, loss_factor = self.pixel_loss.compute_frame_loss(gts, preds, **kw)

        # weight by confidence
        conf_losses = []
        conf_sum = 0
        for i in range(len(losses)):
            conf, log_conf = self.get_conf_log(confs[i][masks[i]])
            conf_sum += conf.mean()
            conf_loss = losses[i] * conf - self.alpha * log_conf
            conf_loss = conf_loss.mean() if conf_loss.numel() > 0 else 0
            conf_losses.append(conf_loss)
        
        conf_losses = torch.stack(conf_losses) * 2.0
        conf_loss_mean = conf_losses.mean()
            

        return conf_loss_mean, dict(conf_loss_1=float(conf_losses[0]), conf_loss2=float(conf_losses[1]), conf_mean=conf_sum/len(losses), **details), loss_factor


class Regr3D_t_ShiftInv (Regr3D_t):
    """ Same than Regr3D but invariant to depth shift.
    """

    def get_all_pts3d_t(self, gts, preds):
        # compute unnormalized points
        gt_pts, pred_pts, gt_factor, pr_factor, masks, monitoring = \
            super().get_all_pts3d_t(gts, preds)
            
        pred_pts_l, pred_pts_r = pred_pts
        gt_zs = [gt_pt[..., 2] for gt_pt in gt_pts]
        
        pred_zs = [pred_pt[..., 2] for pred_pt in pred_pts_l]
        pred_zs.append(pred_pts_r[-1][..., 2])

        # compute median depth
        gt_shift_z = get_joint_pointcloud_depth(gt_zs, masks)[:, None, None]
        pred_shift_z = get_joint_pointcloud_depth(pred_zs, masks)[:, None, None]

        # subtract the median depth
        for i in range(len(gt_pts)):
            gt_pts[i][..., 2] -= gt_shift_z
        
        for i in range(len(pred_pts)):
            for j in range(len(pred_pts[i])):
                pred_pts[i][j][..., 2] -= pred_shift_z

        monitoring = dict(monitoring, gt_shift_z=gt_shift_z.mean().detach(), pred_shift_z=pred_shift_z.mean().detach())
        return gt_pts, pred_pts, gt_factor, pr_factor, masks, monitoring


class Regr3D_t_ScaleInv (Regr3D_t):
    """ Same than Regr3D but invariant to depth shift.
        if gt_scale == True: enforce the prediction to take the same scale than GT
    """

    def get_all_pts3d_t(self, gts, preds):
        # compute depth-normalized points
        gt_pts, pred_pts, gt_factor, pr_factor, masks, monitoring  = super().get_all_pts3d_t(gts, preds)

        # measure scene scale
        
        pred_pts_l, pred_pts_r = pred_pts
        
        pred_pts_all = [pred_pt for pred_pt in pred_pts_l]
        pred_pts_all.append(pred_pts_r[-1])
        
        
        _, gt_scale = get_joint_pointcloud_center_scale(gt_pts, masks)
        _, pred_scale = get_joint_pointcloud_center_scale(pred_pts_all, masks)

        # 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:
            for i in range(len(pred_pts)):
                for j in range(len(pred_pts[i])):
                    pred_pts[i][j] *= gt_scale / pred_scale

        else:
            for i in range(len(pred_pts)):
                for j in range(len(pred_pts[i])):
                    pred_pts[i][j] *= pred_scale / gt_scale
            
            for i in range(len(gt_pts)):
                gt_pts[i] *= gt_scale / pred_scale
        
        monitoring = dict(monitoring, gt_scale=gt_scale.mean(), pred_scale=pred_scale.mean().detach())

        return gt_pts, pred_pts, gt_factor, pr_factor, masks, monitoring


class Regr3D_t_ScaleShiftInv (Regr3D_t_ScaleInv, Regr3D_t_ShiftInv):
    # calls Regr3D_ShiftInv first, then Regr3D_ScaleInv
    pass