



"""simple docstring"""






from typing import Callable, Dict, Optional, Tuple

import torch
from torch import nn
from torch.distributions import (
    AffineTransform,
    Distribution,
    Independent,
    NegativeBinomial,
    Normal,
    StudentT,
    TransformedDistribution,
)






class 	__snake_case	(     _lowercase):




  def __init__(			self					:						Any						, __lowerCAmelCase					:						Distribution						, __lowerCAmelCase					:						List[Any]=None						, __lowerCAmelCase					:						Union[str, Any]=None						, __lowerCAmelCase					:						Dict=0		):






         """simple docstring"""






         _lowerCamelCase						:    List[Any]         =			1.0 if scale is None else scale
         _lowerCamelCase						:    Union[str, Any]         =			0.0 if loc is None else loc

         super().__init__(__lowerCAmelCase						, [AffineTransform(loc=self.loc						, scale=self.scale						, event_dim=__lowerCAmelCase		)]		)




  @property
  def 						SCREAMING_SNAKE_CASE    (			self					:						Dict		):






         """simple docstring"""






         return self.base_dist.mean * self.scale + self.loc




  @property
  def 						SCREAMING_SNAKE_CASE    (			self					:						List[Any]		):






         """simple docstring"""






         return self.base_dist.variance * self.scale**2




  @property
  def 						SCREAMING_SNAKE_CASE    (			self					:						Dict		):






         """simple docstring"""






         return self.variance.sqrt()






class 	__snake_case	(     nn.Module):




  def __init__(			self					:						Any						, __lowerCAmelCase					:						int						, __lowerCAmelCase					:						Dict[str, int]						, __lowerCAmelCase					:						Callable[..., Tuple[torch.Tensor]]						, **__lowerCAmelCase					:						Any		):






         """simple docstring"""






         super().__init__(**__lowerCAmelCase		)
         _lowerCamelCase						:    str         =			args_dim
         _lowerCamelCase						:    Dict         =			nn.ModuleList([nn.Linear(__lowerCAmelCase						, __lowerCAmelCase		) for dim in args_dim.values()]		)
         _lowerCamelCase						:    List[Any]         =			domain_map




  def 						SCREAMING_SNAKE_CASE    (			self					:						Optional[Any]						, __lowerCAmelCase					:						torch.Tensor		):






         """simple docstring"""






         _lowerCamelCase						:    Optional[Any]         =			[proj(__lowerCAmelCase		) for proj in self.proj]

         return self.domain_map(*__lowerCAmelCase		)






class 	__snake_case	(     nn.Module):




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






         """simple docstring"""






         super().__init__()
         _lowerCamelCase						:    str         =			function




  def 						SCREAMING_SNAKE_CASE    (			self					:						Union[str, Any]						, __lowerCAmelCase					:						Optional[int]						, *__lowerCAmelCase					:						int		):






         """simple docstring"""






         return self.function(__lowerCAmelCase						, *__lowerCAmelCase		)






class 	__snake_case	:
  snake_case__    :       type
  snake_case__    :       int
  snake_case__    :       Dict[str, int]




  def __init__(			self					:						int						, __lowerCAmelCase					:						int = 1		):






         """simple docstring"""






         _lowerCamelCase						:    str         =			dim
         _lowerCamelCase						:    Optional[int]         =			{k: dim * self.args_dim[k] for k in self.args_dim}




  def 						SCREAMING_SNAKE_CASE    (			self					:						Any						, __lowerCAmelCase					:						Any		):






         """simple docstring"""






         if self.dim == 1:
                return self.distribution_class(*__lowerCAmelCase		)
         else:
                return Independent(self.distribution_class(*__lowerCAmelCase		)						, 1		)




  def 						SCREAMING_SNAKE_CASE    (			self					:						str						, __lowerCAmelCase					:						Tuple						, __lowerCAmelCase					:						Optional[torch.Tensor] = None						, __lowerCAmelCase					:						Optional[torch.Tensor] = None						, ):






         """simple docstring"""






         _lowerCamelCase						:    int         =			self._base_distribution(__lowerCAmelCase		)
         if loc is None and scale is None:
                return distr
         else:
                return AffineTransformed(__lowerCAmelCase						, loc=__lowerCAmelCase						, scale=__lowerCAmelCase						, event_dim=self.event_dim		)




  @property
  def 						SCREAMING_SNAKE_CASE    (			self					:						int		):






         """simple docstring"""






         return () if self.dim == 1 else (self.dim,)




  @property
  def 						SCREAMING_SNAKE_CASE    (			self					:						Optional[int]		):






         """simple docstring"""






         return len(self.event_shape		)




  @property
  def 						SCREAMING_SNAKE_CASE    (			self					:						Any		):






         """simple docstring"""






         return 0.0




  def 						SCREAMING_SNAKE_CASE    (			self					:						Any						, __lowerCAmelCase					:						int		):






         """simple docstring"""






         return ParameterProjection(
             in_features=__lowerCAmelCase						, args_dim=self.args_dim						, domain_map=LambdaLayer(self.domain_map		)						, )




  def 						SCREAMING_SNAKE_CASE    (			self					:						str						, *__lowerCAmelCase					:						torch.Tensor		):






         """simple docstring"""






         raise NotImplementedError()




  @staticmethod
  def 						SCREAMING_SNAKE_CASE    (			__lowerCAmelCase					:						torch.Tensor		):






         """simple docstring"""






         return (x + torch.sqrt(torch.square(__lowerCAmelCase		) + 4.0		)) / 2.0






class 	__snake_case	(     _lowercase):
  snake_case__    :       Dict[str, int]        =						{"df": 1, "loc": 1, "scale": 1}
  snake_case__    :       type        =						StudentT




  @classmethod
  def 						SCREAMING_SNAKE_CASE    (			cls					:						Optional[int]						, __lowerCAmelCase					:						torch.Tensor						, __lowerCAmelCase					:						torch.Tensor						, __lowerCAmelCase					:						torch.Tensor		):






         """simple docstring"""






         _lowerCamelCase						:    List[str]         =			cls.squareplus(__lowerCAmelCase		).clamp_min(torch.finfo(scale.dtype		).eps		)
         _lowerCamelCase						:    Dict         =			2.0 + cls.squareplus(__lowerCAmelCase		)
         return df.squeeze(-1		), loc.squeeze(-1		), scale.squeeze(-1		)






class 	__snake_case	(     _lowercase):
  snake_case__    :       Dict[str, int]        =						{"loc": 1, "scale": 1}
  snake_case__    :       type        =						Normal




  @classmethod
  def 						SCREAMING_SNAKE_CASE    (			cls					:						Optional[int]						, __lowerCAmelCase					:						torch.Tensor						, __lowerCAmelCase					:						torch.Tensor		):






         """simple docstring"""






         _lowerCamelCase						:    Union[str, Any]         =			cls.squareplus(__lowerCAmelCase		).clamp_min(torch.finfo(scale.dtype		).eps		)
         return loc.squeeze(-1		), scale.squeeze(-1		)






class 	__snake_case	(     _lowercase):
  snake_case__    :       Dict[str, int]        =						{"total_count": 1, "logits": 1}
  snake_case__    :       type        =						NegativeBinomial




  @classmethod
  def 						SCREAMING_SNAKE_CASE    (			cls					:						Optional[int]						, __lowerCAmelCase					:						torch.Tensor						, __lowerCAmelCase					:						torch.Tensor		):






         """simple docstring"""






         _lowerCamelCase						:    Optional[int]         =			cls.squareplus(__lowerCAmelCase		)
         return total_count.squeeze(-1		), logits.squeeze(-1		)




  def 						SCREAMING_SNAKE_CASE    (			self					:						Optional[Any]						, __lowerCAmelCase					:						List[str]		):






         """simple docstring"""






         _lowerCamelCase      ,   _lowerCamelCase						:    Union[str, Any]         =			distr_args
         if self.dim == 1:
                return self.distribution_class(total_count=__lowerCAmelCase						, logits=__lowerCAmelCase		)
         else:
                return Independent(self.distribution_class(total_count=__lowerCAmelCase						, logits=__lowerCAmelCase		)						, 1		)




  def 						SCREAMING_SNAKE_CASE    (			self					:						Tuple						, __lowerCAmelCase					:						List[Any]						, __lowerCAmelCase					:						Optional[torch.Tensor] = None						, __lowerCAmelCase					:						Optional[torch.Tensor] = None		):






         """simple docstring"""






         _lowerCamelCase      ,   _lowerCamelCase						:    Any         =			distr_args

         if scale is not None:
                # See scaling property of Gamma.
                logits += scale.log()

         return self._base_distribution((total_count, logits)		)







"""simple docstring"""






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

lowerCAmelCase__      =    logging.get_logger(__name__)






@add_end_docstrings(_lowercase)
class 	__snake_case	(     _lowercase):




  def __init__(			self					:						Any						, **__lowerCAmelCase					:						Union[str, Any]		):






         """simple docstring"""






         super().__init__(**__lowerCAmelCase		)

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

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




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






         """simple docstring"""






         if "text_queries" in kwargs:
                _lowerCamelCase						:    List[Any]         =			kwargs.pop('''text_queries'''		)

         if isinstance(__lowerCAmelCase						, (str, Image.Image)		):
                _lowerCamelCase						:    Optional[int]         =			{'''image''': image, '''candidate_labels''': candidate_labels}
         else:
                _lowerCamelCase						:    List[Any]         =			image
         _lowerCamelCase						:    List[str]         =			super().__call__(__lowerCAmelCase						, **__lowerCAmelCase		)
         return results




  def 						SCREAMING_SNAKE_CASE    (			self					:						List[Any]						, **__lowerCAmelCase					:						int		):






         """simple docstring"""






         _lowerCamelCase						:    int         =			{}
         if "threshold" in kwargs:
                _lowerCamelCase						:    Optional[Any]         =			kwargs['''threshold''']
         if "top_k" in kwargs:
                _lowerCamelCase						:    int         =			kwargs['''top_k''']
         return {}, {}, postprocess_params




  def 						SCREAMING_SNAKE_CASE    (			self					:						List[Any]						, __lowerCAmelCase					:						Union[str, Any]		):






         """simple docstring"""






         _lowerCamelCase						:    int         =			load_image(inputs['''image''']		)
         _lowerCamelCase						:    Optional[Any]         =			inputs['''candidate_labels''']
         if isinstance(__lowerCAmelCase						, __lowerCAmelCase		):
                _lowerCamelCase						:    int         =			candidate_labels.split(''','''		)

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




  def 						SCREAMING_SNAKE_CASE    (			self					:						Any						, __lowerCAmelCase					:						List[Any]		):






         """simple docstring"""






         _lowerCamelCase						:    Optional[Any]         =			model_inputs.pop('''target_size'''		)
         _lowerCamelCase						:    List[Any]         =			model_inputs.pop('''candidate_label'''		)
         _lowerCamelCase						:    Dict         =			model_inputs.pop('''is_last'''		)

         _lowerCamelCase						:    str         =			self.model(**__lowerCAmelCase		)

         _lowerCamelCase						:    Optional[Any]         =			{'''target_size''': target_size, '''candidate_label''': candidate_label, '''is_last''': is_last, **outputs}
         return model_outputs




  def 						SCREAMING_SNAKE_CASE    (			self					:						Optional[Any]						, __lowerCAmelCase					:						Any						, __lowerCAmelCase					:						Union[str, Any]=0.1						, __lowerCAmelCase					:						Optional[Any]=None		):






         """simple docstring"""






         _lowerCamelCase						:    str         =			[]
         for model_output in model_outputs:
                _lowerCamelCase						:    Any         =			model_output['''candidate_label''']
                _lowerCamelCase						:    Union[str, Any]         =			BaseModelOutput(__lowerCAmelCase		)
                _lowerCamelCase						:    Tuple         =			self.image_processor.post_process_object_detection(
                    outputs=__lowerCAmelCase						, threshold=__lowerCAmelCase						, target_sizes=model_output['''target_size''']		)[0]

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

                       _lowerCamelCase						:    Optional[Any]         =			{'''score''': score, '''label''': label, '''box''': box}
                       results.append(__lowerCAmelCase		)

         _lowerCamelCase						:    int         =			sorted(__lowerCAmelCase						, key=lambda __lowerCAmelCase		: x["score"]						, reverse=__lowerCAmelCase		)
         if top_k:
                _lowerCamelCase						:    Dict         =			results[:top_k]

         return results




  def 						SCREAMING_SNAKE_CASE    (			self					:						Any						, __lowerCAmelCase					:						"torch.Tensor"		):






         """simple docstring"""






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


