import torch import numpy as np from vidar.utils.write import write_image from vidar.utils.viz import viz_depth, viz_normals from vidar.geometry.camera import Camera def torch_quantile( input, q, dim = None, keepdim: bool = False, *, interpolation: str = "nearest", out: torch.Tensor = None, ) -> torch.Tensor: """Better torch.quantile for one SCALAR quantile. Using torch.kthvalue. Better than torch.quantile because: - No 2**24 input size limit (pytorch/issues/67592), - Much faster, at least on big input sizes. Arguments: input (torch.Tensor): See torch.quantile. q (float): See torch.quantile. Supports only scalar input currently. dim (int | None): See torch.quantile. keepdim (bool): See torch.quantile. Supports only False currently. interpolation: {"nearest", "lower", "higher"} See torch.quantile. out (torch.Tensor | None): See torch.quantile. Supports only None currently. """ # https://github.com/pytorch/pytorch/issues/64947 # Sanitization: q try: q = float(q) assert 0 <= q <= 1 except Exception: raise ValueError(f"Only scalar input 0<=q<=1 is currently supported (got {q})!") # Handle dim=None case if dim_was_none := dim is None: dim = 0 input = input.reshape((-1,) + (1,) * (input.ndim - 1)) # Set interpolation method if interpolation == "nearest": inter = round elif interpolation == "lower": inter = floor elif interpolation == "higher": inter = ceil else: raise ValueError( "Supported interpolations currently are {'nearest', 'lower', 'higher'} " f"(got '{interpolation}')!" ) # Validate out parameter if out is not None: raise ValueError(f"Only None value is currently supported for out (got {out})!") # Compute k-th value k = inter(q * (input.shape[dim] - 1)) + 1 out = torch.kthvalue(input, k, dim, keepdim=True, out=out)[0] # Handle keepdim and dim=None cases if keepdim: return out if dim_was_none: return out.squeeze() else: return out.squeeze(dim) return out def filter_by_quantile(loss_tensor, valid_range, hard_max=None): if hard_max is not None: loss_tensor = loss_tensor.clamp(max=hard_max) quantile_thresh = torch_quantile(loss_tensor.detach(), valid_range) if hard_max is not None: quantile_thresh = min(quantile_thresh, hard_max) quantile_mask = loss_tensor <= quantile_thresh return loss_tensor[quantile_mask] # def calculate_normal(p0, p1, p2, pad, return_mask=True): # normals = torch.cross(p1 - p0, p2 - p0, 1) # normals = normals / normals.norm(dim=1, keepdim=True) # normals[torch.isnan(normals)] = 0.0 # normals = torch.nn.functional.pad(normals, pad, mode='constant', value=0.0) # if return_mask: # mask = calculate_normal_mask(p0, p1, p2, pad) # return normals, mask # else: # return normals # def calculate_normal_mask(p0, p1, p2, pad): # mask = (p0.sum(1, keepdim=True) != 0) & \ # (p1.sum(1, keepdim=True) != 0) & \ # (p2.sum(1, keepdim=True) != 0) # return torch.nn.functional.pad(mask, pad, mode='constant', value=0) # def calculate_normals(depth, camera): # points = camera.reconstruct_depth_map( # depth, euclidean=False, to_world=False) # p0 = points[:, :, :-1 , :-1] # p1 = points[:, :, 1: , :-1] # p2 = points[:, :, :-1 , 1: ] # normals1, mask1 = calculate_normal(p0, p1, p2, [0, 1, 0, 1]) # p0 = points[:, :, 1: , 1: ] # p1 = points[:, :, :-1 , 1: ] # p2 = points[:, :, 1: , :-1] # normals2, mask2 = calculate_normal(p0, p1, p2, [1, 0, 1, 0]) # p0 = points[:, :, 1: , :-1] # p1 = points[:, :, :-1 , :-1] # p2 = points[:, :, 1: , 1: ] # normals3, mask3 = calculate_normal(p0, p1, p2, [0, 1, 1, 0]) # p0 = points[:, :, :-1 , 1: ] # p1 = points[:, :, 1: , 1: ] # p2 = points[:, :, :-1 , :-1] # normals4, mask4 = calculate_normal(p0, p1, p2, [1, 0, 0, 1]) # normals = torch.stack([normals1, normals2, normals3, normals4], -1) # mask = torch.stack([mask1, mask2, mask3, mask4], -1) # return normals, mask def grad(x1, x2, pad): diff = x1 - x2 diff = torch.nn.functional.pad(diff, pad, mode='constant', value=0) mask = (x1.sum(1, keepdim=True) != 0) & (x2.sum(1, keepdim=True) != 0) mask = torch.nn.functional.pad(mask, pad, mode='constant', value=0) return diff, mask def grad_all(x): x1, x2 = x[..., 1:, :], x[..., :-1, :] diff_h1, mask_h1 = grad(x1, x2, (0, 0, 1, 0)) x1, x2 = x[..., :, 1:], x[..., :, :-1] diff_w1, mask_w1 = grad(x1, x2, (1, 0, 0, 0)) x1, x2 = x[..., :-1, :], x[..., 1:, :] diff_h2, mask_h2 = grad(x1, x2, (0, 0, 0, 1)) x1, x2 = x[..., :, :-1], x[..., :, 1:] diff_w2, mask_w2 = grad(x1, x2, (1, 0, 0, 0)) diff = torch.stack([diff_h1, diff_h2, diff_w1, diff_w2], dim=-1) mask = torch.stack([mask_h1, mask_h2, mask_w1, mask_w2], dim=-1) return diff, mask def calculate_conf(loss, conf, mask, gamma, alpha): loss = gamma * loss * conf[mask] - alpha * torch.log(conf[mask]) return loss def grad_mask(mask): return mask[..., :-1, :-1] & mask[..., 1:, :-1] & mask[..., :-1, 1:] & mask[..., 1:, 1:] def quantile(data, q, floor=None, ceil=None): thr = [] for d in data: flat = d.view(-1) if floor is not None: flat = flat[flat > floor] if ceil is not None: flat = flat[flat < ceil] thr.append(torch.quantile(flat, q)) return torch.tensor(thr, device=data.device, dtype=data.dtype) class LossDepth(torch.nn.Module): def __init__(self): super().__init__() self.cosine_similarity = torch.nn.CosineSimilarity(dim=-1) self.thr_multiplier = 1.0 self.weights = { 'l1': 1.0, 'absrel': 0.1, 'ratio': 0.01, 'normals': 0.01, 'grad': 0.01, } for key, val in dict(**self.weights).items(): self.weights[f'{key}_thr'] = self.thr_multiplier * val self.max_range = 200.0 self.grad_clamp = {'min': 0.0, 'max': self.max_range / 10} def l1_loss(self, pred, gt): return torch.abs(pred - gt).mean(-1) / self.max_range def absrel_loss(self, pred, gt): return (torch.abs(pred - gt) / gt).mean(-1) def ratio_loss(self, pred, gt): ratio = torch.abs(pred / gt) idx_inv = ratio > 1.0 ratio[idx_inv] = 1.0 / ratio[idx_inv] return (1.0 - ratio).mean(-1) def cossim_loss(self, pred, gt): loss = self.cosine_similarity(pred, gt) mask = ~torch.isnan(loss) return (1.0 - loss[mask]) / 2 def grad_loss(self, pred, gt): return torch.abs(pred - gt).clamp(**self.grad_clamp) def points_loss(self, pred, gt, camera): gt = camera.reconstruct_depth_map(gt, euclidean=False, to_world=False) pred = camera.reconstruct_depth_map(pred, euclidean=False, to_world=False) def forward(self, pred_depth, gt_depth, gt_intrinsics, gt_rgb): # Camera for normals calculation pred_depth = pred_depth.float() gt_depth = gt_depth.float() gt_intrinsics = gt_intrinsics.float() gt_rgb = gt_rgb.float() gt_camera = Camera(K=gt_intrinsics, hw=gt_depth) # Valid depth maps mask_gt_depth = gt_depth > 0 # Calculate gradients gt_depth_grad, mask_gt_depth_grad = grad_all(gt_depth) pred_depth_grad, _ = grad_all(pred_depth) max_gt_depth_grad = gt_depth_grad.abs().max(-1)[0].squeeze(1) max_pred_depth_grad = pred_depth_grad.abs().max(-1)[0].squeeze(1) mask_max_gt_depth_grad = mask_gt_depth_grad.float().abs().min(-1)[0].bool().squeeze(1) # Calculate normals gt_normals, mask_gt_normals = calculate_normals(gt_depth, gt_camera) pred_normals, _ = calculate_normals(pred_depth, gt_camera) # Store visuals to return visuals = { 'gt_rgb': gt_rgb, 'gt_depth': gt_depth, 'mask_gt_depth': mask_gt_depth, 'pred_depth': pred_depth, 'gt_normals': gt_normals[..., 0], 'mask_gt_normals': mask_gt_normals[..., 0], 'pred_normals': pred_normals[..., 0], 'max_gt_depth_grad': max_gt_depth_grad.unsqueeze(-1), 'mask_max_gt_depth_grad': mask_max_gt_depth_grad.unsqueeze(-1), 'max_pred_depth_grad': max_pred_depth_grad.unsqueeze(-1), } viz_gt_rgb = visuals['gt_rgb'][0].permute(1, 2, 0).detach().cpu().numpy() viz_gt_depth = viz_depth(visuals['gt_depth'][0]) viz_pred_depth = viz_depth(visuals['pred_depth'][0]) viz_gt_normals = viz_normals(visuals['gt_normals'][0]) viz_pred_normals = viz_normals(visuals['pred_normals'][0]) viz_max_gt_depth_grad = np.repeat(visuals['max_gt_depth_grad'][0].detach().cpu().numpy(), 3, -1) viz_max_pred_depth_grad = np.repeat(visuals['max_pred_depth_grad'][0].detach().cpu().numpy(), 3, -1) viz_max_gt_depth_grad /= viz_max_gt_depth_grad.max() viz_max_gt_depth_grad = viz_max_gt_depth_grad.clip(max=1.0) viz_max_pred_depth_grad /= viz_max_gt_depth_grad.max() viz_max_pred_depth_grad = viz_max_pred_depth_grad.clip(max=1.0) viz_mask_gt_depth = visuals['mask_gt_depth'][0] viz_mask_gt_normals = visuals['mask_gt_normals'][0] viz_mask_max_gt_depth_grad = np.repeat(visuals['mask_max_gt_depth_grad'][0].detach().cpu().numpy(), 3, -1) visuals = { 'gt_rgb': viz_gt_rgb, 'gt_depth': viz_gt_depth, 'mask_gt_depth': viz_mask_gt_depth, 'pred_depth': viz_pred_depth, 'gt_normals': viz_gt_normals, 'mask_gt_normals': viz_mask_gt_normals, 'pred_normals': viz_pred_normals, 'max_gt_depth_grad': viz_max_gt_depth_grad, 'mask_max_gt_depth_grad': viz_mask_max_gt_depth_grad, 'max_pred_depth_grad': viz_max_pred_depth_grad, } col0 = np.concatenate([viz_gt_rgb,viz_gt_rgb], axis=0) col1 = np.concatenate([viz_gt_depth,viz_pred_depth], axis=0) col2 = np.concatenate([viz_gt_normals,viz_pred_normals], axis=0) col3 = np.concatenate([viz_max_gt_depth_grad,viz_max_pred_depth_grad], axis=0) merge = np.concatenate([col0,col1,col3,col2], axis=1) visuals['merge'] = merge gt_points = gt_camera.reconstruct_depth_map(gt_depth, euclidean=False, to_world=False) visuals['gt_points'] = torch.cat([ gt_points[0].view(3, -1).permute(1, 0), gt_rgb[0].view(3, -1).permute(1, 0) * 255, ], -1).detach().cpu().numpy() pred_points = gt_camera.reconstruct_depth_map(pred_depth, euclidean=False, to_world=False) visuals['pred_points'] = torch.cat([ pred_points[0].view(3, -1).permute(1, 0), gt_rgb[0].view(3, -1).permute(1, 0) * 255, ], -1).detach().cpu().numpy() # write_image('viz/gt_rgb.png', viz_gt_rgb) # write_image('viz/gt_depth.png', viz_gt_depth) # write_image('viz/mask_gt_depth.png', viz_mask_gt_depth) # write_image('viz/pred_depth.png', viz_pred_depth) # write_image('viz/gt_normals.png', viz_gt_normals) # write_image('viz/mask_gt_normals.png', viz_mask_gt_normals) # write_image('viz/pred_normals.png', viz_pred_normals) # write_image('viz/max_gt_depth_grad.png', viz_max_gt_depth_grad) # write_image('viz/mask_max_gt_depth_grad.png', viz_mask_max_gt_depth_grad) # write_image('viz/max_pred_depth_grad.png', viz_max_pred_depth_grad) # write_image('viz/merge.png', merge) # Permute tensors for loss calculation gt_depth = gt_depth.permute(0, 2, 3, 1) pred_depth = pred_depth.permute(0, 2, 3, 1) gt_normals = gt_normals.permute(0, 2, 3, 4, 1) pred_normals = pred_normals.permute(0, 2, 3, 4, 1) gt_depth_grad = gt_depth_grad.permute(0, 2, 3, 4, 1) pred_depth_grad = pred_depth_grad.permute(0, 2, 3, 4, 1) # Squeeze masks mask_gt_depth = mask_gt_depth.squeeze(1) mask_gt_depth_grad = mask_gt_depth_grad.squeeze(1) mask_gt_normals = mask_gt_normals.squeeze(1) # Loop over each sample losses = { 'l1': [], 'absrel': [], 'ratio': [], 'normals': [], 'grad': [], 'l1_thr': [], 'absrel_thr': [], 'ratio_thr': [], 'normals_thr': [], 'grad_thr': [], } for b in range(gt_depth.shape[0]): gt_depth_b = gt_depth[b] pred_depth_b = pred_depth[b] mask_gt_depth_b = mask_gt_depth[b] gt_depth_grad_b = gt_depth_grad[b] pred_depth_grad_b = pred_depth_grad[b] mask_gt_depth_grad_b = mask_gt_depth_grad[b] gt_normals_b = gt_normals[b] pred_normals_b = pred_normals[b] mask_gt_normals_b = mask_gt_normals[b] valid_gt_depth_b = gt_depth_b[mask_gt_depth_b] valid_pred_depth_b = pred_depth_b[mask_gt_depth_b] valid_gt_normals_b = gt_normals_b[mask_gt_normals_b] valid_pred_normals_b = pred_normals_b[mask_gt_normals_b] valid_gt_depth_grad_b = gt_depth_grad_b[mask_gt_depth_grad_b] valid_pred_depth_grad_b = pred_depth_grad_b[mask_gt_depth_grad_b] # Get threshold quantile max_gt_depth_grad_b = max_gt_depth_grad[b] thr_max_gt_depth_grad_b = torch_quantile(max_gt_depth_grad_b[max_gt_depth_grad_b > 0], 0.7) # Get threshold gradient mask mask_thr_gt_depth_grad_b = gt_depth_grad_b > thr_max_gt_depth_grad_b valid_thr_gt_depth_grad_b = gt_depth_grad_b[mask_thr_gt_depth_grad_b] valid_thr_pred_depth_grad_b = pred_depth_grad_b[mask_thr_gt_depth_grad_b] # Get threshold depth mask mask_thr_max_gt_depth_grad_b = max_gt_depth_grad_b > thr_max_gt_depth_grad_b mask_thr_max_gt_depth_b = mask_thr_max_gt_depth_grad_b * mask_gt_depth_b valid_thr_gt_depth_b = gt_depth_b[mask_thr_max_gt_depth_b] valid_thr_pred_depth_b = pred_depth_b[mask_thr_max_gt_depth_b] # Get threshold normals mask mask_thr_max_gt_normals_b = mask_thr_gt_depth_grad_b[..., 0] * mask_gt_normals_b valid_thr_gt_normals_b = gt_normals_b[mask_thr_max_gt_normals_b] valid_thr_pred_normals_b = pred_normals_b[mask_thr_max_gt_normals_b] # Depth losses l1_loss = self.l1_loss(valid_pred_depth_b, valid_gt_depth_b) losses['l1'].append(l1_loss.mean()) absrel_loss = self.absrel_loss(valid_pred_depth_b, valid_gt_depth_b) losses['absrel'].append(absrel_loss.mean()) ratio_loss = self.ratio_loss(valid_pred_depth_b, valid_gt_depth_b) losses['ratio'].append(ratio_loss.mean()) # Threshold Depth losses l1_thr_loss = self.l1_loss(valid_thr_pred_depth_b, valid_thr_gt_depth_b) losses['l1_thr'].append(l1_thr_loss.mean()) absrel_thr_loss = self.absrel_loss(valid_thr_pred_depth_b, valid_thr_gt_depth_b) losses['absrel_thr'].append(absrel_thr_loss.mean()) ratio_thr_loss = self.ratio_loss(valid_thr_pred_depth_b, valid_thr_gt_depth_b) losses['ratio_thr'].append(ratio_thr_loss.mean()) # Normals losses normals_loss = self.cossim_loss(valid_pred_normals_b, valid_gt_normals_b) losses['normals'].append(normals_loss.mean()) # Normals Threshold losses normals_thr_loss = self.cossim_loss(valid_thr_pred_normals_b, valid_thr_gt_normals_b) losses['normals_thr'].append(normals_thr_loss.mean()) # Gradient losses grad_loss = self.grad_loss(valid_pred_depth_grad_b, valid_gt_depth_grad_b) losses['grad'].append(grad_loss.mean()) # Threshold Gradient losses grad_thr_loss = self.grad_loss(valid_thr_pred_depth_grad_b, valid_thr_gt_depth_grad_b) losses['grad_thr'].append(grad_thr_loss.mean()) loss = [] losses_avg = {} for key, val in losses.items(): losses_avg[f'loss_{key}'] = torch.stack(val).mean() loss.append(self.weights[key] * losses_avg[f'loss_{key}']) losses_avg['loss'] = torch.stack(loss).sum() return losses_avg, visuals