facebook/map-anything
132
1# Copyright (c) Meta Platforms, Inc. and affiliates.2#3# This source code is licensed under the Apache License, Version 2.04# found in the LICENSE file in the root directory of this source tree.5 6"""7Multi-view geometric losses for training 3D reconstruction models.8 9References: DUSt3R & MASt3R10"""11 12import math13from copy import copy, deepcopy14 15import einops as ein16import torch17import torch.nn as nn18 19from mapanything.utils.geometry import (20 angle_diff_vec3,21 apply_log_to_norm,22 closed_form_pose_inverse,23 convert_ray_dirs_depth_along_ray_pose_trans_quats_to_pointmap,24 geotrf,25 normalize_multiple_pointclouds,26 quaternion_inverse,27 quaternion_multiply,28 quaternion_to_rotation_matrix,29 transform_pose_using_quats_and_trans_2_to_1,30)31 32 33def get_loss_terms_and_details(34 losses_dict, valid_masks, self_name, n_views, flatten_across_image_only35):36 """37 Helper function to generate loss terms and details for different loss types.38 39 Args:40 losses_dict (dict): Dictionary mapping loss types to their values.41 Format: {42 'loss_type': {43 'values': list_of_loss_tensors or single_tensor,44 'use_mask': bool,45 'is_multi_view': bool46 }47 }48 valid_masks (list): List of valid masks for each view.49 self_name (str): Name of the loss class.50 n_views (int): Number of views.51 flatten_across_image_only (bool): Whether flattening was done across image only.52 53 Returns:54 tuple: (loss_terms, details) where loss_terms is a list of tuples (loss, mask, type)55 and details is a dictionary of loss details.56 """57 loss_terms = []58 details = {}59 60 for loss_type, loss_info in losses_dict.items():61 values = loss_info["values"]62 use_mask = loss_info["use_mask"]63 is_multi_view = loss_info["is_multi_view"]64 if is_multi_view:65 # Handle multi-view losses (list of tensors)66 view_loss_details = []67 for i in range(n_views):68 mask = valid_masks[i] if use_mask else None69 loss_terms.append((values[i], mask, loss_type))70 71 # Add details for individual view72 if not flatten_across_image_only or not use_mask:73 values_after_masking = values[i]74 else:75 values_after_masking = values[i][mask]76 77 if values_after_masking.numel() > 0:78 view_loss_detail = float(values_after_masking.mean())79 if view_loss_detail > 0:80 details[f"{self_name}_{loss_type}_view{i + 1}"] = (81 view_loss_detail82 )83 view_loss_details.append(view_loss_detail)84 # Add average across views85 if len(view_loss_details) > 0:86 details[f"{self_name}_{loss_type}_avg"] = sum(view_loss_details) / len(87 view_loss_details88 )89 else:90 # Handle single tensor losses91 if values is not None:92 loss_terms.append((values, None, loss_type))93 if values.numel() > 0:94 loss_detail = float(values.mean())95 if loss_detail > 0:96 details[f"{self_name}_{loss_type}"] = loss_detail97 98 return loss_terms, details99 100 101def _smooth(err: torch.FloatTensor, beta: float = 0.0) -> torch.FloatTensor:102 if beta == 0:103 return err104 else:105 return torch.where(err < beta, 0.5 * err.square() / beta, err - 0.5 * beta)106 107 108def compute_normal_loss(points, gt_points, mask):109 """110 Compute the normal loss between the predicted and ground truth points.111 References:112 https://github.com/microsoft/MoGe/blob/a8c37341bc0325ca99b9d57981cc3bb2bd3e255b/moge/train/losses.py#L205113 114 Args:115 points (torch.Tensor): Predicted points. Shape: (..., H, W, 3).116 gt_points (torch.Tensor): Ground truth points. Shape: (..., H, W, 3).117 mask (torch.Tensor): Mask indicating valid points. Shape: (..., H, W).118 119 Returns:120 torch.Tensor: Normal loss.121 """122 height, width = points.shape[-3:-1]123 124 leftup, rightup, leftdown, rightdown = (125 points[..., :-1, :-1, :],126 points[..., :-1, 1:, :],127 points[..., 1:, :-1, :],128 points[..., 1:, 1:, :],129 )130 upxleft = torch.cross(rightup - rightdown, leftdown - rightdown, dim=-1)131 leftxdown = torch.cross(leftup - rightup, rightdown - rightup, dim=-1)132 downxright = torch.cross(leftdown - leftup, rightup - leftup, dim=-1)133 rightxup = torch.cross(rightdown - leftdown, leftup - leftdown, dim=-1)134 135 gt_leftup, gt_rightup, gt_leftdown, gt_rightdown = (136 gt_points[..., :-1, :-1, :],137 gt_points[..., :-1, 1:, :],138 gt_points[..., 1:, :-1, :],139 gt_points[..., 1:, 1:, :],140 )141 gt_upxleft = torch.cross(142 gt_rightup - gt_rightdown, gt_leftdown - gt_rightdown, dim=-1143 )144 gt_leftxdown = torch.cross(145 gt_leftup - gt_rightup, gt_rightdown - gt_rightup, dim=-1146 )147 gt_downxright = torch.cross(gt_leftdown - gt_leftup, gt_rightup - gt_leftup, dim=-1)148 gt_rightxup = torch.cross(149 gt_rightdown - gt_leftdown, gt_leftup - gt_leftdown, dim=-1150 )151 152 mask_leftup, mask_rightup, mask_leftdown, mask_rightdown = (153 mask[..., :-1, :-1],154 mask[..., :-1, 1:],155 mask[..., 1:, :-1],156 mask[..., 1:, 1:],157 )158 mask_upxleft = mask_rightup & mask_leftdown & mask_rightdown159 mask_leftxdown = mask_leftup & mask_rightdown & mask_rightup160 mask_downxright = mask_leftdown & mask_rightup & mask_leftup161 mask_rightxup = mask_rightdown & mask_leftup & mask_leftdown162 163 MIN_ANGLE, MAX_ANGLE, BETA_RAD = math.radians(1), math.radians(90), math.radians(3)164 165 loss = (166 mask_upxleft167 * _smooth(168 angle_diff_vec3(upxleft, gt_upxleft).clamp(MIN_ANGLE, MAX_ANGLE),169 beta=BETA_RAD,170 )171 + mask_leftxdown172 * _smooth(173 angle_diff_vec3(leftxdown, gt_leftxdown).clamp(MIN_ANGLE, MAX_ANGLE),174 beta=BETA_RAD,175 )176 + mask_downxright177 * _smooth(178 angle_diff_vec3(downxright, gt_downxright).clamp(MIN_ANGLE, MAX_ANGLE),179 beta=BETA_RAD,180 )181 + mask_rightxup182 * _smooth(183 angle_diff_vec3(rightxup, gt_rightxup).clamp(MIN_ANGLE, MAX_ANGLE),184 beta=BETA_RAD,185 )186 )187 188 total_valid_mask = mask_upxleft | mask_leftxdown | mask_downxright | mask_rightxup189 valid_count = total_valid_mask.sum()190 if valid_count > 0:191 loss = loss.sum() / (valid_count * (4 * max(points.shape[-3:-1])))192 else:193 loss = 0 * loss.sum()194 195 return loss196 197 198def compute_gradient_loss(prediction, gt_target, mask):199 """200 Compute the gradient loss between the prediction and GT target at valid points.201 References:202 https://docs.nerf.studio/_modules/nerfstudio/model_components/losses.html#GradientLoss203 https://github.com/autonomousvision/monosdf/blob/main/code/model/loss.py204 205 Args:206 prediction (torch.Tensor): Predicted scene representation. Shape: (B, H, W, C).207 gt_target (torch.Tensor): Ground truth scene representation. Shape: (B, H, W, C).208 mask (torch.Tensor): Mask indicating valid points. Shape: (B, H, W).209 """210 # Expand mask to match number of channels in prediction211 mask = mask[..., None].expand(-1, -1, -1, prediction.shape[-1])212 summed_mask = torch.sum(mask, (1, 2, 3))213 214 # Compute the gradient of the prediction and GT target215 diff = prediction - gt_target216 diff = torch.mul(mask, diff)217 218 # Gradient in x direction219 grad_x = torch.abs(diff[:, :, 1:] - diff[:, :, :-1])220 mask_x = torch.mul(mask[:, :, 1:], mask[:, :, :-1])221 grad_x = torch.mul(mask_x, grad_x)222 223 # Gradient in y direction224 grad_y = torch.abs(diff[:, 1:, :] - diff[:, :-1, :])225 mask_y = torch.mul(mask[:, 1:, :], mask[:, :-1, :])226 grad_y = torch.mul(mask_y, grad_y)227 228 # Clamp the outlier gradients229 grad_x = grad_x.clamp(max=100)230 grad_y = grad_y.clamp(max=100)231 232 # Compute the total loss233 image_loss = torch.sum(grad_x, (1, 2, 3)) + torch.sum(grad_y, (1, 2, 3))234 num_valid_pixels = torch.sum(summed_mask)235 if num_valid_pixels > 0:236 image_loss = torch.sum(image_loss) / num_valid_pixels237 else:238 image_loss = 0 * torch.sum(image_loss)239 240 return image_loss241 242 243def compute_gradient_matching_loss(prediction, gt_target, mask, scales=4):244 """245 Compute the multi-scale gradient matching loss between the prediction and GT target at valid points.246 This loss biases discontinuities to be sharp and to coincide with discontinuities in the ground truth.247 More info in MiDAS: https://arxiv.org/pdf/1907.01341.pdf; Equation 11248 References:249 https://docs.nerf.studio/_modules/nerfstudio/model_components/losses.html#GradientLoss250 https://github.com/autonomousvision/monosdf/blob/main/code/model/loss.py251 252 Args:253 prediction (torch.Tensor): Predicted scene representation. Shape: (B, H, W, C).254 gt_target (torch.Tensor): Ground truth scene representation. Shape: (B, H, W, C).255 mask (torch.Tensor): Mask indicating valid points. Shape: (B, H, W).256 scales (int): Number of scales to compute the loss at. Default: 4.257 """258 # Define total loss259 total_loss = 0.0260 261 # Compute the gradient loss at different scales262 for scale in range(scales):263 step = pow(2, scale)264 grad_loss = compute_gradient_loss(265 prediction[:, ::step, ::step],266 gt_target[:, ::step, ::step],267 mask[:, ::step, ::step],268 )269 total_loss += grad_loss270 271 return total_loss272 273 274def Sum(*losses_and_masks):275 """276 Aggregates multiple losses into a single loss value or returns the original losses.277 278 Args:279 *losses_and_masks: Variable number of tuples, each containing (loss, mask, rep_type)280 - loss: Tensor containing loss values281 - mask: Mask indicating valid pixels/regions282 - rep_type: String indicating the type of representation (e.g., 'pts3d', 'depth')283 284 Returns:285 If the first loss has dimensions > 0:286 Returns the original list of (loss, mask, rep_type) tuples287 Otherwise:288 Returns a scalar tensor that is the sum of all loss values289 """290 loss, mask, rep_type = losses_and_masks[0]291 if loss.ndim > 0:292 # we are actually returning the loss for every pixels293 return losses_and_masks294 else:295 # we are returning the global loss296 for loss2, mask2, rep_type2 in losses_and_masks[1:]:297 loss = loss + loss2298 return loss299 300 301class BaseCriterion(nn.Module):302 "Base Criterion to support different reduction methods"303 304 def __init__(self, reduction="mean"):305 super().__init__()306 self.reduction = reduction307 308 309class LLoss(BaseCriterion):310 "L-norm loss"311 312 def forward(self, a, b, **kwargs):313 assert a.shape == b.shape and a.ndim >= 2 and 1 <= a.shape[-1] <= 4, (314 f"Bad shape = {a.shape}"315 )316 dist = self.distance(a, b, **kwargs)317 assert dist.ndim == a.ndim - 1 # one dimension less318 if self.reduction == "none":319 return dist320 if self.reduction == "sum":321 return dist.sum()322 if self.reduction == "mean":323 return dist.mean() if dist.numel() > 0 else dist.new_zeros(())324 raise ValueError(f"bad {self.reduction=} mode")325 326 def distance(self, a, b, **kwargs):327 raise NotImplementedError()328 329 330class L1Loss(LLoss):331 "L1 distance"332 333 def distance(self, a, b, **kwargs):334 return torch.abs(a - b).sum(dim=-1)335 336 337class L2Loss(LLoss):338 "Euclidean (L2 Norm) distance"339 340 def distance(self, a, b, **kwargs):341 return torch.norm(a - b, dim=-1)342 343 344class GenericLLoss(LLoss):345 "Criterion that supports different L-norms"346 347 def distance(self, a, b, loss_type, **kwargs):348 if loss_type == "l1":349 # L1 distance350 return torch.abs(a - b).sum(dim=-1)351 elif loss_type == "l2":352 # Euclidean (L2 norm) distance353 return torch.norm(a - b, dim=-1)354 else:355 raise ValueError(356 f"Unsupported loss type: {loss_type}. Supported types are 'l1' and 'l2'."357 )358 359 360class FactoredLLoss(LLoss):361 "Criterion that supports different L-norms for the factored loss functions"362 363 def __init__(364 self,365 reduction="mean",366 points_loss_type="l2",367 depth_loss_type="l1",368 ray_directions_loss_type="l1",369 pose_quats_loss_type="l1",370 pose_trans_loss_type="l1",371 scale_loss_type="l1",372 ):373 super().__init__(reduction)374 self.points_loss_type = points_loss_type375 self.depth_loss_type = depth_loss_type376 self.ray_directions_loss_type = ray_directions_loss_type377 self.pose_quats_loss_type = pose_quats_loss_type378 self.pose_trans_loss_type = pose_trans_loss_type379 self.scale_loss_type = scale_loss_type380 381 def _distance(self, a, b, loss_type):382 if loss_type == "l1":383 # L1 distance384 return torch.abs(a - b).sum(dim=-1)385 elif loss_type == "l2":386 # Euclidean (L2 norm) distance387 return torch.norm(a - b, dim=-1)388 else:389 raise ValueError(f"Unsupported loss type: {loss_type}.")390 391 def distance(self, a, b, factor, **kwargs):392 if factor == "points":393 return self._distance(a, b, self.points_loss_type)394 elif factor == "depth":395 return self._distance(a, b, self.depth_loss_type)396 elif factor == "ray_directions":397 return self._distance(a, b, self.ray_directions_loss_type)398 elif factor == "pose_quats":399 return self._distance(a, b, self.pose_quats_loss_type)400 elif factor == "pose_trans":401 return self._distance(a, b, self.pose_trans_loss_type)402 elif factor == "scale":403 return self._distance(a, b, self.scale_loss_type)404 else:405 raise ValueError(f"Unsupported factor type: {factor}.")406 407 408class RobustRegressionLoss(LLoss):409 """410 Generalized Robust Loss introduced in https://arxiv.org/abs/1701.03077.411 """412 413 def __init__(self, alpha=0.5, scaling_c=0.25, reduction="mean"):414 """415 Initialize the Robust Regression Loss.416 417 Args:418 alpha (float): Shape parameter controlling the robustness of the loss.419 Lower values make the loss more robust to outliers. Default: 0.5.420 scaling_c (float): Scale parameter controlling the transition between421 quadratic and robust behavior. Default: 0.1.422 reduction (str): Specifies the reduction to apply to the output:423 'none' | 'mean' | 'sum'. Default: 'mean'.424 """425 super().__init__(reduction)426 self.alpha = alpha427 self.scaling_c = scaling_c428 429 def distance(self, a, b, **kwargs):430 error_scaled = torch.sum(((a - b) / self.scaling_c) ** 2, dim=-1)431 robust_loss = (abs(self.alpha - 2) / self.alpha) * (432 torch.pow((error_scaled / abs(self.alpha - 2)) + 1, self.alpha / 2) - 1433 )434 return robust_loss435 436 437class BCELoss(BaseCriterion):438 """Binary Cross Entropy loss"""439 440 def forward(self, predicted_logits, reference_mask):441 """442 Args:443 predicted_logits: (B, H, W) tensor of predicted logits for the mask444 reference_mask: (B, H, W) tensor of reference mask445 446 Returns:447 loss: scalar tensor of the BCE loss448 """449 bce_loss = torch.nn.functional.binary_cross_entropy_with_logits(450 predicted_logits, reference_mask.float()451 )452 453 return bce_loss454 455 456class Criterion(nn.Module):457 """458 Base class for all criterion modules that wrap a BaseCriterion.459 460 This class serves as a wrapper around BaseCriterion objects, providing461 additional functionality like naming and reduction mode control.462 463 Args:464 criterion (BaseCriterion): The base criterion to wrap.465 """466 467 def __init__(self, criterion=None):468 super().__init__()469 assert isinstance(criterion, BaseCriterion), (470 f"{criterion} is not a proper criterion!"471 )472 self.criterion = copy(criterion)473 474 def get_name(self):475 """476 Returns a string representation of this criterion.477 478 Returns:479 str: A string containing the class name and the wrapped criterion.480 """481 return f"{type(self).__name__}({self.criterion})"482 483 def with_reduction(self, mode="none"):484 """485 Creates a deep copy of this criterion with the specified reduction mode.486 487 This method recursively sets the reduction mode for this criterion and488 any chained MultiLoss criteria.489 490 Args:491 mode (str): The reduction mode to set. Default: "none".492 493 Returns:494 Criterion: A new criterion with the specified reduction mode.495 """496 res = loss = deepcopy(self)497 while loss is not None:498 assert isinstance(loss, Criterion)499 loss.criterion.reduction = mode # make it return the loss for each sample500 loss = loss._loss2 # we assume loss is a Multiloss501 return res502 503 504class MultiLoss(nn.Module):505 """506 Base class for combinable loss functions with automatic tracking of individual loss values.507 508 This class enables easy combination of multiple loss functions through arithmetic operations:509 loss = MyLoss1() + 0.1*MyLoss2()510 511 The combined loss functions maintain their individual weights and the forward pass512 automatically computes and aggregates all losses while tracking individual loss values.513 514 Usage:515 Inherit from this class and override get_name() and compute_loss() methods.516 517 Attributes:518 _alpha (float): Weight multiplier for this loss component.519 _loss2 (MultiLoss): Reference to the next loss in the chain, if any.520 """521 522 def __init__(self):523 """Initialize the MultiLoss with default weight of 1 and no chained loss."""524 super().__init__()525 self._alpha = 1526 self._loss2 = None527 528 def compute_loss(self, *args, **kwargs):529 """530 Compute the loss value for this specific loss component.531 532 Args:533 *args: Variable length argument list.534 **kwargs: Arbitrary keyword arguments.535 536 Returns:537 torch.Tensor or tuple: Either the loss tensor or a tuple of (loss, details_dict).538 539 Raises:540 NotImplementedError: This method must be implemented by subclasses.541 """542 raise NotImplementedError()543 544 def get_name(self):545 """546 Get the name of this loss component.547 548 Returns:549 str: The name of the loss.550 551 Raises:552 NotImplementedError: This method must be implemented by subclasses.553 """554 raise NotImplementedError()555 556 def __mul__(self, alpha):557 """558 Multiply the loss by a scalar weight.559 560 Args:561 alpha (int or float): The weight to multiply the loss by.562 563 Returns:564 MultiLoss: A new loss object with the updated weight.565 566 Raises:567 AssertionError: If alpha is not a number.568 """569 assert isinstance(alpha, (int, float))570 res = copy(self)571 res._alpha = alpha572 return res573 574 __rmul__ = __mul__ # Support both loss*alpha and alpha*loss575 576 def __add__(self, loss2):577 """578 Add another loss to this loss, creating a chain of losses.579 580 Args:581 loss2 (MultiLoss): Another loss to add to this one.582 583 Returns:584 MultiLoss: A new loss object representing the combined losses.585 586 Raises:587 AssertionError: If loss2 is not a MultiLoss.588 """589 assert isinstance(loss2, MultiLoss)590 res = cur = copy(self)591 # Find the end of the chain592 while cur._loss2 is not None:593 cur = cur._loss2594 cur._loss2 = loss2595 return res596 597 def __repr__(self):598 """599 Create a string representation of the loss, including weights and chained losses.600 601 Returns:602 str: String representation of the loss.603 """604 name = self.get_name()605 if self._alpha != 1:606 name = f"{self._alpha:g}*{name}"607 if self._loss2:608 name = f"{name} + {self._loss2}"609 return name610 611 def forward(self, *args, **kwargs):612 """613 Compute the weighted loss and aggregate with any chained losses.614 615 Args:616 *args: Variable length argument list.617 **kwargs: Arbitrary keyword arguments.618 619 Returns:620 tuple: A tuple containing:621 - torch.Tensor: The total weighted loss.622 - dict: Details about individual loss components.623 """624 loss = self.compute_loss(*args, **kwargs)625 if isinstance(loss, tuple):626 loss, details = loss627 elif loss.ndim == 0:628 details = {self.get_name(): float(loss)}629 else:630 details = {}631 loss = loss * self._alpha632 633 if self._loss2:634 loss2, details2 = self._loss2(*args, **kwargs)635 loss = loss + loss2636 details |= details2637 638 return loss, details639 640 641class NonAmbiguousMaskLoss(Criterion, MultiLoss):642 """643 Loss on non-ambiguous mask prediction logits.644 """645 646 def __init__(self, criterion):647 super().__init__(criterion)648 649 def compute_loss(self, batch, preds, **kw):650 """651 Args:652 batch: list of dicts with the gt data653 preds: list of dicts with the predictions654 655 Returns:656 loss: Sum class of the lossses for N-views and the loss details657 """658 # Init loss list to keep track of individual losses for each view659 loss_list = []660 mask_loss_details = {}661 mask_loss_total = 0662 self_name = type(self).__name__663 664 # Loop over the views665 for view_idx, (gt, pred) in enumerate(zip(batch, preds)):666 # Get the GT non-ambiguous masks667 gt_non_ambiguous_mask = gt["non_ambiguous_mask"]668 669 # Get the predicted non-ambiguous mask logits670 pred_non_ambiguous_mask_logits = pred["non_ambiguous_mask_logits"]671 672 # Compute the loss for the current view673 loss = self.criterion(pred_non_ambiguous_mask_logits, gt_non_ambiguous_mask)674 675 # Add the loss to the list676 loss_list.append((loss, None, "non_ambiguous_mask"))677 678 # Add the loss details to the dictionary679 mask_loss_details[f"{self_name}_mask_view{view_idx + 1}"] = float(loss)680 mask_loss_total += float(loss)681 682 # Compute the average loss across all views683 mask_loss_details[f"{self_name}_mask_avg"] = mask_loss_total / len(batch)684 685 return Sum(*loss_list), (mask_loss_details | {})686 687 688class ConfLoss(MultiLoss):689 """690 Applies confidence-weighted regression loss using model-predicted confidence values.691 692 The confidence-weighted loss has the form:693 conf_loss = raw_loss * conf - alpha * log(conf)694 695 Where:696 - raw_loss is the original per-pixel loss697 - conf is the predicted confidence (higher values = higher confidence)698 - alpha is a hyperparameter controlling the regularization strength699 700 This loss can be selectively applied to specific loss components in factored and multi-view settings.701 """702 703 def __init__(self, pixel_loss, alpha=1, loss_set_indices=None):704 """705 Args:706 pixel_loss (MultiLoss): The pixel-level regression loss to be used.707 alpha (float): Hyperparameter controlling the confidence regularization strength.708 loss_set_indices (list or None): Indices of the loss sets to apply confidence weighting to.709 Each index selects a specific loss set across all views (with the same rep_type).710 If None, defaults to [0] which applies to the first loss set only.711 """712 super().__init__()713 assert alpha > 0714 self.alpha = alpha715 self.pixel_loss = pixel_loss.with_reduction("none")716 self.loss_set_indices = [0] if loss_set_indices is None else loss_set_indices717 718 def get_name(self):719 return f"ConfLoss({self.pixel_loss})"720 721 def get_conf_log(self, x):722 return x, torch.log(x)723 724 def compute_loss(self, batch, preds, **kw):725 # Init loss list and details726 total_loss = 0727 conf_loss_details = {}728 running_avg_dict = {}729 self_name = type(self.pixel_loss).__name__730 n_views = len(batch)731 732 # Compute per-pixel loss for each view733 losses, pixel_loss_details = self.pixel_loss(batch, preds, **kw)734 735 # Select specific loss sets based on indices736 selected_losses = []737 processed_indices = set()738 for idx in self.loss_set_indices:739 start_idx = idx * n_views740 end_idx = min((idx + 1) * n_views, len(losses))741 selected_losses.extend(losses[start_idx:end_idx])742 processed_indices.update(range(start_idx, end_idx))743 744 # Process selected losses with confidence weighting745 for loss_idx, (loss, msk, rep_type) in enumerate(selected_losses):746 view_idx = loss_idx % n_views # Map to corresponding view index747 748 if loss.numel() == 0:749 # print(f"NO VALID VALUES in loss idx {loss_idx} (Rep Type: {rep_type}, Num Views: {n_views})", force=True)750 continue751 752 # Get the confidence and log confidence753 if (754 hasattr(self.pixel_loss, "flatten_across_image_only")755 and self.pixel_loss.flatten_across_image_only756 ):757 # Reshape confidence to match the flattened dimensions758 conf_reshaped = preds[view_idx]["conf"].view(759 preds[view_idx]["conf"].shape[0], -1760 )761 conf, log_conf = self.get_conf_log(conf_reshaped[msk])762 loss = loss[msk]763 else:764 conf, log_conf = self.get_conf_log(preds[view_idx]["conf"][msk])765 766 # Weight the loss by the confidence767 conf_loss = loss * conf - self.alpha * log_conf768 769 # Only add to total loss and store details if there are valid elements770 if conf_loss.numel() > 0:771 conf_loss = conf_loss.mean()772 total_loss = total_loss + conf_loss773 774 # Store details775 conf_loss_details[776 f"{self_name}_{rep_type}_conf_loss_view{view_idx + 1}"777 ] = float(conf_loss)778 779 # Initialize or update running average directly780 avg_key = f"{self_name}_{rep_type}_conf_loss_avg"781 if avg_key not in conf_loss_details:782 conf_loss_details[avg_key] = float(conf_loss)783 running_avg_dict[784 f"{self_name}_{rep_type}_conf_loss_valid_views"785 ] = 1786 else:787 valid_views = (788 running_avg_dict[789 f"{self_name}_{rep_type}_conf_loss_valid_views"790 ]791 + 1792 )793 running_avg_dict[794 f"{self_name}_{rep_type}_conf_loss_valid_views"795 ] = valid_views796 conf_loss_details[avg_key] += (797 float(conf_loss) - conf_loss_details[avg_key]798 ) / valid_views799 800 # Add unmodified losses for sets not in selected_losses801 for idx, (loss, msk, rep_type) in enumerate(losses):802 if idx not in processed_indices:803 if msk is not None:804 loss_after_masking = loss[msk]805 else:806 loss_after_masking = loss807 if loss_after_masking.numel() > 0:808 loss_mean = loss_after_masking.mean()809 else:810 # print(f"NO VALID VALUES in loss idx {idx} (Rep Type: {rep_type}, Num Views: {n_views})", force=True)811 loss_mean = 0812 total_loss = total_loss + loss_mean813 814 return total_loss, dict(**conf_loss_details, **pixel_loss_details)815 816 817class ExcludeTopNPercentPixelLoss(MultiLoss):818 """819 Pixel-level regression loss where for each instance in a batch the top N% of per-pixel loss values are ignored820 for the mean loss computation.821 Allows selecting which pixel-level regression loss sets to apply the exclusion to.822 """823 824 def __init__(825 self,826 pixel_loss,827 top_n_percent=5,828 apply_to_real_data_only=True,829 loss_set_indices=None,830 ):831 """832 Args:833 pixel_loss (MultiLoss): The pixel-level regression loss to be used.834 top_n_percent (float): The percentage of top per-pixel loss values to ignore. Range: [0, 100]. Default: 5.835 apply_to_real_data_only (bool): Whether to apply the loss only to real world data. Default: True.836 loss_set_indices (list or None): Indices of the loss sets to apply the exclusion to.837 Each index selects a specific loss set across all views (with the same rep_type).838 If None, defaults to [0] which applies to the first loss set only.839 """840 super().__init__()841 self.pixel_loss = pixel_loss.with_reduction("none")842 self.top_n_percent = top_n_percent843 self.bottom_n_percent = 100 - top_n_percent844 self.apply_to_real_data_only = apply_to_real_data_only845 self.loss_set_indices = [0] if loss_set_indices is None else loss_set_indices846 847 def get_name(self):848 return f"ExcludeTopNPercentPixelLoss({self.pixel_loss})"849 850 def keep_bottom_n_percent(self, tensor, mask, bottom_n_percent):851 """852 Function to compute the mask for keeping the bottom n percent of per-pixel loss values.853 854 Args:855 tensor (torch.Tensor): The tensor containing the per-pixel loss values.856 Shape: (B, N) where B is the batch size and N is the number of total pixels.857 mask (torch.Tensor): The mask indicating valid pixels. Shape: (B, N).858 859 Returns:860 torch.Tensor: Flattened tensor containing the bottom n percent of per-pixel loss values.861 """862 B, N = tensor.shape863 864 # Calculate the number of valid elements (where mask is True)865 num_valid = mask.sum(dim=1)866 867 # Calculate the number of elements to keep (n% of valid elements)868 num_keep = (num_valid * bottom_n_percent / 100).long()869 870 # Create a mask for the bottom n% elements871 keep_mask = torch.arange(N, device=tensor.device).unsqueeze(872 0873 ) < num_keep.unsqueeze(1)874 875 # Create a tensor with inf where mask is False876 masked_tensor = torch.where(877 mask, tensor, torch.tensor(float("inf"), device=tensor.device)878 )879 880 # Sort the masked tensor along the N dimension881 sorted_tensor, _ = torch.sort(masked_tensor, dim=1, descending=False)882 883 # Get the bottom n% elements884 bottom_n_percent_elements = sorted_tensor[keep_mask]885 886 return bottom_n_percent_elements887 888 def compute_loss(self, batch, preds, **kw):889 # Compute per-pixel loss890 losses, details = self.pixel_loss(batch, preds, **kw)891 n_views = len(batch)892 893 # Select specific loss sets based on indices894 selected_losses = []895 processed_indices = set()896 for idx in self.loss_set_indices:897 start_idx = idx * n_views898 end_idx = min((idx + 1) * n_views, len(losses))899 selected_losses.extend(losses[start_idx:end_idx])900 processed_indices.update(range(start_idx, end_idx))901 902 # Initialize total loss903 total_loss = 0.0904 loss_details = {}905 running_avg_dict = {}906 self_name = type(self.pixel_loss).__name__907 908 # Process selected losses with top N percent exclusion909 for loss_idx, (loss, msk, rep_type) in enumerate(selected_losses):910 view_idx = loss_idx % n_views # Map to corresponding view index911 912 if loss.numel() == 0:913 # print(f"NO VALID VALUES in loss idx {loss_idx} (Rep Type: {rep_type}, Num Views: {n_views})", force=True)914 continue915 916 # Create empty list for current view's aggregated tensors917 aggregated_losses = []918 919 if self.apply_to_real_data_only:920 # Get the synthetic and real world data mask921 synthetic_mask = batch[view_idx]["is_synthetic"]922 real_data_mask = ~batch[view_idx]["is_synthetic"]923 else:924 # Apply the filtering to all data925 synthetic_mask = torch.zeros_like(batch[view_idx]["is_synthetic"])926 real_data_mask = torch.ones_like(batch[view_idx]["is_synthetic"])927 928 # Process synthetic data929 if synthetic_mask.any():930 synthetic_loss = loss[synthetic_mask]931 synthetic_msk = msk[synthetic_mask]932 aggregated_losses.append(synthetic_loss[synthetic_msk])933 934 # Process real data935 if real_data_mask.any():936 real_loss = loss[real_data_mask]937 real_msk = msk[real_data_mask]938 real_bottom_n_percent_loss = self.keep_bottom_n_percent(939 real_loss, real_msk, self.bottom_n_percent940 )941 aggregated_losses.append(real_bottom_n_percent_loss)942 943 # Compute view loss944 view_loss = torch.cat(aggregated_losses, dim=0)945 946 # Only add to total loss and store details if there are valid elements947 if view_loss.numel() > 0:948 view_loss = view_loss.mean()949 total_loss = total_loss + view_loss950 951 # Store details952 loss_details[953 f"{self_name}_{rep_type}_bot{self.bottom_n_percent}%_loss_view{view_idx + 1}"954 ] = float(view_loss)955 956 # Initialize or update running average directly957 avg_key = f"{self_name}_{rep_type}_bot{self.bottom_n_percent}%_loss_avg"958 if avg_key not in loss_details:959 loss_details[avg_key] = float(view_loss)960 running_avg_dict[961 f"{self_name}_{rep_type}_bot{self.bottom_n_percent}%_valid_views"962 ] = 1963 else:964 valid_views = (965 running_avg_dict[966 f"{self_name}_{rep_type}_bot{self.bottom_n_percent}%_valid_views"967 ]968 + 1969 )970 running_avg_dict[971 f"{self_name}_{rep_type}_bot{self.bottom_n_percent}%_valid_views"972 ] = valid_views973 loss_details[avg_key] += (974 float(view_loss) - loss_details[avg_key]975 ) / valid_views976 977 # Add unmodified losses for sets not in selected_losses978 for idx, (loss, msk, rep_type) in enumerate(losses):979 if idx not in processed_indices:980 if msk is not None:981 loss_after_masking = loss[msk]982 else:983 loss_after_masking = loss984 if loss_after_masking.numel() > 0:985 loss_mean = loss_after_masking.mean()986 else:987 # print(f"NO VALID VALUES in loss idx {idx} (Rep Type: {rep_type}, Num Views: {n_views})", force=True)988 loss_mean = 0989 total_loss = total_loss + loss_mean990 991 return total_loss, dict(**loss_details, **details)992 993 994class ConfAndExcludeTopNPercentPixelLoss(MultiLoss):995 """996 Combined loss that applies ConfLoss to one set of pixel-level regression losses997 and ExcludeTopNPercentPixelLoss to another set of pixel-level regression losses.998 """999 1000 def __init__(1001 self,1002 pixel_loss,1003 conf_alpha=1,1004 top_n_percent=5,1005 apply_to_real_data_only=True,1006 conf_loss_set_indices=None,1007 exclude_loss_set_indices=None,1008 ):1009 """1010 Args:1011 pixel_loss (MultiLoss): The pixel-level regression loss to be used.1012 conf_alpha (float): Alpha parameter for ConfLoss. Default: 1.1013 top_n_percent (float): Percentage of top per-pixel loss values to ignore. Range: [0, 100]. Default: 5.1014 apply_to_real_data_only (bool): Whether to apply the exclude loss only to real world data. Default: True.1015 conf_loss_set_indices (list or None): Indices of the loss sets to apply confidence weighting to.1016 Each index selects a specific loss set across all views (with the same rep_type).1017 If None, defaults to [0] which applies to the first loss set only.1018 exclude_loss_set_indices (list or None): Indices of the loss sets to apply top N percent exclusion to.1019 Each index selects a specific loss set across all views (with the same rep_type).1020 If None, defaults to [1] which applies to the second loss set only.1021 """1022 super().__init__()1023 self.pixel_loss = pixel_loss.with_reduction("none")1024 assert conf_alpha > 01025 self.conf_alpha = conf_alpha1026 self.top_n_percent = top_n_percent1027 self.bottom_n_percent = 100 - top_n_percent1028 self.apply_to_real_data_only = apply_to_real_data_only1029 self.conf_loss_set_indices = (1030 [0] if conf_loss_set_indices is None else conf_loss_set_indices1031 )1032 self.exclude_loss_set_indices = (1033 [1] if exclude_loss_set_indices is None else exclude_loss_set_indices1034 )1035 1036 def get_name(self):1037 return f"ConfAndExcludeTopNPercentPixelLoss({self.pixel_loss})"1038 1039 def get_conf_log(self, x):1040 return x, torch.log(x)1041 1042 def keep_bottom_n_percent(self, tensor, mask, bottom_n_percent):1043 """1044 Function to compute the mask for keeping the bottom n percent of per-pixel loss values.1045 """1046 B, N = tensor.shape1047 1048 # Calculate the number of valid elements (where mask is True)1049 num_valid = mask.sum(dim=1)1050 1051 # Calculate the number of elements to keep (n% of valid elements)1052 num_keep = (num_valid * bottom_n_percent / 100).long()1053 1054 # Create a mask for the bottom n% elements1055 keep_mask = torch.arange(N, device=tensor.device).unsqueeze(1056 01057 ) < num_keep.unsqueeze(1)1058 1059 # Create a tensor with inf where mask is False1060 masked_tensor = torch.where(1061 mask, tensor, torch.tensor(float("inf"), device=tensor.device)1062 )1063 1064 # Sort the masked tensor along the N dimension1065 sorted_tensor, _ = torch.sort(masked_tensor, dim=1, descending=False)1066 1067 # Get the bottom n% elements1068 bottom_n_percent_elements = sorted_tensor[keep_mask]1069 1070 return bottom_n_percent_elements1071 1072 def compute_loss(self, batch, preds, **kw):1073 # Compute per-pixel loss1074 losses, pixel_loss_details = self.pixel_loss(batch, preds, **kw)1075 n_views = len(batch)1076 1077 # Select specific loss sets for confidence weighting1078 conf_selected_losses = []1079 conf_processed_indices = set()1080 for idx in self.conf_loss_set_indices:1081 start_idx = idx * n_views1082 end_idx = min((idx + 1) * n_views, len(losses))1083 conf_selected_losses.extend(losses[start_idx:end_idx])1084 conf_processed_indices.update(range(start_idx, end_idx))1085 1086 # Select specific loss sets for top N percent exclusion1087 exclude_selected_losses = []1088 exclude_processed_indices = set()1089 for idx in self.exclude_loss_set_indices:1090 start_idx = idx * n_views1091 end_idx = min((idx + 1) * n_views, len(losses))1092 exclude_selected_losses.extend(losses[start_idx:end_idx])1093 exclude_processed_indices.update(range(start_idx, end_idx))1094 1095 # Initialize total loss and details1096 total_loss = 01097 loss_details = {}1098 running_avg_dict = {}1099 self_name = type(self.pixel_loss).__name__1100 1101 # Process selected losses with confidence weighting1102 for loss_idx, (loss, msk, rep_type) in enumerate(conf_selected_losses):1103 view_idx = loss_idx % n_views # Map to corresponding view index1104 1105 if loss.numel() == 0:1106 # print(f"NO VALID VALUES in loss idx {loss_idx} (Rep Type: {rep_type}, Num Views: {n_views}) for conf loss", force=True)1107 continue1108 1109 # Get the confidence and log confidence1110 if (1111 hasattr(self.pixel_loss, "flatten_across_image_only")1112 and self.pixel_loss.flatten_across_image_only1113 ):1114 # Reshape confidence to match the flattened dimensions1115 conf_reshaped = preds[view_idx]["conf"].view(1116 preds[view_idx]["conf"].shape[0], -11117 )1118 conf, log_conf = self.get_conf_log(conf_reshaped[msk])1119 loss = loss[msk]1120 else:1121 conf, log_conf = self.get_conf_log(preds[view_idx]["conf"][msk])1122 1123 # Weight the loss by the confidence1124 conf_loss = loss * conf - self.conf_alpha * log_conf1125 1126 # Only add to total loss and store details if there are valid elements1127 if conf_loss.numel() > 0:1128 conf_loss = conf_loss.mean()1129 total_loss = total_loss + conf_loss1130 1131 # Store details1132 loss_details[f"{self_name}_{rep_type}_conf_loss_view{view_idx + 1}"] = (1133 float(conf_loss)1134 )1135 1136 # Initialize or update running average directly1137 avg_key = f"{self_name}_{rep_type}_conf_loss_avg"1138 if avg_key not in loss_details:1139 loss_details[avg_key] = float(conf_loss)1140 running_avg_dict[1141 f"{self_name}_{rep_type}_conf_loss_valid_views"1142 ] = 11143 else:1144 valid_views = (1145 running_avg_dict[1146 f"{self_name}_{rep_type}_conf_loss_valid_views"1147 ]1148 + 11149 )1150 running_avg_dict[1151 f"{self_name}_{rep_type}_conf_loss_valid_views"1152 ] = valid_views1153 loss_details[avg_key] += (1154 float(conf_loss) - loss_details[avg_key]1155 ) / valid_views1156 1157 # Process selected losses with top N percent exclusion1158 for loss_idx, (loss, msk, rep_type) in enumerate(exclude_selected_losses):1159 view_idx = loss_idx % n_views # Map to corresponding view index1160 1161 if loss.numel() == 0:1162 # print(f"NO VALID VALUES in loss idx {loss_idx} (Rep Type: {rep_type}, Num Views: {n_views}) for exclude loss", force=True)1163 continue1164 1165 # Create empty list for current view's aggregated tensors1166 aggregated_losses = []1167 1168 if self.apply_to_real_data_only:1169 # Get the synthetic and real world data mask1170 synthetic_mask = batch[view_idx]["is_synthetic"]1171 real_data_mask = ~batch[view_idx]["is_synthetic"]1172 else:1173 # Apply the filtering to all data1174 synthetic_mask = torch.zeros_like(batch[view_idx]["is_synthetic"])1175 real_data_mask = torch.ones_like(batch[view_idx]["is_synthetic"])1176 1177 # Process synthetic data1178 if synthetic_mask.any():1179 synthetic_loss = loss[synthetic_mask]1180 synthetic_msk = msk[synthetic_mask]1181 aggregated_losses.append(synthetic_loss[synthetic_msk])1182 1183 # Process real data1184 if real_data_mask.any():1185 real_loss = loss[real_data_mask]1186 real_msk = msk[real_data_mask]1187 real_bottom_n_percent_loss = self.keep_bottom_n_percent(1188 real_loss, real_msk, self.bottom_n_percent1189 )1190 aggregated_losses.append(real_bottom_n_percent_loss)1191 1192 # Compute view loss1193 view_loss = torch.cat(aggregated_losses, dim=0)1194 1195 # Only add to total loss and store details if there are valid elements1196 if view_loss.numel() > 0:1197 view_loss = view_loss.mean()1198 total_loss = total_loss + view_loss1199 1200 # Store details