"""simple docstring"""






import inspect
import unittest

from transformers import is_torch_available, is_vision_available
from transformers.models.auto import get_values
from transformers.testing_utils import require_torch, slow, torch_device

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 MODEL_MAPPING, PoolFormerConfig, PoolFormerForImageClassification, PoolFormerModel
      from transformers.models.poolformer.modeling_poolformer import POOLFORMER_PRETRAINED_MODEL_ARCHIVE_LIST


if is_vision_available():
      from PIL import Image

      from transformers import PoolFormerImageProcessor
class     snake_case	(			lowercase_    ):




     """simple docstring"""

     def 		__lowerCAmelCase						(    self       :     Any      ):
         UpperCAmelCase__					=							self.config_class(**self.inputs_dict      )
         self.parent.assertTrue(hasattr(lowerCamelCase__   ,'hidden_sizes'      )      )
         self.parent.assertTrue(hasattr(lowerCamelCase__   ,'num_encoder_blocks'      )      )






class     snake_case	:




     """simple docstring"""

     def __init__(    self       :     str   ,lowerCamelCase__       :     Tuple   ,lowerCamelCase__       :     str=13   ,lowerCamelCase__       :     Optional[Any]=64   ,lowerCamelCase__       :     Optional[int]=3   ,lowerCamelCase__       :     Union[str, Any]=4   ,lowerCamelCase__       :     List[Any]=[2, 2, 2, 2]   ,lowerCamelCase__       :     int=[8, 4, 2, 1]   ,lowerCamelCase__       :     Any=[16, 32, 64, 128]   ,lowerCamelCase__       :     Optional[Any]=[1, 4, 8, 16]   ,lowerCamelCase__       :     Dict=False   ,lowerCamelCase__       :     Optional[Any]=True   ,lowerCamelCase__       :     Any="gelu"   ,lowerCamelCase__       :     Optional[int]=0.1   ,lowerCamelCase__       :     Tuple=0.0_2   ,lowerCamelCase__       :     Tuple=3   ,lowerCamelCase__       :     Tuple=None   ,):
         UpperCAmelCase__					=							parent
         UpperCAmelCase__					=							batch_size
         UpperCAmelCase__					=							image_size
         UpperCAmelCase__					=							num_channels
         UpperCAmelCase__					=							num_encoder_blocks
         UpperCAmelCase__					=							sr_ratios
         UpperCAmelCase__					=							depths
         UpperCAmelCase__					=							hidden_sizes
         UpperCAmelCase__					=							downsampling_rates
         UpperCAmelCase__					=							is_training
         UpperCAmelCase__					=							use_labels
         UpperCAmelCase__					=							hidden_act
         UpperCAmelCase__					=							hidden_dropout_prob
         UpperCAmelCase__					=							initializer_range
         UpperCAmelCase__					=							num_labels
         UpperCAmelCase__					=							scope







     def 		__lowerCAmelCase						(    self       :     Optional[Any]      ):
         UpperCAmelCase__					=							floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]      )

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

         UpperCAmelCase__					=							PoolFormerConfig(
             image_size=self.image_size   ,num_channels=self.num_channels   ,num_encoder_blocks=self.num_encoder_blocks   ,depths=self.depths   ,hidden_sizes=self.hidden_sizes   ,hidden_act=self.hidden_act   ,hidden_dropout_prob=self.hidden_dropout_prob   ,initializer_range=self.initializer_range   ,)

         return config, pixel_values, labels







     def 		__lowerCAmelCase						(    self       :     Optional[int]   ,lowerCamelCase__       :     Optional[Any]   ,lowerCamelCase__       :     Any   ,lowerCamelCase__       :     Optional[int]      ):
         UpperCAmelCase__					=							PoolFormerModel(config=lowerCamelCase__      )
         model.to(lowerCamelCase__      )
         model.eval()
         UpperCAmelCase__					=							model(lowerCamelCase__      )
         UpperCAmelCase__					=							UpperCAmelCase__					=							self.image_size // 3_2.0
         self.parent.assertEqual(
             result.last_hidden_state.shape   ,(self.batch_size, self.hidden_sizes[-1], expected_height, expected_width)      )







     def 		__lowerCAmelCase						(    self       :     Dict      ):
         UpperCAmelCase__					=							self.prepare_config_and_inputs()
         UpperCAmelCase__				,     UpperCAmelCase__				,     UpperCAmelCase__					=							config_and_inputs
         UpperCAmelCase__					=							{'pixel_values': pixel_values}
         return config, inputs_dict






@require_torch
class     snake_case	(			lowercase_				,   lowercase_				,   unittest.TestCase    ):




     """simple docstring"""

     snake_case__       =			(PoolFormerModel, PoolFormerForImageClassification) if is_torch_available() else ()
     snake_case__       =			(
         {"""feature-extraction""": PoolFormerModel, """image-classification""": PoolFormerForImageClassification}
         if is_torch_available()
         else {}
     )

     snake_case__       =			False
     snake_case__       =			False
     snake_case__       =			False
     snake_case__       =			False
     snake_case__       =			False







     def 		__lowerCAmelCase						(    self       :     List[Any]      ):
         UpperCAmelCase__					=							PoolFormerModelTester(self      )
         UpperCAmelCase__					=							PoolFormerConfigTester(self   ,config_class=lowerCamelCase__      )







     def 		__lowerCAmelCase						(    self       :     Union[str, Any]      ):
         self.config_tester.run_common_tests()







     def 		__lowerCAmelCase						(    self       :     str      ):
         UpperCAmelCase__					=							self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_model(*lowerCamelCase__      )







     @unittest.skip('PoolFormer does not use inputs_embeds'      )
     def 		__lowerCAmelCase						(    self       :     Any      ):
         pass







     @unittest.skip('PoolFormer does not have get_input_embeddings method and get_output_embeddings methods'      )
     def 		__lowerCAmelCase						(    self       :     Union[str, Any]      ):
         pass







     def 		__lowerCAmelCase						(    self       :     List[str]      ):







         def check_hidden_states_output(lowerCamelCase__       :     str   ,lowerCamelCase__       :     Tuple   ,lowerCamelCase__       :     List[str]      ):
             UpperCAmelCase__					=							model_class(lowerCamelCase__      )
             model.to(lowerCamelCase__      )
             model.eval()

             with torch.no_grad():
                 UpperCAmelCase__					=							model(**self._prepare_for_class(lowerCamelCase__   ,lowerCamelCase__      )      )

             UpperCAmelCase__					=							outputs.hidden_states

             UpperCAmelCase__					=							self.model_tester.num_encoder_blocks
             self.assertEqual(len(lowerCamelCase__      )   ,lowerCamelCase__      )

             # verify the first hidden states (first block)
             self.assertListEqual(
                 list(hidden_states[0].shape[-3:]      )   ,[
                     self.model_tester.hidden_sizes[0],
                     self.model_tester.image_size // 4,
                     self.model_tester.image_size // 4,
                 ]   ,)

         UpperCAmelCase__				,     UpperCAmelCase__					=							self.model_tester.prepare_config_and_inputs_for_common()

         for model_class in self.all_model_classes:
             UpperCAmelCase__					=							True
             check_hidden_states_output(lowerCamelCase__   ,lowerCamelCase__   ,lowerCamelCase__      )

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

             check_hidden_states_output(lowerCamelCase__   ,lowerCamelCase__   ,lowerCamelCase__      )







     def 		__lowerCAmelCase						(    self       :     int      ):







         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:
             if model_class in get_values(lowerCamelCase__      ):
                 continue
             UpperCAmelCase__					=							model_class(lowerCamelCase__      )
             model.to(lowerCamelCase__      )
             model.train()
             UpperCAmelCase__					=							self._prepare_for_class(lowerCamelCase__   ,lowerCamelCase__   ,return_labels=lowerCamelCase__      )
             UpperCAmelCase__					=							model(**lowerCamelCase__      ).loss
             loss.backward()







     def 		__lowerCAmelCase						(    self       :     Optional[Any]      ):
         UpperCAmelCase__				,     UpperCAmelCase__					=							self.model_tester.prepare_config_and_inputs_for_common()

         for model_class in self.all_model_classes:
             UpperCAmelCase__					=							model_class(lowerCamelCase__      )
             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]   ,lowerCamelCase__      )







     @slow
     def 		__lowerCAmelCase						(    self       :     Any      ):







         for model_name in POOLFORMER_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
             UpperCAmelCase__					=							PoolFormerModel.from_pretrained(lowerCamelCase__      )
             self.assertIsNotNone(lowerCamelCase__      )






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







@require_torch
class     snake_case	(			unittest.TestCase    ):




     """simple docstring"""

     @slow
     def 		__lowerCAmelCase						(    self       :     List[str]      ):
         UpperCAmelCase__					=							PoolFormerImageProcessor()
         UpperCAmelCase__					=							PoolFormerForImageClassification.from_pretrained('sail/poolformer_s12'      ).to(lowerCamelCase__      )

         UpperCAmelCase__					=							image_processor(images=prepare_img()   ,return_tensors='pt'      ).to(lowerCamelCase__      )

         # forward pass
         with torch.no_grad():
             UpperCAmelCase__					=							model(**lowerCamelCase__      )

         # verify the logits
         UpperCAmelCase__					=							torch.Size((1, 1_000)      )
         self.assertEqual(outputs.logits.shape   ,lowerCamelCase__      )

         UpperCAmelCase__					=							torch.tensor([-0.6_1_1_3, 0.1_6_8_5, -0.0_4_9_2]      ).to(lowerCamelCase__      )
         self.assertTrue(torch.allclose(outputs.logits[0, :3]   ,lowerCamelCase__   ,atol=1e-4      )      )




"""simple docstring"""






from typing import SupportsIndex

import numpy as np
from scipy.ndimage import convolve



def   a_						( lowerCamelCase ,      lowerCamelCase ,      lowerCamelCase	):
				UpperCAmelCase__					=							np.stack((horizontal_flow, vertical_flow) ,      2	)

				# Create a grid of all pixel coordinates and subtract the flow to get the
				# target pixels coordinates
				UpperCAmelCase__					=							np.stack(
				    np.meshgrid(np.arange(0 ,      image.shape[1]	) ,      np.arange(0 ,      image.shape[0]	)	) ,      2	)
				UpperCAmelCase__					=							np.round(grid - flow	).astype(np.intaa	)

				# Find the locations outside of the original image
				UpperCAmelCase__					=							(grid < 0) | (grid >= np.array([image.shape[1], image.shape[0]]	))
				UpperCAmelCase__					=							0

				UpperCAmelCase__					=							image[grid[:, :, 1], grid[:, :, 0]]

				# Set pixels at invalid locations to 0
				UpperCAmelCase__					=							0

				return warped



def   a_						( lowerCamelCase ,      lowerCamelCase ,      lowerCamelCase ,      lowerCamelCase = None ,      ):
				if alpha is None:
								UpperCAmelCase__					=							0.1

				# Initialize flow
				UpperCAmelCase__					=							np.zeros_like(lowerCamelCase	)
				UpperCAmelCase__					=							np.zeros_like(lowerCamelCase	)

				# Prepare kernels for the calculation of the derivatives and the average velocity
				UpperCAmelCase__					=							np.array([[-1, 1], [-1, 1]]	) * 0.25
				UpperCAmelCase__					=							np.array([[-1, -1], [1, 1]]	) * 0.25
				UpperCAmelCase__					=							np.array([[1, 1], [1, 1]]	) * 0.25
				UpperCAmelCase__					=							np.array(
				    [[1 / 1_2, 1 / 6, 1 / 1_2], [1 / 6, 0, 1 / 6], [1 / 1_2, 1 / 6, 1 / 1_2]]	)

				# Iteratively refine the flow
				for _ in range(lowerCamelCase	):
								UpperCAmelCase__					=							warp(lowerCamelCase ,      lowerCamelCase ,      lowerCamelCase	)
								UpperCAmelCase__					=							convolve(lowerCamelCase ,      lowerCamelCase	) + convolve(lowerCamelCase ,      lowerCamelCase	)
								UpperCAmelCase__					=							convolve(lowerCamelCase ,      lowerCamelCase	) + convolve(lowerCamelCase ,      lowerCamelCase	)
								UpperCAmelCase__					=							convolve(lowerCamelCase ,      lowerCamelCase	) + convolve(lowerCamelCase ,      -kernel_t	)

								UpperCAmelCase__					=							convolve(lowerCamelCase ,      lowerCamelCase	)
								UpperCAmelCase__					=							convolve(lowerCamelCase ,      lowerCamelCase	)

								# This updates the flow as proposed in the paper (Step 12)
								UpperCAmelCase__					=							(
								    derivative_x * avg_horizontal_velocity
								    + derivative_y * avg_vertical_velocity
								    + derivative_t
								)
								UpperCAmelCase__					=							update / (alpha**2 + derivative_x**2 + derivative_y**2)

								UpperCAmelCase__					=							avg_horizontal_velocity - derivative_x * update
								UpperCAmelCase__					=							avg_vertical_velocity - derivative_y * update

				return horizontal_flow, vertical_flow


if __name__ == "__main__":
						import doctest

						doctest.testmod()



