





import warnings
from typing import List

import numpy as np

from ...processing_utils import ProcessorMixin
from ...tokenization_utils_base import BatchEncoding
from ...utils import is_flax_available, is_tf_available, is_torch_available



class 							_UpperCAmelCase       (       _UpperCamelCase       ):

  """simple docstring"""


  a_              =					["""image_processor""", """tokenizer"""]
  a_              =					"""OwlViTImageProcessor"""
  a_              =					("""CLIPTokenizer""", """CLIPTokenizerFast""")






  def __init__(							self					:						Optional[int]   ,    lowerCAmelCase_					:						int=None   ,    lowerCAmelCase_					:						Optional[int]=None   ,    **lowerCAmelCase_					:						Optional[Any]	)				->		Optional[int]:
     __lowerCAmelCase				      =      None
     if "feature_extractor" in kwargs:
        warnings.warn(
            'The `feature_extractor` argument is deprecated and will be removed in v5, use `image_processor`'
            ' instead.'   ,    lowerCAmelCase_   ,    )
        __lowerCAmelCase				      =      kwargs.pop('feature_extractor'	)

     __lowerCAmelCase				      =      image_processor if image_processor is not None else feature_extractor
     if image_processor is None:
        raise ValueError('You need to specify an `image_processor`.'	)
     if tokenizer is None:
        raise ValueError('You need to specify a `tokenizer`.'	)

     super().__init__(lowerCAmelCase_   ,    lowerCAmelCase_	)






  def __call__(							self					:						List[Any]   ,    lowerCAmelCase_					:						List[Any]=None   ,    lowerCAmelCase_					:						Tuple=None   ,    lowerCAmelCase_					:						List[Any]=None   ,    lowerCAmelCase_					:						Union[str, Any]="max_length"   ,    lowerCAmelCase_					:						int="np"   ,    **lowerCAmelCase_					:						str	)				->		Tuple:

     if text is None and query_images is None and images is None:
        raise ValueError(
            'You have to specify at least one text or query image or image. All three cannot be none.'	)

     if text is not None:
        if isinstance(lowerCAmelCase_   ,    lowerCAmelCase_	) or (isinstance(lowerCAmelCase_   ,    lowerCAmelCase_	) and not isinstance(text[0]   ,    lowerCAmelCase_	)):
           __lowerCAmelCase				      =      [self.tokenizer(lowerCAmelCase_   ,    padding=lowerCAmelCase_   ,    return_tensors=lowerCAmelCase_   ,    **lowerCAmelCase_	)]

        elif isinstance(lowerCAmelCase_   ,    lowerCAmelCase_	) and isinstance(text[0]   ,    lowerCAmelCase_	):
           __lowerCAmelCase				      =      []

           # Maximum number of queries across batch
           __lowerCAmelCase				      =      max([len(lowerCAmelCase_	) for t in text]	)

           # Pad all batch samples to max number of text queries
           for t in text:
              if len(lowerCAmelCase_	) != max_num_queries:
                 __lowerCAmelCase				      =      t + [' '] * (max_num_queries - len(lowerCAmelCase_	))

              __lowerCAmelCase				      =      self.tokenizer(lowerCAmelCase_   ,    padding=lowerCAmelCase_   ,    return_tensors=lowerCAmelCase_   ,    **lowerCAmelCase_	)
              encodings.append(lowerCAmelCase_	)
        else:
           raise TypeError('Input text should be a string, a list of strings or a nested list of strings'	)

        if return_tensors == "np":
           __lowerCAmelCase				      =      np.concatenate([encoding['input_ids'] for encoding in encodings]   ,    axis=0	)
           __lowerCAmelCase				      =      np.concatenate([encoding['attention_mask'] for encoding in encodings]   ,    axis=0	)

        elif return_tensors == "jax" and is_flax_available():
           import jax.numpy as jnp

           __lowerCAmelCase				      =      jnp.concatenate([encoding['input_ids'] for encoding in encodings]   ,    axis=0	)
           __lowerCAmelCase				      =      jnp.concatenate([encoding['attention_mask'] for encoding in encodings]   ,    axis=0	)

        elif return_tensors == "pt" and is_torch_available():
           import torch

           __lowerCAmelCase				      =      torch.cat([encoding['input_ids'] for encoding in encodings]   ,    dim=0	)
           __lowerCAmelCase				      =      torch.cat([encoding['attention_mask'] for encoding in encodings]   ,    dim=0	)

        elif return_tensors == "tf" and is_tf_available():
           import tensorflow as tf

           __lowerCAmelCase				      =      tf.stack([encoding['input_ids'] for encoding in encodings]   ,    axis=0	)
           __lowerCAmelCase				      =      tf.stack([encoding['attention_mask'] for encoding in encodings]   ,    axis=0	)

        else:
           raise ValueError('Target return tensor type could not be returned'	)

        __lowerCAmelCase				      =      BatchEncoding()
        __lowerCAmelCase				      =      input_ids
        __lowerCAmelCase				      =      attention_mask

     if query_images is not None:
        __lowerCAmelCase				      =      BatchEncoding()
        __lowerCAmelCase				      =      self.image_processor(
            lowerCAmelCase_   ,    return_tensors=lowerCAmelCase_   ,    **lowerCAmelCase_	).pixel_values
        __lowerCAmelCase				      =      query_pixel_values

     if images is not None:
        __lowerCAmelCase				      =      self.image_processor(lowerCAmelCase_   ,    return_tensors=lowerCAmelCase_   ,    **lowerCAmelCase_	)

     if text is not None and images is not None:
        __lowerCAmelCase				      =      image_features.pixel_values
        return encoding
     elif query_images is not None and images is not None:
        __lowerCAmelCase				      =      image_features.pixel_values
        return encoding
     elif text is not None or query_images is not None:
        return encoding
     else:
        return BatchEncoding(data=dict(**lowerCAmelCase_	)   ,    tensor_type=lowerCAmelCase_	)






  def 							lowercase						(							self					:						str   ,    *lowerCAmelCase_					:						Union[str, Any]   ,    **lowerCAmelCase_					:						List[str]	)				->		List[str]:
     return self.image_processor.post_process(*lowerCAmelCase_   ,    **lowerCAmelCase_	)






  def 							lowercase						(							self					:						Union[str, Any]   ,    *lowerCAmelCase_					:						Union[str, Any]   ,    **lowerCAmelCase_					:						Tuple	)				->		Optional[Any]:
     return self.image_processor.post_process_object_detection(*lowerCAmelCase_   ,    **lowerCAmelCase_	)






  def 							lowercase						(							self					:						Union[str, Any]   ,    *lowerCAmelCase_					:						Any   ,    **lowerCAmelCase_					:						List[str]	)				->		Dict:
     return self.image_processor.post_process_image_guided_detection(*lowerCAmelCase_   ,    **lowerCAmelCase_	)






  def 							lowercase						(							self					:						str   ,    *lowerCAmelCase_					:						Union[str, Any]   ,    **lowerCAmelCase_					:						Optional[Any]	)				->		Optional[Any]:
     return self.tokenizer.batch_decode(*lowerCAmelCase_   ,    **lowerCAmelCase_	)






  def 							lowercase						(							self					:						Union[str, Any]   ,    *lowerCAmelCase_					:						int   ,    **lowerCAmelCase_					:						List[str]	)				->		Dict:
     return self.tokenizer.decode(*lowerCAmelCase_   ,    **lowerCAmelCase_	)






  @property
  def 							lowercase						(							self					:						Dict	)				->		int:
     warnings.warn(
         '`feature_extractor_class` is deprecated and will be removed in v5. Use `image_processor_class` instead.'   ,    lowerCAmelCase_   ,    )
     return self.image_processor_class






  @property
  def 							lowercase						(							self					:						Union[str, Any]	)				->		List[Any]:
     warnings.warn(
         '`feature_extractor` is deprecated and will be removed in v5. Use `image_processor` instead.'   ,    lowerCAmelCase_   ,    )
     return self.image_processor









import re

import jax.numpy as jnp
from flax.traverse_util import flatten_dict, unflatten_dict
from jax.random import PRNGKey

from ..utils import logging


lowerCamelCase      :       int    =						logging.get_logger(__name__)





def    __lowerCAmelCase   (    __snake_case							):
     __lowerCAmelCase             =       r"\w+[.]\d+"
     __lowerCAmelCase             =       re.findall(__snake_case       ,			__snake_case							)
     for pat in pats:
          __lowerCAmelCase             =       key.replace(__snake_case       ,			"_".join(pat.split("."							)							)							)
     return key





def    __lowerCAmelCase   (    __snake_case       ,			__snake_case       ,			__snake_case							):
     __lowerCAmelCase             =       pt_tuple_key[:-1] + ("scale",)
     if (
         any("norm" in str_ for str_ in pt_tuple_key							)
         and (pt_tuple_key[-1] == "bias")
         and (pt_tuple_key[:-1] + ("bias",) not in random_flax_state_dict)
         and (pt_tuple_key[:-1] + ("scale",) in random_flax_state_dict)
     ):
          __lowerCAmelCase             =       pt_tuple_key[:-1] + ("scale",)
          return renamed_pt_tuple_key, pt_tensor
     elif pt_tuple_key[-1] in ["weight", "gamma"] and pt_tuple_key[:-1] + ("scale",) in random_flax_state_dict:
          __lowerCAmelCase             =       pt_tuple_key[:-1] + ("scale",)
          return renamed_pt_tuple_key, pt_tensor

     # embedding
     if pt_tuple_key[-1] == "weight" and pt_tuple_key[:-1] + ("embedding",) in random_flax_state_dict:
          __lowerCAmelCase             =       pt_tuple_key[:-1] + ("embedding",)
          return renamed_pt_tuple_key, pt_tensor

     # conv layer
     __lowerCAmelCase             =       pt_tuple_key[:-1] + ("kernel",)
     if pt_tuple_key[-1] == "weight" and pt_tensor.ndim == 4:
          __lowerCAmelCase             =       pt_tensor.transpose(2       ,			3       ,			1       ,			0							)
          return renamed_pt_tuple_key, pt_tensor

     # linear layer
     __lowerCAmelCase             =       pt_tuple_key[:-1] + ("kernel",)
     if pt_tuple_key[-1] == "weight":
          __lowerCAmelCase             =       pt_tensor.T
          return renamed_pt_tuple_key, pt_tensor

     # old PyTorch layer norm weight
     __lowerCAmelCase             =       pt_tuple_key[:-1] + ("weight",)
     if pt_tuple_key[-1] == "gamma":
          return renamed_pt_tuple_key, pt_tensor

     # old PyTorch layer norm bias
     __lowerCAmelCase             =       pt_tuple_key[:-1] + ("bias",)
     if pt_tuple_key[-1] == "beta":
          return renamed_pt_tuple_key, pt_tensor

     return pt_tuple_key, pt_tensor





def    __lowerCAmelCase   (    __snake_case       ,			__snake_case       ,			__snake_case=42							):
     # Step 1: Convert pytorch tensor to numpy
     __lowerCAmelCase             =       {k: v.numpy() for k, v in pt_state_dict.items()}

     # Step 2: Since the model is stateless, get random Flax params
     __lowerCAmelCase             =       flax_model.init_weights(PRNGKey(__snake_case							)							)

     __lowerCAmelCase             =       flatten_dict(__snake_case							)
     __lowerCAmelCase             =       {}

     # Need to change some parameters name to match Flax names
     for pt_key, pt_tensor in pt_state_dict.items():
          __lowerCAmelCase             =       rename_key(__snake_case							)
          __lowerCAmelCase             =       tuple(renamed_pt_key.split("."							)							)

          # Correctly rename weight parameters
          __lowerCAmelCase      ,			__lowerCAmelCase             =       rename_key_and_reshape_tensor(__snake_case       ,			__snake_case       ,			__snake_case							)

          if flax_key in random_flax_state_dict:
               if flax_tensor.shape != random_flax_state_dict[flax_key].shape:
                    raise ValueError(
                        F"""PyTorch checkpoint seems to be incorrect. Weight {pt_key} was expected to be of shape """
                        F"""{random_flax_state_dict[flax_key].shape}, but is {flax_tensor.shape}."""							)

        # also add unexpected weight so that warning is thrown
          __lowerCAmelCase             =       jnp.asarray(__snake_case							)

     return unflatten_dict(__snake_case							)






