





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_       (     UpperCamelCase__    ):


							def      _snake_case  (							self    :Dict		)   ->   int:



									"""simple docstring"""



									SCREAMING_SNAKE_CASE__			=      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_       :


							def __init__(							self    :List[str]							,				__A    :int							,				__A    :List[Any]=13							,				__A    :Tuple=32							,				__A    :Optional[Any]=2							,				__A    :Optional[Any]=3							,				__A    :Tuple=640							,				__A    :Optional[int]=4							,				__A    :Optional[int]="silu"							,				__A    :Any=3							,				__A    :List[str]=32							,				__A    :List[Any]=0.1							,				__A    :Any=0.1							,				__A    :Union[str, Any]=0.1							,				__A    :str=0.0_2							,				__A    :List[Any]=True							,				__A    :Tuple=True							,				__A    :Optional[int]=10							,				__A    :Dict=None							,				)   ->   str:



									"""simple docstring"""



									SCREAMING_SNAKE_CASE__			=      parent
									SCREAMING_SNAKE_CASE__			=      batch_size
									SCREAMING_SNAKE_CASE__			=      image_size
									SCREAMING_SNAKE_CASE__			=      patch_size
									SCREAMING_SNAKE_CASE__			=      num_channels
									SCREAMING_SNAKE_CASE__			=      last_hidden_size
									SCREAMING_SNAKE_CASE__			=      num_attention_heads
									SCREAMING_SNAKE_CASE__			=      hidden_act
									SCREAMING_SNAKE_CASE__			=      conv_kernel_size
									SCREAMING_SNAKE_CASE__			=      output_stride
									SCREAMING_SNAKE_CASE__			=      hidden_dropout_prob
									SCREAMING_SNAKE_CASE__			=      attention_probs_dropout_prob
									SCREAMING_SNAKE_CASE__			=      classifier_dropout_prob
									SCREAMING_SNAKE_CASE__			=      use_labels
									SCREAMING_SNAKE_CASE__			=      is_training
									SCREAMING_SNAKE_CASE__			=      num_labels
									SCREAMING_SNAKE_CASE__			=      initializer_range
									SCREAMING_SNAKE_CASE__			=      scope


							def      _snake_case  (							self    :str		)   ->   int:



									"""simple docstring"""



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

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

									SCREAMING_SNAKE_CASE__			=      self.get_config()

									return config, pixel_values, labels, pixel_labels


							def      _snake_case  (							self    :Tuple		)   ->   int:



									"""simple docstring"""



									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    :int							,				__A    :int							,				__A    :List[Any]							,				__A    :Optional[Any]							,				__A    :Dict		)   ->   Optional[int]:



									"""simple docstring"""



									SCREAMING_SNAKE_CASE__			=      MobileViTModel(config=__A		)
									model.to(__A		)
									model.eval()
									SCREAMING_SNAKE_CASE__			=      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    :int							,				__A    :str							,				__A    :str							,				__A    :List[str]							,				__A    :Any		)   ->   str:



									"""simple docstring"""



									SCREAMING_SNAKE_CASE__			=      self.num_labels
									SCREAMING_SNAKE_CASE__			=      MobileViTForImageClassification(__A		)
									model.to(__A		)
									model.eval()
									SCREAMING_SNAKE_CASE__			=      model(__A							,				labels=__A		)
									self.parent.assertEqual(result.logits.shape							,				(self.batch_size, self.num_labels)		)


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



									"""simple docstring"""



									SCREAMING_SNAKE_CASE__			=      self.num_labels
									SCREAMING_SNAKE_CASE__			=      MobileViTForSemanticSegmentation(__A		)
									model.to(__A		)
									model.eval()
									SCREAMING_SNAKE_CASE__			=      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,
									    )							,				)
									SCREAMING_SNAKE_CASE__			=      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    :Dict		)   ->   List[Any]:



									"""simple docstring"""



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




@require_torch
class   UpperCamelCase_       (     UpperCamelCase__				,  UpperCamelCase__				,  unittest.TestCase    ):
							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    :str		)   ->   Union[str, Any]:



									"""simple docstring"""



									SCREAMING_SNAKE_CASE__			=      MobileViTModelTester(self		)
									SCREAMING_SNAKE_CASE__			=      MobileViTConfigTester(self							,				config_class=__A							,				has_text_modality=__A		)


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



									"""simple docstring"""



									self.config_tester.run_common_tests()


							@unittest.skip(reason="""MobileViT does not use inputs_embeds"""		)
							def      _snake_case  (							self    :Tuple		)   ->   str:



									"""simple docstring"""



									pass


							@unittest.skip(reason="""MobileViT does not support input and output embeddings"""		)
							def      _snake_case  (							self    :Dict		)   ->   Dict:



									"""simple docstring"""



									pass


							@unittest.skip(reason="""MobileViT does not output attentions"""		)
							def      _snake_case  (							self    :Dict		)   ->   Optional[int]:



									"""simple docstring"""



									pass


							def      _snake_case  (							self    :Optional[Any]		)   ->   Any:



									"""simple docstring"""



									SCREAMING_SNAKE_CASE__	,						SCREAMING_SNAKE_CASE__			=      self.model_tester.prepare_config_and_inputs_for_common()

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

											SCREAMING_SNAKE_CASE__			=      ["""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    :str		)   ->   Any:



									"""simple docstring"""



									pass


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



									"""simple docstring"""



									SCREAMING_SNAKE_CASE__			=      self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_model(*__A		)


							def      _snake_case  (							self    :Any		)   ->   str:



									"""simple docstring"""



									def check_hidden_states_output(__A    :Union[str, Any]							,				__A    :Any							,				__A    :Any		):
											SCREAMING_SNAKE_CASE__			=      model_class(__A		)
											model.to(__A		)
											model.eval()

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

											SCREAMING_SNAKE_CASE__			=      outputs.hidden_states

											SCREAMING_SNAKE_CASE__			=      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.
											SCREAMING_SNAKE_CASE__			=      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		)

									SCREAMING_SNAKE_CASE__	,						SCREAMING_SNAKE_CASE__			=      self.model_tester.prepare_config_and_inputs_for_common()

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

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

											check_hidden_states_output(__A							,				__A							,				__A		)


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



									"""simple docstring"""



									SCREAMING_SNAKE_CASE__			=      self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_for_image_classification(*__A		)


							def      _snake_case  (							self    :Any		)   ->   Optional[Any]:



									"""simple docstring"""



									SCREAMING_SNAKE_CASE__			=      self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_for_semantic_segmentation(*__A		)


							@slow
							def      _snake_case  (							self    :List[Any]		)   ->   Union[str, Any]:



									"""simple docstring"""



									for model_name in MOBILEVIT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
											SCREAMING_SNAKE_CASE__			=      MobileViTModel.from_pretrained(__A		)
											self.assertIsNotNone(__A		)






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




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


							@cached_property
							def      _snake_case  (							self    :Optional[int]		)   ->   Union[str, Any]:



									"""simple docstring"""



									return MobileViTImageProcessor.from_pretrained("""apple/mobilevit-xx-small"""		) if is_vision_available() else None


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



									"""simple docstring"""



									SCREAMING_SNAKE_CASE__			=      MobileViTForImageClassification.from_pretrained("""apple/mobilevit-xx-small"""		).to(__A		)

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

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

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

									SCREAMING_SNAKE_CASE__			=      torch.tensor([-1.9_3_6_4, -1.2_3_2_7, -0.4_6_5_3]		).to(__A		)

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


							@slow
							def      _snake_case  (							self    :int		)   ->   str:



									"""simple docstring"""



									SCREAMING_SNAKE_CASE__			=      MobileViTForSemanticSegmentation.from_pretrained("""apple/deeplabv3-mobilevit-xx-small"""		)
									SCREAMING_SNAKE_CASE__			=      model.to(__A		)

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

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

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

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

									SCREAMING_SNAKE_CASE__			=      torch.tensor(
									    [
									        [[6.9_7_1_3, 6.9_7_8_6, 7.2_4_2_2], [7.2_8_9_3, 7.2_8_2_5, 7.4_4_4_6], [7.6_5_8_0, 7.8_7_9_7, 7.9_4_2_0]],
									        [[-1_0.6_8_6_9, -1_0.3_2_5_0, -1_0.3_4_7_1], [-1_0.4_2_2_8, -9.9_8_6_8, -9.7_1_3_2], [-1_1.0_4_0_5, -1_1.0_2_2_1, -1_0.7_3_1_8]],
									        [[-3.3_0_8_9, -2.8_5_3_9, -2.6_7_4_0], [-3.2_7_0_6, -2.5_6_2_1, -2.5_1_0_8], [-3.2_5_3_4, -2.6_6_1_5, -2.6_6_5_1]],
									    ]							,				device=__A							,				)

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


							@slow
							def      _snake_case  (							self    :List[Any]		)   ->   Optional[Any]:



									"""simple docstring"""



									SCREAMING_SNAKE_CASE__			=      MobileViTForSemanticSegmentation.from_pretrained("""apple/deeplabv3-mobilevit-xx-small"""		)
									SCREAMING_SNAKE_CASE__			=      model.to(__A		)

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

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

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

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

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

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






import copy
from dataclasses import dataclass, field
from typing import ClassVar, Dict

from ..features import ClassLabel, Features, Image
from .base import TaskTemplate




@dataclass(frozen=UpperCamelCase__    )
class   UpperCamelCase_       (     UpperCamelCase__    ):
							lowerCamelCase_           =					field(default="image-classification"				,  metadata={"include_in_asdict_even_if_is_default": True}    )
							lowerCamelCase_           =					Features({"image": Image()}    )
							lowerCamelCase_           =					Features({"labels": ClassLabel}    )
							lowerCamelCase_           =					"image"
							lowerCamelCase_           =					"labels"


							def      _snake_case  (							self    :List[str]							,				__A    :Tuple		)   ->   Tuple:



									"""simple docstring"""



									if self.label_column not in features:
											raise ValueError(f'''Column {self.label_column} is not present in features.'''		)
									if not isinstance(features[self.label_column]							,				__A		):
											raise ValueError(f'''Column {self.label_column} is not a ClassLabel.'''		)
									SCREAMING_SNAKE_CASE__			=      copy.deepcopy(self		)
									SCREAMING_SNAKE_CASE__			=      self.label_schema.copy()
									SCREAMING_SNAKE_CASE__			=      features[self.label_column]
									SCREAMING_SNAKE_CASE__			=      label_schema
									return task_template


							@property
							def      _snake_case  (							self    :Dict		)   ->   Dict[str, str]:



									"""simple docstring"""



									return {
									    self.image_column: "image",
									    self.label_column: "labels",
									}