


import collections
import inspect
import unittest

from transformers import FocalNetConfig
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_backbone_common import BackboneTesterMixin
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 (
		    FocalNetBackbone,
		    FocalNetForImageClassification,
		    FocalNetForMaskedImageModeling,
		    FocalNetModel,
		)
		from transformers.models.focalnet.modeling_focalnet import FOCALNET_PRETRAINED_MODEL_ARCHIVE_LIST

if is_vision_available():
		from PIL import Image

		from transformers import AutoImageProcessor



class 						_UpperCAmelCase      :






		def __init__(						self				:     List[str]   ,  UpperCAmelCase				:     Optional[int]   ,  UpperCAmelCase				:     Optional[Any]=13   ,  UpperCAmelCase				:     List[str]=32   ,  UpperCAmelCase				:     str=2   ,  UpperCAmelCase				:     Dict=3   ,  UpperCAmelCase				:     Dict=16   ,  UpperCAmelCase				:     str=[32, 64, 1_28]   ,  UpperCAmelCase				:     Optional[Any]=[1, 2, 1]   ,  UpperCAmelCase				:     Any=[2, 2, 4]   ,  UpperCAmelCase				:     str=2   ,  UpperCAmelCase				:     Optional[int]=2.0   ,  UpperCAmelCase				:     Optional[int]=True   ,  UpperCAmelCase				:     Dict=0.0   ,  UpperCAmelCase				:     Dict=0.0   ,  UpperCAmelCase				:     Union[str, Any]=0.1   ,  UpperCAmelCase				:     int="gelu"   ,  UpperCAmelCase				:     Union[str, Any]=False   ,  UpperCAmelCase				:     Dict=True   ,  UpperCAmelCase				:     Optional[int]=0.02   ,  UpperCAmelCase				:     str=1E-5   ,  UpperCAmelCase				:     Tuple=True   ,  UpperCAmelCase				:     str=None   ,  UpperCAmelCase				:     Optional[Any]=True   ,  UpperCAmelCase				:     Any=10   ,  UpperCAmelCase				:     int=8   ,  UpperCAmelCase				:     List[Any]=["stage1", "stage2"]   ,  UpperCAmelCase				:     Optional[Any]=[1, 2]   ,  ):
				SCREAMING_SNAKE_CASE_				:int				     =						parent
				SCREAMING_SNAKE_CASE_				:str				     =						batch_size
				SCREAMING_SNAKE_CASE_				:Union[str, Any]				     =						image_size
				SCREAMING_SNAKE_CASE_				:Tuple				     =						patch_size
				SCREAMING_SNAKE_CASE_				:Dict				     =						num_channels
				SCREAMING_SNAKE_CASE_				:Any				     =						embed_dim
				SCREAMING_SNAKE_CASE_				:Optional[int]				     =						hidden_sizes
				SCREAMING_SNAKE_CASE_				:Optional[Any]				     =						depths
				SCREAMING_SNAKE_CASE_				:List[Any]				     =						num_heads
				SCREAMING_SNAKE_CASE_				:int				     =						window_size
				SCREAMING_SNAKE_CASE_				:List[str]				     =						mlp_ratio
				SCREAMING_SNAKE_CASE_				:List[Any]				     =						qkv_bias
				SCREAMING_SNAKE_CASE_				:Dict				     =						hidden_dropout_prob
				SCREAMING_SNAKE_CASE_				:Dict				     =						attention_probs_dropout_prob
				SCREAMING_SNAKE_CASE_				:Dict				     =						drop_path_rate
				SCREAMING_SNAKE_CASE_				:Dict				     =						hidden_act
				SCREAMING_SNAKE_CASE_				:Any				     =						use_absolute_embeddings
				SCREAMING_SNAKE_CASE_				:int				     =						patch_norm
				SCREAMING_SNAKE_CASE_				:Any				     =						layer_norm_eps
				SCREAMING_SNAKE_CASE_				:Optional[int]				     =						initializer_range
				SCREAMING_SNAKE_CASE_				:List[Any]				     =						is_training
				SCREAMING_SNAKE_CASE_				:Optional[int]				     =						scope
				SCREAMING_SNAKE_CASE_				:int				     =						use_labels
				SCREAMING_SNAKE_CASE_				:int				     =						type_sequence_label_size
				SCREAMING_SNAKE_CASE_				:Optional[int]				     =						encoder_stride
				SCREAMING_SNAKE_CASE_				:Optional[int]				     =						out_features
				SCREAMING_SNAKE_CASE_				:Union[str, Any]				     =						out_indices






		def _snake_case		(						self				:     Dict):
				SCREAMING_SNAKE_CASE_				:List[str]				     =						floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size])

				SCREAMING_SNAKE_CASE_				:int				     =						None
				if self.use_labels:
						SCREAMING_SNAKE_CASE_				:Optional[Any]				     =						ids_tensor([self.batch_size]   ,  self.type_sequence_label_size)

				SCREAMING_SNAKE_CASE_				:List[str]				     =						self.get_config()

				return config, pixel_values, labels






		def _snake_case		(						self				:     Union[str, Any]):
				return FocalNetConfig(
				    image_size=self.image_size   ,  patch_size=self.patch_size   ,  num_channels=self.num_channels   ,  embed_dim=self.embed_dim   ,  hidden_sizes=self.hidden_sizes   ,  depths=self.depths   ,  num_heads=self.num_heads   ,  window_size=self.window_size   ,  mlp_ratio=self.mlp_ratio   ,  qkv_bias=self.qkv_bias   ,  hidden_dropout_prob=self.hidden_dropout_prob   ,  attention_probs_dropout_prob=self.attention_probs_dropout_prob   ,  drop_path_rate=self.drop_path_rate   ,  hidden_act=self.hidden_act   ,  use_absolute_embeddings=self.use_absolute_embeddings   ,  path_norm=self.patch_norm   ,  layer_norm_eps=self.layer_norm_eps   ,  initializer_range=self.initializer_range   ,  encoder_stride=self.encoder_stride   ,  out_features=self.out_features   ,  out_indices=self.out_indices   ,  )






		def _snake_case		(						self				:     Union[str, Any]   ,  UpperCAmelCase				:     Optional[int]   ,  UpperCAmelCase				:     Optional[int]   ,  UpperCAmelCase				:     Optional[Any]):
				SCREAMING_SNAKE_CASE_				:Tuple				     =						FocalNetModel(config=UpperCAmelCase)
				model.to(UpperCAmelCase)
				model.eval()
				SCREAMING_SNAKE_CASE_				:Any				     =						model(UpperCAmelCase)

				SCREAMING_SNAKE_CASE_				:List[str]				     =						((config.image_size // config.patch_size) ** 2) // (4 ** (len(config.depths) - 1))
				SCREAMING_SNAKE_CASE_				:Optional[int]				     =						int(config.embed_dim * 2 ** (len(config.depths) - 1))

				self.parent.assertEqual(result.last_hidden_state.shape   ,  (self.batch_size, expected_seq_len, expected_dim))






		def _snake_case		(						self				:     Optional[Any]   ,  UpperCAmelCase				:     Optional[Any]   ,  UpperCAmelCase				:     List[str]   ,  UpperCAmelCase				:     Dict):
				SCREAMING_SNAKE_CASE_				:Optional[int]				     =						FocalNetBackbone(config=UpperCAmelCase)
				model.to(UpperCAmelCase)
				model.eval()
				SCREAMING_SNAKE_CASE_				:Any				     =						model(UpperCAmelCase)

				# verify feature maps
				self.parent.assertEqual(len(result.feature_maps)   ,  len(config.out_features))
				self.parent.assertListEqual(list(result.feature_maps[0].shape)   ,  [self.batch_size, self.image_size, 8, 8])

				# verify channels
				self.parent.assertEqual(len(model.channels)   ,  len(config.out_features))
				self.parent.assertListEqual(model.channels   ,  config.hidden_sizes[:-1])

				# verify backbone works with out_features=None
				SCREAMING_SNAKE_CASE_				:int				     =						None
				SCREAMING_SNAKE_CASE_				:int				     =						FocalNetBackbone(config=UpperCAmelCase)
				model.to(UpperCAmelCase)
				model.eval()
				SCREAMING_SNAKE_CASE_				:Tuple				     =						model(UpperCAmelCase)

				# verify feature maps
				self.parent.assertEqual(len(result.feature_maps)   ,  1)
				self.parent.assertListEqual(list(result.feature_maps[0].shape)   ,  [self.batch_size, self.image_size * 2, 4, 4])

				# verify channels
				self.parent.assertEqual(len(model.channels)   ,  1)
				self.parent.assertListEqual(model.channels   ,  [config.hidden_sizes[-1]])






		def _snake_case		(						self				:     Optional[Any]   ,  UpperCAmelCase				:     List[Any]   ,  UpperCAmelCase				:     str   ,  UpperCAmelCase				:     str):
				SCREAMING_SNAKE_CASE_				:List[str]				     =						FocalNetForMaskedImageModeling(config=UpperCAmelCase)
				model.to(UpperCAmelCase)
				model.eval()
				SCREAMING_SNAKE_CASE_				:Optional[int]				     =						model(UpperCAmelCase)
				self.parent.assertEqual(
				    result.reconstruction.shape   ,  (self.batch_size, self.num_channels, self.image_size, self.image_size))

				# test greyscale images
				SCREAMING_SNAKE_CASE_				:Any				     =						1
				SCREAMING_SNAKE_CASE_				:Tuple				     =						FocalNetForMaskedImageModeling(UpperCAmelCase)
				model.to(UpperCAmelCase)
				model.eval()

				SCREAMING_SNAKE_CASE_				:str				     =						floats_tensor([self.batch_size, 1, self.image_size, self.image_size])
				SCREAMING_SNAKE_CASE_				:Optional[Any]				     =						model(UpperCAmelCase)
				self.parent.assertEqual(result.reconstruction.shape   ,  (self.batch_size, 1, self.image_size, self.image_size))






		def _snake_case		(						self				:     Any   ,  UpperCAmelCase				:     Any   ,  UpperCAmelCase				:     Optional[int]   ,  UpperCAmelCase				:     List[str]):
				SCREAMING_SNAKE_CASE_				:List[str]				     =						self.type_sequence_label_size
				SCREAMING_SNAKE_CASE_				:Dict				     =						FocalNetForImageClassification(UpperCAmelCase)
				model.to(UpperCAmelCase)
				model.eval()
				SCREAMING_SNAKE_CASE_				:Optional[int]				     =						model(UpperCAmelCase   ,  labels=UpperCAmelCase)
				self.parent.assertEqual(result.logits.shape   ,  (self.batch_size, self.type_sequence_label_size))

				# test greyscale images
				SCREAMING_SNAKE_CASE_				:Optional[int]				     =						1
				SCREAMING_SNAKE_CASE_				:Union[str, Any]				     =						FocalNetForImageClassification(UpperCAmelCase)
				model.to(UpperCAmelCase)
				model.eval()

				SCREAMING_SNAKE_CASE_				:Dict				     =						floats_tensor([self.batch_size, 1, self.image_size, self.image_size])
				SCREAMING_SNAKE_CASE_				:Any				     =						model(UpperCAmelCase)
				self.parent.assertEqual(result.logits.shape   ,  (self.batch_size, self.type_sequence_label_size))







		def _snake_case		(						self				:     str):
				SCREAMING_SNAKE_CASE_				:Optional[Any]				     =						self.prepare_config_and_inputs()

				SCREAMING_SNAKE_CASE_     ,				SCREAMING_SNAKE_CASE_     ,				SCREAMING_SNAKE_CASE_				:Optional[int]				     =						config_and_inputs
				SCREAMING_SNAKE_CASE_				:Optional[Any]				     =						{"pixel_values": pixel_values}
				return config, inputs_dict




@require_torch
class 						_UpperCAmelCase      (     lowercase		,		lowercase		,		unittest.TestCase							):
		lowerCamelCase_	:			List[Any]       			=  (
		    (
		        FocalNetModel,
		        FocalNetForImageClassification,
		        FocalNetForMaskedImageModeling,
		        FocalNetBackbone,
		    )
		    if is_torch_available()
		    else ()
		)
		lowerCamelCase_	:			Any       			=  (
		    {"""feature-extraction""": FocalNetModel, """image-classification""": FocalNetForImageClassification}
		    if is_torch_available()
		    else {}
		)
		lowerCamelCase_	:			Tuple       			=  False

		lowerCamelCase_	:			Any       			=  False
		lowerCamelCase_	:			List[Any]       			=  False
		lowerCamelCase_	:			Optional[int]       			=  False
		lowerCamelCase_	:			str       			=  False






		def _snake_case		(						self				:     Union[str, Any]):
				SCREAMING_SNAKE_CASE_				:int				     =						FocalNetModelTester(self)
				SCREAMING_SNAKE_CASE_				:str				     =						ConfigTester(self   ,  config_class=UpperCAmelCase   ,  embed_dim=37   ,  has_text_modality=UpperCAmelCase)






		def _snake_case		(						self				:     Optional[int]):
				self.create_and_test_config_common_properties()
				self.config_tester.create_and_test_config_to_json_string()
				self.config_tester.create_and_test_config_to_json_file()
				self.config_tester.create_and_test_config_from_and_save_pretrained()
				self.config_tester.create_and_test_config_with_num_labels()
				self.config_tester.check_config_can_be_init_without_params()
				self.config_tester.check_config_arguments_init()






		def _snake_case		(						self				:     Dict):
				return






		def _snake_case		(						self				:     int):
				SCREAMING_SNAKE_CASE_				:Optional[Any]				     =						self.model_tester.prepare_config_and_inputs()
				self.model_tester.create_and_check_model(*UpperCAmelCase)






		def _snake_case		(						self				:     Optional[Any]):
				SCREAMING_SNAKE_CASE_				:Union[str, Any]				     =						self.model_tester.prepare_config_and_inputs()
				self.model_tester.create_and_check_backbone(*UpperCAmelCase)






		def _snake_case		(						self				:     str):
				SCREAMING_SNAKE_CASE_				:List[str]				     =						self.model_tester.prepare_config_and_inputs()
				self.model_tester.create_and_check_for_masked_image_modeling(*UpperCAmelCase)






		def _snake_case		(						self				:     Dict):
				SCREAMING_SNAKE_CASE_				:Any				     =						self.model_tester.prepare_config_and_inputs()
				self.model_tester.create_and_check_for_image_classification(*UpperCAmelCase)






		@unittest.skip(reason="FocalNet does not use inputs_embeds")
		def _snake_case		(						self				:     str):
				pass






		@unittest.skip(reason="FocalNet does not use feedforward chunking")
		def _snake_case		(						self				:     Tuple):
				pass






		def _snake_case		(						self				:     int):
				SCREAMING_SNAKE_CASE_     ,				SCREAMING_SNAKE_CASE_				:List[Any]				     =						self.model_tester.prepare_config_and_inputs_for_common()

				for model_class in self.all_model_classes[:-1]:
						SCREAMING_SNAKE_CASE_				:Tuple				     =						model_class(UpperCAmelCase)
						self.assertIsInstance(model.get_input_embeddings()   ,  (nn.Module))
						SCREAMING_SNAKE_CASE_				:Union[str, Any]				     =						model.get_output_embeddings()
						self.assertTrue(x is None or isinstance(UpperCAmelCase   ,  nn.Linear))






		def _snake_case		(						self				:     Optional[int]):
				SCREAMING_SNAKE_CASE_     ,				SCREAMING_SNAKE_CASE_				:Dict				     =						self.model_tester.prepare_config_and_inputs_for_common()

				for model_class in self.all_model_classes[:-1]:
						SCREAMING_SNAKE_CASE_				:int				     =						model_class(UpperCAmelCase)
						SCREAMING_SNAKE_CASE_				:Optional[int]				     =						inspect.signature(model.forward)
						# signature.parameters is an OrderedDict => so arg_names order is deterministic
						SCREAMING_SNAKE_CASE_				:Optional[int]				     =						[*signature.parameters.keys()]

						SCREAMING_SNAKE_CASE_				:Optional[int]				     =						["pixel_values"]
						self.assertListEqual(arg_names[:1]   ,  UpperCAmelCase)






		def _snake_case		(						self				:     Any   ,  UpperCAmelCase				:     Optional[int]   ,  UpperCAmelCase				:     Optional[Any]   ,  UpperCAmelCase				:     Optional[int]   ,  UpperCAmelCase				:     Optional[int]):
				SCREAMING_SNAKE_CASE_				:Tuple				     =						model_class(UpperCAmelCase)
				model.to(UpperCAmelCase)
				model.eval()

				with torch.no_grad():
						SCREAMING_SNAKE_CASE_				:Optional[Any]				     =						model(**self._prepare_for_class(UpperCAmelCase   ,  UpperCAmelCase))

				SCREAMING_SNAKE_CASE_				:Any				     =						outputs.hidden_states

				SCREAMING_SNAKE_CASE_				:Dict				     =						getattr(
				    self.model_tester   ,  "expected_num_hidden_layers"   ,  len(self.model_tester.depths) + 1)
				self.assertEqual(len(UpperCAmelCase)   ,  UpperCAmelCase)

				# FocalNet has a different seq_length
				SCREAMING_SNAKE_CASE_				:Dict				     =						(
				    config.patch_size
				    if isinstance(config.patch_size   ,  collections.abc.Iterable)
				    else (config.patch_size, config.patch_size)
				)

				SCREAMING_SNAKE_CASE_				:List[str]				     =						(image_size[1] // patch_size[1]) * (image_size[0] // patch_size[0])

				self.assertListEqual(
				    list(hidden_states[0].shape[-2:])   ,  [num_patches, self.model_tester.embed_dim]   ,  )

				SCREAMING_SNAKE_CASE_				:Optional[Any]				     =						outputs.reshaped_hidden_states
				self.assertEqual(len(UpperCAmelCase)   ,  UpperCAmelCase)

				SCREAMING_SNAKE_CASE_     ,				SCREAMING_SNAKE_CASE_     ,				SCREAMING_SNAKE_CASE_     ,				SCREAMING_SNAKE_CASE_				:Union[str, Any]				     =						reshaped_hidden_states[0].shape
				SCREAMING_SNAKE_CASE_				:List[Any]				     =						(
				    reshaped_hidden_states[0].view(UpperCAmelCase   ,  UpperCAmelCase   ,  height * width).permute(0   ,  2   ,  1)
				)
				self.assertListEqual(
				    list(reshaped_hidden_states.shape[-2:])   ,  [num_patches, self.model_tester.embed_dim]   ,  )






		def _snake_case		(						self				:     Any):
				SCREAMING_SNAKE_CASE_     ,				SCREAMING_SNAKE_CASE_				:Tuple				     =						self.model_tester.prepare_config_and_inputs_for_common()

				SCREAMING_SNAKE_CASE_				:List[str]				     =						(
				    self.model_tester.image_size
				    if isinstance(self.model_tester.image_size   ,  collections.abc.Iterable)
				    else (self.model_tester.image_size, self.model_tester.image_size)
				)

				for model_class in self.all_model_classes[:-1]:
						SCREAMING_SNAKE_CASE_				:Dict				     =						True
						self.check_hidden_states_output(UpperCAmelCase   ,  UpperCAmelCase   ,  UpperCAmelCase   ,  UpperCAmelCase)

						# check that output_hidden_states also work using config
						del inputs_dict["output_hidden_states"]
						SCREAMING_SNAKE_CASE_				:Union[str, Any]				     =						True

						self.check_hidden_states_output(UpperCAmelCase   ,  UpperCAmelCase   ,  UpperCAmelCase   ,  UpperCAmelCase)






		def _snake_case		(						self				:     Union[str, Any]):
				SCREAMING_SNAKE_CASE_     ,				SCREAMING_SNAKE_CASE_				:Optional[int]				     =						self.model_tester.prepare_config_and_inputs_for_common()
				SCREAMING_SNAKE_CASE_				:Tuple				     =						3

				SCREAMING_SNAKE_CASE_				:List[Any]				     =						(
				    self.model_tester.image_size
				    if isinstance(self.model_tester.image_size   ,  collections.abc.Iterable)
				    else (self.model_tester.image_size, self.model_tester.image_size)
				)
				SCREAMING_SNAKE_CASE_				:Tuple				     =						(
				    config.patch_size
				    if isinstance(config.patch_size   ,  collections.abc.Iterable)
				    else (config.patch_size, config.patch_size)
				)

				SCREAMING_SNAKE_CASE_				:Any				     =						image_size[0] + patch_size[0] - (image_size[0] % patch_size[0])
				SCREAMING_SNAKE_CASE_				:List[Any]				     =						image_size[1] + patch_size[1] - (image_size[1] % patch_size[1])

				for model_class in self.all_model_classes[:-1]:
						SCREAMING_SNAKE_CASE_				:Optional[int]				     =						True
						self.check_hidden_states_output(UpperCAmelCase   ,  UpperCAmelCase   ,  UpperCAmelCase   ,  (padded_height, padded_width))

						# check that output_hidden_states also work using config
						del inputs_dict["output_hidden_states"]
						SCREAMING_SNAKE_CASE_				:Dict				     =						True
						self.check_hidden_states_output(UpperCAmelCase   ,  UpperCAmelCase   ,  UpperCAmelCase   ,  (padded_height, padded_width))






		@slow
		def _snake_case		(						self				:     str):
				for model_name in FOCALNET_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
						SCREAMING_SNAKE_CASE_				:str				     =						FocalNetModel.from_pretrained(UpperCAmelCase)
						self.assertIsNotNone(UpperCAmelCase)







		def _snake_case		(						self				:     List[Any]):
				SCREAMING_SNAKE_CASE_     ,				SCREAMING_SNAKE_CASE_				:Dict				     =						self.model_tester.prepare_config_and_inputs_for_common()

				SCREAMING_SNAKE_CASE_				:Union[str, Any]				     =						_config_zero_init(UpperCAmelCase)
				for model_class in self.all_model_classes:
						SCREAMING_SNAKE_CASE_				:Any				     =						model_class(config=UpperCAmelCase)
						for name, param in model.named_parameters():
								if "embeddings" not in name and 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"   ,  )




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






		@cached_property
		def _snake_case		(						self				:     Any):
				# TODO update organization
				return AutoImageProcessor.from_pretrained("microsoft/focalnet-tiny") if is_vision_available() else None







		@slow
		def _snake_case		(						self				:     Dict):
				SCREAMING_SNAKE_CASE_				:List[str]				     =						FocalNetForImageClassification.from_pretrained("microsoft/focalnet-tiny").to(UpperCAmelCase)
				SCREAMING_SNAKE_CASE_				:Tuple				     =						self.default_image_processor

				SCREAMING_SNAKE_CASE_				:Dict				     =						Image.open("./tests/fixtures/tests_samples/COCO/000000039769.png")
				SCREAMING_SNAKE_CASE_				:Any				     =						image_processor(images=UpperCAmelCase   ,  return_tensors="pt").to(UpperCAmelCase)

				# forward pass
				with torch.no_grad():
						SCREAMING_SNAKE_CASE_				:List[Any]				     =						model(**UpperCAmelCase)

				# verify the logits
				SCREAMING_SNAKE_CASE_				:Optional[Any]				     =						torch.Size((1, 10_00))
				self.assertEqual(outputs.logits.shape   ,  UpperCAmelCase)
				SCREAMING_SNAKE_CASE_				:Optional[int]				     =						torch.tensor([0.2166, -0.4368, 0.2191]).to(UpperCAmelCase)
				self.assertTrue(torch.allclose(outputs.logits[0, :3]   ,  UpperCAmelCase   ,  atol=1E-4))
				self.assertTrue(outputs.logits.argmax(dim=-1).item()   ,  2_81)




@require_torch
class 						_UpperCAmelCase      (     lowercase		,		unittest.TestCase							):
		lowerCamelCase_	:			str       			=  (FocalNetBackbone,) if is_torch_available() else ()
		lowerCamelCase_	:			int       			=  FocalNetConfig

		lowerCamelCase_	:			Dict       			=  False






		def _snake_case		(						self				:     Dict):
				SCREAMING_SNAKE_CASE_				:int				     =						FocalNetModelTester(self)










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


SCREAMING_SNAKE_CASE__     					=			logging.get_logger(__name__)


def     lowercase	(    a     , a ):



		'''simple docstring'''




		SCREAMING_SNAKE_CASE_				:Dict				     =						b.T
		SCREAMING_SNAKE_CASE_				:str				     =						np.sum(np.square(a )     , axis=1 )
		SCREAMING_SNAKE_CASE_				:List[Any]				     =						np.sum(np.square(a )     , axis=0 )
		SCREAMING_SNAKE_CASE_				:Optional[int]				     =						np.matmul(a     , a )
		SCREAMING_SNAKE_CASE_				:Any				     =						aa[:, None] - 2 * ab + ba[None, :]
		return d


def     lowercase	(    a     , a ):



		'''simple docstring'''




		SCREAMING_SNAKE_CASE_				:str				     =						x.reshape(-1     , 3 )
		SCREAMING_SNAKE_CASE_				:str				     =						squared_euclidean_distance(a     , a )
		return np.argmin(a     , axis=1 )



class 						_UpperCAmelCase      (     lowercase							):
		lowerCamelCase_	:			List[str]       			=  ["""pixel_values"""]






		def __init__(						self				:     Dict   ,  UpperCAmelCase				:     Optional[Union[List[List[int]], np.ndarray]] = None   ,  UpperCAmelCase				:     bool = True   ,  UpperCAmelCase				:     Dict[str, int] = None   ,  UpperCAmelCase				:     PILImageResampling = PILImageResampling.BILINEAR   ,  UpperCAmelCase				:     bool = True   ,  UpperCAmelCase				:     bool = True   ,  **UpperCAmelCase				:     Union[str, Any]   ,  ):
				super().__init__(**UpperCAmelCase)
				SCREAMING_SNAKE_CASE_				:Union[str, Any]				     =						size if size is not None else {"height": 2_56, "width": 2_56}
				SCREAMING_SNAKE_CASE_				:Optional[int]				     =						get_size_dict(UpperCAmelCase)
				SCREAMING_SNAKE_CASE_				:int				     =						np.array(UpperCAmelCase) if clusters is not None else None
				SCREAMING_SNAKE_CASE_				:List[Any]				     =						do_resize
				SCREAMING_SNAKE_CASE_				:Dict				     =						size
				SCREAMING_SNAKE_CASE_				:Optional[Any]				     =						resample
				SCREAMING_SNAKE_CASE_				:Dict				     =						do_normalize
				SCREAMING_SNAKE_CASE_				:Tuple				     =						do_color_quantize






		def _snake_case		(						self				:     Union[str, Any]   ,  UpperCAmelCase				:     np.ndarray   ,  UpperCAmelCase				:     Dict[str, int]   ,  UpperCAmelCase				:     PILImageResampling = PILImageResampling.BILINEAR   ,  UpperCAmelCase				:     Optional[Union[str, ChannelDimension]] = None   ,  **UpperCAmelCase				:     List[Any]   ,  ):
				SCREAMING_SNAKE_CASE_				:str				     =						get_size_dict(UpperCAmelCase)
				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(
				    UpperCAmelCase   ,  size=(size["height"], size["width"])   ,  resample=UpperCAmelCase   ,  data_format=UpperCAmelCase   ,  **UpperCAmelCase)






		def _snake_case		(						self				:     Union[str, Any]   ,  UpperCAmelCase				:     np.ndarray   ,  UpperCAmelCase				:     Optional[Union[str, ChannelDimension]] = None   ,  ):
				SCREAMING_SNAKE_CASE_				:int				     =						rescale(image=UpperCAmelCase   ,  scale=1 / 127.5   ,  data_format=UpperCAmelCase)
				SCREAMING_SNAKE_CASE_				:List[str]				     =						image - 1
				return image







		def _snake_case		(						self				:     List[str]   ,  UpperCAmelCase				:     ImageInput   ,  UpperCAmelCase				:     bool = None   ,  UpperCAmelCase				:     Dict[str, int] = None   ,  UpperCAmelCase				:     PILImageResampling = None   ,  UpperCAmelCase				:     bool = None   ,  UpperCAmelCase				:     Optional[bool] = None   ,  UpperCAmelCase				:     Optional[Union[List[List[int]], np.ndarray]] = None   ,  UpperCAmelCase				:     Optional[Union[str, TensorType]] = None   ,  UpperCAmelCase				:     Optional[Union[str, ChannelDimension]] = ChannelDimension.FIRST   ,  **UpperCAmelCase				:     Optional[Any]   ,  ):
				SCREAMING_SNAKE_CASE_				:Any				     =						do_resize if do_resize is not None else self.do_resize
				SCREAMING_SNAKE_CASE_				:Dict				     =						size if size is not None else self.size
				SCREAMING_SNAKE_CASE_				:Union[str, Any]				     =						get_size_dict(UpperCAmelCase)
				SCREAMING_SNAKE_CASE_				:Optional[int]				     =						resample if resample is not None else self.resample
				SCREAMING_SNAKE_CASE_				:Dict				     =						do_normalize if do_normalize is not None else self.do_normalize
				SCREAMING_SNAKE_CASE_				:List[str]				     =						do_color_quantize if do_color_quantize is not None else self.do_color_quantize
				SCREAMING_SNAKE_CASE_				:Tuple				     =						clusters if clusters is not None else self.clusters
				SCREAMING_SNAKE_CASE_				:Tuple				     =						np.array(UpperCAmelCase)

				SCREAMING_SNAKE_CASE_				:Dict				     =						make_list_of_images(UpperCAmelCase)

				if not valid_images(UpperCAmelCase):
						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.
				SCREAMING_SNAKE_CASE_				:Union[str, Any]				     =						[to_numpy_array(UpperCAmelCase) for image in images]

				if do_resize:
						SCREAMING_SNAKE_CASE_				:Tuple				     =						[self.resize(image=UpperCAmelCase   ,  size=UpperCAmelCase   ,  resample=UpperCAmelCase) for image in images]

				if do_normalize:
						SCREAMING_SNAKE_CASE_				:int				     =						[self.normalize(image=UpperCAmelCase) for image in images]

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

						# flatten to (batch_size, height*width)
						SCREAMING_SNAKE_CASE_				:Any				     =						images.shape[0]
						SCREAMING_SNAKE_CASE_				:Dict				     =						images.reshape(UpperCAmelCase   ,  -1)

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

				SCREAMING_SNAKE_CASE_				:str				     =						{"input_ids": images}
				return BatchFeature(data=UpperCAmelCase   ,  tensor_type=UpperCAmelCase)






