class_docstr: Head for predicting masks inside bounding boxes.
  
  The architecture is informed by https://arxiv.org/abs/2104.00613.
  
  The head takes the following inputs:
   * Predicted boxes and image features (output tokens of the image backbone)
     for each box as inputs.
   * The input image, for extracting additional low-level image features.
   * During training, the ground-truth boxes, to select which predicted boxes
     to predict masks for.
  
  The head performs the following steps:
    1. Apply a small ResNet to the input image to extract low-level features.
    2. Apply ROIAlign to the ResNet features to get per-box low-level features.
    3. Merge per-box low-level features with the image features coming from the
       main image backbone.
    4. Apply an Hourglass network to the per-box features, to merge low- and
       high-level features. Applying a relatively large/deep per-box network was
       found to be useful especially for novel classes in
       https://arxiv.org/abs/2104.00613. The outputs of the Hourglass network
       are the final segmentation masks.
  
  Attributes:
    mask_size: Integer specifying the width and height of the predicted masks.
    roi_align_num_parallel: Number of boxes to call roi_align on in parallel.
      Larger values are faster but consume more memory.
    stop_box_gradients: Whether to stop the box gradients from flowing back to
      the main model.
    stop_image_gradients: Whether to stop the image feature gradients from
      flowing back to the main model.
    num_training_boxes: If set, only the top predicted boxes by IoU with ground-
      truth boxes will be used during training. This speeds up training because
      most predicted boxes will not be matched to true boxes during training,
      and predicting masks for these unmatched boxes takes time but provides no
      training signal.
    num_mlp_layers_backbone_features: How many MLP layers to apply to the
      backbone image features before merging them with the low-level image
      features.
    image_resnet_width: Width multiplier of the low-level image ResNet.
    image_resnet_depth: Depth spec of the low-level image ResNet.
    mask_resnet_width: Width multiplier of the per-box Hourglass network.
    mask_resnet_depth: Depth spec of the per-box Hourglass network.
    add_image_coords: Whether to add image-centric x/y-coordinate maps to the
      low-level features.
    add_mask_coords: Whether to add mask-centric x/y-coordinate maps to the
      per-box features.
    resnet_out_width_mult: Width multiplier for the low-level features.
    backbone_out_width_mult: Width multiplier for the image backbone features.
class_docstr_tok: ["Head","for","predicting","masks","inside","bounding","boxes",".","The","architecture","is","informed","by","https",":","\/\/arxiv.org\/abs\/2104.00613",".","The","head","takes","the","following","inputs",":","*","Predicted","boxes","and","image","features","(","output","tokens","of","the","image","backbone",")","for","each","box","as","inputs",".","*","The","input","image",",","for","extracting","additional","low-level","image","features",".","*","During","training",",","the","ground-truth","boxes",",","to","select","which","predicted","boxes","to","predict","masks","for",".","The","head","performs","the","following","steps",":","1",".","Apply","a","small","ResNet","to","the","input","image","to","extract","low-level","features",".","2",".","Apply","ROIAlign","to","the","ResNet","features","to","get","per-box","low-level","features",".","3",".","Merge","per-box","low-level","features","with","the","image","features","coming","from","the","main","image","backbone",".","4",".","Apply","an","Hourglass","network","to","the","per-box","features",",","to","merge","low-","and","high-level","features",".","Applying","a","relatively","large\/deep","per-box","network","was","found","to","be","useful","especially","for","novel","classes","in","https",":","\/\/arxiv.org\/abs\/2104.00613",".","The","outputs","of","the","Hourglass","network","are","the","final","segmentation","masks",".","Attributes",":","mask_size",":","Integer","specifying","the","width","and","height","of","the","predicted","masks",".","roi_align_num_parallel",":","Number","of","boxes","to","call","roi_align","on","in","parallel",".","Larger","values","are","faster","but","consume","more","memory",".","stop_box_gradients",":","Whether","to","stop","the","box","gradients","from","flowing","back","to","the","main","model",".","stop_image_gradients",":","Whether","to","stop","the","image","feature","gradients","from","flowing","back","to","the","main","model",".","num_training_boxes",":","If","set",",","only","the","top","predicted","boxes","by","IoU","with","ground-","truth","boxes","will","be","used","during","training",".","This","speeds","up","training","because","most","predicted","boxes","will","not","be","matched","to","true","boxes","during","training",",","and","predicting","masks","for","these","unmatched","boxes","takes","time","but","provides","no","training","signal",".","num_mlp_layers_backbone_features",":","How","many","MLP","layers","to","apply","to","the","backbone","image","features","before","merging","them","with","the","low-level","image","features",".","image_resnet_width",":","Width","multiplier","of","the","low-level","image","ResNet",".","image_resnet_depth",":","Depth","spec","of","the","low-level","image","ResNet",".","mask_resnet_width",":","Width","multiplier","of","the","per-box","Hourglass","network",".","mask_resnet_depth",":","Depth","spec","of","the","per-box","Hourglass","network",".","add_image_coords",":","Whether","to","add","image-centric","x\/y-coordinate","maps","to","the","low-level","features",".","add_mask_coords",":","Whether","to","add","mask-centric","x\/y-coordinate","maps","to","the","per-box","features",".","resnet_out_width_mult",":","Width","multiplier","for","the","low-level","features",".","backbone_out_width_mult",":","Width","multiplier","for","the","image","backbone","features","."]
docstr: Forward pass of the mask head.
  
  Args:
    image: [B, H, W, C] input image.
    image_backbone_features: [B, num_boxes, D] image backbone output tokens.
    boxes: [B, num_boxes, 4] predicted boxes.
    true_boxes: [B, num_true_boxes, 4] ground-truth bounding boxes. Only used
      during training.
  
  Returns:
    [B, num_masks, mask_width, mask_width] array of segmentation masks.
  
  Raises:
    ValueError if true boxes are provided but gradients are not stopped.
    ValueError if roi_align_num_parallel does not evenly divide the number of
      boxes.
docstr_tok: ["Forward","pass","of","the","mask","head",".","Args",":","image",":","[","B",",","H",",","W",",","C","]","input","image",".","image_backbone_features",":","[","B",",","num_boxes",",","D","]","image","backbone","output","tokens",".","boxes",":","[","B",",","num_boxes",",","4","]","predicted","boxes",".","true_boxes",":","[","B",",","num_true_boxes",",","4","]","ground-truth","bounding","boxes",".","Only","used","during","training",".","Returns",":","[","B",",","num_masks",",","mask_width",",","mask_width","]","array","of","segmentation","masks",".","Raises",":","ValueError","if","true","boxes","are","provided","but","gradients","are","not","stopped",".","ValueError","if","roi_align_num_parallel","does","not","evenly","divide","the","number","of","boxes","."]
code: def __call__(
      self,
      image: jnp.ndarray,
      image_backbone_features: jnp.ndarray,
      boxes: jnp.ndarray,
      *,
      true_boxes: Optional[jnp.ndarray] = None,
  ) -> jnp.ndarray:
      """Forward pass of the mask head.
  
      Args:
        image: [B, H, W, C] input image.
        image_backbone_features: [B, num_boxes, D] image backbone output tokens.
        boxes: [B, num_boxes, 4] predicted boxes.
        true_boxes: [B, num_true_boxes, 4] ground-truth bounding boxes. Only used
          during training.
  
      Returns:
        [B, num_masks, mask_width, mask_width] array of segmentation masks.
  
      Raises:
        ValueError if true boxes are provided but gradients are not stopped.
        ValueError if roi_align_num_parallel does not evenly divide the number of
          boxes.
      """
      if not self.stop_box_gradients and true_boxes is not None:
          raise ValueError(
              "stop_box_gradients must be true when using true boxes."
          )
  
      if self.stop_box_gradients:
          boxes = jax.lax.stop_gradient(boxes)
      if self.stop_image_gradients:
          image_backbone_features = jax.lax.stop_gradient(
              image_backbone_features
          )
  
      # If true boxes are available, only compute masks for the predictions that
      # overlap most with the true boxes (the others won't get matched anyway).
      orig_num_boxes = boxes.shape[-2]
      top_k_indices = None
      if true_boxes is not None:
          (
              boxes,
              image_backbone_features,
              top_k_indices,
          ) = self._gather_top_boxes(
              boxes, image_backbone_features, true_boxes
          )
  
      # Apply a ResNet to get low-level image features:
      _, out = ResNet(
          num_classes=0,
          width=self.image_resnet_width,
          depth=self.image_resnet_depth,
          name="image_resnet",
      )(image)
      resnet_features = {
          k: v
          for k, v in out.items()
          if k == "stem" or k.startswith("stage")
      }
  
      # Reduce number of channels for each stage and resize:
      b, h, w, c = resnet_features["stem"].shape
      for name, feature in resnet_features.items():
          if name.startswith("stage"):
              feature = nn.Dense(c)(feature)
              resnet_features[name] = jax.image.resize(
                  feature, (b, h, w, c), method="linear"
              )
  
      # Concatenate and project to manageable size:
      resnet_features = jnp.concatenate(
          list(resnet_features.values()), axis=-1
      )
      resnet_features = nn.Dense(c * self.resnet_out_width_mult)(
          resnet_features
      )
  
      if self.add_image_coords:
          resnet_features = self._add_coord_channels(resnet_features)
  
      # Get feature map for each box using RoIAlign:
      roi_resnet_features = self._roi_align_image_features(
          resnet_features, boxes
      )
  
      # Process and project backbone features:
      for _ in range(self.num_mlp_layers_backbone_features):
          image_backbone_features = vit.MlpBlock()(
              image_backbone_features
          )
      image_backbone_features = nn.Dense(
          c * self.backbone_out_width_mult
      )(image_backbone_features)
  
      # Concatenate image features with maps created by replicating the
      # backbone output features in space:
      b, num_instances, h, w, _ = roi_resnet_features.shape
      backbone_feature_map = jnp.broadcast_to(
          image_backbone_features[:, :, None, None, :],
          (
              b,
              num_instances,
              self.mask_size,
              self.mask_size,
              c * self.backbone_out_width_mult,
          ),
      )
      roi_features = jnp.concatenate(
          [roi_resnet_features, backbone_feature_map], axis=-1
      )
  
      # Apply per-mask Hourglass network:
      roi_features = jnp.reshape(
          roi_features,
          (b * num_instances, h, w, roi_features.shape[-1]),
      )
      if self.add_mask_coords:
          roi_features = self._add_coord_channels(roi_features)
      pred_masks_batch, _ = HourglassNetwork(
          num_classes=1,
          width=self.mask_resnet_width,
          depth=self.mask_resnet_depth,
          name="mask_hourglass",
      )(roi_features)
      pred_masks = jnp.reshape(
          pred_masks_batch, (b, num_instances, h, w)
      )
  
      # If we're only predicting masks for the best boxes, scatter them back to
      # full size to align with the predicted boxes:
      if top_k_indices is not None:
          pred_masks = self._scatter_top_masks(
              pred_masks, top_k_indices, orig_num_boxes
          )
  
      return pred_masks
  
imports: import abc
  import functools
  from typing import Any, Callable, Dict, Optional, Tuple, Union, Sequence
  from absl import logging
  from big_vision.models import bit
  from big_vision.models import vit
  import flax
  import flax.linen
  import jax
  import jax.numpy
  import ml_collections
  import numpy
  from scenic.model_lib.base_models import box_utils
  from scenic.projects.owl_vit import utils
  from scenic.projects.owl_vit.clip import layers
  from scenic.projects.owl_vit.clip import model
label_desc: Use image node_id 5 for calling the BoxMaskHead obj's underlying member method code with example usage: obj.__call__(image, image_backbone_features, boxes) and returns: pred_masks