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: Applies TextZeroShotDetectionModule on the input.
  
  Args:
    inputs: Images [batch_size, height, width, 3].
    text_queries: Queries to score boxes on. Queries starting with 0 stand for
      padding [batch_size=b, num_queries=q, max_query_length=l].
    train: Whether it is training.
    true_boxes: For filtering mask head predictions during training.
    debug: Unused.
  
  Returns:
    Outputs dict with items:
      pred_logits: Class logits [b, num_patches, num_queries].
      pred_boxes: Predicted bounding boxes [b, num_patches, 4].
      feature_map: Image embeddings 2d feature map [b, sp, sp, img_emb_dim].
docstr_tok: ["Applies","TextZeroShotDetectionModule","on","the","input",".","Args",":","inputs",":","Images","[","batch_size",",","height",",","width",",","3","]",".","text_queries",":","Queries","to","score","boxes","on",".","Queries","starting","with","0","stand","for","padding","[","batch_size=b",",","num_queries=q",",","max_query_length=l","]",".","train",":","Whether","it","is","training",".","true_boxes",":","For","filtering","mask","head","predictions","during","training",".","debug",":","Unused",".","Returns",":","Outputs","dict","with","items",":","pred_logits",":","Class","logits","[","b",",","num_patches",",","num_queries","]",".","pred_boxes",":","Predicted","bounding","boxes","[","b",",","num_patches",",","4","]",".","feature_map",":","Image","embeddings","2d","feature","map","[","b",",","sp",",","sp",",","img_emb_dim","]","."]
code: def __call__(
      self,
      inputs: jnp.ndarray,
      text_queries: jnp.ndarray,
      train: bool,
      *,
      true_boxes: Optional[jnp.ndarray] = None,
      debug: bool = False,
  ) -> Mapping[str, Any]:
      """Applies TextZeroShotDetectionModule on the input.
  
      Args:
        inputs: Images [batch_size, height, width, 3].
        text_queries: Queries to score boxes on. Queries starting with 0 stand for
          padding [batch_size=b, num_queries=q, max_query_length=l].
        train: Whether it is training.
        true_boxes: For filtering mask head predictions during training.
        debug: Unused.
  
      Returns:
        Outputs dict with items:
          pred_logits: Class logits [b, num_patches, num_queries].
          pred_boxes: Predicted bounding boxes [b, num_patches, 4].
          feature_map: Image embeddings 2d feature map [b, sp, sp, img_emb_dim].
      """
      del debug
      if not train and true_boxes is not None:
          raise ValueError(
              "True boxes should only be supplied during training."
          )
  
      keep_tokens = None
  
      # Embed images:
      feature_map = self.image_embedder(inputs, train)
      b, h, w, d = feature_map.shape
      image_features = jnp.reshape(feature_map, (b, h * w, d))
  
      # Embed queries:
      query_embeddings = self.text_embedder(text_queries, train)
      # If first token is 0, then this is a padding query [b, q].
      query_mask = (text_queries[..., 0] > 0).astype(jnp.float32)
  
      outputs = {
          "feature_map": feature_map,
          "query_embeddings": query_embeddings,
      }
  
      # Get objectness scores:
      if self.objectness_head_configs is not None:
          outputs.update(self.objectness_predictor(image_features))
  
      # During training, sample top tokens by objectness:
      num_instances = image_features.shape[-2]
      top_k = self.body_configs.get("objectness_top_k", num_instances)
      if train and (0 < top_k < num_instances):
          if "objectness_logits" not in outputs:
              raise ValueError(
                  "Need objectness head to sample by objectness."
              )
          outputs["objectness_logits"], keep_tokens = jax.lax.top_k(
              outputs["objectness_logits"],
              k=self.body_configs.objectness_top_k,
          )
          image_features = jnp.take_along_axis(
              image_features, keep_tokens[..., None], axis=-2
          )
  
      # Classification [b, num_patches, num_queries]:
      outputs.update(
          self.class_predictor(
              image_features, query_embeddings, query_mask
          )
      )
  
      # Predict boxes:
      outputs.update(
          self.box_predictor(
              image_features=image_features,
              feature_map=feature_map,
              keep_image_tokens=keep_tokens,
          )
      )
  
      # Predict masks:
      if self.mask_head_configs is not None:
          outputs.update(
              self.mask_predictor(
                  inputs,
                  image_features,
                  outputs["pred_boxes"],
                  true_boxes=true_boxes,
              )
          )
  
      return outputs
  
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 10 for calling the TextZeroShotDetectionModule obj's underlying member method code with example usage: obj.__call__(inputs, text_queries, train) and returns: outputs