


import inspect
import unittest

from transformers import MobileViTConfig
from transformers.testing_utils import require_torch, 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, floats_tensor, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
  import torch

  from transformers import MobileViTForImageClassification, MobileViTForSemanticSegmentation, MobileViTModel
  from transformers.models.mobilevit.modeling_mobilevit import MOBILEVIT_PRETRAINED_MODEL_ARCHIVE_LIST


if is_vision_available():
  from PIL import Image

  from transformers import MobileViTImageProcessor


class     __UpperCAmelCase							(       __A							):


     """simple docstring"""


     def 	snake_case_  (	self    ):
        __a					=  self.config_class(**self.inputs_dict    )
        self.parent.assertTrue(hasattr(__A  ,       """hidden_sizes"""    )    )
        self.parent.assertTrue(hasattr(__A  ,       """neck_hidden_sizes"""    )    )
        self.parent.assertTrue(hasattr(__A  ,       """num_attention_heads"""    )    )







class     __UpperCAmelCase							:


     """simple docstring"""


     def __init__(	self  ,       __A  ,       __A=13  ,       __A=32  ,       __A=2  ,       __A=3  ,       __A=640  ,       __A=4  ,       __A="silu"  ,       __A=3  ,       __A=32  ,       __A=0.1  ,       __A=0.1  ,       __A=0.1  ,       __A=0.02  ,       __A=True  ,       __A=True  ,       __A=10  ,       __A=None  ,       ):
        __a					=  parent
        __a					=  batch_size
        __a					=  image_size
        __a					=  patch_size
        __a					=  num_channels
        __a					=  last_hidden_size
        __a					=  num_attention_heads
        __a					=  hidden_act
        __a					=  conv_kernel_size
        __a					=  output_stride
        __a					=  hidden_dropout_prob
        __a					=  attention_probs_dropout_prob
        __a					=  classifier_dropout_prob
        __a					=  use_labels
        __a					=  is_training
        __a					=  num_labels
        __a					=  initializer_range
        __a					=  scope


     def 	snake_case_  (	self    ):
        __a					=  floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]    )

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

        __a					=  self.get_config()

        return config, pixel_values, labels, pixel_labels


     def 	snake_case_  (	self    ):
        return MobileViTConfig(
            image_size=self.image_size  ,       patch_size=self.patch_size  ,       num_channels=self.num_channels  ,       num_attention_heads=self.num_attention_heads  ,       hidden_act=self.hidden_act  ,       conv_kernel_size=self.conv_kernel_size  ,       output_stride=self.output_stride  ,       hidden_dropout_prob=self.hidden_dropout_prob  ,       attention_probs_dropout_prob=self.attention_probs_dropout_prob  ,       classifier_dropout_prob=self.classifier_dropout_prob  ,       initializer_range=self.initializer_range  ,       )


     def 	snake_case_  (	self  ,       __A  ,       __A  ,       __A  ,       __A    ):
        __a					=  MobileViTModel(config=__A    )
        model.to(__A    )
        model.eval()
        __a					=  model(__A    )
        self.parent.assertEqual(
            result.last_hidden_state.shape  ,       (
                self.batch_size,
                self.last_hidden_size,
                self.image_size // self.output_stride,
                self.image_size // self.output_stride,
            )  ,       )


     def 	snake_case_  (	self  ,       __A  ,       __A  ,       __A  ,       __A    ):
        __a					=  self.num_labels
        __a					=  MobileViTForImageClassification(__A    )
        model.to(__A    )
        model.eval()
        __a					=  model(__A  ,       labels=__A    )
        self.parent.assertEqual(result.logits.shape  ,       (self.batch_size, self.num_labels)    )


     def 	snake_case_  (	self  ,       __A  ,       __A  ,       __A  ,       __A    ):
        __a					=  self.num_labels
        __a					=  MobileViTForSemanticSegmentation(__A    )
        model.to(__A    )
        model.eval()
        __a					=  model(__A    )
        self.parent.assertEqual(
            result.logits.shape  ,       (
                self.batch_size,
                self.num_labels,
                self.image_size // self.output_stride,
                self.image_size // self.output_stride,
            )  ,       )
        __a					=  model(__A  ,       labels=__A    )
        self.parent.assertEqual(
            result.logits.shape  ,       (
                self.batch_size,
                self.num_labels,
                self.image_size // self.output_stride,
                self.image_size // self.output_stride,
            )  ,       )


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







@require_torch
class     __UpperCAmelCase							(       __A ,					__A ,					unittest.TestCase							):


     """simple docstring"""
     _lowerCamelCase   =      (
         (MobileViTModel, MobileViTForImageClassification, MobileViTForSemanticSegmentation)
         if is_torch_available()
         else ()
     )
     _lowerCamelCase   =      (
         {
             """feature-extraction""": MobileViTModel,
             """image-classification""": MobileViTForImageClassification,
             """image-segmentation""": MobileViTForSemanticSegmentation,
         }
         if is_torch_available()
         else {}
     )

     _lowerCamelCase   =      False
     _lowerCamelCase   =      False
     _lowerCamelCase   =      False
     _lowerCamelCase   =      False


     def 	snake_case_  (	self    ):
        __a					=  MobileViTModelTester(self    )
        __a					=  MobileViTConfigTester(self  ,       config_class=__A  ,       has_text_modality=__A    )


     def 	snake_case_  (	self    ):
        self.config_tester.run_common_tests()


     @unittest.skip(reason="""MobileViT does not use inputs_embeds"""    )
     def 	snake_case_  (	self    ):
        pass


     @unittest.skip(reason="""MobileViT does not support input and output embeddings"""    )
     def 	snake_case_  (	self    ):
        pass


     @unittest.skip(reason="""MobileViT does not output attentions"""    )
     def 	snake_case_  (	self    ):
        pass


     def 	snake_case_  (	self    ):
        __a	,    __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.forward    )
           # signature.parameters is an OrderedDict => so arg_names order is deterministic
           __a					=  [*signature.parameters.keys()]

           __a					=  ["""pixel_values"""]
           self.assertListEqual(arg_names[:1]  ,       __A    )


     @unittest.skip("""Will be fixed soon by reducing the size of the model used for common tests."""    )
     def 	snake_case_  (	self    ):
        pass


     def 	snake_case_  (	self    ):
        __a					=  self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_model(*__A    )


     def 	snake_case_  (	self    ):
        def check_hidden_states_output(__A  ,       __A  ,       __A    ):
           __a					=  model_class(__A    )
           model.to(__A    )
           model.eval()

           with torch.no_grad():
              __a					=  model(**self._prepare_for_class(__A  ,       __A    )    )

           __a					=  outputs.hidden_states

           __a					=  5
           self.assertEqual(len(__A    )  ,       __A    )

           # MobileViT's feature maps are of shape (batch_size, num_channels, height, width)
           # with the width and height being successively divided by 2.
           __a					=  2
           for i in range(len(__A    )    ):
              self.assertListEqual(
                  list(hidden_states[i].shape[-2:]    )  ,       [self.model_tester.image_size // divisor, self.model_tester.image_size // divisor]  ,       )
              divisor *= 2

           self.assertEqual(self.model_tester.output_stride  ,       divisor // 2    )

        __a	,    __a					=  self.model_tester.prepare_config_and_inputs_for_common()

        for model_class in self.all_model_classes:
           __a					=  True
           check_hidden_states_output(__A  ,       __A  ,       __A    )

           # check that output_hidden_states also work using config
           del inputs_dict["output_hidden_states"]
           __a					=  True

           check_hidden_states_output(__A  ,       __A  ,       __A    )


     def 	snake_case_  (	self    ):
        __a					=  self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_for_image_classification(*__A    )


     def 	snake_case_  (	self    ):
        __a					=  self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_for_semantic_segmentation(*__A    )


     @slow
     def 	snake_case_  (	self    ):
        for model_name in MOBILEVIT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
           __a					=  MobileViTModel.from_pretrained(__A    )
           self.assertIsNotNone(__A    )







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


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


     """simple docstring"""


     @cached_property
     def 	snake_case_  (	self    ):
        return MobileViTImageProcessor.from_pretrained("""apple/mobilevit-xx-small"""    ) if is_vision_available() else None


     @slow
     def 	snake_case_  (	self    ):
        __a					=  MobileViTForImageClassification.from_pretrained("""apple/mobilevit-xx-small"""    ).to(__A    )

        __a					=  self.default_image_processor
        __a					=  prepare_img()
        __a					=  image_processor(images=__A  ,       return_tensors="""pt"""    ).to(__A    )

        # forward pass
        with torch.no_grad():
           __a					=  model(**__A    )

        # verify the logits
        __a					=  torch.Size((1, 1000)    )
        self.assertEqual(outputs.logits.shape  ,       __A    )

        __a					=  torch.tensor([-1.9364, -1.2327, -0.4653]    ).to(__A    )

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


     @slow
     def 	snake_case_  (	self    ):
        __a					=  MobileViTForSemanticSegmentation.from_pretrained("""apple/deeplabv3-mobilevit-xx-small"""    )
        __a					=  model.to(__A    )

        __a					=  MobileViTImageProcessor.from_pretrained("""apple/deeplabv3-mobilevit-xx-small"""    )

        __a					=  prepare_img()
        __a					=  image_processor(images=__A  ,       return_tensors="""pt"""    ).to(__A    )

        # forward pass
        with torch.no_grad():
           __a					=  model(**__A    )
        __a					=  outputs.logits

        # verify the logits
        __a					=  torch.Size((1, 21, 32, 32)    )
        self.assertEqual(logits.shape  ,       __A    )

        __a					=  torch.tensor(
            [
                [[6.9713, 6.9786, 7.2422], [7.2893, 7.2825, 7.4446], [7.6580, 7.8797, 7.9420]],
                [[-10.6869, -10.3250, -10.3471], [-10.4228, -9.9868, -9.7132], [-11.0405, -11.0221, -10.7318]],
                [[-3.3089, -2.8539, -2.6740], [-3.2706, -2.5621, -2.5108], [-3.2534, -2.6615, -2.6651]],
            ]  ,       device=__A  ,       )

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


     @slow
     def 	snake_case_  (	self    ):
        __a					=  MobileViTForSemanticSegmentation.from_pretrained("""apple/deeplabv3-mobilevit-xx-small"""    )
        __a					=  model.to(__A    )

        __a					=  MobileViTImageProcessor.from_pretrained("""apple/deeplabv3-mobilevit-xx-small"""    )

        __a					=  prepare_img()
        __a					=  image_processor(images=__A  ,       return_tensors="""pt"""    ).to(__A    )

        # forward pass
        with torch.no_grad():
           __a					=  model(**__A    )

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

        __a					=  image_processor.post_process_semantic_segmentation(outputs=__A  ,       target_sizes=[(50, 60)]    )
        __a					=  torch.Size((50, 60)    )
        self.assertEqual(segmentation[0].shape  ,       __A    )

        __a					=  image_processor.post_process_semantic_segmentation(outputs=__A    )
        __a					=  torch.Size((32, 32)    )
        self.assertEqual(segmentation[0].shape  ,       __A    )







class     __UpperCAmelCase							:


     """simple docstring"""


     def __init__(	self  ,       __A    ):
        __a					=  set_counts
        __a					=  max(__A    )
        __a					=  len(__A    )
        __a					=  [1] * num_sets
        __a					=  list(range(__A    )    )


     def 	snake_case_  (	self  ,       __A  ,       __A    ):
        __a					=  self.get_parent(__A    )
        __a					=  self.get_parent(__A    )

        if src_parent == dst_parent:
           return False

        if self.ranks[dst_parent] >= self.ranks[src_parent]:
           self.set_counts[dst_parent] += self.set_counts[src_parent]
           __a					=  0
           __a					=  dst_parent
           if self.ranks[dst_parent] == self.ranks[src_parent]:
              self.ranks[dst_parent] += 1
           __a					=  self.set_counts[dst_parent]
        else:
           self.set_counts[src_parent] += self.set_counts[dst_parent]
           __a					=  0
           __a					=  src_parent
           __a					=  self.set_counts[src_parent]

        __a					=  max(self.max_set  ,       __A    )
        return True


     def 	snake_case_  (	self  ,       __A    ):
        if self.parents[disj_set] == disj_set:
           return disj_set
        __a					=  self.get_parent(self.parents[disj_set]    )
        return self.parents[disj_set]



