class_docstr: Text-query-based OWL-ViT model.
  
  This module computes joint text and image embeddings which are then
  used for localized prediction of bounding boxes and classes.
  
  Attributes:
    body_configs: Configurations of the image-text module.
    objectness_head_configs: Configurations for the (optional) objectness head.
    mask_head_configs: Configurations for the (optional) mask head.
    normalize: Whether to normalize the output of the model and the
      label_embeddings before computing the class logits.
    box_bias: Type of box bias - one of 'location', 'size' or 'both'.
class_docstr_tok: ["Text-query-based","OWL-ViT","model",".","This","module","computes","joint","text","and","image","embeddings","which","are","then","used","for","localized","prediction","of","bounding","boxes","and","classes",".","Attributes",":","body_configs",":","Configurations","of","the","image-text","module",".","objectness_head_configs",":","Configurations","for","the","(","optional",")","objectness","head",".","mask_head_configs",":","Configurations","for","the","(","optional",")","mask","head",".","normalize",":","Whether","to","normalize","the","output","of","the","model","and","the","label_embeddings","before","computing","the","class","logits",".","box_bias",":","Type","of","box","bias","-","one","of","'location","'",",","'size","'","or","'both","'","."]
docstr: Predicts bounding boxes from image features.
  
  Args:
    image_features: Features extracted from the image, flattened into a 1d
      sequence of tokens.
    feature_map: A 2d spatial re-arrangement of image_features.
    keep_image_tokens: If keep_image_tokens is not None, this indicates that
      image_features is a subset of tokens of the full grid. keep_image_tokens
      then contains the 1d indices of the kept tokens within the full token
      sequence. In that case, feature_map will contain dummy values at the
      dropped locations.
  
  Returns:
    List of predicted boxes (cxcywh normalized to 0, 1) nested within
      a dictionary.
docstr_tok: ["Predicts","bounding","boxes","from","image","features",".","Args",":","image_features",":","Features","extracted","from","the","image",",","flattened","into","a","1d","sequence","of","tokens",".","feature_map",":","A","2d","spatial","re-arrangement","of","image_features",".","keep_image_tokens",":","If","keep_image_tokens","is","not","None",",","this","indicates","that","image_features","is","a","subset","of","tokens","of","the","full","grid",".","keep_image_tokens","then","contains","the","1d","indices","of","the","kept","tokens","within","the","full","token","sequence",".","In","that","case",",","feature_map","will","contain","dummy","values","at","the","dropped","locations",".","Returns",":","List","of","predicted","boxes","(","cxcywh","normalized","to","0",",","1",")","nested","within","a","dictionary","."]
code: def box_predictor(
      self,
      *,
      image_features: jnp.ndarray,
      feature_map: jnp.ndarray,
      keep_image_tokens: Optional[jnp.ndarray] = None,
  ) -> Dict[str, jnp.ndarray]:
      """Predicts bounding boxes from image features.
  
      Args:
        image_features: Features extracted from the image, flattened into a 1d
          sequence of tokens.
        feature_map: A 2d spatial re-arrangement of image_features.
        keep_image_tokens: If keep_image_tokens is not None, this indicates that
          image_features is a subset of tokens of the full grid. keep_image_tokens
          then contains the 1d indices of the kept tokens within the full token
          sequence. In that case, feature_map will contain dummy values at the
          dropped locations.
  
      Returns:
        List of predicted boxes (cxcywh normalized to 0, 1) nested within
          a dictionary.
      """
      # Bounding box detection head [b, num_patches, 4].
      pred_boxes = self._box_head(image_features)
  
      # We compute the location of each token on the grid and use it to compute
      # a bias for the bbox prediction, i.e., each token is biased towards
      # predicting its location on the grid as the center.
      box_bias = utils.compute_box_bias(
          feature_map=feature_map, kind=self.box_bias
      )
  
      if keep_image_tokens is not None:
          box_bias = jnp.take_along_axis(
              box_bias[None, ...], keep_image_tokens[..., None], axis=-2
          )
  
      pred_boxes += box_bias
      pred_boxes = nn.sigmoid(pred_boxes)
      return {"pred_boxes": pred_boxes}
  
imports: import copy
  from typing import Any, Dict, List, Mapping, Optional
  import flax.linen
  from flax.training import checkpoints
  import jax
  import jax.numpy
  import ml_collections
  from scenic.projects.owl_vit import layers
  from scenic.projects.owl_vit import matching_base_models
  from scenic.projects.owl_vit import utils
  from scenic.projects.owl_vit.clip import model
  from scenic.projects.owl_vit.clip import tokenizer
label_desc: Use image node_id 5 for calling the TextZeroShotDetectionModule obj's underlying member method code with example usage: obj.box_predictor() and returns: dict