





import numpy as np
import torch
from torch.utils.data import Dataset, IterableDataset

from ..utils.generic import ModelOutput




class      UpperCamelCase (     UpperCamelCase_   ):

    def __init__(			self :      str							,    snake_case__ :      Optional[int]							,    snake_case__ :      Optional[int]							,    snake_case__ :      Optional[Any]    ):

     """simple docstring"""






     SCREAMING_SNAKE_CASE								=			dataset
     SCREAMING_SNAKE_CASE								=			process
     SCREAMING_SNAKE_CASE								=			params

    def __len__(			self :      Tuple    ):

     """simple docstring"""






     return len(self.dataset    )




    def __getitem__(			self :      List[str]							,    snake_case__ :      Optional[Any]    ):

     """simple docstring"""






     SCREAMING_SNAKE_CASE								=			self.dataset[i]
     SCREAMING_SNAKE_CASE								=			self.process(UpperCamelCase__							,    **self.params    )
     return processed




class      UpperCamelCase (     UpperCamelCase_   ):

    def __init__(			self :      Dict							,    snake_case__ :      Union[str, Any]							,    snake_case__ :      Optional[int]							,    snake_case__ :      Dict							,    snake_case__ :      List[Any]=None    ):

     """simple docstring"""






     SCREAMING_SNAKE_CASE								=			loader
     SCREAMING_SNAKE_CASE								=			infer
     SCREAMING_SNAKE_CASE								=			params
     if loader_batch_size == 1:
      # Let's spare some time by deactivating altogether
      SCREAMING_SNAKE_CASE								=			None
     SCREAMING_SNAKE_CASE								=			loader_batch_size

     # Internal bookkeeping
     SCREAMING_SNAKE_CASE								=			None
     SCREAMING_SNAKE_CASE								=			None

    def __len__(			self :      int    ):

     """simple docstring"""






     return len(self.loader    )

    def __iter__(			self :      List[str]    ):

     """simple docstring"""






     SCREAMING_SNAKE_CASE								=			iter(self.loader    )
     return self

    def UpperCamelCase			(			self :      List[str]    ):

     """simple docstring"""






     if isinstance(self._loader_batch_data							,    torch.Tensor    ):
      # Batch data is simple tensor, just fetch the slice
      SCREAMING_SNAKE_CASE								=			self._loader_batch_data[self._loader_batch_index]
     else:
      # Batch data is assumed to be BaseModelOutput (or dict)
      SCREAMING_SNAKE_CASE								=			{}
      for k, element in self._loader_batch_data.items():
       if isinstance(UpperCamelCase__							,    UpperCamelCase__    ):
        # Convert ModelOutput to tuple first
        SCREAMING_SNAKE_CASE								=			element.to_tuple()
        if isinstance(element[0]							,    torch.Tensor    ):
         SCREAMING_SNAKE_CASE								=			tuple(el[self._loader_batch_index].unsqueeze(0    ) for el in element    )
        elif isinstance(element[0]							,    np.ndarray    ):
         SCREAMING_SNAKE_CASE								=			tuple(np.expand_dims(el[self._loader_batch_index]							,    0    ) for el in element    )
        continue
       if k in {"hidden_states", "past_key_values", "attentions"} and isinstance(UpperCamelCase__							,    UpperCamelCase__    ):
        # Those are stored as lists of tensors so need specific unbatching.
        if isinstance(element[0]							,    torch.Tensor    ):
         SCREAMING_SNAKE_CASE								=			tuple(el[self._loader_batch_index].unsqueeze(0    ) for el in element    )
        elif isinstance(element[0]							,    np.ndarray    ):
         SCREAMING_SNAKE_CASE								=			tuple(np.expand_dims(el[self._loader_batch_index]							,    0    ) for el in element    )
        continue
       if element is None:
        # This can happen for optional data that get passed around
        SCREAMING_SNAKE_CASE								=			None
       elif isinstance(element[self._loader_batch_index]							,    torch.Tensor    ):
        # Take correct batch data, but make it looked like batch_size=1
        # For compatibility with other methods within transformers

        SCREAMING_SNAKE_CASE								=			element[self._loader_batch_index].unsqueeze(0    )
       elif isinstance(element[self._loader_batch_index]							,    np.ndarray    ):
        # Take correct batch data, but make it looked like batch_size=1
        # For compatibility with other methods within transformers
        SCREAMING_SNAKE_CASE								=			np.expand_dims(element[self._loader_batch_index]							,    0    )
       else:
        # This is typically a list, so no need to `unsqueeze`.
        SCREAMING_SNAKE_CASE								=			element[self._loader_batch_index]
            # Recreate the element by reusing the original class to make it look
            # batch_size=1
      SCREAMING_SNAKE_CASE								=			self._loader_batch_data.__class__(UpperCamelCase__    )
     self._loader_batch_index += 1
     return result




    def UpperCamelCase			(			self :      Union[str, Any]    ):

     """simple docstring"""






     if self._loader_batch_index is not None and self._loader_batch_index < self.loader_batch_size:
      # We are currently unrolling a batch so we just need to return
      # the current item within a batch
      return self.loader_batch_item()

     # We're out of items within a batch
     SCREAMING_SNAKE_CASE								=			next(self.iterator    )
     SCREAMING_SNAKE_CASE								=			self.infer(UpperCamelCase__							,    **self.params    )
     # We now have a batch of "inferred things".
     if self.loader_batch_size is not None:
      # Try to infer the size of the batch
      if isinstance(UpperCamelCase__							,    torch.Tensor    ):
       SCREAMING_SNAKE_CASE								=			processed
      else:
       SCREAMING_SNAKE_CASE								=			list(processed.keys()    )[0]
       SCREAMING_SNAKE_CASE								=			processed[key]
      if isinstance(UpperCamelCase__							,    UpperCamelCase__    ):
       SCREAMING_SNAKE_CASE								=			len(UpperCamelCase__    )
      else:
       SCREAMING_SNAKE_CASE								=			first_tensor.shape[0]
      if 0 < observed_batch_size < self.loader_batch_size:
       # could be last batch so we can't unroll as many
       # elements.
       SCREAMING_SNAKE_CASE								=			observed_batch_size
      # Setting internal index to unwrap the batch
      SCREAMING_SNAKE_CASE								=			processed
      SCREAMING_SNAKE_CASE								=			0
      return self.loader_batch_item()
     else:
      # We're not unrolling batches
      return processed




class      UpperCamelCase (     UpperCamelCase_   ):

    def __init__(			self :      str							,    snake_case__ :      List[Any]							,    snake_case__ :      Optional[int]							,    snake_case__ :      Union[str, Any]							,    snake_case__ :      Union[str, Any]=None    ):

     """simple docstring"""






     super().__init__(UpperCamelCase__							,    UpperCamelCase__							,    UpperCamelCase__    )

    def __iter__(			self :      Union[str, Any]    ):

     """simple docstring"""






     SCREAMING_SNAKE_CASE								=			iter(self.loader    )
     SCREAMING_SNAKE_CASE								=			None
     return self




    def UpperCamelCase			(			self :      List[str]    ):

     """simple docstring"""






     if self.subiterator is None:
      SCREAMING_SNAKE_CASE								=			self.infer(next(self.iterator    )							,    **self.params    )
     try:
      # Try to return next item
      SCREAMING_SNAKE_CASE								=			next(self.subiterator    )
     except StopIteration:
      # When a preprocess iterator ends, we can start lookig at the next item
      # ChunkIterator will keep feeding until ALL elements of iterator
      # all have created their subiterator and have been iterating against.
      #
      # Another way to look at it, is we're basically flattening lists of lists
      # into a single list, but with generators
      SCREAMING_SNAKE_CASE								=			self.infer(next(self.iterator    )							,    **self.params    )
      SCREAMING_SNAKE_CASE								=			next(self.subiterator    )
     return processed




class      UpperCamelCase (     UpperCamelCase_   ):

    def __iter__(			self :      Tuple    ):

     """simple docstring"""






     SCREAMING_SNAKE_CASE								=			iter(self.loader    )
     return self




    def UpperCamelCase			(			self :      List[str]    ):

     """simple docstring"""






     SCREAMING_SNAKE_CASE								=			False
     SCREAMING_SNAKE_CASE								=			[]
     if self._loader_batch_index is not None and self._loader_batch_index < self.loader_batch_size:
      while self._loader_batch_index < self.loader_batch_size:
       SCREAMING_SNAKE_CASE								=			self.loader_batch_item()
       SCREAMING_SNAKE_CASE								=			item.pop('is_last'    )
       accumulator.append(UpperCamelCase__    )
       if is_last:
        return accumulator

     while not is_last:
      SCREAMING_SNAKE_CASE								=			self.infer(next(self.iterator    )							,    **self.params    )
      if self.loader_batch_size is not None:
       if isinstance(UpperCamelCase__							,    torch.Tensor    ):
        SCREAMING_SNAKE_CASE								=			processed
       else:
        SCREAMING_SNAKE_CASE								=			list(processed.keys()    )[0]
        SCREAMING_SNAKE_CASE								=			processed[key]
       if isinstance(UpperCamelCase__							,    UpperCamelCase__    ):
        SCREAMING_SNAKE_CASE								=			len(UpperCamelCase__    )
       else:
        SCREAMING_SNAKE_CASE								=			first_tensor.shape[0]
       if 0 < observed_batch_size < self.loader_batch_size:
        # could be last batch so we can't unroll as many
        # elements.
        SCREAMING_SNAKE_CASE								=			observed_batch_size
       SCREAMING_SNAKE_CASE								=			processed
       SCREAMING_SNAKE_CASE								=			0
       while self._loader_batch_index < self.loader_batch_size:
        SCREAMING_SNAKE_CASE								=			self.loader_batch_item()
        SCREAMING_SNAKE_CASE								=			item.pop('is_last'    )
        accumulator.append(UpperCamelCase__    )
        if is_last:
         return accumulator
      else:
       SCREAMING_SNAKE_CASE								=			processed
       SCREAMING_SNAKE_CASE								=			item.pop('is_last'    )
       accumulator.append(UpperCamelCase__    )
     return accumulator




class      UpperCamelCase (     UpperCamelCase_   ):

    def __init__(			self :      Optional[Any]							,    snake_case__ :      Dataset							,    snake_case__ :      str    ):

     """simple docstring"""






     SCREAMING_SNAKE_CASE								=			dataset
     SCREAMING_SNAKE_CASE								=			key

    def __len__(			self :      Optional[int]    ):

     """simple docstring"""






     return len(self.dataset    )




    def __getitem__(			self :      Dict							,    snake_case__ :      Tuple    ):

     """simple docstring"""






     return self.dataset[i][self.key]




class      UpperCamelCase (     UpperCamelCase_   ):

    def __init__(			self :      List[Any]							,    snake_case__ :      Dataset							,    snake_case__ :      str							,    snake_case__ :      str    ):

     """simple docstring"""






     SCREAMING_SNAKE_CASE								=			dataset
     SCREAMING_SNAKE_CASE								=			keya
     SCREAMING_SNAKE_CASE								=			keya

    def __len__(			self :      List[str]    ):

     """simple docstring"""






     return len(self.dataset    )




    def __getitem__(			self :      Union[str, Any]							,    snake_case__ :      Any    ):

     """simple docstring"""






     return {"text": self.dataset[i][self.keya], "text_pair": self.dataset[i][self.keya]}




from typing import Any, Dict, List, Union

from ..utils import add_end_docstrings, is_torch_available, is_vision_available, logging, requires_backends
from .base import PIPELINE_INIT_ARGS, ChunkPipeline


if is_vision_available():
      from PIL import Image

      from ..image_utils import load_image

if is_torch_available():
      import torch

      from transformers.modeling_outputs import BaseModelOutput

      from ..models.auto.modeling_auto import MODEL_FOR_ZERO_SHOT_OBJECT_DETECTION_MAPPING

a_    :						Dict    		=      logging.get_logger(__name__)




@add_end_docstrings(SCREAMING_SNAKE_CASE   )
class      UpperCamelCase (     SCREAMING_SNAKE_CASE   ):

    def __init__(			self :      Any							,    **snake_case__ :      Optional[int]    ):

     """simple docstring"""






     super().__init__(**snake_case__    )

     if self.framework == "tf":
      raise ValueError(F"""The {self.__class__} is only available in PyTorch."""    )

     requires_backends(self							,    'vision'    )
     self.check_model_type(snake_case__    )

    def __call__(			self :      List[Any]							,    snake_case__ :      Union[str, "Image.Image", List[Dict[str, Any]]]							,    snake_case__ :      Union[str, List[str]] = None							,    **snake_case__ :      Union[str, Any]							,    ):

     """simple docstring"""






     if "text_queries" in kwargs:
      SCREAMING_SNAKE_CASE								=			kwargs.pop('text_queries'    )

     if isinstance(snake_case__							,    (str, Image.Image)    ):
      SCREAMING_SNAKE_CASE								=			{'image': image, 'candidate_labels': candidate_labels}
     else:
      SCREAMING_SNAKE_CASE								=			image
     SCREAMING_SNAKE_CASE								=			super().__call__(snake_case__							,    **snake_case__    )
     return results

    def UpperCamelCase			(			self :      Union[str, Any]							,    **snake_case__ :      Optional[Any]    ):

     """simple docstring"""






     SCREAMING_SNAKE_CASE								=			{}
     if "threshold" in kwargs:
      SCREAMING_SNAKE_CASE								=			kwargs['threshold']
     if "top_k" in kwargs:
      SCREAMING_SNAKE_CASE								=			kwargs['top_k']
     return {}, {}, postprocess_params

    def UpperCamelCase			(			self :      List[Any]							,    snake_case__ :      Optional[int]    ):

     """simple docstring"""






     SCREAMING_SNAKE_CASE								=			load_image(inputs['image']    )
     SCREAMING_SNAKE_CASE								=			inputs['candidate_labels']
     if isinstance(snake_case__							,    snake_case__    ):
      SCREAMING_SNAKE_CASE								=			candidate_labels.split(','    )

     SCREAMING_SNAKE_CASE								=			torch.tensor([[image.height, image.width]]							,    dtype=torch.intaa    )
     for i, candidate_label in enumerate(snake_case__    ):
      SCREAMING_SNAKE_CASE								=			self.tokenizer(snake_case__							,    return_tensors=self.framework    )
      SCREAMING_SNAKE_CASE								=			self.image_processor(snake_case__							,    return_tensors=self.framework    )
      yield {
          "is_last": i == len(snake_case__    ) - 1,
          "target_size": target_size,
          "candidate_label": candidate_label,
          **text_inputs,
          **image_features,
      }

    def UpperCamelCase			(			self :      Any							,    snake_case__ :      Dict    ):

     """simple docstring"""






     SCREAMING_SNAKE_CASE								=			model_inputs.pop('target_size'    )
     SCREAMING_SNAKE_CASE								=			model_inputs.pop('candidate_label'    )
     SCREAMING_SNAKE_CASE								=			model_inputs.pop('is_last'    )

     SCREAMING_SNAKE_CASE								=			self.model(**snake_case__    )

     SCREAMING_SNAKE_CASE								=			{'target_size': target_size, 'candidate_label': candidate_label, 'is_last': is_last, **outputs}
     return model_outputs

    def UpperCamelCase			(			self :      Union[str, Any]							,    snake_case__ :      Tuple							,    snake_case__ :      str=0.1							,    snake_case__ :      Union[str, Any]=None    ):

     """simple docstring"""






     SCREAMING_SNAKE_CASE								=			[]
     for model_output in model_outputs:
      SCREAMING_SNAKE_CASE								=			model_output['candidate_label']
      SCREAMING_SNAKE_CASE								=			BaseModelOutput(snake_case__    )
      SCREAMING_SNAKE_CASE								=			self.image_processor.post_process_object_detection(
          outputs=snake_case__							,    threshold=snake_case__							,    target_sizes=model_output['target_size']    )[0]

      for index in outputs["scores"].nonzero():
       SCREAMING_SNAKE_CASE								=			outputs['scores'][index].item()
       SCREAMING_SNAKE_CASE								=			self._get_bounding_box(outputs['boxes'][index][0]    )

       SCREAMING_SNAKE_CASE								=			{'score': score, 'label': label, 'box': box}
       results.append(snake_case__    )

     SCREAMING_SNAKE_CASE								=			sorted(snake_case__							,    key=lambda snake_case__    : x["score"]							,    reverse=snake_case__    )
     if top_k:
      SCREAMING_SNAKE_CASE								=			results[:top_k]

     return results




    def UpperCamelCase			(			self :      List[Any]							,    snake_case__ :      "torch.Tensor"    ):

     """simple docstring"""






     if self.framework != "pt":
      raise ValueError('The ZeroShotObjectDetectionPipeline is only available in PyTorch.'    )
     SCREAMING_SNAKE_CASE						,	SCREAMING_SNAKE_CASE						,	SCREAMING_SNAKE_CASE						,	SCREAMING_SNAKE_CASE								=			box.int().tolist()
     SCREAMING_SNAKE_CASE								=			{
         'xmin': xmin,
         'ymin': ymin,
         'xmax': xmax,
         'ymax': ymax,
     }
     return bbox

