

from typing import List, Union

from ..utils import (
    add_end_docstrings,
    is_tf_available,
    is_torch_available,
    is_vision_available,
    logging,
    requires_backends,
)
from .base import PIPELINE_INIT_ARGS, Pipeline


if is_vision_available():
  from PIL import Image

  from ..image_utils import load_image

if is_tf_available():
  import tensorflow as tf

  from ..models.auto.modeling_tf_auto import TF_MODEL_FOR_IMAGE_CLASSIFICATION_MAPPING
  from ..tf_utils import stable_softmax

if is_torch_available():
  from ..models.auto.modeling_auto import MODEL_FOR_IMAGE_CLASSIFICATION_MAPPING

lowercase_             =     logging.get_logger(__name__)





@add_end_docstrings(SCREAMING_SNAKE_CASE							)
class 				__lowerCAmelCase  (      SCREAMING_SNAKE_CASE							):




  def __init__(       self							,     *lowerCAmelCase							,     **lowerCAmelCase      )						->				Dict:
        '''simple docstring'''






        super().__init__(*lowerCAmelCase							,     **lowerCAmelCase      )
        requires_backends(self							,     'vision'      )
        self.check_model_type(
            TF_MODEL_FOR_IMAGE_CLASSIFICATION_MAPPING
            if self.framework == 'tf'
            else MODEL_FOR_IMAGE_CLASSIFICATION_MAPPING      )




  def 						A__  (       self							,     lowerCAmelCase=None      )						->				int:
        '''simple docstring'''






        _lowercase			 ={}
        if top_k is not None:
              _lowercase			 =top_k
        return {}, {}, postprocess_params




  def __call__(       self							,     lowerCAmelCase							,     **lowerCAmelCase      )						->				Dict:
        '''simple docstring'''






        return super().__call__(lowerCAmelCase							,     **lowerCAmelCase      )




  def 						A__  (       self							,     lowerCAmelCase      )						->				Optional[int]:
        '''simple docstring'''






        _lowercase			 =load_image(lowerCAmelCase      )
        _lowercase			 =self.image_processor(images=lowerCAmelCase							,     return_tensors=self.framework      )
        return model_inputs




  def 						A__  (       self							,     lowerCAmelCase      )						->				Tuple:
        '''simple docstring'''






        _lowercase			 =self.model(**lowerCAmelCase      )
        return model_outputs




  def 						A__  (       self							,     lowerCAmelCase							,     lowerCAmelCase=5      )						->				str:
        '''simple docstring'''






        if top_k > self.model.config.num_labels:
              _lowercase			 =self.model.config.num_labels

        if self.framework == "pt":
              _lowercase			 =model_outputs.logits.softmax(-1      )[0]
              _lowercase , _lowercase			 =probs.topk(lowerCAmelCase      )
        elif self.framework == "tf":
              _lowercase			 =stable_softmax(model_outputs.logits							,     axis=-1      )[0]
              _lowercase			 =tf.math.top_k(lowerCAmelCase							,     k=lowerCAmelCase      )
              _lowercase , _lowercase			 =topk.values.numpy(), topk.indices.numpy()
        else:
              raise ValueError(F'''Unsupported framework: {self.framework}'''      )

        _lowercase			 =scores.tolist()
        _lowercase			 =ids.tolist()
        return [{"score": score, "label": self.model.config.idalabel[_id]} for score, _id in zip(lowerCAmelCase							,     lowerCAmelCase      )]





from typing import TYPE_CHECKING

from ...utils import (
    OptionalDependencyNotAvailable,
    _LazyModule,
    is_flax_available,
    is_torch_available,
    is_vision_available,
)


lowercase_             =     {'configuration_beit': ['BEIT_PRETRAINED_CONFIG_ARCHIVE_MAP', 'BeitConfig', 'BeitOnnxConfig']}

try:
  if not is_vision_available():
    raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
  pass
else:
  lowercase_             =     ['BeitFeatureExtractor']
  lowercase_             =     ['BeitImageProcessor']

try:
  if not is_torch_available():
    raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
  pass
else:
  lowercase_             =     [
      'BEIT_PRETRAINED_MODEL_ARCHIVE_LIST',
      'BeitForImageClassification',
      'BeitForMaskedImageModeling',
      'BeitForSemanticSegmentation',
      'BeitModel',
      'BeitPreTrainedModel',
  ]


try:
  if not is_flax_available():
    raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
  pass
else:
  lowercase_             =     [
      'FlaxBeitForImageClassification',
      'FlaxBeitForMaskedImageModeling',
      'FlaxBeitModel',
      'FlaxBeitPreTrainedModel',
  ]

if TYPE_CHECKING:
  from .configuration_beit import BEIT_PRETRAINED_CONFIG_ARCHIVE_MAP, BeitConfig, BeitOnnxConfig

  try:
    if not is_vision_available():
      raise OptionalDependencyNotAvailable()
  except OptionalDependencyNotAvailable:
    pass
  else:
    from .feature_extraction_beit import BeitFeatureExtractor
    from .image_processing_beit import BeitImageProcessor

  try:
    if not is_torch_available():
      raise OptionalDependencyNotAvailable()
  except OptionalDependencyNotAvailable:
    pass
  else:
    from .modeling_beit import (
        BEIT_PRETRAINED_MODEL_ARCHIVE_LIST,
        BeitForImageClassification,
        BeitForMaskedImageModeling,
        BeitForSemanticSegmentation,
        BeitModel,
        BeitPreTrainedModel,
    )

  try:
    if not is_flax_available():
      raise OptionalDependencyNotAvailable()
  except OptionalDependencyNotAvailable:
    pass
  else:
    from .modeling_flax_beit import (
        FlaxBeitForImageClassification,
        FlaxBeitForMaskedImageModeling,
        FlaxBeitModel,
        FlaxBeitPreTrainedModel,
    )


else:
  import sys

  lowercase_             =     _LazyModule(__name__, globals()['__file__'], _import_structure, module_spec=__spec__)


