docstr: Calculate metrics for the current stage.
docstr_tok: ["Calculate","metrics","for","the","current","stage","."]
code: def _calculate_stage(
      self,
      raymarched: RendererOutput,
      ray_bundle: ImplicitronRayBundle,
      image_rgb: Optional[torch.Tensor] = None,
      depth_map: Optional[torch.Tensor] = None,
      fg_probability: Optional[torch.Tensor] = None,
      mask_crop: Optional[torch.Tensor] = None,
      keys_prefix: str = "loss_",
      **kwargs,
  ) -> Dict[str, Any]:
      """
      Calculate metrics for the current stage.
      """
      # TODO: extract functions
  
      # reshape from B x ... x DIM to B x DIM x -1 x 1
      image_rgb_pred, fg_probability_pred, depth_map_pred = [
          _reshape_nongrid_var(x)
          for x in [
              raymarched.features,
              raymarched.masks,
              raymarched.depths,
          ]
      ]
      xys = ray_bundle.xys
  
      # If ray_bundle is packed than we can sample images in padded state to lower
      # memory requirements. Instead of having one image for every element in
      # ray_bundle we can than have one image per unique sampled camera.
      if ray_bundle.is_packed():
          xys, first_idxs, num_inputs = ray_bundle.get_padded_xys()
  
      # reshape the sampling grid as well
      # TODO: we can get rid of the singular dimension here and in _reshape_nongrid_var
      # now that we use rend_utils.ndc_grid_sample
      xys = xys.reshape(xys.shape[0], -1, 1, 2)
  
      # closure with the given xys
      def sample_full(tensor, mode):
          if tensor is None:
              return tensor
          return rend_utils.ndc_grid_sample(tensor, xys, mode=mode)
  
      def sample_packed(tensor, mode):
          if tensor is None:
              return tensor
  
          # select images that corespond to sampled cameras if raybundle is packed
          tensor = tensor[ray_bundle.camera_ids]
          if ray_bundle.is_packed():
              # select images that corespond to sampled cameras if raybundle is packed
              tensor = tensor[ray_bundle.camera_ids]
          result = rend_utils.ndc_grid_sample(tensor, xys, mode=mode)
          return padded_to_packed(
              result, first_idxs, num_inputs, max_size_dim=2
          )[
              :, :, None
          ]  # the result is [n_rays_total_training, 3, 1, 1]
  
      sample = sample_packed if ray_bundle.is_packed() else sample_full
  
      # eval all results in this size
      image_rgb = sample(image_rgb, mode="bilinear")
      depth_map = sample(depth_map, mode="nearest")
      fg_probability = sample(fg_probability, mode="nearest")
      mask_crop = sample(mask_crop, mode="nearest")
      if mask_crop is None and image_rgb_pred is not None:
          mask_crop = torch.ones_like(image_rgb_pred[:, :1])
      if mask_crop is None and depth_map_pred is not None:
          mask_crop = torch.ones_like(depth_map_pred[:, :1])
  
      metrics = {}
      if image_rgb is not None and image_rgb_pred is not None:
          metrics.update(
              _rgb_metrics(
                  image_rgb,
                  image_rgb_pred,
                  fg_probability,
                  fg_probability_pred,
                  mask_crop,
              )
          )
  
      if fg_probability_pred is not None:
          metrics["mask_beta_prior"] = utils.beta_prior(
              fg_probability_pred
          )
      if fg_probability is not None and fg_probability_pred is not None:
          metrics["mask_neg_iou"] = utils.neg_iou_loss(
              fg_probability_pred, fg_probability, mask=mask_crop
          )
          metrics["mask_bce"] = utils.calc_bce(
              fg_probability_pred, fg_probability, mask=mask_crop
          )
  
      if depth_map is not None and depth_map_pred is not None:
          assert mask_crop is not None
          _, abs_ = utils.eval_depth(
              depth_map_pred,
              depth_map,
              get_best_scale=True,
              mask=mask_crop,
              crop=0,
          )
          metrics["depth_abs"] = abs_.mean()
  
          if fg_probability is not None:
              mask = fg_probability * mask_crop
              _, abs_ = utils.eval_depth(
                  depth_map_pred,
                  depth_map,
                  get_best_scale=True,
                  mask=mask,
                  crop=0,
              )
              metrics["depth_abs_fg"] = abs_.mean()
  
      # regularizers
      grad_theta = raymarched.aux.get("grad_theta")
      if grad_theta is not None:
          metrics["eikonal"] = _get_eikonal_loss(grad_theta)
  
      density_grid = raymarched.aux.get("density_grid")
      if density_grid is not None:
          metrics["density_tv"] = _get_grid_tv_loss(density_grid)
  
      if depth_map_pred is not None:
          metrics["depth_neg_penalty"] = _get_depth_neg_penalty_loss(
              depth_map_pred
          )
  
      if keys_prefix is not None:
          metrics = {(keys_prefix + k): v for k, v in metrics.items()}
  
      return metrics
  
imports: import warnings
  from typing import Any, Dict, Optional
  import torch
  from pytorch3d.implicitron.models.renderer.ray_sampler import ImplicitronRayBundle
  from pytorch3d.implicitron.tools import metric_utils
  from pytorch3d.implicitron.tools.config import registry, ReplaceableBase
  from pytorch3d.ops import padded_to_packed
  from pytorch3d.renderer import utils
  from .renderer.base import RendererOutput
label_desc: Use image node_id 2 for calling the ViewMetrics obj's underlying member method code with example usage: obj._calculate_stage(raymarched, ray_bundle, image_rgb, depth_map, fg_probability, mask_crop, keys_prefix) and returns: metrics, rend_utils, padded_to_packed, tensor, tensor