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 (cropped) segmentation masks from the image features.
  
  Args:
    image: Input image, for extracting low-level image features.
    image_tokens: High-level features from the image embedder.
    boxes: Predicted bounding boxes corresponding to the image tokens.
    true_boxes: For filtering mask head predictions during training.
  
  Returns:
    A dictionary containing the predicted segmentation masks. The mask at
      index i corresponds to the predicted box in `pred_boxes` at index i.
docstr_tok: ["Predicts","(","cropped",")","segmentation","masks","from","the","image","features",".","Args",":","image",":","Input","image",",","for","extracting","low-level","image","features",".","image_tokens",":","High-level","features","from","the","image","embedder",".","boxes",":","Predicted","bounding","boxes","corresponding","to","the","image","tokens",".","true_boxes",":","For","filtering","mask","head","predictions","during","training",".","Returns",":","A","dictionary","containing","the","predicted","segmentation","masks",".","The","mask","at","index","i","corresponds","to","the","predicted","box","in","`","pred_boxes","`","at","index","i","."]
code: def mask_predictor(
      self, image, image_tokens, boxes, *, true_boxes=None
  ) -> Dict[str, jnp.ndarray]:
      """Predicts (cropped) segmentation masks from the image features.
  
      Args:
        image: Input image, for extracting low-level image features.
        image_tokens: High-level features from the image embedder.
        boxes: Predicted bounding boxes corresponding to the image tokens.
        true_boxes: For filtering mask head predictions during training.
  
      Returns:
        A dictionary containing the predicted segmentation masks. The mask at
          index i corresponds to the predicted box in `pred_boxes` at index i.
      """
      if self.mask_head_configs is None:
          raise ValueError(
              "Must pass mask_head_configs to use mask head."
          )
      pred_masks = self._mask_head(
          image, image_tokens, boxes, true_boxes=true_boxes
      )
      batch_size = image_tokens.shape[0]
      mask_size = self.mask_head_configs.mask_size
      return {
          "pred_masks": jnp.reshape(
              pred_masks, (batch_size, -1, mask_size, mask_size)
          )
      }
  
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 7 for calling the TextZeroShotDetectionModule obj's underlying member method code with example usage: obj.mask_predictor(image, image_tokens, boxes) and returns: dict