
'''simple docstring'''





import collections
import inspect
import unittest

from transformers import DinatConfig
from transformers.testing_utils import require_natten, 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 DinatBackbone, DinatForImageClassification, DinatModel
 from transformers.models.dinat.modeling_dinat import DINAT_PRETRAINED_MODEL_ARCHIVE_LIST

if is_vision_available():
 from PIL import Image

 from transformers import AutoImageProcessor






class 	UpperCamelCase__					:






 def __init__(     self					:List[str]		,			_A					:Dict		,			_A					:int=13		,			_A					:Optional[Any]=64		,			_A					:Tuple=4		,			_A					:Union[str, Any]=3		,			_A					:Union[str, Any]=16		,			_A					:Any=[1, 2, 1]		,			_A					:List[str]=[2, 4, 8]		,			_A					:Dict=3		,			_A					:Any=[[3], [1, 2], [1]]		,			_A					:Any=2.0		,			_A					:Optional[int]=True		,			_A					:Optional[int]=0.0		,			_A					:Any=0.0		,			_A					:Any=0.1		,			_A					:List[str]="gelu"		,			_A					:Optional[Any]=True		,			_A					:Optional[Any]=0.02		,			_A					:List[Any]=1E-5		,			_A					:List[str]=True		,			_A					:Union[str, Any]=None		,			_A					:Optional[int]=True		,			_A					:Tuple=10		,			_A					:List[Any]=["stage1", "stage2"]		,			_A					:List[str]=[1, 2]		,			)			->     Union[str, Any]:







       '''simple docstring'''



       __A    		=     parent
       __A    		=     batch_size
       __A    		=     image_size
       __A    		=     patch_size
       __A    		=     num_channels
       __A    		=     embed_dim
       __A    		=     depths
       __A    		=     num_heads
       __A    		=     kernel_size
       __A    		=     dilations
       __A    		=     mlp_ratio
       __A    		=     qkv_bias
       __A    		=     hidden_dropout_prob
       __A    		=     attention_probs_dropout_prob
       __A    		=     drop_path_rate
       __A    		=     hidden_act
       __A    		=     patch_norm
       __A    		=     layer_norm_eps
       __A    		=     initializer_range
       __A    		=     is_training
       __A    		=     scope
       __A    		=     use_labels
       __A    		=     num_labels
       __A    		=     out_features
       __A    		=     out_indices






 def 		lowercase_  (     self					:Any     )			->     Tuple:







       '''simple docstring'''



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

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

       __A    		=     self.get_config()

       return config, pixel_values, labels






 def 		lowercase_  (     self					:Dict     )			->     Optional[int]:







       '''simple docstring'''



       return DinatConfig(
           num_labels=self.num_labels		,			image_size=self.image_size		,			patch_size=self.patch_size		,			num_channels=self.num_channels		,			embed_dim=self.embed_dim		,			depths=self.depths		,			num_heads=self.num_heads		,			kernel_size=self.kernel_size		,			dilations=self.dilations		,			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		,			patch_norm=self.patch_norm		,			layer_norm_eps=self.layer_norm_eps		,			initializer_range=self.initializer_range		,			out_features=self.out_features		,			out_indices=self.out_indices		,			)






 def 		lowercase_  (     self					:str		,			_A					:Optional[Any]		,			_A					:Optional[int]		,			_A					:Union[str, Any]     )			->     Union[str, Any]:







       '''simple docstring'''



       __A    		=     DinatModel(config=__UpperCamelCase     )
       model.to(__UpperCamelCase     )
       model.eval()
       __A    		=     model(__UpperCamelCase     )

       __A    		=     __A    		=     (config.image_size // config.patch_size) // (2 ** (len(config.depths     ) - 1))
       __A    		=     int(config.embed_dim * 2 ** (len(config.depths     ) - 1)     )

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






 def 		lowercase_  (     self					:List[Any]		,			_A					:str		,			_A					:int		,			_A					:Union[str, Any]     )			->     Optional[Any]:







       '''simple docstring'''



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

       # test greyscale images
       __A    		=     1
       __A    		=     DinatForImageClassification(__UpperCamelCase     )
       model.to(__UpperCamelCase     )
       model.eval()

       __A    		=     floats_tensor([self.batch_size, 1, self.image_size, self.image_size]     )
       __A    		=     model(__UpperCamelCase     )
       self.parent.assertEqual(result.logits.shape		,			(self.batch_size, self.num_labels)     )






 def 		lowercase_  (     self					:Tuple		,			_A					:List[Any]		,			_A					:Tuple		,			_A					:Dict     )			->     Tuple:







       '''simple docstring'''



       __A    		=     DinatBackbone(config=__UpperCamelCase     )
       model.to(__UpperCamelCase     )
       model.eval()
       __A    		=     model(__UpperCamelCase     )

       # verify hidden states
       self.parent.assertEqual(len(result.feature_maps     )		,			len(config.out_features     )     )
       self.parent.assertListEqual(list(result.feature_maps[0].shape     )		,			[self.batch_size, model.channels[0], 16, 16]     )

       # verify channels
       self.parent.assertEqual(len(model.channels     )		,			len(config.out_features     )     )

       # verify backbone works with out_features=None
       __A    		=     None
       __A    		=     DinatBackbone(config=__UpperCamelCase     )
       model.to(__UpperCamelCase     )
       model.eval()
       __A    		=     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, model.channels[-1], 4, 4]     )

       # verify channels
       self.parent.assertEqual(len(model.channels     )		,			1     )






 def 		lowercase_  (     self					:int     )			->     List[Any]:







       '''simple docstring'''



       __A    		=     self.prepare_config_and_inputs()
       __A		, __A		, __A    		=     config_and_inputs
       __A    		=     {'pixel_values': pixel_values}
       return config, inputs_dict






@require_natten
@require_torch
class 	UpperCamelCase__					(      UpperCAmelCase_ ,       UpperCAmelCase_ ,       unittest.TestCase):
 UpperCAmelCase__      :      Union[str, Any]  			=   (
     (
         DinatModel,
         DinatForImageClassification,
         DinatBackbone,
     )
     if is_torch_available()
     else ()
 )
 UpperCAmelCase__      :      Dict  			=   (
     {'feature-extraction': DinatModel, 'image-classification': DinatForImageClassification}
     if is_torch_available()
     else {}
 )
 UpperCAmelCase__      :      List[str]  			=   False

 UpperCAmelCase__      :      Any  			=   False
 UpperCAmelCase__      :      str  			=   False
 UpperCAmelCase__      :      Tuple  			=   False
 UpperCAmelCase__      :      List[str]  			=   False






 def 		lowercase_  (     self					:Dict     )			->     List[Any]:







       '''simple docstring'''



       __A    		=     DinatModelTester(self     )
       __A    		=     ConfigTester(self		,			config_class=__UpperCamelCase		,			embed_dim=37     )






 def 		lowercase_  (     self					:Optional[int]     )			->     Optional[int]:







       '''simple docstring'''



       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 		lowercase_  (     self					:int     )			->     Optional[int]:







       '''simple docstring'''



       return






 def 		lowercase_  (     self					:int     )			->     Union[str, Any]:







       '''simple docstring'''



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






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







       '''simple docstring'''



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






 def 		lowercase_  (     self					:Optional[int]     )			->     str:







       '''simple docstring'''



       __A    		=     self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_backbone(*__UpperCamelCase     )






 @unittest.skip(reason='Dinat does not use inputs_embeds'     )
 def 		lowercase_  (     self					:Tuple     )			->     int:







       '''simple docstring'''



       pass






 @unittest.skip(reason='Dinat does not use feedforward chunking'     )
 def 		lowercase_  (     self					:Any     )			->     Optional[Any]:







       '''simple docstring'''



       pass






 def 		lowercase_  (     self					:Optional[int]     )			->     List[str]:







       '''simple docstring'''



       __A		, __A    		=     self.model_tester.prepare_config_and_inputs_for_common()

       for model_class in self.all_model_classes:
             __A    		=     model_class(__UpperCamelCase     )
             self.assertIsInstance(model.get_input_embeddings()		,			(nn.Module)     )
             __A    		=     model.get_output_embeddings()
             self.assertTrue(x is None or isinstance(__UpperCamelCase		,			nn.Linear     )     )






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







       '''simple docstring'''



       __A		, __A    		=     self.model_tester.prepare_config_and_inputs_for_common()

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






 def 		lowercase_  (     self					:Dict     )			->     int:







       '''simple docstring'''



       self.skipTest('Dinat\'s attention operation is handled entirely by NATTEN.'     )






 def 		lowercase_  (     self					:List[str]		,			_A					:List[str]		,			_A					:List[Any]		,			_A					:Union[str, Any]		,			_A					:Any     )			->     Dict:







       '''simple docstring'''



       __A    		=     model_class(__UpperCamelCase     )
       model.to(__UpperCamelCase     )
       model.eval()

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

       __A    		=     outputs.hidden_states

       __A    		=     getattr(
           self.model_tester		,			'expected_num_hidden_layers'		,			len(self.model_tester.depths     ) + 1     )
       self.assertEqual(len(__UpperCamelCase     )		,			__UpperCamelCase     )

       # Dinat has a different seq_length
       __A    		=     (
           config.patch_size
           if isinstance(config.patch_size		,			collections.abc.Iterable     )
           else (config.patch_size, config.patch_size)
       )

       __A    		=     image_size[0] // patch_size[0]
       __A    		=     image_size[1] // patch_size[1]

       self.assertListEqual(
           list(hidden_states[0].shape[-3:]     )		,			[height, width, self.model_tester.embed_dim]		,			)

       if model_class.__name__ != "DinatBackbone":
             __A    		=     outputs.reshaped_hidden_states
             self.assertEqual(len(__UpperCamelCase     )		,			__UpperCamelCase     )

             __A		, __A		, __A		, __A    		=     reshaped_hidden_states[0].shape
             __A    		=     (
                 reshaped_hidden_states[0].view(__UpperCamelCase		,			__UpperCamelCase		,			__UpperCamelCase		,			__UpperCamelCase     ).permute(0		,			2		,			3		,			1     )
             )
             self.assertListEqual(
                 list(reshaped_hidden_states.shape[-3:]     )		,			[height, width, self.model_tester.embed_dim]		,			)






 def 		lowercase_  (     self					:int     )			->     Any:







       '''simple docstring'''



       __A		, __A    		=     self.model_tester.prepare_config_and_inputs_for_common()

       __A    		=     (
           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:
             __A    		=     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"]
             __A    		=     True

             self.check_hidden_states_output(__UpperCamelCase		,			__UpperCamelCase		,			__UpperCamelCase		,			__UpperCamelCase     )






 @slow
 def 		lowercase_  (     self					:Optional[int]     )			->     str:







       '''simple docstring'''



       for model_name in DINAT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
             __A    		=     DinatModel.from_pretrained(__UpperCamelCase     )
             self.assertIsNotNone(__UpperCamelCase     )






 def 		lowercase_  (     self					:Dict     )			->     List[Any]:







       '''simple docstring'''



       __A		, __A    		=     self.model_tester.prepare_config_and_inputs_for_common()

       __A    		=     _config_zero_init(__UpperCamelCase     )
       for model_class in self.all_model_classes:
             __A    		=     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_natten
@require_vision
@require_torch
class 	UpperCamelCase__					(      unittest.TestCase):






 @cached_property
 def 		lowercase_  (     self					:Dict     )			->     Optional[int]:







       '''simple docstring'''



       return AutoImageProcessor.from_pretrained('shi-labs/dinat-mini-in1k-224'     ) if is_vision_available() else None






 @slow
 def 		lowercase_  (     self					:Tuple     )			->     List[Any]:







       '''simple docstring'''



       __A    		=     DinatForImageClassification.from_pretrained('shi-labs/dinat-mini-in1k-224'     ).to(__UpperCamelCase     )
       __A    		=     self.default_image_processor

       __A    		=     Image.open('./tests/fixtures/tests_samples/COCO/000000039769.png'     )
       __A    		=     image_processor(images=__UpperCamelCase		,			return_tensors='pt'     ).to(__UpperCamelCase     )

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

       # verify the logits
       __A    		=     torch.Size((1, 1_000)     )
       self.assertEqual(outputs.logits.shape		,			__UpperCamelCase     )
       __A    		=     torch.tensor([-0.1_545, -0.7_667, 0.4_642]     ).to(__UpperCamelCase     )
       self.assertTrue(torch.allclose(outputs.logits[0, :3]		,			__UpperCamelCase		,			atol=1E-4     )     )






@require_torch
@require_natten
class 	UpperCamelCase__					(      unittest.TestCase ,       UpperCAmelCase_):
 UpperCAmelCase__      :      Dict  			=   (DinatBackbone,) if is_torch_available() else ()
 UpperCAmelCase__      :      Optional[Any]  			=   DinatConfig






 def 		lowercase_  (     self					:Dict     )			->     Optional[int]:







       '''simple docstring'''



       __A    		=     DinatModelTester(self     )








'''simple docstring'''





import math
import os

import cva
import numpy as np

a__						:				List[Any]           =    255.0





def     snake_case							(						UpperCAmelCase ,					UpperCAmelCase		)->     float:







      """simple docstring"""
      __A    		=     np.mean((original - contrast) ** 2		)
      if mse == 0:
            return 1_0_0

      return 2_0 * math.logaa(PIXEL_MAX / math.sqrt(UpperCAmelCase		)		)





def     snake_case							(						)->     None:







      """simple docstring"""
      __A    		=     os.path.dirname(os.path.realpath(UpperCAmelCase		)		)
      # Loading images (original image and compressed image)
      __A    		=     cva.imread(os.path.join(UpperCAmelCase ,					'image_data/original_image.png'		)		)
      __A    		=     cva.imread(os.path.join(UpperCAmelCase ,					'image_data/compressed_image.png'		) ,					1		)

      __A    		=     cva.imread(os.path.join(UpperCAmelCase ,					'image_data/PSNR-example-base.png'		)		)
      __A    		=     cva.imread(
          os.path.join(UpperCAmelCase ,					'image_data/PSNR-example-comp-10.jpg'		) ,					1		)

      # Value expected: 29.73dB
      print('-- First Test --'		)
      print(f'PSNR value is {peak_signal_to_noise_ratio(UpperCAmelCase ,					UpperCAmelCase		)} dB'		)

      # # Value expected: 31.53dB (Wikipedia Example)
      print('\n-- Second Test --'		)
      print(f'PSNR value is {peak_signal_to_noise_ratio(UpperCAmelCase ,					UpperCAmelCase		)} dB'		)


if __name__ == "__main__":
 main()






