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: Gathers the top instances based on IoU with ground-truth boxes.
docstr_tok: ["Gathers","the","top","instances","based","on","IoU","with","ground-truth","boxes","."]
code: def _gather_top_boxes(
      self, boxes, image_backbone_features, true_boxes
  ):
      """Gathers the top instances based on IoU with ground-truth boxes."""
      top_k_indices = None
      if self.num_training_boxes is not None and true_boxes is not None:
          iou_mat, _ = box_utils.box_iou(
              boxes1=box_utils.box_cxcywh_to_xyxy(boxes),
              boxes2=box_utils.box_cxcywh_to_xyxy(true_boxes),
              all_pairs=True,
          )
          max_iou = jnp.max(iou_mat, axis=-1)
          _, top_k_indices = jax.lax.top_k(
              max_iou, self.num_training_boxes or true_boxes.shape[-2]
          )
          gather = jax.vmap(lambda arr, idx: arr[idx])
          boxes = gather(boxes, top_k_indices)
          image_backbone_features = gather(
              image_backbone_features, top_k_indices
          )
      return boxes, image_backbone_features, top_k_indices
  
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 1 for calling the BoxMaskHead obj's underlying member method code with example usage: obj._gather_top_boxes(boxes, image_backbone_features, true_boxes) and returns: boxes,  image_backbone_features,  top_k_indices