class_docstr: Wraps an OWL-ViT FLAX model for convenient inference.
  
  All public methods apply to a single example and take and return Numpy arrays.
  
  Attributes:
    config: ConfigDict with model configuration.
    module: OWL-ViT Flax module.
    variables: Variable dict to be used with module.apply.
docstr: Extracts image features in the region of `desired_boxes_yxyx`.
  
  
  This works by taking the features of a bounding box with high IOU
  to the desired_boxes_yxyx and returning the corresponding embedding
  features.
  
  Args:
    query_image: Image showing the example object to be used as query.
    query_box_yxyx: A single bounding box around the example object, in the
      TensorFlow format (y_min, x_min, y_max, x_max), normalized to [0, 1],
      for which to extract a query embedding.
  
  Returns:
    Queryable features and index of the predicted box whose features were
    selected.
docstr_tok: ["Extracts","image","features","in","the","region","of","`","desired_boxes_yxyx","`",".","This","works","by","taking","the","features","of","a","bounding","box","with","high","IOU","to","the","desired_boxes_yxyx","and","returning","the","corresponding","embedding","features",".","Args",":","query_image",":","Image","showing","the","example","object","to","be","used","as","query",".","query_box_yxyx",":","A","single","bounding","box","around","the","example","object",",","in","the","TensorFlow","format","(","y_min",",","x_min",",","y_max",",","x_max",")",",","normalized","to","[","0",",","1","]",",","for","which","to","extract","a","query","embedding",".","Returns",":","Queryable","features","and","index","of","the","predicted","box","whose","features","were","selected","."]
code: def embed_image_query(
      self,
      query_image: np.ndarray,
      query_box_yxyx: Tuple[float, float, float, float],
  ) -> Tuple[np.ndarray, np.ndarray]:
      """Extracts image features in the region of `desired_boxes_yxyx`.
  
  
      This works by taking the features of a bounding box with high IOU
      to the desired_boxes_yxyx and returning the corresponding embedding
      features.
  
      Args:
        query_image: Image showing the example object to be used as query.
        query_box_yxyx: A single bounding box around the example object, in the
          TensorFlow format (y_min, x_min, y_max, x_max), normalized to [0, 1],
          for which to extract a query embedding.
  
      Returns:
        Queryable features and index of the predicted box whose features were
        selected.
      """
      _, class_embeddings, pred_boxes = self.embed_image(query_image)
      ious = box_utils.box_iou(
          np.array(query_box_yxyx)[None, ...],
          box_utils.box_cxcywh_to_yxyx(pred_boxes, np),
          np_backbone=np,
      )[0][0]
  
      # If there are no overlapping boxes, fall back to generalized IoU:
      if np.all(ious == 0.0):
          ious = box_utils.generalized_box_iou(
              np.array(query_box_yxyx)[None, ...],
              box_utils.box_cxcywh_to_yxyx(pred_boxes, np),
              np_backbone=np,
          )[0]
  
      # Use an adaptive threshold such that all boxes within 80% of the best IoU
      # are included:
      iou_thresh = np.max(ious) * 0.8
  
      # Select class_embeddings that are above the IoU threshold:
      selected_inds = (ious >= iou_thresh).nonzero()[0]
      assert selected_inds.size
      selected_embeddings = class_embeddings[selected_inds]
  
      # Due to the DETR style bipartite matching loss, only one embedding
      # feature for each object is "good" and the rest are "background." To find
      # the one "good" feature we use the heuristic that it should be dissimilar
      # to the mean embedding.
      mean_embedding = np.mean(class_embeddings, axis=0)
      mean_sim = np.einsum(
          "d,id->i", mean_embedding, selected_embeddings
      )
  
      # Find box with lowest overall similarity:
      best_box_ind = selected_inds[np.argmin(mean_sim)]
  
      return class_embeddings[best_box_ind], best_box_ind
  
imports: import dataclasses
  import functools
  from typing import Any, Dict, Tuple
  from flax import 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.notebooks import numpy_cache
  from scipy import special
  from skimage import transform
  import tensorflow
label_desc: Use image node_id 7 for calling the Model obj's underlying member method code with example usage: obj.embed_image_query(query_image, query_box_yxyx) and returns: class_embeddings,  best_box_ind