





'''simple docstring'''

from __future__ import annotations

import inspect
import unittest

import numpy as np

from transformers import DeiTConfig
from transformers.testing_utils import require_tf, require_vision, slow
from transformers.utils import cached_property, is_tf_available, is_vision_available

from ...test_configuration_common import ConfigTester
from ...test_modeling_tf_common import TFModelTesterMixin, floats_tensor, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_tf_available():
  import tensorflow as tf

  from transformers import (
      TFDeiTForImageClassification,
      TFDeiTForImageClassificationWithTeacher,
      TFDeiTForMaskedImageModeling,
      TFDeiTModel,
  )
  from transformers.models.deit.modeling_tf_deit import TF_DEIT_PRETRAINED_MODEL_ARCHIVE_LIST


if is_vision_available():
  from PIL import Image

  from transformers import DeiTImageProcessor






class    UpperCAmelCase	:
 def __init__(				self      :Union[str, Any]						,       lowercase_      :Union[str, Any]						,       lowercase_      :Union[str, Any]=13						,       lowercase_      :Any=30						,       lowercase_      :Optional[int]=2						,       lowercase_      :Tuple=3						,       lowercase_      :str=True						,       lowercase_      :str=True						,       lowercase_      :List[str]=32						,       lowercase_      :str=2						,       lowercase_      :Optional[int]=4						,       lowercase_      :str=37						,       lowercase_      :Union[str, Any]="gelu"						,       lowercase_      :Any=0.1						,       lowercase_      :Optional[Any]=0.1						,       lowercase_      :Optional[int]=10						,       lowercase_      :Tuple=0.0_2						,       lowercase_      :Optional[int]=3						,       lowercase_      :Tuple=None						,       lowercase_      :str=2						,       )->							Union[str, Any]:
     A__				=			parent
     A__				=			batch_size
     A__				=			image_size
     A__				=			patch_size
     A__				=			num_channels
     A__				=			is_training
     A__				=			use_labels
     A__				=			hidden_size
     A__				=			num_hidden_layers
     A__				=			num_attention_heads
     A__				=			intermediate_size
     A__				=			hidden_act
     A__				=			hidden_dropout_prob
     A__				=			attention_probs_dropout_prob
     A__				=			type_sequence_label_size
     A__				=			initializer_range
     A__				=			scope
     A__				=			encoder_stride

     # in DeiT, the seq length equals the number of patches + 2 (we add 2 for the [CLS] and distilation tokens)
     A__				=			(image_size // patch_size) ** 2
     A__				=			num_patches + 2
 def      UpperCAmelCase_							(				self      :str							)->							Any:
     A__				=			floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]							)

     A__				=			None
     if self.use_labels:
         A__				=			ids_tensor([self.batch_size]						,       self.type_sequence_label_size							)

     A__				=			self.get_config()

     return config, pixel_values, labels
 def      UpperCAmelCase_							(				self      :str							)->							List[Any]:
     return DeiTConfig(
         image_size=self.image_size						,       patch_size=self.patch_size						,       num_channels=self.num_channels						,       hidden_size=self.hidden_size						,       num_hidden_layers=self.num_hidden_layers						,       num_attention_heads=self.num_attention_heads						,       intermediate_size=self.intermediate_size						,       hidden_act=self.hidden_act						,       hidden_dropout_prob=self.hidden_dropout_prob						,       attention_probs_dropout_prob=self.attention_probs_dropout_prob						,       is_decoder=_a						,       initializer_range=self.initializer_range						,       encoder_stride=self.encoder_stride						,       )
 def      UpperCAmelCase_							(				self      :Dict						,       lowercase_      :Optional[int]						,       lowercase_      :Optional[Any]						,       lowercase_      :Optional[Any]							)->							List[str]:
     A__				=			TFDeiTModel(config=_a							)
     A__				=			model(_a							)
     self.parent.assertEqual(result.last_hidden_state.shape						,       (self.batch_size, self.seq_length, self.hidden_size)							)
 def      UpperCAmelCase_							(				self      :int						,       lowercase_      :int						,       lowercase_      :Any						,       lowercase_      :List[str]							)->							Union[str, Any]:
     A__				=			TFDeiTForMaskedImageModeling(config=_a							)
     A__				=			model(_a							)
     self.parent.assertEqual(
         result.reconstruction.shape						,       (self.batch_size, self.num_channels, self.image_size, self.image_size)							)

     # test greyscale images
     A__				=			1
     A__				=			TFDeiTForMaskedImageModeling(_a							)

     A__				=			floats_tensor([self.batch_size, 1, self.image_size, self.image_size]							)
     A__				=			model(_a							)
     self.parent.assertEqual(result.reconstruction.shape						,       (self.batch_size, 1, self.image_size, self.image_size)							)
 def      UpperCAmelCase_							(				self      :List[Any]						,       lowercase_      :Optional[Any]						,       lowercase_      :Optional[int]						,       lowercase_      :str							)->							Optional[Any]:
     A__				=			self.type_sequence_label_size
     A__				=			TFDeiTForImageClassification(_a							)
     A__				=			model(_a						,       labels=_a							)
     self.parent.assertEqual(result.logits.shape						,       (self.batch_size, self.type_sequence_label_size)							)

     # test greyscale images
     A__				=			1
     A__				=			TFDeiTForImageClassification(_a							)

     A__				=			floats_tensor([self.batch_size, 1, self.image_size, self.image_size]							)
     A__				=			model(_a						,       labels=_a							)
     self.parent.assertEqual(result.logits.shape						,       (self.batch_size, self.type_sequence_label_size)							)



 def      UpperCAmelCase_							(				self      :Dict							)->							Optional[Any]:
     A__				=			self.prepare_config_and_inputs()
     A__				=			config_and_inputs
     A__				=			{"pixel_values": pixel_values}
     return config, inputs_dict






@require_tf
class    UpperCAmelCase	(       a_			,       a_			,       unittest.TestCase				):
 __lowercase   		=				(
     (
         TFDeiTModel,
         TFDeiTForImageClassification,
         TFDeiTForImageClassificationWithTeacher,
         TFDeiTForMaskedImageModeling,
     )
     if is_tf_available()
     else ()
 )
 __lowercase   		=				(
     {
         """feature-extraction""": TFDeiTModel,
         """image-classification""": (TFDeiTForImageClassification, TFDeiTForImageClassificationWithTeacher),
     }
     if is_tf_available()
     else {}
 )

 __lowercase   		=				False
 __lowercase   		=				False
 __lowercase   		=				False
 __lowercase   		=				False
 def      UpperCAmelCase_							(				self      :int							)->							Any:
     A__				=			TFDeiTModelTester(self							)
     A__				=			ConfigTester(self						,       config_class=_a						,       has_text_modality=_a						,       hidden_size=37							)
 def      UpperCAmelCase_							(				self      :Optional[Any]							)->							Any:
     self.config_tester.run_common_tests()
 @unittest.skip(reason="DeiT does not use inputs_embeds"							)
 def      UpperCAmelCase_							(				self      :List[str]							)->							Optional[int]:
     pass
 def      UpperCAmelCase_							(				self      :str							)->							str:
     A__				=			self.model_tester.prepare_config_and_inputs_for_common()

     for model_class in self.all_model_classes:
         A__				=			model_class(_a							)
         self.assertIsInstance(model.get_input_embeddings()						,       (tf.keras.layers.Layer)							)
         A__				=			model.get_output_embeddings()
         self.assertTrue(x is None or isinstance(_a						,       tf.keras.layers.Dense							)							)
 def      UpperCAmelCase_							(				self      :int							)->							str:
     A__				=			self.model_tester.prepare_config_and_inputs_for_common()

     for model_class in self.all_model_classes:
         A__				=			model_class(_a							)
         A__				=			inspect.signature(model.call							)
         # signature.parameters is an OrderedDict => so arg_names order is deterministic
         A__				=			[*signature.parameters.keys()]

         A__				=			["pixel_values"]
         self.assertListEqual(arg_names[:1]						,       _a							)
 def      UpperCAmelCase_							(				self      :Optional[Any]							)->							int:
     A__				=			self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_model(*_a							)
 def      UpperCAmelCase_							(				self      :Dict							)->							int:
     A__				=			self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_for_masked_image_modeling(*_a							)
 def      UpperCAmelCase_							(				self      :List[str]							)->							str:
     A__				=			self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_for_image_classification(*_a							)
 def      UpperCAmelCase_							(				self      :Any						,       lowercase_      :Optional[Any]						,       lowercase_      :Dict						,       lowercase_      :Tuple=False							)->							Dict:
     A__				=			super()._prepare_for_class(_a						,       _a						,       return_labels=_a							)

     if return_labels:
         if "labels" in inputs_dict and "labels" not in inspect.signature(model_class.call							).parameters:
             del inputs_dict["labels"]

     return inputs_dict



 @slow
 def      UpperCAmelCase_							(				self      :Dict							)->							Any:
     for model_name in TF_DEIT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
         A__				=			TFDeiTModel.from_pretrained(_a							)
         self.assertIsNotNone(_a							)





def 						UpperCamelCase				(							):
    A__				=			Image.open("./tests/fixtures/tests_samples/COCO/000000039769.png"       )
    return image






@require_tf
@require_vision
class    UpperCAmelCase	(       unittest.TestCase				):
 @cached_property
 def      UpperCAmelCase_							(				self      :Any							)->							Dict:
     return (
         DeiTImageProcessor.from_pretrained("facebook/deit-base-distilled-patch16-224"							)
         if is_vision_available()
         else None
     )



 @slow
 def      UpperCAmelCase_							(				self      :int							)->							Optional[int]:
     A__				=			TFDeiTForImageClassificationWithTeacher.from_pretrained("facebook/deit-base-distilled-patch16-224"							)

     A__				=			self.default_image_processor
     A__				=			prepare_img()
     A__				=			image_processor(images=_a						,       return_tensors="tf"							)

     # forward pass
     A__				=			model(**_a							)

     # verify the logits
     A__				=			tf.TensorShape((1, 10_00)							)
     self.assertEqual(outputs.logits.shape						,       _a							)

     A__				=			tf.constant([-1.0_2_6_6, 0.1_9_1_2, -1.2_8_6_1]							)

     self.assertTrue(np.allclose(outputs.logits[0, :3]						,       _a						,       atol=1E-4							)							)





"""simple docstring"""




from typing import Dict, List, Optional, Tuple, Union

import numpy as np

from ...image_processing_utils import BaseImageProcessor, BatchFeature, get_size_dict
from ...image_transforms import (
    center_crop,
    flip_channel_order,
    get_resize_output_image_size,
    rescale,
    resize,
    to_channel_dimension_format,
)
from ...image_utils import (
    ChannelDimension,
    ImageInput,
    PILImageResampling,
    make_list_of_images,
    to_numpy_array,
    valid_images,
)
from ...utils import TensorType, is_torch_available, is_torch_tensor, is_vision_available, logging


if is_vision_available():
   import PIL

if is_torch_available():
   import torch


_A			:       List[Any]   		=       logging.get_logger(__name__)

class 		a__						(      a_						):
   __lowerCAmelCase            =				["""pixel_values"""]
   def __init__(			self	,    _a = True	,    _a = None	,    _a = PILImageResampling.BILINEAR	,    _a = True	,    _a = 1 / 255	,    _a = True	,    _a = None	,    _a = True	,    **_a	,    ):
          super().__init__(**_a				)
          lowercase    :				Optional[Any]       		=							size if size is not None else {"shortest_edge": 224}
          lowercase    :				List[Any]       		=							get_size_dict(_a	,    default_to_square=_a				)
          lowercase    :				str       		=							crop_size if crop_size is not None else {"height": 256, "width": 256}
          lowercase    :				List[str]       		=							get_size_dict(_a	,    param_name="crop_size"				)

          lowercase    :				int       		=							do_resize
          lowercase    :				Optional[int]       		=							size
          lowercase    :				str       		=							resample
          lowercase    :				List[Any]       		=							do_rescale
          lowercase    :				Union[str, Any]       		=							rescale_factor
          lowercase    :				Optional[int]       		=							do_center_crop
          lowercase    :				Union[str, Any]       		=							crop_size
          lowercase    :				Optional[Any]       		=							do_flip_channel_order
   def 			__magic_name__			(			self	,    _a	,    _a	,    _a = PIL.Image.BILINEAR	,    _a = None	,    **_a	,    ):
          lowercase    :				List[Any]       		=							get_size_dict(_a	,    default_to_square=_a				)
          if "shortest_edge" not in size:
                 raise ValueError(f"""The `size` dictionary must contain the key `shortest_edge`. Got {size.keys()}"""				)
          lowercase    :				Union[str, Any]       		=							get_resize_output_image_size(_a	,    size=size["shortest_edge"]	,    default_to_square=_a				)
          return resize(_a	,    size=_a	,    resample=_a	,    data_format=_a	,    **_a				)
   def 			__magic_name__			(			self	,    _a	,    _a	,    _a = None	,    **_a	,    ):
          lowercase    :				str       		=							get_size_dict(_a				)
          if "height" not in size or "width" not in size:
                 raise ValueError(f"""The `size` dictionary must contain the keys `height` and `width`. Got {size.keys()}"""				)
          return center_crop(_a	,    size=(size["height"], size["width"])	,    data_format=_a	,    **_a				)
   def 			__magic_name__			(			self	,    _a	,    _a	,    _a = None	,    **_a	,    ):
          return rescale(_a	,    scale=_a	,    data_format=_a	,    **_a				)
   def 			__magic_name__			(			self	,    _a	,    _a = None				):
          return flip_channel_order(_a	,    data_format=_a				)
   def 			__magic_name__			(			self	,    _a	,    _a = None	,    _a = None	,    _a = None	,    _a = None	,    _a = None	,    _a = None	,    _a = None	,    _a = None	,    _a = None	,    _a = ChannelDimension.FIRST	,    **_a	,    ):
          lowercase    :				Optional[Any]       		=							do_resize if do_resize is not None else self.do_resize
          lowercase    :				Tuple       		=							resample if resample is not None else self.resample
          lowercase    :				List[str]       		=							do_rescale if do_rescale is not None else self.do_rescale
          lowercase    :				List[Any]       		=							rescale_factor if rescale_factor is not None else self.rescale_factor
          lowercase    :				Union[str, Any]       		=							do_center_crop if do_center_crop is not None else self.do_center_crop
          lowercase    :				Optional[int]       		=							(
              do_flip_channel_order if do_flip_channel_order is not None else self.do_flip_channel_order
          )

          lowercase    :				str       		=							size if size is not None else self.size
          lowercase    :				Any       		=							get_size_dict(_a	,    default_to_square=_a				)
          lowercase    :				int       		=							crop_size if crop_size is not None else self.crop_size
          lowercase    :				Any       		=							get_size_dict(_a	,    param_name="crop_size"				)

          lowercase    :				int       		=							make_list_of_images(_a				)

          if not valid_images(_a				):
                 raise ValueError(
                     "Invalid image type. Must be of type PIL.Image.Image, numpy.ndarray, "
                     "torch.Tensor, tf.Tensor or jax.ndarray."				)

          if do_resize and size is None:
                 raise ValueError("Size must be specified if do_resize is True."				)

          if do_rescale and rescale_factor is None:
                 raise ValueError("Rescale factor must be specified if do_rescale is True."				)

          if do_center_crop and crop_size is None:
                 raise ValueError("Crop size must be specified if do_center_crop is True."				)

          # All transformations expect numpy arrays.
          lowercase    :				Any       		=							[to_numpy_array(_a				) for image in images]

          if do_resize:
                 lowercase    :				Optional[int]       		=							[self.resize(image=_a	,    size=_a	,    resample=_a				) for image in images]

          if do_center_crop:
                 lowercase    :				str       		=							[self.center_crop(image=_a	,    size=_a				) for image in images]

          if do_rescale:
                 lowercase    :				Union[str, Any]       		=							[self.rescale(image=_a	,    scale=_a				) for image in images]

          # the pretrained checkpoints assume images are BGR, not RGB
          if do_flip_channel_order:
                 lowercase    :				int       		=							[self.flip_channel_order(image=_a				) for image in images]

          lowercase    :				int       		=							[to_channel_dimension_format(_a	,    _a				) for image in images]

          lowercase    :				Optional[Any]       		=							{"pixel_values": images}
          return BatchFeature(data=_a	,    tensor_type=_a				)

   def 			__magic_name__			(			self	,    _a	,    _a = None				):
          lowercase    :				Optional[int]       		=							outputs.logits

          # Resize logits and compute semantic segmentation maps
          if target_sizes is not None:
                 if len(_a				) != len(_a				):
                        raise ValueError(
                            "Make sure that you pass in as many target sizes as the batch dimension of the logits"				)

                 if is_torch_tensor(_a				):
                        lowercase    :				Tuple       		=							target_sizes.numpy()

                 lowercase    :				List[Any]       		=							[]

                 for idx in range(len(_a				)				):
                        lowercase    :				Dict       		=							torch.nn.functional.interpolate(
                            logits[idx].unsqueeze(dim=0				)	,    size=target_sizes[idx]	,    mode="bilinear"	,    align_corners=_a				)
                        lowercase    :				Tuple       		=							resized_logits[0].argmax(dim=0				)
                        semantic_segmentation.append(_a				)
          else:
                 lowercase    :				str       		=							logits.argmax(dim=1				)
                 lowercase    :				Dict       		=							[semantic_segmentation[i] for i in range(semantic_segmentation.shape[0]				)]

          return semantic_segmentation


