class_docstr: Base model for object detection with matching including a mask loss.
  
  The masks are predicted one for each instance and must be in direct
  correspondence with the bounding boxes (both predicted and ground truth).
  
  In practice this is typically used to predict cropped masks within the
  predicted bounding boxes, similar to what has been done with, for example,
  mask RCNN. Right now the cropping must occur outside this class.
class_docstr_tok: ["Base","model","for","object","detection","with","matching","including","a","mask","loss",".","The","masks","are","predicted","one","for","each","instance","and","must","be","in","direct","correspondence","with","the","bounding","boxes","(","both","predicted","and","ground","truth",")",".","In","practice","this","is","typically","used","to","predict","cropped","masks","within","the","predicted","bounding","boxes",",","similar","to","what","has","been","done","with",",","for","example",",","mask","RCNN",".","Right","now","the","cropping","must","occur","outside","this","class","."]
docstr: Mask losses - pixelwise cross entropy a la mask-RCNN.
  
  Note that it is assumed all masks, both predicted and ground truth, have
  been resized to a fixed height and width (i.e., allowing distortion of the
  aspect ratio).
  
  Args:
    outputs: Model predictions. For the purpose of this loss, outputs
      must have key 'pred_masks'. outputs['pred_masks'] is a nd-array of the
      predicted masks with (batch, num-instances, h, w) shape, type float32
      in range [0, 1].
    batch: Has keys 'inputs', 'batch_mask' and, 'label' (ground truth).
      batch['label'] is a dict. For the purpose of this loss, batch['label']
      must have key 'masks', which the value has the same format
      as outputs['pred_masks']. Additionally in batch['label'], key
      'valid_masks'. This is to decide which masks are invalid
      and need to be ignored. Invalid masks have label 0. If
      batch['batch_mask'] is provided it is used to weight the loss for
      different images in the current batch of examples.
    indices: Matcher output which conveys source to target pairing of objects.
  
  Returns:
    loss: Has keys 'loss_mask'. This is mask losses averaged over the batch.
      Therefore it has shape [].
    metrics: Has keys 'loss_mask'. This is the metrics psumed over the batch.
      Therefore it has shape [].
docstr_tok: ["Mask","losses","-","pixelwise","cross","entropy","a","la","mask-RCNN",".","Note","that","it","is","assumed","all","masks",",","both","predicted","and","ground","truth",",","have","been","resized","to","a","fixed","height","and","width","(","i.e.",",","allowing","distortion","of","the","aspect","ratio",")",".","Args",":","outputs",":","Model","predictions",".","For","the","purpose","of","this","loss",",","outputs","must","have","key","'pred_masks","'",".","outputs","[","'pred_masks","'","]","is","a","nd-array","of","the","predicted","masks","with","(","batch",",","num-instances",",","h",",","w",")","shape",",","type","float32","in","range","[","0",",","1","]",".","batch",":","Has","keys","'inputs","'",",","'batch_mask","'","and",",","'label","'","(","ground","truth",")",".","batch","[","'label","'","]","is","a","dict",".","For","the","purpose","of","this","loss",",","batch","[","'label","'","]","must","have","key","'masks","'",",","which","the","value","has","the","same","format","as","outputs","[","'pred_masks","'","]",".","Additionally","in","batch","[","'label","'","]",",","key","'valid_masks","'",".","This","is","to","decide","which","masks","are","invalid","and","need","to","be","ignored",".","Invalid","masks","have","label","0",".","If","batch","[","'batch_mask","'","]","is","provided","it","is","used","to","weight","the","loss","for","different","images","in","the","current","batch","of","examples",".","indices",":","Matcher","output","which","conveys","source","to","target","pairing","of","objects",".","Returns",":","loss",":","Has","keys","'loss_mask","'",".","This","is","mask","losses","averaged","over","the","batch",".","Therefore","it","has","shape","[","]",".","metrics",":","Has","keys","'loss_mask","'",".","This","is","the","metrics","psumed","over","the","batch",".","Therefore","it","has","shape","[","]","."]
code: def masks_losses_and_metrics(
      self, outputs: ArrayDict, batch: ArrayDict, indices: jnp.ndarray
  ) -> Tuple[ArrayDict, MetricsDict]:
      """Mask losses - pixelwise cross entropy a la mask-RCNN.
  
      Note that it is assumed all masks, both predicted and ground truth, have
      been resized to a fixed height and width (i.e., allowing distortion of the
      aspect ratio).
  
      Args:
        outputs: Model predictions. For the purpose of this loss, outputs
          must have key 'pred_masks'. outputs['pred_masks'] is a nd-array of the
          predicted masks with (batch, num-instances, h, w) shape, type float32
          in range [0, 1].
        batch: Has keys 'inputs', 'batch_mask' and, 'label' (ground truth).
          batch['label'] is a dict. For the purpose of this loss, batch['label']
          must have key 'masks', which the value has the same format
          as outputs['pred_masks']. Additionally in batch['label'], key
          'valid_masks'. This is to decide which masks are invalid
          and need to be ignored. Invalid masks have label 0. If
          batch['batch_mask'] is provided it is used to weight the loss for
          different images in the current batch of examples.
        indices: Matcher output which conveys source to target pairing of objects.
  
      Returns:
        loss: Has keys 'loss_mask'. This is mask losses averaged over the batch.
          Therefore it has shape [].
        metrics: Has keys 'loss_mask'. This is the metrics psumed over the batch.
          Therefore it has shape [].
      """
      assert "pred_masks" in outputs
      assert "label" in batch
  
      targets = batch["label"]
      assert "masks" in targets
      assert "valid_masks" in targets
      assert (
          targets["valid_masks"].shape[1] == targets["masks"].shape[1]
      )
      assert (
          targets["valid_masks"].shape[1] == targets["boxes"].shape[1]
      )
      losses, metrics = {}, {}
      batch_weights = batch.get("batch_mask")
  
      pred_masks = model_utils.simple_gather(
          outputs["pred_masks"], indices[:, 0]
      )
      tgt_masks = model_utils.simple_gather(
          targets["masks"], indices[:, 1]
      )
      tgt_not_padding = model_utils.simple_gather(
          targets["valid_masks"], indices[:, 1]
      )
  
      src_boxes = model_utils.simple_gather(
          outputs["pred_boxes"], indices[:, 0]
      )
      tgt_boxes = model_utils.simple_gather(
          targets["boxes"], indices[:, 1]
      )
  
      if self.config.weight_mask_loss_by_iou:
          src_boxes_xyxy = box_utils.box_cxcywh_to_xyxy(src_boxes)
          tgt_boxes_xyxy = box_utils.box_cxcywh_to_xyxy(tgt_boxes)
          box_iou, _ = box_utils.box_iou(
              src_boxes_xyxy, tgt_boxes_xyxy, all_pairs=False
          )
          loss_weight = jax.lax.stop_gradient(box_iou)
      else:
          loss_weight = 1.0
  
      if self.config.move_true_to_pred_mask:
          tgt_masks = jax.vmap(jax.vmap(move_true_to_pred_mask))(
              true_mask=tgt_masks,
              true_box=tgt_boxes,
              pred_box=jax.lax.stop_gradient(src_boxes),
          )
  
      unnormalized_ce_loss_mask = (
          model_utils.weighted_unnormalized_sigmoid_cross_entropy(
              logits=pred_masks, multi_hot_targets=tgt_masks
          )
      )
      # Take the mean across the spatial height, width of the masks.
      unnormalized_ce_loss_mask = jnp.mean(
          unnormalized_ce_loss_mask, axis=(-1, -2)
      )
      unnormalized_ce_loss_mask *= (
          loss_weight  # Weigh loss by box overlap.
      )
      denom = tgt_not_padding.sum(axis=1)
      if batch_weights is not None:
          denom *= batch_weights
          unnormalized_ce_loss_mask = model_utils.apply_weights(
              unnormalized_ce_loss_mask, batch_weights
          )
      unnormalized_ce_loss_mask *= tgt_not_padding
      norm_type = self.config.get("normalization")
  
      if norm_type != "per_example":
          denom = jnp.maximum(
              jax.lax.pmean(denom.sum(), axis_name="batch"), 1
          )
          normalized_ce_loss_mask = (
              unnormalized_ce_loss_mask.sum() / denom
          )
      else:
          denom = jnp.maximum(denom, 1.0)
          normalized_ce_loss_mask = (
              unnormalized_ce_loss_mask.sum(axis=1) / denom
          ).mean()
  
      losses["loss_mask"] = normalized_ce_loss_mask.sum()
      metrics["loss_mask"] = (normalized_ce_loss_mask.sum(), 1.0)
  
      # Sum metrics and normalizers over all replicas.
      for k, v in metrics.items():
          metrics[k] = model_utils.psum_metric_normalizer(
              v
          )  # pytype: disable=wrong-arg-types  # jax-ndarray
      return losses, metrics
  
imports: import abc
  import functools
  from typing import Any, Callable, Dict, Optional, Tuple
  import jax
  import jax.numpy
  import ml_collections
  from scenic.model_lib import matchers
  from scenic.model_lib.base_models import base_model
  from scenic.model_lib.base_models import box_utils
  from scenic.model_lib.base_models import model_utils
  from scenic.projects.owl_vit import losses
label_desc: Use image node_id 3 for calling the ObjectDetectionModelWithMasks obj's underlying member method code with example usage: obj.masks_losses_and_metrics(outputs, batch, indices) and returns: losses,  metrics