



"""simple docstring"""





import inspect
import unittest

from datasets import load_dataset
from packaging import version

from transformers import BeitConfig
from transformers.models.auto import get_values
from transformers.testing_utils import require_torch, require_torch_multi_gpu, require_vision, slow, torch_device
from transformers.utils import cached_property, is_torch_available, is_vision_available

from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, _config_zero_init, floats_tensor, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
  import torch
  from torch import nn

  from transformers import (
      MODEL_MAPPING,
      BeitForImageClassification,
      BeitForMaskedImageModeling,
      BeitForSemanticSegmentation,
      BeitModel,
  )
  from transformers.models.beit.modeling_beit import BEIT_PRETRAINED_MODEL_ARCHIVE_LIST


if is_vision_available():
  import PIL
  from PIL import Image

  from transformers import BeitImageProcessor
class 				lowercase__       :







       '''simple docstring'''


       def __init__(							self      :  Optional[Any]		,    _UpperCAmelCase      :  str		,    _UpperCAmelCase      :  str=100		,    _UpperCAmelCase      :  Optional[Any]=13		,    _UpperCAmelCase      :  List[str]=30		,    _UpperCAmelCase      :  List[str]=2		,    _UpperCAmelCase      :  Union[str, Any]=3		,    _UpperCAmelCase      :  Dict=True		,    _UpperCAmelCase      :  Dict=True		,    _UpperCAmelCase      :  Optional[int]=32		,    _UpperCAmelCase      :  List[str]=4		,    _UpperCAmelCase      :  Dict=4		,    _UpperCAmelCase      :  List[str]=37		,    _UpperCAmelCase      :  Tuple="gelu"		,    _UpperCAmelCase      :  Optional[int]=0.1		,    _UpperCAmelCase      :  str=0.1		,    _UpperCAmelCase      :  Any=10		,    _UpperCAmelCase      :  Optional[Any]=0.02		,    _UpperCAmelCase      :  str=3		,    _UpperCAmelCase      :  int=None		,    _UpperCAmelCase      :  Optional[Any]=[0, 1, 2, 3]		,    )			->						Optional[Any]:


             '''simple docstring'''




             UpperCAmelCase_      			=							parent
             UpperCAmelCase_      			=							100
             UpperCAmelCase_      			=							batch_size
             UpperCAmelCase_      			=							image_size
             UpperCAmelCase_      			=							patch_size
             UpperCAmelCase_      			=							num_channels
             UpperCAmelCase_      			=							is_training
             UpperCAmelCase_      			=							use_labels
             UpperCAmelCase_      			=							hidden_size
             UpperCAmelCase_      			=							num_hidden_layers
             UpperCAmelCase_      			=							num_attention_heads
             UpperCAmelCase_      			=							intermediate_size
             UpperCAmelCase_      			=							hidden_act
             UpperCAmelCase_      			=							hidden_dropout_prob
             UpperCAmelCase_      			=							attention_probs_dropout_prob
             UpperCAmelCase_      			=							type_sequence_label_size
             UpperCAmelCase_      			=							initializer_range
             UpperCAmelCase_      			=							scope
             UpperCAmelCase_      			=							out_indices
             UpperCAmelCase_      			=							num_labels

             # in BeiT, the seq length equals the number of patches + 1 (we add 1 for the [CLS] token)
             UpperCAmelCase_      			=							(image_size // patch_size) ** 2
             UpperCAmelCase_      			=							num_patches + 1



       def 					lowercase__    (							self      :  Optional[int]   )			->						List[Any]:


             '''simple docstring'''




             UpperCAmelCase_      			=							floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]   )

             UpperCAmelCase_      			=							None
             UpperCAmelCase_      			=							None
             if self.use_labels:
                   UpperCAmelCase_      			=							ids_tensor([self.batch_size]		,    self.type_sequence_label_size   )
                   UpperCAmelCase_      			=							ids_tensor([self.batch_size, self.image_size, self.image_size]		,    self.num_labels   )

             UpperCAmelCase_      			=							self.get_config()

             return config, pixel_values, labels, pixel_labels



       def 					lowercase__    (							self      :  Optional[int]   )			->						str:


             '''simple docstring'''




             return BeitConfig(
                 vocab_size=self.vocab_size		,    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=_UpperCAmelCase		,    initializer_range=self.initializer_range		,    out_indices=self.out_indices		,    )



       def 					lowercase__    (							self      :  List[str]		,    _UpperCAmelCase      :  List[str]		,    _UpperCAmelCase      :  Any		,    _UpperCAmelCase      :  List[str]		,    _UpperCAmelCase      :  List[str]   )			->						int:


             '''simple docstring'''




             UpperCAmelCase_      			=							BeitModel(config=_UpperCAmelCase   )
             model.to(_UpperCAmelCase   )
             model.eval()
             UpperCAmelCase_      			=							model(_UpperCAmelCase   )
             self.parent.assertEqual(result.last_hidden_state.shape		,    (self.batch_size, self.seq_length, self.hidden_size)   )



       def 					lowercase__    (							self      :  List[Any]		,    _UpperCAmelCase      :  Union[str, Any]		,    _UpperCAmelCase      :  Union[str, Any]		,    _UpperCAmelCase      :  str		,    _UpperCAmelCase      :  List[str]   )			->						str:


             '''simple docstring'''




             UpperCAmelCase_      			=							BeitForMaskedImageModeling(config=_UpperCAmelCase   )
             model.to(_UpperCAmelCase   )
             model.eval()
             UpperCAmelCase_      			=							model(_UpperCAmelCase   )
             self.parent.assertEqual(result.logits.shape		,    (self.batch_size, self.seq_length - 1, self.vocab_size)   )



       def 					lowercase__    (							self      :  List[Any]		,    _UpperCAmelCase      :  Any		,    _UpperCAmelCase      :  Union[str, Any]		,    _UpperCAmelCase      :  Union[str, Any]		,    _UpperCAmelCase      :  str   )			->						int:


             '''simple docstring'''




             UpperCAmelCase_      			=							self.type_sequence_label_size
             UpperCAmelCase_      			=							BeitForImageClassification(_UpperCAmelCase   )
             model.to(_UpperCAmelCase   )
             model.eval()
             UpperCAmelCase_      			=							model(_UpperCAmelCase		,    labels=_UpperCAmelCase   )
             self.parent.assertEqual(result.logits.shape		,    (self.batch_size, self.type_sequence_label_size)   )

             # test greyscale images
             UpperCAmelCase_      			=							1
             UpperCAmelCase_      			=							BeitForImageClassification(_UpperCAmelCase   )
             model.to(_UpperCAmelCase   )
             model.eval()

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



       def 					lowercase__    (							self      :  str		,    _UpperCAmelCase      :  Dict		,    _UpperCAmelCase      :  str		,    _UpperCAmelCase      :  Union[str, Any]		,    _UpperCAmelCase      :  Any   )			->						Optional[int]:


             '''simple docstring'''




             UpperCAmelCase_      			=							self.num_labels
             UpperCAmelCase_      			=							BeitForSemanticSegmentation(_UpperCAmelCase   )
             model.to(_UpperCAmelCase   )
             model.eval()
             UpperCAmelCase_      			=							model(_UpperCAmelCase   )
             self.parent.assertEqual(
                 result.logits.shape		,    (self.batch_size, self.num_labels, self.image_size * 2, self.image_size * 2)   )
             UpperCAmelCase_      			=							model(_UpperCAmelCase		,    labels=_UpperCAmelCase   )
             self.parent.assertEqual(
                 result.logits.shape		,    (self.batch_size, self.num_labels, self.image_size * 2, self.image_size * 2)   )







       def 					lowercase__    (							self      :  Dict   )			->						Any:


             '''simple docstring'''




             UpperCAmelCase_      			=							self.prepare_config_and_inputs()
             UpperCAmelCase_			,   UpperCAmelCase_			,   UpperCAmelCase_			,   UpperCAmelCase_      			=							config_and_inputs
             UpperCAmelCase_      			=							{"pixel_values": pixel_values}
             return config, inputs_dict



@require_torch
class 				lowercase__       (   SCREAMING_SNAKE_CASE    , SCREAMING_SNAKE_CASE    , unittest.TestCase				):







       '''simple docstring'''


       UpperCamelCase						=					(
           (BeitModel, BeitForImageClassification, BeitForMaskedImageModeling, BeitForSemanticSegmentation)
           if is_torch_available()
           else ()
       )
       UpperCamelCase						=					(
           {
               '''feature-extraction''': BeitModel,
               '''image-classification''': BeitForImageClassification,
               '''image-segmentation''': BeitForSemanticSegmentation,
           }
           if is_torch_available()
           else {}
       )

       UpperCamelCase						=					False
       UpperCamelCase						=					False
       UpperCamelCase						=					False



       def 					lowercase__    (							self      :  List[Any]   )			->						List[Any]:


             '''simple docstring'''




             UpperCAmelCase_      			=							BeitModelTester(self   )
             UpperCAmelCase_      			=							ConfigTester(self		,    config_class=_UpperCAmelCase		,    has_text_modality=_UpperCAmelCase		,    hidden_size=37   )



       def 					lowercase__    (							self      :  Dict   )			->						Dict:


             '''simple docstring'''




             self.config_tester.run_common_tests()



       @unittest.skip(reason="BEiT does not use inputs_embeds"   )
       def 					lowercase__    (							self      :  str   )			->						List[Any]:


             '''simple docstring'''




             pass



       @require_torch_multi_gpu
       @unittest.skip(reason="BEiT has some layers using `add_module` which doesn't work well with `nn.DataParallel`"   )
       def 					lowercase__    (							self      :  Optional[int]   )			->						Any:


             '''simple docstring'''




             pass



       def 					lowercase__    (							self      :  str   )			->						Any:


             '''simple docstring'''




             UpperCAmelCase_			,   UpperCAmelCase_      			=							self.model_tester.prepare_config_and_inputs_for_common()

             for model_class in self.all_model_classes:
                   UpperCAmelCase_      			=							model_class(_UpperCAmelCase   )
                   self.assertIsInstance(model.get_input_embeddings()		,    (nn.Module)   )
                   UpperCAmelCase_      			=							model.get_output_embeddings()
                   self.assertTrue(x is None or isinstance(_UpperCAmelCase		,    nn.Linear   )   )



       def 					lowercase__    (							self      :  Optional[int]   )			->						Optional[int]:


             '''simple docstring'''




             UpperCAmelCase_			,   UpperCAmelCase_      			=							self.model_tester.prepare_config_and_inputs_for_common()

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

                   UpperCAmelCase_      			=							["pixel_values"]
                   self.assertListEqual(arg_names[:1]		,    _UpperCAmelCase   )



       def 					lowercase__    (							self      :  List[str]   )			->						Tuple:


             '''simple docstring'''




             UpperCAmelCase_      			=							self.model_tester.prepare_config_and_inputs()
             self.model_tester.create_and_check_model(*_UpperCAmelCase   )



       def 					lowercase__    (							self      :  Optional[int]   )			->						Dict:


             '''simple docstring'''




             UpperCAmelCase_      			=							self.model_tester.prepare_config_and_inputs()
             self.model_tester.create_and_check_for_masked_lm(*_UpperCAmelCase   )



       def 					lowercase__    (							self      :  List[Any]   )			->						List[Any]:


             '''simple docstring'''




             UpperCAmelCase_      			=							self.model_tester.prepare_config_and_inputs()
             self.model_tester.create_and_check_for_image_classification(*_UpperCAmelCase   )



       def 					lowercase__    (							self      :  Tuple   )			->						Tuple:


             '''simple docstring'''




             UpperCAmelCase_      			=							self.model_tester.prepare_config_and_inputs()
             self.model_tester.create_and_check_for_semantic_segmentation(*_UpperCAmelCase   )



       def 					lowercase__    (							self      :  Union[str, Any]   )			->						Tuple:


             '''simple docstring'''




             if not self.model_tester.is_training:
                   return

             UpperCAmelCase_			,   UpperCAmelCase_      			=							self.model_tester.prepare_config_and_inputs_for_common()
             UpperCAmelCase_      			=							True

             for model_class in self.all_model_classes:
                   # we don't test BeitForMaskedImageModeling
                   if model_class in [*get_values(_UpperCAmelCase   ), BeitForMaskedImageModeling]:
                         continue

                   UpperCAmelCase_      			=							model_class(_UpperCAmelCase   )
                   model.to(_UpperCAmelCase   )
                   model.train()
                   UpperCAmelCase_      			=							self._prepare_for_class(_UpperCAmelCase		,    _UpperCAmelCase		,    return_labels=_UpperCAmelCase   )
                   UpperCAmelCase_      			=							model(**_UpperCAmelCase   ).loss
                   loss.backward()



       def 					lowercase__    (							self      :  str   )			->						List[str]:


             '''simple docstring'''




             UpperCAmelCase_			,   UpperCAmelCase_      			=							self.model_tester.prepare_config_and_inputs_for_common()
             if not self.model_tester.is_training:
                   return

             UpperCAmelCase_      			=							False
             UpperCAmelCase_      			=							True

             for model_class in self.all_model_classes:
                   # we don't test BeitForMaskedImageModeling
                   if (
                       model_class in [*get_values(_UpperCAmelCase   ), BeitForMaskedImageModeling]
                       or not model_class.supports_gradient_checkpointing
                   ):
                         continue

                   UpperCAmelCase_      			=							model_class(_UpperCAmelCase   )
                   model.gradient_checkpointing_enable()
                   model.to(_UpperCAmelCase   )
                   model.train()
                   UpperCAmelCase_      			=							self._prepare_for_class(_UpperCAmelCase		,    _UpperCAmelCase		,    return_labels=_UpperCAmelCase   )
                   UpperCAmelCase_      			=							model(**_UpperCAmelCase   ).loss
                   loss.backward()



       def 					lowercase__    (							self      :  Tuple   )			->						Optional[Any]:


             '''simple docstring'''




             UpperCAmelCase_			,   UpperCAmelCase_      			=							self.model_tester.prepare_config_and_inputs_for_common()

             UpperCAmelCase_      			=							_config_zero_init(_UpperCAmelCase   )
             for model_class in self.all_model_classes:
                   UpperCAmelCase_      			=							model_class(config=_UpperCAmelCase   )
                   for name, param in model.named_parameters():
                         # we skip lambda parameters as these require special initial values
                         # determined by config.layer_scale_init_value
                         if "lambda" in name:
                               continue
                         if param.requires_grad:
                               self.assertIn(
                                   ((param.data.mean() * 1e9).round() / 1e9).item()		,    [0.0, 1.0]		,    msg=F"""Parameter {name} of model {model_class} seems not properly initialized"""		,    )







       @slow
       def 					lowercase__    (							self      :  List[Any]   )			->						List[str]:


             '''simple docstring'''




             for model_name in BEIT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
                   UpperCAmelCase_      			=							BeitModel.from_pretrained(_UpperCAmelCase   )
                   self.assertIsNotNone(_UpperCAmelCase   )



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


@require_torch
@require_vision
class 				lowercase__       (   unittest.TestCase				):







       '''simple docstring'''


       @cached_property
       def 					lowercase__    (							self      :  List[Any]   )			->						Tuple:


             '''simple docstring'''




             return BeitImageProcessor.from_pretrained("microsoft/beit-base-patch16-224"   ) if is_vision_available() else None



       @slow
       def 					lowercase__    (							self      :  Any   )			->						Dict:


             '''simple docstring'''




             UpperCAmelCase_      			=							BeitForMaskedImageModeling.from_pretrained("microsoft/beit-base-patch16-224-pt22k"   ).to(_UpperCAmelCase   )

             UpperCAmelCase_      			=							self.default_image_processor
             UpperCAmelCase_      			=							prepare_img()
             UpperCAmelCase_      			=							image_processor(images=_UpperCAmelCase		,    return_tensors="pt"   ).pixel_values.to(_UpperCAmelCase   )

             # prepare bool_masked_pos
             UpperCAmelCase_      			=							torch.ones((1, 196)		,    dtype=torch.bool   ).to(_UpperCAmelCase   )

             # forward pass
             with torch.no_grad():
                   UpperCAmelCase_      			=							model(pixel_values=_UpperCAmelCase		,    bool_masked_pos=_UpperCAmelCase   )
             UpperCAmelCase_      			=							outputs.logits

             # verify the logits
             UpperCAmelCase_      			=							torch.Size((1, 196, 8192)   )
             self.assertEqual(logits.shape		,    _UpperCAmelCase   )

             UpperCAmelCase_      			=							torch.tensor(
                 [[-3.2437, 0.5072, -13.9174], [-3.2456, 0.4948, -13.9401], [-3.2033, 0.5121, -13.8550]]   ).to(_UpperCAmelCase   )

             self.assertTrue(torch.allclose(logits[bool_masked_pos][:3, :3]		,    _UpperCAmelCase		,    atol=1e-2   )   )



       @slow
       def 					lowercase__    (							self      :  Dict   )			->						int:


             '''simple docstring'''




             UpperCAmelCase_      			=							BeitForImageClassification.from_pretrained("microsoft/beit-base-patch16-224"   ).to(_UpperCAmelCase   )

             UpperCAmelCase_      			=							self.default_image_processor
             UpperCAmelCase_      			=							prepare_img()
             UpperCAmelCase_      			=							image_processor(images=_UpperCAmelCase		,    return_tensors="pt"   ).to(_UpperCAmelCase   )

             # forward pass
             with torch.no_grad():
                   UpperCAmelCase_      			=							model(**_UpperCAmelCase   )
             UpperCAmelCase_      			=							outputs.logits

             # verify the logits
             UpperCAmelCase_      			=							torch.Size((1, 1000)   )
             self.assertEqual(logits.shape		,    _UpperCAmelCase   )

             UpperCAmelCase_      			=							torch.tensor([-1.2385, -1.0987, -1.0108]   ).to(_UpperCAmelCase   )

             self.assertTrue(torch.allclose(logits[0, :3]		,    _UpperCAmelCase		,    atol=1e-4   )   )

             UpperCAmelCase_      			=							281
             self.assertEqual(logits.argmax(-1   ).item()		,    _UpperCAmelCase   )



       @slow
       def 					lowercase__    (							self      :  Any   )			->						Optional[int]:


             '''simple docstring'''




             UpperCAmelCase_      			=							BeitForImageClassification.from_pretrained("microsoft/beit-large-patch16-224-pt22k-ft22k"   ).to(
                 _UpperCAmelCase   )

             UpperCAmelCase_      			=							self.default_image_processor
             UpperCAmelCase_      			=							prepare_img()
             UpperCAmelCase_      			=							image_processor(images=_UpperCAmelCase		,    return_tensors="pt"   ).to(_UpperCAmelCase   )

             # forward pass
             with torch.no_grad():
                   UpperCAmelCase_      			=							model(**_UpperCAmelCase   )
             UpperCAmelCase_      			=							outputs.logits

             # verify the logits
             UpperCAmelCase_      			=							torch.Size((1, 21841)   )
             self.assertEqual(logits.shape		,    _UpperCAmelCase   )

             UpperCAmelCase_      			=							torch.tensor([1.6881, -0.2787, 0.5901]   ).to(_UpperCAmelCase   )

             self.assertTrue(torch.allclose(logits[0, :3]		,    _UpperCAmelCase		,    atol=1e-4   )   )

             UpperCAmelCase_      			=							2396
             self.assertEqual(logits.argmax(-1   ).item()		,    _UpperCAmelCase   )



       @slow
       def 					lowercase__    (							self      :  Tuple   )			->						str:


             '''simple docstring'''




             UpperCAmelCase_      			=							BeitForSemanticSegmentation.from_pretrained("microsoft/beit-base-finetuned-ade-640-640"   )
             UpperCAmelCase_      			=							model.to(_UpperCAmelCase   )

             UpperCAmelCase_      			=							BeitImageProcessor(do_resize=_UpperCAmelCase		,    size=640		,    do_center_crop=_UpperCAmelCase   )

             UpperCAmelCase_      			=							load_dataset("hf-internal-testing/fixtures_ade20k"		,    split="test"   )
             UpperCAmelCase_      			=							Image.open(ds[0]["file"]   )
             UpperCAmelCase_      			=							image_processor(images=_UpperCAmelCase		,    return_tensors="pt"   ).to(_UpperCAmelCase   )

             # forward pass
             with torch.no_grad():
                   UpperCAmelCase_      			=							model(**_UpperCAmelCase   )
             UpperCAmelCase_      			=							outputs.logits

             # verify the logits
             UpperCAmelCase_      			=							torch.Size((1, 150, 160, 160)   )
             self.assertEqual(logits.shape		,    _UpperCAmelCase   )

             UpperCAmelCase_      			=							version.parse(PIL.__version__   ) < version.parse("9.0.0"   )

             if is_pillow_less_than_a:
                   UpperCAmelCase_      			=							torch.tensor(
                       [
                           [[-4.9225, -2.3954, -3.0522], [-2.8822, -1.0046, -1.7561], [-2.9549, -1.3228, -2.1347]],
                           [[-5.8168, -3.4129, -4.0778], [-3.8651, -2.2214, -3.0277], [-3.8356, -2.4643, -3.3535]],
                           [[-0.0078, 3.9952, 4.0754], [2.9856, 4.6944, 5.0035], [3.2413, 4.7813, 4.9969]],
                       ]		,    device=_UpperCAmelCase		,    )
             else:
                   UpperCAmelCase_      			=							torch.tensor(
                       [
                           [[-4.8960, -2.3688, -3.0355], [-2.8478, -0.9836, -1.7418], [-2.9449, -1.3332, -2.1456]],
                           [[-5.8081, -3.4124, -4.1006], [-3.8561, -2.2081, -3.0323], [-3.8365, -2.4601, -3.3669]],
                           [[-0.0309, 3.9868, 4.0540], [2.9640, 4.6877, 4.9976], [3.2081, 4.7690, 4.9942]],
                       ]		,    device=_UpperCAmelCase		,    )

             self.assertTrue(torch.allclose(logits[0, :3, :3, :3]		,    _UpperCAmelCase		,    atol=1e-4   )   )







       @slow
       def 					lowercase__    (							self      :  int   )			->						int:


             '''simple docstring'''




             UpperCAmelCase_      			=							BeitForSemanticSegmentation.from_pretrained("microsoft/beit-base-finetuned-ade-640-640"   )
             UpperCAmelCase_      			=							model.to(_UpperCAmelCase   )

             UpperCAmelCase_      			=							BeitImageProcessor(do_resize=_UpperCAmelCase		,    size=640		,    do_center_crop=_UpperCAmelCase   )

             UpperCAmelCase_      			=							load_dataset("hf-internal-testing/fixtures_ade20k"		,    split="test"   )
             UpperCAmelCase_      			=							Image.open(ds[0]["file"]   )
             UpperCAmelCase_      			=							image_processor(images=_UpperCAmelCase		,    return_tensors="pt"   ).to(_UpperCAmelCase   )

             # forward pass
             with torch.no_grad():
                   UpperCAmelCase_      			=							model(**_UpperCAmelCase   )

             UpperCAmelCase_      			=							outputs.logits.detach().cpu()

             UpperCAmelCase_      			=							image_processor.post_process_semantic_segmentation(outputs=_UpperCAmelCase		,    target_sizes=[(500, 300)]   )
             UpperCAmelCase_      			=							torch.Size((500, 300)   )
             self.assertEqual(segmentation[0].shape		,    _UpperCAmelCase   )

             UpperCAmelCase_      			=							image_processor.post_process_semantic_segmentation(outputs=_UpperCAmelCase   )
             UpperCAmelCase_      			=							torch.Size((160, 160)   )
             self.assertEqual(segmentation[0].shape		,    _UpperCAmelCase   )









"""simple docstring"""





import inspect
import unittest

import numpy as np

from tests.test_modeling_common import floats_tensor
from transformers import MaskaFormerConfig, is_torch_available, is_vision_available
from transformers.testing_utils import require_torch, require_torch_multi_gpu, require_vision, slow, torch_device
from transformers.utils import cached_property

from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
  import torch

  from transformers import MaskaFormerForUniversalSegmentation, MaskaFormerModel

  if is_vision_available():
    from transformers import MaskaFormerImageProcessor

if is_vision_available():
  from PIL import Image
class 				lowercase__       :







       '''simple docstring'''


       def __init__(							self      :  Any		,    _UpperCAmelCase      :  str		,    _UpperCAmelCase      :  Union[str, Any]=2		,    _UpperCAmelCase      :  Any=True		,    _UpperCAmelCase      :  int=False		,    _UpperCAmelCase      :  str=10		,    _UpperCAmelCase      :  int=3		,    _UpperCAmelCase      :  Optional[int]=32 * 8		,    _UpperCAmelCase      :  str=32 * 8		,    _UpperCAmelCase      :  Any=4		,    _UpperCAmelCase      :  List[Any]=64		,    )			->						str:


             '''simple docstring'''




             UpperCAmelCase_      			=							parent
             UpperCAmelCase_      			=							batch_size
             UpperCAmelCase_      			=							is_training
             UpperCAmelCase_      			=							use_auxiliary_loss
             UpperCAmelCase_      			=							num_queries
             UpperCAmelCase_      			=							num_channels
             UpperCAmelCase_      			=							min_size
             UpperCAmelCase_      			=							max_size
             UpperCAmelCase_      			=							num_labels
             UpperCAmelCase_      			=							hidden_dim
             UpperCAmelCase_      			=							hidden_dim



       def 					lowercase__    (							self      :  int   )			->						Optional[Any]:


             '''simple docstring'''




             UpperCAmelCase_      			=							floats_tensor([self.batch_size, self.num_channels, self.min_size, self.max_size]   ).to(
                 _UpperCAmelCase   )

             UpperCAmelCase_      			=							torch.ones([self.batch_size, self.min_size, self.max_size]		,    device=_UpperCAmelCase   )

             UpperCAmelCase_      			=							(
                 torch.rand([self.batch_size, self.num_labels, self.min_size, self.max_size]		,    device=_UpperCAmelCase   ) > 0.5
             ).float()
             UpperCAmelCase_      			=							(torch.rand((self.batch_size, self.num_labels)		,    device=_UpperCAmelCase   ) > 0.5).long()

             UpperCAmelCase_      			=							self.get_config()
             return config, pixel_values, pixel_mask, mask_labels, class_labels



       def 					lowercase__    (							self      :  Any   )			->						Tuple:


             '''simple docstring'''




             UpperCAmelCase_      			=							MaskaFormerConfig(
                 hidden_size=self.hidden_dim		,    )
             UpperCAmelCase_      			=							self.num_queries
             UpperCAmelCase_      			=							self.num_labels

             UpperCAmelCase_      			=							[1, 1, 1, 1]
             UpperCAmelCase_      			=							self.num_channels

             UpperCAmelCase_      			=							64
             UpperCAmelCase_      			=							128
             UpperCAmelCase_      			=							self.hidden_dim
             UpperCAmelCase_      			=							self.hidden_dim
             UpperCAmelCase_      			=							self.hidden_dim
             return config



       def 					lowercase__    (							self      :  Dict   )			->						Dict:


             '''simple docstring'''




             UpperCAmelCase_			,   UpperCAmelCase_			,   UpperCAmelCase_			,   UpperCAmelCase_			,   UpperCAmelCase_      			=							self.prepare_config_and_inputs()
             UpperCAmelCase_      			=							{"pixel_values": pixel_values, "pixel_mask": pixel_mask}
             return config, inputs_dict



       def 					lowercase__    (							self      :  Optional[Any]		,    _UpperCAmelCase      :  str		,    _UpperCAmelCase      :  Union[str, Any]   )			->						str:


             '''simple docstring'''




             UpperCAmelCase_      			=							output.encoder_hidden_states
             UpperCAmelCase_      			=							output.pixel_decoder_hidden_states
             UpperCAmelCase_      			=							output.transformer_decoder_hidden_states

             self.parent.assertTrue(len(_UpperCAmelCase   )		,    len(config.backbone_config.depths   )   )
             self.parent.assertTrue(len(_UpperCAmelCase   )		,    len(config.backbone_config.depths   )   )
             self.parent.assertTrue(len(_UpperCAmelCase   )		,    config.decoder_layers   )



       def 					lowercase__    (							self      :  Tuple		,    _UpperCAmelCase      :  List[Any]		,    _UpperCAmelCase      :  Any		,    _UpperCAmelCase      :  Union[str, Any]		,    _UpperCAmelCase      :  Optional[int]=False   )			->						str:


             '''simple docstring'''




             with torch.no_grad():
                   UpperCAmelCase_      			=							MaskaFormerModel(config=_UpperCAmelCase   )
                   model.to(_UpperCAmelCase   )
                   model.eval()

                   UpperCAmelCase_      			=							model(pixel_values=_UpperCAmelCase		,    pixel_mask=_UpperCAmelCase   )
                   UpperCAmelCase_      			=							model(_UpperCAmelCase		,    output_hidden_states=_UpperCAmelCase   )

             self.parent.assertEqual(
                 output.transformer_decoder_last_hidden_state.shape		,    (self.batch_size, self.num_queries, self.hidden_dim)		,    )
             # let's ensure the other two hidden state exists
             self.parent.assertTrue(output.pixel_decoder_last_hidden_state is not None   )
             self.parent.assertTrue(output.encoder_last_hidden_state is not None   )

             if output_hidden_states:
                   self.check_output_hidden_state(_UpperCAmelCase		,    _UpperCAmelCase   )







       def 					lowercase__    (							self      :  Optional[Any]		,    _UpperCAmelCase      :  Optional[int]		,    _UpperCAmelCase      :  Dict		,    _UpperCAmelCase      :  Optional[int]		,    _UpperCAmelCase      :  Union[str, Any]		,    _UpperCAmelCase      :  int   )			->						Optional[Any]:


             '''simple docstring'''




             UpperCAmelCase_      			=							MaskaFormerForUniversalSegmentation(config=_UpperCAmelCase   )
             model.to(_UpperCAmelCase   )
             model.eval()

             def comm_check_on_output(_UpperCAmelCase      :  List[Any]   ):
                   # let's still check that all the required stuff is there
                   self.parent.assertTrue(result.transformer_decoder_last_hidden_state is not None   )
                   self.parent.assertTrue(result.pixel_decoder_last_hidden_state is not None   )
                   self.parent.assertTrue(result.encoder_last_hidden_state is not None   )
                   # okay, now we need to check the logits shape
                   # due to the encoder compression, masks have a //4 spatial size
                   self.parent.assertEqual(
                       result.masks_queries_logits.shape		,    (self.batch_size, self.num_queries, self.min_size // 4, self.max_size // 4)		,    )
                   # + 1 for null class
                   self.parent.assertEqual(
                       result.class_queries_logits.shape		,    (self.batch_size, self.num_queries, self.num_labels + 1)   )

             with torch.no_grad():
                   UpperCAmelCase_      			=							model(pixel_values=_UpperCAmelCase		,    pixel_mask=_UpperCAmelCase   )
                   UpperCAmelCase_      			=							model(_UpperCAmelCase   )

                   comm_check_on_output(_UpperCAmelCase   )

                   UpperCAmelCase_      			=							model(
                       pixel_values=_UpperCAmelCase		,    pixel_mask=_UpperCAmelCase		,    mask_labels=_UpperCAmelCase		,    class_labels=_UpperCAmelCase   )

             comm_check_on_output(_UpperCAmelCase   )

             self.parent.assertTrue(result.loss is not None   )
             self.parent.assertEqual(result.loss.shape		,    torch.Size([1]   )   )



@require_torch
class 				lowercase__       (   SCREAMING_SNAKE_CASE    , SCREAMING_SNAKE_CASE    , unittest.TestCase				):







       '''simple docstring'''


       UpperCamelCase						=					(MaskaFormerModel, MaskaFormerForUniversalSegmentation) if is_torch_available() else ()
       UpperCamelCase						=					{'''feature-extraction''': MaskaFormerModel} if is_torch_available() else {}

       UpperCamelCase						=					False
       UpperCamelCase						=					False
       UpperCamelCase						=					False
       UpperCamelCase						=					False



       def 					lowercase__    (							self      :  Tuple   )			->						Tuple:


             '''simple docstring'''




             UpperCAmelCase_      			=							MaskaFormerModelTester(self   )
             UpperCAmelCase_      			=							ConfigTester(self		,    config_class=_UpperCAmelCase		,    has_text_modality=_UpperCAmelCase   )



       def 					lowercase__    (							self      :  List[str]   )			->						Dict:


             '''simple docstring'''




             self.config_tester.run_common_tests()



       def 					lowercase__    (							self      :  Dict   )			->						str:


             '''simple docstring'''




             UpperCAmelCase_			,   UpperCAmelCase_      			=							self.model_tester.prepare_config_and_inputs_for_common()
             self.model_tester.create_and_check_maskaformer_model(_UpperCAmelCase		,    **_UpperCAmelCase		,    output_hidden_states=_UpperCAmelCase   )



       def 					lowercase__    (							self      :  Optional[int]   )			->						Dict:


             '''simple docstring'''




             UpperCAmelCase_      			=							self.model_tester.prepare_config_and_inputs()
             self.model_tester.create_and_check_maskaformer_instance_segmentation_head_model(*_UpperCAmelCase   )



       @unittest.skip(reason="Mask2Former does not use inputs_embeds"   )
       def 					lowercase__    (							self      :  Tuple   )			->						Any:


             '''simple docstring'''




             pass



       @unittest.skip(reason="Mask2Former does not have a get_input_embeddings method"   )
       def 					lowercase__    (							self      :  Tuple   )			->						Optional[int]:


             '''simple docstring'''




             pass



       @unittest.skip(reason="Mask2Former is not a generative model"   )
       def 					lowercase__    (							self      :  str   )			->						Dict:


             '''simple docstring'''




             pass



       @unittest.skip(reason="Mask2Former does not use token embeddings"   )
       def 					lowercase__    (							self      :  Union[str, Any]   )			->						Union[str, Any]:


             '''simple docstring'''




             pass



       @require_torch_multi_gpu
       @unittest.skip(
           reason="Mask2Former has some layers using `add_module` which doesn't work well with `nn.DataParallel`"   )
       def 					lowercase__    (							self      :  List[Any]   )			->						Dict:


             '''simple docstring'''




             pass



       @unittest.skip("Will be fixed soon by reducing the size of the model used for common tests."   )
       def 					lowercase__    (							self      :  List[str]   )			->						List[str]:


             '''simple docstring'''




             pass



       def 					lowercase__    (							self      :  Tuple   )			->						str:


             '''simple docstring'''




             UpperCAmelCase_			,   UpperCAmelCase_      			=							self.model_tester.prepare_config_and_inputs_for_common()

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

                   UpperCAmelCase_      			=							["pixel_values"]
                   self.assertListEqual(arg_names[:1]		,    _UpperCAmelCase   )



       @slow
       def 					lowercase__    (							self      :  List[Any]   )			->						Dict:


             '''simple docstring'''




             for model_name in ["facebook/mask2former-swin-small-coco-instance"]:
                   UpperCAmelCase_      			=							MaskaFormerModel.from_pretrained(_UpperCAmelCase   )
                   self.assertIsNotNone(_UpperCAmelCase   )



       def 					lowercase__    (							self      :  List[str]   )			->						List[str]:


             '''simple docstring'''




             UpperCAmelCase_      			=							(self.model_tester.min_size,) * 2
             UpperCAmelCase_      			=							{
                 "pixel_values": torch.randn((2, 3, *size)		,    device=_UpperCAmelCase   ),
                 "mask_labels": torch.randn((2, 10, *size)		,    device=_UpperCAmelCase   ),
                 "class_labels": torch.zeros(2		,    10		,    device=_UpperCAmelCase   ).long(),
             }
             UpperCAmelCase_      			=							self.model_tester.get_config()

             UpperCAmelCase_      			=							MaskaFormerForUniversalSegmentation(_UpperCAmelCase   ).to(_UpperCAmelCase   )
             UpperCAmelCase_      			=							model(**_UpperCAmelCase   )
             self.assertTrue(outputs.loss is not None   )



       def 					lowercase__    (							self      :  Union[str, Any]   )			->						str:


             '''simple docstring'''




             UpperCAmelCase_			,   UpperCAmelCase_      			=							self.model_tester.prepare_config_and_inputs_for_common()
             self.model_tester.create_and_check_maskaformer_model(_UpperCAmelCase		,    **_UpperCAmelCase		,    output_hidden_states=_UpperCAmelCase   )



       def 					lowercase__    (							self      :  Union[str, Any]   )			->						List[str]:


             '''simple docstring'''




             UpperCAmelCase_			,   UpperCAmelCase_      			=							self.model_tester.prepare_config_and_inputs_for_common()

             for model_class in self.all_model_classes:
                   UpperCAmelCase_      			=							model_class(_UpperCAmelCase   ).to(_UpperCAmelCase   )
                   UpperCAmelCase_      			=							model(**_UpperCAmelCase		,    output_attentions=_UpperCAmelCase   )
                   self.assertTrue(outputs.attentions is not None   )



       def 					lowercase__    (							self      :  Union[str, Any]   )			->						List[str]:


             '''simple docstring'''




             if not self.model_tester.is_training:
                   return

             UpperCAmelCase_      			=							self.all_model_classes[1]
             UpperCAmelCase_			,   UpperCAmelCase_			,   UpperCAmelCase_			,   UpperCAmelCase_			,   UpperCAmelCase_      			=							self.model_tester.prepare_config_and_inputs()

             UpperCAmelCase_      			=							model_class(_UpperCAmelCase   )
             model.to(_UpperCAmelCase   )
             model.train()

             UpperCAmelCase_      			=							model(_UpperCAmelCase		,    mask_labels=_UpperCAmelCase		,    class_labels=_UpperCAmelCase   ).loss
             loss.backward()







       def 					lowercase__    (							self      :  str   )			->						Optional[int]:


             '''simple docstring'''




             UpperCAmelCase_      			=							self.all_model_classes[1]
             UpperCAmelCase_			,   UpperCAmelCase_			,   UpperCAmelCase_			,   UpperCAmelCase_			,   UpperCAmelCase_      			=							self.model_tester.prepare_config_and_inputs()
             UpperCAmelCase_      			=							True
             UpperCAmelCase_      			=							True

             UpperCAmelCase_      			=							model_class(_UpperCAmelCase   ).to(_UpperCAmelCase   )
             model.train()

             UpperCAmelCase_      			=							model(_UpperCAmelCase		,    mask_labels=_UpperCAmelCase		,    class_labels=_UpperCAmelCase   )

             UpperCAmelCase_      			=							outputs.encoder_hidden_states[0]
             encoder_hidden_states.retain_grad()

             UpperCAmelCase_      			=							outputs.pixel_decoder_hidden_states[0]
             pixel_decoder_hidden_states.retain_grad()

             UpperCAmelCase_      			=							outputs.transformer_decoder_hidden_states[0]
             transformer_decoder_hidden_states.retain_grad()

             UpperCAmelCase_      			=							outputs.attentions[0]
             attentions.retain_grad()

             outputs.loss.backward(retain_graph=_UpperCAmelCase   )

             self.assertIsNotNone(encoder_hidden_states.grad   )
             self.assertIsNotNone(pixel_decoder_hidden_states.grad   )
             self.assertIsNotNone(transformer_decoder_hidden_states.grad   )
             self.assertIsNotNone(attentions.grad   )



lowerCamelCase       	=						1e-4
def        a__     (				):
      UpperCAmelCase_      			=							Image.open("./tests/fixtures/tests_samples/COCO/000000039769.png"				)
      return image


@require_vision
@slow
class 				lowercase__       (   unittest.TestCase				):







       '''simple docstring'''


       @cached_property
       def 					lowercase__    (							self      :  Optional[Any]   )			->						int:


             '''simple docstring'''




             return "facebook/mask2former-swin-small-coco-instance"



       @cached_property
       def 					lowercase__    (							self      :  Dict   )			->						Optional[int]:


             '''simple docstring'''




             return MaskaFormerImageProcessor.from_pretrained(self.model_checkpoints   ) if is_vision_available() else None



       def 					lowercase__    (							self      :  Optional[Any]   )			->						Union[str, Any]:


             '''simple docstring'''




             UpperCAmelCase_      			=							MaskaFormerModel.from_pretrained(self.model_checkpoints   ).to(_UpperCAmelCase   )
             UpperCAmelCase_      			=							self.default_image_processor
             UpperCAmelCase_      			=							prepare_img()
             UpperCAmelCase_      			=							image_processor(_UpperCAmelCase		,    return_tensors="pt"   ).to(_UpperCAmelCase   )
             UpperCAmelCase_      			=							inputs["pixel_values"].shape
             # check size is divisible by 32
             self.assertTrue((inputs_shape[-1] % 32) == 0 and (inputs_shape[-2] % 32) == 0   )
             # check size
             self.assertEqual(_UpperCAmelCase		,    (1, 3, 384, 384)   )

             with torch.no_grad():
                   UpperCAmelCase_      			=							model(**_UpperCAmelCase   )

             UpperCAmelCase_      			=							torch.tensor(
                 [[-0.2790, -1.0717, -1.1668], [-0.5128, -0.3128, -0.4987], [-0.5832, 0.1971, -0.0197]]   ).to(_UpperCAmelCase   )
             self.assertTrue(
                 torch.allclose(
                     outputs.encoder_last_hidden_state[0, 0, :3, :3]		,    _UpperCAmelCase		,    atol=_UpperCAmelCase   )   )

             UpperCAmelCase_      			=							torch.tensor(
                 [[0.8973, 1.1847, 1.1776], [1.1934, 1.5040, 1.5128], [1.1153, 1.4486, 1.4951]]   ).to(_UpperCAmelCase   )
             self.assertTrue(
                 torch.allclose(
                     outputs.pixel_decoder_last_hidden_state[0, 0, :3, :3]		,    _UpperCAmelCase		,    atol=_UpperCAmelCase   )   )

             UpperCAmelCase_      			=							torch.tensor(
                 [[2.1152, 1.7000, -0.8603], [1.5808, 1.8004, -0.9353], [1.6043, 1.7495, -0.5999]]   ).to(_UpperCAmelCase   )
             self.assertTrue(
                 torch.allclose(
                     outputs.transformer_decoder_last_hidden_state[0, :3, :3]		,    _UpperCAmelCase		,    atol=_UpperCAmelCase   )   )



       def 					lowercase__    (							self      :  int   )			->						List[str]:


             '''simple docstring'''




             UpperCAmelCase_      			=							MaskaFormerForUniversalSegmentation.from_pretrained(self.model_checkpoints   ).to(_UpperCAmelCase   ).eval()
             UpperCAmelCase_      			=							self.default_image_processor
             UpperCAmelCase_      			=							prepare_img()
             UpperCAmelCase_      			=							image_processor(_UpperCAmelCase		,    return_tensors="pt"   ).to(_UpperCAmelCase   )
             UpperCAmelCase_      			=							inputs["pixel_values"].shape
             # check size is divisible by 32
             self.assertTrue((inputs_shape[-1] % 32) == 0 and (inputs_shape[-2] % 32) == 0   )
             # check size
             self.assertEqual(_UpperCAmelCase		,    (1, 3, 384, 384)   )

             with torch.no_grad():
                   UpperCAmelCase_      			=							model(**_UpperCAmelCase   )
             # masks_queries_logits
             UpperCAmelCase_      			=							outputs.masks_queries_logits
             self.assertEqual(
                 masks_queries_logits.shape		,    (1, model.config.num_queries, inputs_shape[-2] // 4, inputs_shape[-1] // 4)   )
             UpperCAmelCase_      			=							[
                 [-8.7839, -9.0056, -8.8121],
                 [-7.4104, -7.0313, -6.5401],
                 [-6.6105, -6.3427, -6.4675],
             ]
             UpperCAmelCase_      			=							torch.tensor(_UpperCAmelCase   ).to(_UpperCAmelCase   )
             self.assertTrue(torch.allclose(masks_queries_logits[0, 0, :3, :3]		,    _UpperCAmelCase		,    atol=_UpperCAmelCase   )   )
             # class_queries_logits
             UpperCAmelCase_      			=							outputs.class_queries_logits
             self.assertEqual(class_queries_logits.shape		,    (1, model.config.num_queries, model.config.num_labels + 1)   )
             UpperCAmelCase_      			=							torch.tensor(
                 [
                     [1.8324, -8.0835, -4.1922],
                     [0.8450, -9.0050, -3.6053],
                     [0.3045, -7.7293, -3.0275],
                 ]   ).to(_UpperCAmelCase   )
             self.assertTrue(torch.allclose(outputs.class_queries_logits[0, :3, :3]		,    _UpperCAmelCase		,    atol=_UpperCAmelCase   )   )







       def 					lowercase__    (							self      :  List[Any]   )			->						Union[str, Any]:


             '''simple docstring'''




             UpperCAmelCase_      			=							MaskaFormerForUniversalSegmentation.from_pretrained(self.model_checkpoints   ).to(_UpperCAmelCase   ).eval()
             UpperCAmelCase_      			=							self.default_image_processor

             UpperCAmelCase_      			=							image_processor(
                 [np.zeros((3, 800, 1333)   ), np.zeros((3, 800, 1333)   )]		,    segmentation_maps=[np.zeros((384, 384)   ).astype(np.floataa   ), np.zeros((384, 384)   ).astype(np.floataa   )]		,    return_tensors="pt"		,    )

             UpperCAmelCase_      			=							inputs["pixel_values"].to(_UpperCAmelCase   )
             UpperCAmelCase_      			=							[el.to(_UpperCAmelCase   ) for el in inputs["mask_labels"]]
             UpperCAmelCase_      			=							[el.to(_UpperCAmelCase   ) for el in inputs["class_labels"]]

             with torch.no_grad():
                   UpperCAmelCase_      			=							model(**_UpperCAmelCase   )

             self.assertTrue(outputs.loss is not None   )




