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 the class head to the image features.
  
  Args:
    image_features: Feature tokens extracted by the image embedder.
    query_embeddings: Optional list of text (or image) embeddings. If no
      embeddings are provided, no logits will be computed and only the class
      embeddings for the image will be returned.
    query_mask: Must be provided with query_embeddings. A mask indicating
      which query embeddings are valid.
  
  Returns:
    A dictionary containing the class_embeddings and the pred_logits if
      query_embeddings and query_mask are provided.
docstr_tok: ["Applies","the","class","head","to","the","image","features",".","Args",":","image_features",":","Feature","tokens","extracted","by","the","image","embedder",".","query_embeddings",":","Optional","list","of","text","(","or","image",")","embeddings",".","If","no","embeddings","are","provided",",","no","logits","will","be","computed","and","only","the","class","embeddings","for","the","image","will","be","returned",".","query_mask",":","Must","be","provided","with","query_embeddings",".","A","mask","indicating","which","query","embeddings","are","valid",".","Returns",":","A","dictionary","containing","the","class_embeddings","and","the","pred_logits","if","query_embeddings","and","query_mask","are","provided","."]
code: def class_predictor(
      self,
      image_features: jnp.ndarray,
      query_embeddings: Optional[jnp.ndarray] = None,
      query_mask: Optional[jnp.ndarray] = None,
  ) -> Dict[str, jnp.ndarray]:
      """Applies the class head to the image features.
  
      Args:
        image_features: Feature tokens extracted by the image embedder.
        query_embeddings: Optional list of text (or image) embeddings. If no
          embeddings are provided, no logits will be computed and only the class
          embeddings for the image will be returned.
        query_mask: Must be provided with query_embeddings. A mask indicating
          which query embeddings are valid.
  
      Returns:
        A dictionary containing the class_embeddings and the pred_logits if
          query_embeddings and query_mask are provided.
      """
      return self._class_head(
          image_features, query_embeddings, query_mask
      )
  
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 6 for calling the TextZeroShotDetectionModule obj's underlying member method code with example usage: obj.class_predictor(image_features, query_embeddings, query_mask) and returns: self