

"""simple docstring"""




import inspect
import unittest

from transformers import SegformerConfig, 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,
          SegformerForImageClassification,
          SegformerForSemanticSegmentation,
          SegformerModel,
      )
      from transformers.models.segformer.modeling_segformer import SEGFORMER_PRETRAINED_MODEL_ARCHIVE_LIST


if is_vision_available():
      from PIL import Image

      from transformers import SegformerImageProcessor







class      __lowercase			(    lowercase_					):

  '''simple docstring'''





  def    __lowercase       (  self		:			Dict    ):






   '''simple docstring'''
   UpperCAmelCase__					:  Optional[Any]					   =					self.config_class(**self.inputs_dict    )
   self.parent.assertTrue(hasattr(A					,"""hidden_sizes"""    )    )
   self.parent.assertTrue(hasattr(A					,"""num_attention_heads"""    )    )
   self.parent.assertTrue(hasattr(A					,"""num_encoder_blocks"""    )    )







class      __lowercase			:

  '''simple docstring'''





  def __init__(  self		:			Optional[Any]					,A		:			Dict					,A		:			List[Any]=13					,A		:			List[Any]=64					,A		:			Optional[Any]=3					,A		:			Union[str, Any]=4					,A		:			Any=[2, 2, 2, 2]					,A		:			str=[8, 4, 2, 1]					,A		:			Dict=[16, 32, 64, 128]					,A		:			Any=[1, 4, 8, 16]					,A		:			List[str]=[1, 2, 4, 8]					,A		:			Dict=True					,A		:			Any=True					,A		:			List[str]="gelu"					,A		:			str=0.1					,A		:			Any=0.1					,A		:			Union[str, Any]=0.0_2					,A		:			Optional[Any]=3					,A		:			Optional[int]=None					,):






   '''simple docstring'''
   UpperCAmelCase__					:  Tuple					   =					parent
   UpperCAmelCase__					:  List[str]					   =					batch_size
   UpperCAmelCase__					:  Tuple					   =					image_size
   UpperCAmelCase__					:  Dict					   =					num_channels
   UpperCAmelCase__					:  Any					   =					num_encoder_blocks
   UpperCAmelCase__					:  List[str]					   =					sr_ratios
   UpperCAmelCase__					:  Union[str, Any]					   =					depths
   UpperCAmelCase__					:  Optional[Any]					   =					hidden_sizes
   UpperCAmelCase__					:  Any					   =					downsampling_rates
   UpperCAmelCase__					:  str					   =					num_attention_heads
   UpperCAmelCase__					:  List[Any]					   =					is_training
   UpperCAmelCase__					:  Union[str, Any]					   =					use_labels
   UpperCAmelCase__					:  Tuple					   =					hidden_act
   UpperCAmelCase__					:  List[Any]					   =					hidden_dropout_prob
   UpperCAmelCase__					:  Optional[Any]					   =					attention_probs_dropout_prob
   UpperCAmelCase__					:  List[str]					   =					initializer_range
   UpperCAmelCase__					:  Tuple					   =					num_labels
   UpperCAmelCase__					:  Dict					   =					scope

  def    __lowercase       (  self		:			Any    ):






   '''simple docstring'''
   UpperCAmelCase__					:  Union[str, Any]					   =					floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]    )

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

   UpperCAmelCase__					:  str					   =					self.get_config()
   return config, pixel_values, labels

  def    __lowercase       (  self		:			List[str]    ):






   '''simple docstring'''
   return SegformerConfig(
       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					,num_attention_heads=self.num_attention_heads					,hidden_act=self.hidden_act					,hidden_dropout_prob=self.hidden_dropout_prob					,attention_probs_dropout_prob=self.attention_probs_dropout_prob					,initializer_range=self.initializer_range					,)

  def    __lowercase       (  self		:			Union[str, Any]					,A		:			List[str]					,A		:			int					,A		:			Any    ):






   '''simple docstring'''
   UpperCAmelCase__					:  Dict					   =					SegformerModel(config=A    )
   model.to(A    )
   model.eval()
   UpperCAmelCase__					:  List[Any]					   =					model(A    )
   UpperCAmelCase__					:  Dict					   =					self.image_size // (self.downsampling_rates[-1] * 2)
   self.parent.assertEqual(
       result.last_hidden_state.shape					,(self.batch_size, self.hidden_sizes[-1], expected_height, expected_width)    )

  def    __lowercase       (  self		:			str					,A		:			List[str]					,A		:			List[Any]					,A		:			Optional[int]    ):






   '''simple docstring'''
   UpperCAmelCase__					:  Tuple					   =					self.num_labels
   UpperCAmelCase__					:  Tuple					   =					SegformerForSemanticSegmentation(A    )
   model.to(A    )
   model.eval()
   UpperCAmelCase__					:  List[Any]					   =					model(A    )
   self.parent.assertEqual(
       result.logits.shape					,(self.batch_size, self.num_labels, self.image_size // 4, self.image_size // 4)    )
   UpperCAmelCase__					:  List[str]					   =					model(A					,labels=A    )
   self.parent.assertEqual(
       result.logits.shape					,(self.batch_size, self.num_labels, self.image_size // 4, self.image_size // 4)    )
   self.parent.assertGreater(result.loss					,0.0    )

  def    __lowercase       (  self		:			Tuple					,A		:			Optional[Any]					,A		:			int					,A		:			Optional[Any]    ):






   '''simple docstring'''
   UpperCAmelCase__					:  List[str]					   =					1
   UpperCAmelCase__					:  List[Any]					   =					SegformerForSemanticSegmentation(config=A    )
   model.to(A    )
   model.eval()
   UpperCAmelCase__					:  List[Any]					   =					torch.randint(0					,1					,(self.batch_size, self.image_size, self.image_size)    ).to(A    )
   UpperCAmelCase__					:  int					   =					model(A					,labels=A    )
   self.parent.assertGreater(result.loss					,0.0    )




  def    __lowercase       (  self		:			int    ):






   '''simple docstring'''
   UpperCAmelCase__					:  int					   =					self.prepare_config_and_inputs()
   UpperCAmelCase__   ,      UpperCAmelCase__   ,      UpperCAmelCase__					:  str					   =					config_and_inputs
   UpperCAmelCase__					:  int					   =					{"""pixel_values""": pixel_values}
   return config, inputs_dict







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

  '''simple docstring'''





  snake_case_				       =       (
      (
          SegformerModel,
          SegformerForSemanticSegmentation,
          SegformerForImageClassification,
      )
      if is_torch_available()
      else ()
  )
  snake_case_				       =       (
      {
          '''feature-extraction''': SegformerModel,
          '''image-classification''': SegformerForImageClassification,
          '''image-segmentation''': SegformerForSemanticSegmentation,
      }
      if is_torch_available()
      else {}
  )

  snake_case_				       =       True
  snake_case_				       =       False
  snake_case_				       =       False
  snake_case_				       =       False

  def    __lowercase       (  self		:			Dict    ):






   '''simple docstring'''
   UpperCAmelCase__					:  Dict					   =					SegformerModelTester(self    )
   UpperCAmelCase__					:  Optional[int]					   =					SegformerConfigTester(self					,config_class=A    )

  def    __lowercase       (  self		:			Dict    ):






   '''simple docstring'''
   self.config_tester.run_common_tests()

  def    __lowercase       (  self		:			Tuple    ):






   '''simple docstring'''
   UpperCAmelCase__					:  Dict					   =					self.model_tester.prepare_config_and_inputs()
   self.model_tester.create_and_check_model(*A    )

  def    __lowercase       (  self		:			Optional[Any]    ):






   '''simple docstring'''
   UpperCAmelCase__					:  List[str]					   =					self.model_tester.prepare_config_and_inputs()
   self.model_tester.create_and_check_for_binary_image_segmentation(*A    )

  def    __lowercase       (  self		:			Tuple    ):






   '''simple docstring'''
   UpperCAmelCase__					:  Dict					   =					self.model_tester.prepare_config_and_inputs()
   self.model_tester.create_and_check_for_image_segmentation(*A    )

  @unittest.skip("""SegFormer does not use inputs_embeds"""    )
  def    __lowercase       (  self		:			Tuple    ):






   '''simple docstring'''
   pass

  @unittest.skip("""SegFormer does not have get_input_embeddings method and get_output_embeddings methods"""    )
  def    __lowercase       (  self		:			Dict    ):






   '''simple docstring'''
   pass

  def    __lowercase       (  self		:			str    ):






   '''simple docstring'''
   UpperCAmelCase__   ,      UpperCAmelCase__					:  Any					   =					self.model_tester.prepare_config_and_inputs_for_common()

   for model_class in self.all_model_classes:
    UpperCAmelCase__					:  List[str]					   =					model_class(A    )
    UpperCAmelCase__					:  Tuple					   =					inspect.signature(model.forward    )
    # signature.parameters is an OrderedDict => so arg_names order is deterministic
    UpperCAmelCase__					:  Any					   =					[*signature.parameters.keys()]

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

  def    __lowercase       (  self		:			Union[str, Any]    ):






   '''simple docstring'''
   UpperCAmelCase__   ,      UpperCAmelCase__					:  Union[str, Any]					   =					self.model_tester.prepare_config_and_inputs_for_common()
   UpperCAmelCase__					:  Union[str, Any]					   =					True

   for model_class in self.all_model_classes:
    UpperCAmelCase__					:  Optional[int]					   =					True
    UpperCAmelCase__					:  List[Any]					   =					False
    UpperCAmelCase__					:  Dict					   =					True
    UpperCAmelCase__					:  Optional[int]					   =					model_class(A    )
    model.to(A    )
    model.eval()
    with torch.no_grad():
     UpperCAmelCase__					:  Union[str, Any]					   =					model(**self._prepare_for_class(A					,A    )    )
    UpperCAmelCase__					:  str					   =					outputs.attentions

    UpperCAmelCase__					:  Any					   =					sum(self.model_tester.depths    )
    self.assertEqual(len(A    )					,A    )

    # check that output_attentions also work using config
    del inputs_dict["output_attentions"]
    UpperCAmelCase__					:  int					   =					True
    UpperCAmelCase__					:  List[str]					   =					model_class(A    )
    model.to(A    )
    model.eval()
    with torch.no_grad():
     UpperCAmelCase__					:  Optional[int]					   =					model(**self._prepare_for_class(A					,A    )    )
    UpperCAmelCase__					:  str					   =					outputs.attentions

    self.assertEqual(len(A    )					,A    )

    # verify the first attentions (first block, first layer)
    UpperCAmelCase__					:  Any					   =					(self.model_tester.image_size // 4) ** 2
    UpperCAmelCase__					:  str					   =					(self.model_tester.image_size // (4 * self.model_tester.sr_ratios[0])) ** 2
    self.assertListEqual(
        list(attentions[0].shape[-3:]    )					,[self.model_tester.num_attention_heads[0], expected_seq_len, expected_reduced_seq_len]					,)

    # verify the last attentions (last block, last layer)
    UpperCAmelCase__					:  List[Any]					   =					(self.model_tester.image_size // 32) ** 2
    UpperCAmelCase__					:  List[str]					   =					(self.model_tester.image_size // (32 * self.model_tester.sr_ratios[-1])) ** 2
    self.assertListEqual(
        list(attentions[-1].shape[-3:]    )					,[self.model_tester.num_attention_heads[-1], expected_seq_len, expected_reduced_seq_len]					,)
    UpperCAmelCase__					:  Dict					   =					len(A    )

    # Check attention is always last and order is fine
    UpperCAmelCase__					:  List[str]					   =					True
    UpperCAmelCase__					:  str					   =					True
    UpperCAmelCase__					:  List[Any]					   =					model_class(A    )
    model.to(A    )
    model.eval()
    with torch.no_grad():
     UpperCAmelCase__					:  Union[str, Any]					   =					model(**self._prepare_for_class(A					,A    )    )

    self.assertEqual(out_len + 1					,len(A    )    )

    UpperCAmelCase__					:  Tuple					   =					outputs.attentions

    self.assertEqual(len(A    )					,A    )
    # verify the first attentions (first block, first layer)
    UpperCAmelCase__					:  Optional[Any]					   =					(self.model_tester.image_size // 4) ** 2
    UpperCAmelCase__					:  int					   =					(self.model_tester.image_size // (4 * self.model_tester.sr_ratios[0])) ** 2
    self.assertListEqual(
        list(self_attentions[0].shape[-3:]    )					,[self.model_tester.num_attention_heads[0], expected_seq_len, expected_reduced_seq_len]					,)

  def    __lowercase       (  self		:			Optional[int]    ):






   '''simple docstring'''
   def check_hidden_states_output(A		:			Tuple					,A		:			Optional[Any]					,A		:			Tuple    ):
    UpperCAmelCase__					:  Any					   =					model_class(A    )
    model.to(A    )
    model.eval()

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

    UpperCAmelCase__					:  List[Any]					   =					outputs.hidden_states

    UpperCAmelCase__					:  List[str]					   =					self.model_tester.num_encoder_blocks
    self.assertEqual(len(A    )					,A    )

    # 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__					:  Tuple					   =					self.model_tester.prepare_config_and_inputs_for_common()

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

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

    check_hidden_states_output(A					,A					,A    )

  def    __lowercase       (  self		:			Union[str, Any]    ):






   '''simple docstring'''
   if not self.model_tester.is_training:
    return

   UpperCAmelCase__   ,      UpperCAmelCase__					:  List[Any]					   =					self.model_tester.prepare_config_and_inputs_for_common()
   UpperCAmelCase__					:  str					   =					True

   for model_class in self.all_model_classes:
    if model_class in get_values(A    ):
     continue

    UpperCAmelCase__					:  List[str]					   =					model_class(A    )
    model.to(A    )
    model.train()
    UpperCAmelCase__					:  Dict					   =					self._prepare_for_class(A					,A					,return_labels=A    )
    UpperCAmelCase__					:  int					   =					model(**A    ).loss
    loss.backward()

  @unittest.skip("""Will be fixed soon by reducing the size of the model used for common tests."""    )
  def    __lowercase       (  self		:			Optional[Any]    ):






   '''simple docstring'''
   pass




  @slow
  def    __lowercase       (  self		:			Tuple    ):






   '''simple docstring'''
   for model_name in SEGFORMER_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
    UpperCAmelCase__					:  Union[str, Any]					   =					SegformerModel.from_pretrained(A    )
    self.assertIsNotNone(A    )







def   lowerCAmelCase						(					):
 '''simple docstring'''





 UpperCAmelCase__					:  Optional[int]					   =					Image.open("""./tests/fixtures/tests_samples/COCO/000000039769.png"""      )
 return image







@require_torch
class      __lowercase			(    unittest.TestCase					):

  '''simple docstring'''





  @slow
  def    __lowercase       (  self		:			List[str]    ):






   '''simple docstring'''
   # only resize + normalize
   UpperCAmelCase__					:  Optional[Any]					   =					SegformerImageProcessor(
       image_scale=(512, 512)					,keep_ratio=A					,align=A					,do_random_crop=A    )
   UpperCAmelCase__					:  Optional[Any]					   =					SegformerForSemanticSegmentation.from_pretrained("""nvidia/segformer-b0-finetuned-ade-512-512"""    ).to(
       A    )

   UpperCAmelCase__					:  List[str]					   =					prepare_img()
   UpperCAmelCase__					:  Dict					   =					image_processor(images=A					,return_tensors="""pt"""    )
   UpperCAmelCase__					:  Union[str, Any]					   =					encoded_inputs.pixel_values.to(A    )

   with torch.no_grad():
    UpperCAmelCase__					:  str					   =					model(A    )

   UpperCAmelCase__					:  Tuple					   =					torch.Size((1, model.config.num_labels, 128, 128)    )
   self.assertEqual(outputs.logits.shape					,A    )

   UpperCAmelCase__					:  Optional[int]					   =					torch.tensor(
       [
           [[-4.6_3_1_0, -5.5_2_3_2, -6.2_3_5_6], [-5.1_9_2_1, -6.1_4_4_4, -6.5_9_9_6], [-5.4_4_2_4, -6.2_7_9_0, -6.7_5_7_4]],
           [[-1_2.1_3_9_1, -1_3.3_1_2_2, -1_3.9_5_5_4], [-1_2.8_7_3_2, -1_3.9_3_5_2, -1_4.3_5_6_3], [-1_2.9_4_3_8, -1_3.8_2_2_6, -1_4.2_5_1_3]],
           [[-1_2.5_1_3_4, -1_3.4_6_8_6, -1_4.4_9_1_5], [-1_2.8_6_6_9, -1_4.4_3_4_3, -1_4.7_7_5_8], [-1_3.2_5_2_3, -1_4.5_8_1_9, -1_5.0_6_9_4]],
       ]    ).to(A    )
   self.assertTrue(torch.allclose(outputs.logits[0, :3, :3, :3]					,A					,atol=1e-4    )    )

  @slow
  def    __lowercase       (  self		:			Optional[Any]    ):






   '''simple docstring'''
   # only resize + normalize
   UpperCAmelCase__					:  str					   =					SegformerImageProcessor(
       image_scale=(512, 512)					,keep_ratio=A					,align=A					,do_random_crop=A    )
   UpperCAmelCase__					:  str					   =					SegformerForSemanticSegmentation.from_pretrained(
       """nvidia/segformer-b1-finetuned-cityscapes-1024-1024"""    ).to(A    )

   UpperCAmelCase__					:  Tuple					   =					prepare_img()
   UpperCAmelCase__					:  Optional[int]					   =					image_processor(images=A					,return_tensors="""pt"""    )
   UpperCAmelCase__					:  Tuple					   =					encoded_inputs.pixel_values.to(A    )

   with torch.no_grad():
    UpperCAmelCase__					:  int					   =					model(A    )

   UpperCAmelCase__					:  str					   =					torch.Size((1, model.config.num_labels, 128, 128)    )
   self.assertEqual(outputs.logits.shape					,A    )

   UpperCAmelCase__					:  Optional[Any]					   =					torch.tensor(
       [
           [[-1_3.5_7_4_8, -1_3.9_1_1_1, -1_2.6_5_0_0], [-1_4.3_5_0_0, -1_5.3_6_8_3, -1_4.2_3_2_8], [-1_4.7_5_3_2, -1_6.0_4_2_4, -1_5.6_0_8_7]],
           [[-1_7.1_6_5_1, -1_5.8_7_2_5, -1_2.9_6_5_3], [-1_7.2_5_8_0, -1_7.3_7_1_8, -1_4.8_2_2_3], [-1_6.6_0_5_8, -1_6.8_7_8_3, -1_6.7_4_5_2]],
           [[-3.6_4_5_6, -3.0_2_0_9, -1.4_2_0_3], [-3.0_7_9_7, -3.1_9_5_9, -2.0_0_0_0], [-1.8_7_5_7, -1.9_2_1_7, -1.6_9_9_7]],
       ]    ).to(A    )
   self.assertTrue(torch.allclose(outputs.logits[0, :3, :3, :3]					,A					,atol=1e-1    )    )




  @slow
  def    __lowercase       (  self		:			List[str]    ):






   '''simple docstring'''
   # only resize + normalize
   UpperCAmelCase__					:  Dict					   =					SegformerImageProcessor(
       image_scale=(512, 512)					,keep_ratio=A					,align=A					,do_random_crop=A    )
   UpperCAmelCase__					:  List[str]					   =					SegformerForSemanticSegmentation.from_pretrained("""nvidia/segformer-b0-finetuned-ade-512-512"""    ).to(
       A    )

   UpperCAmelCase__					:  int					   =					prepare_img()
   UpperCAmelCase__					:  Any					   =					image_processor(images=A					,return_tensors="""pt"""    )
   UpperCAmelCase__					:  int					   =					encoded_inputs.pixel_values.to(A    )

   with torch.no_grad():
    UpperCAmelCase__					:  int					   =					model(A    )

   UpperCAmelCase__					:  Any					   =					outputs.logits.detach().cpu()

   UpperCAmelCase__					:  Dict					   =					image_processor.post_process_semantic_segmentation(outputs=A					,target_sizes=[(500, 300)]    )
   UpperCAmelCase__					:  int					   =					torch.Size((500, 300)    )
   self.assertEqual(segmentation[0].shape					,A    )

   UpperCAmelCase__					:  Optional[Any]					   =					image_processor.post_process_semantic_segmentation(outputs=A    )
   UpperCAmelCase__					:  Union[str, Any]					   =					torch.Size((128, 128)    )
   self.assertEqual(segmentation[0].shape					,A    )








"""simple docstring"""




def   lowerCAmelCase						(					__UpperCamelCase = 1000      ):
 '''simple docstring'''





 return sum(2 * a * ((a - 1) // 2) for a in range(3							,    n + 1      )      )


if __name__ == "__main__":
      print(solution())





