


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 	snake_case_       (SCREAMING_SNAKE_CASE_				):
					def 			lowerCamelCase__(				self       :Any				) ->   Optional[int]:
												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 	snake_case_       :
					def __init__(				self       :Any      ,__snake_case       :List[str]      ,__snake_case       :Optional[int]=13      ,__snake_case       :int=32      ,__snake_case       :Tuple=2      ,__snake_case       :Tuple=3      ,__snake_case       :Any=6_40      ,__snake_case       :List[str]=4      ,__snake_case       :Any="silu"      ,__snake_case       :int=3      ,__snake_case       :List[str]=32      ,__snake_case       :Optional[int]=0.1      ,__snake_case       :Any=0.1      ,__snake_case       :List[str]=0.1      ,__snake_case       :List[Any]=0.02      ,__snake_case       :List[str]=True      ,__snake_case       :str=True      ,__snake_case       :Optional[int]=10      ,__snake_case       :Optional[Any]=None      ,) ->   Optional[Any]:
												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 			lowerCamelCase__(				self       :Optional[int]				) ->   Optional[int]:
												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 			lowerCamelCase__(				self       :List[str]				) ->   int:
												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 			lowerCamelCase__(				self       :str      ,__snake_case       :Union[str, Any]      ,__snake_case       :Optional[int]      ,__snake_case       :Optional[Any]      ,__snake_case       :List[str]				) ->   Union[str, Any]:
												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 			lowerCamelCase__(				self       :Optional[Any]      ,__snake_case       :str      ,__snake_case       :Dict      ,__snake_case       :Optional[Any]      ,__snake_case       :Any				) ->   int:
												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 			lowerCamelCase__(				self       :Dict      ,__snake_case       :List[Any]      ,__snake_case       :List[str]      ,__snake_case       :str      ,__snake_case       :Optional[int]				) ->   Dict:
												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 			lowerCamelCase__(				self       :Any				) ->   Any:
												a__					=					self.prepare_config_and_inputs()
												a__					=					config_and_inputs
												a__					=					{'pixel_values': pixel_values}
												return config, inputs_dict







@require_torch
class 	snake_case_       (SCREAMING_SNAKE_CASE_     ,						SCREAMING_SNAKE_CASE_     ,						unittest.TestCase				):
					UpperCAmelCase__ :   List[Any]   				=      (
					    (MobileViTModel, MobileViTForImageClassification, MobileViTForSemanticSegmentation)
					    if is_torch_available()
					    else ()
					)
					UpperCAmelCase__ :   Any   				=      (
					    {
					        '''feature-extraction''': MobileViTModel,
					        '''image-classification''': MobileViTForImageClassification,
					        '''image-segmentation''': MobileViTForSemanticSegmentation,
					    }
					    if is_torch_available()
					    else {}
					)

					UpperCAmelCase__ :   str   				=      False
					UpperCAmelCase__ :   Tuple   				=      False
					UpperCAmelCase__ :   Optional[Any]   				=      False
					UpperCAmelCase__ :   List[str]   				=      False
					def 			lowerCamelCase__(				self       :Tuple				) ->   Optional[Any]:
												a__					=					MobileViTModelTester(self				)
												a__					=					MobileViTConfigTester(self      ,config_class=__a      ,has_text_modality=__a				)
					def 			lowerCamelCase__(				self       :Optional[Any]				) ->   List[str]:
												self.config_tester.run_common_tests()
					@unittest.skip(reason='MobileViT does not use inputs_embeds'				)
					def 			lowerCamelCase__(				self       :Any				) ->   Union[str, Any]:
												pass
					@unittest.skip(reason='MobileViT does not support input and output embeddings'				)
					def 			lowerCamelCase__(				self       :Tuple				) ->   List[Any]:
												pass
					@unittest.skip(reason='MobileViT does not output attentions'				)
					def 			lowerCamelCase__(				self       :int				) ->   Union[str, Any]:
												pass
					def 			lowerCamelCase__(				self       :int				) ->   Any:
												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 			lowerCamelCase__(				self       :Optional[int]				) ->   Any:
												pass
					def 			lowerCamelCase__(				self       :Any				) ->   Optional[Any]:
												a__					=					self.model_tester.prepare_config_and_inputs()
												self.model_tester.create_and_check_model(*__a				)
					def 			lowerCamelCase__(				self       :Union[str, Any]				) ->   Tuple:
												def check_hidden_states_output(__snake_case       :Any      ,__snake_case       :Tuple      ,__snake_case       :Optional[Any]				):
																			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__					=					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 			lowerCamelCase__(				self       :Optional[int]				) ->   str:
												a__					=					self.model_tester.prepare_config_and_inputs()
												self.model_tester.create_and_check_for_image_classification(*__a				)
					def 			lowerCamelCase__(				self       :int				) ->   str:
												a__					=					self.model_tester.prepare_config_and_inputs()
												self.model_tester.create_and_check_for_semantic_segmentation(*__a				)



					@slow
					def 			lowerCamelCase__(				self       :Union[str, Any]				) ->   Optional[int]:
												for model_name in MOBILEVIT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
																			a__					=					MobileViTModel.from_pretrained(__a				)
																			self.assertIsNotNone(__a				)






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







@require_torch
@require_vision
class 	snake_case_       (unittest.TestCase				):
					@cached_property
					def 			lowerCamelCase__(				self       :Union[str, Any]				) ->   List[str]:
												return MobileViTImageProcessor.from_pretrained('apple/mobilevit-xx-small'				) if is_vision_available() else None
					@slow
					def 			lowerCamelCase__(				self       :int				) ->   Tuple:
												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, 10_00)				)
												self.assertEqual(outputs.logits.shape      ,__a				)

												a__					=					torch.tensor([-1.93_64, -1.23_27, -0.46_53]				).to(__a				)

												self.assertTrue(torch.allclose(outputs.logits[0, :3]      ,__a      ,atol=1E-4				)				)
					@slow
					def 			lowerCamelCase__(				self       :Any				) ->   Optional[int]:
												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.97_13, 6.97_86, 7.24_22], [7.28_93, 7.28_25, 7.44_46], [7.65_80, 7.87_97, 7.94_20]],
												        [[-10.68_69, -10.32_50, -10.34_71], [-10.42_28, -9.98_68, -9.71_32], [-11.04_05, -11.02_21, -10.73_18]],
												        [[-3.30_89, -2.85_39, -2.67_40], [-3.27_06, -2.56_21, -2.51_08], [-3.25_34, -2.66_15, -2.66_51]],
												    ]      ,device=__a      ,)

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



					@slow
					def 			lowerCamelCase__(				self       :List[str]				) ->   Any:
												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				)








from typing import Dict, List, Optional, Union

import numpy as np

from ...image_processing_utils import BaseImageProcessor, BatchFeature, get_size_dict
from ...image_transforms import 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_vision_available, logging


if is_vision_available():
						import PIL


snake_case							:	str					      =				logging.get_logger(__name__)






def 							__lowercase       (			__lowerCAmelCase							:    Union[str, Any] ,   __lowerCAmelCase							:    Optional[Any]       ):
							a__					=					b.T
							a__					=					np.sum(np.square(__lowerCAmelCase       ) ,   axis=1       )
							a__					=					np.sum(np.square(__lowerCAmelCase       ) ,   axis=0       )
							a__					=					np.matmul(__lowerCAmelCase ,   __lowerCAmelCase       )
							a__					=					aa[:, None] - 2 * ab + ba[None, :]
							return d






def 							__lowercase       (			__lowerCAmelCase							:    List[Any] ,   __lowerCAmelCase							:    str       ):
							a__					=					x.reshape(-1 ,   3       )
							a__					=					squared_euclidean_distance(__lowerCAmelCase ,   __lowerCAmelCase       )
							return np.argmin(__lowerCAmelCase ,   axis=1       )







class 	snake_case_       (lowerCamelCase_				):
					UpperCAmelCase__ :   Optional[Any]   				=      ['''pixel_values''']
					def __init__(				self       :Dict      ,__snake_case       :Optional[Union[List[List[int]], np.ndarray]] = None      ,__snake_case       :bool = True      ,__snake_case       :Dict[str, int] = None      ,__snake_case       :PILImageResampling = PILImageResampling.BILINEAR      ,__snake_case       :bool = True      ,__snake_case       :bool = True      ,**__snake_case       :Optional[int]      ,) ->   None:
												super().__init__(**__snake_case				)
												a__					=					size if size is not None else {'height': 2_56, 'width': 2_56}
												a__					=					get_size_dict(__snake_case				)
												a__					=					np.array(__snake_case				) if clusters is not None else None
												a__					=					do_resize
												a__					=					size
												a__					=					resample
												a__					=					do_normalize
												a__					=					do_color_quantize
					def 			lowerCamelCase__(				self       :Union[str, Any]      ,__snake_case       :np.ndarray      ,__snake_case       :Dict[str, int]      ,__snake_case       :PILImageResampling = PILImageResampling.BILINEAR      ,__snake_case       :Optional[Union[str, ChannelDimension]] = None      ,**__snake_case       :Any      ,) ->   np.ndarray:
												a__					=					get_size_dict(__snake_case				)
												if "height" not in size or "width" not in size:
																			raise ValueError(F'Size dictionary must contain both height and width keys. Got {size.keys()}'				)
												return resize(
												    __snake_case      ,size=(size['height'], size['width'])      ,resample=__snake_case      ,data_format=__snake_case      ,**__snake_case				)
					def 			lowerCamelCase__(				self       :List[Any]      ,__snake_case       :np.ndarray      ,__snake_case       :Optional[Union[str, ChannelDimension]] = None      ,) ->   np.ndarray:
												a__					=					rescale(image=__snake_case      ,scale=1 / 1_27.5      ,data_format=__snake_case				)
												a__					=					image - 1
												return image



					def 			lowerCamelCase__(				self       :Optional[Any]      ,__snake_case       :ImageInput      ,__snake_case       :bool = None      ,__snake_case       :Dict[str, int] = None      ,__snake_case       :PILImageResampling = None      ,__snake_case       :bool = None      ,__snake_case       :Optional[bool] = None      ,__snake_case       :Optional[Union[List[List[int]], np.ndarray]] = None      ,__snake_case       :Optional[Union[str, TensorType]] = None      ,__snake_case       :Optional[Union[str, ChannelDimension]] = ChannelDimension.FIRST      ,**__snake_case       :List[str]      ,) ->   PIL.Image.Image:
												a__					=					do_resize if do_resize is not None else self.do_resize
												a__					=					size if size is not None else self.size
												a__					=					get_size_dict(__snake_case				)
												a__					=					resample if resample is not None else self.resample
												a__					=					do_normalize if do_normalize is not None else self.do_normalize
												a__					=					do_color_quantize if do_color_quantize is not None else self.do_color_quantize
												a__					=					clusters if clusters is not None else self.clusters
												a__					=					np.array(__snake_case				)

												a__					=					make_list_of_images(__snake_case				)

												if not valid_images(__snake_case				):
																			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 or resample is None:
																			raise ValueError('Size and resample must be specified if do_resize is True.'				)

												if do_color_quantize and clusters is None:
																			raise ValueError('Clusters must be specified if do_color_quantize is True.'				)

												# All transformations expect numpy arrays.
												a__					=					[to_numpy_array(__snake_case				) for image in images]

												if do_resize:
																			a__					=					[self.resize(image=__snake_case      ,size=__snake_case      ,resample=__snake_case				) for image in images]

												if do_normalize:
																			a__					=					[self.normalize(image=__snake_case				) for image in images]

												if do_color_quantize:
																			a__					=					[to_channel_dimension_format(__snake_case      ,ChannelDimension.LAST				) for image in images]
																			# color quantize from (batch_size, height, width, 3) to (batch_size, height, width)
																			a__					=					np.array(__snake_case				)
																			a__					=					color_quantize(__snake_case      ,__snake_case				).reshape(images.shape[:-1]				)

																			# flatten to (batch_size, height*width)
																			a__					=					images.shape[0]
																			a__					=					images.reshape(__snake_case      ,-1				)

																			# We need to convert back to a list of images to keep consistent behaviour across processors.
																			a__					=					list(__snake_case				)
												else:
																			a__					=					[to_channel_dimension_format(__snake_case      ,__snake_case				) for image in images]

												a__					=					{'input_ids': images}
												return BatchFeature(data=__snake_case      ,tensor_type=__snake_case				)




