





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 				__lowerCAmelCase       (    UpperCamelCase__):







      def    _lowercase  (    self							)   ->						Optional[Any]:



        '''simple docstring'''




        a__	:					Optional[Any]										=self.config_class(**self.inputs_dict							)
        self.parent.assertTrue(hasattr(lowerCAmelCase__		,      "hidden_sizes"							)							)
        self.parent.assertTrue(hasattr(lowerCAmelCase__		,      "num_attention_heads"							)							)
        self.parent.assertTrue(hasattr(lowerCAmelCase__		,      "num_encoder_blocks"							)							)







class 				__lowerCAmelCase       :







      def __init__(    self		,      lowerCAmelCase__		,      lowerCAmelCase__=1_3		,      lowerCAmelCase__=6_4		,      lowerCAmelCase__=3		,      lowerCAmelCase__=4		,      lowerCAmelCase__=[2, 2, 2, 2]		,      lowerCAmelCase__=[8, 4, 2, 1]		,      lowerCAmelCase__=[1_6, 3_2, 6_4, 1_2_8]		,      lowerCAmelCase__=[1, 4, 8, 1_6]		,      lowerCAmelCase__=[1, 2, 4, 8]		,      lowerCAmelCase__=True		,      lowerCAmelCase__=True		,      lowerCAmelCase__="gelu"		,      lowerCAmelCase__=0.1		,      lowerCAmelCase__=0.1		,      lowerCAmelCase__=0.02		,      lowerCAmelCase__=3		,      lowerCAmelCase__=None		,      )   ->						Union[str, Any]:



        '''simple docstring'''




        a__	:					Union[str, Any]										=parent
        a__	:					List[Any]										=batch_size
        a__	:					List[Any]										=image_size
        a__	:					Any										=num_channels
        a__	:					Union[str, Any]										=num_encoder_blocks
        a__	:					List[str]										=sr_ratios
        a__	:					List[str]										=depths
        a__	:					Any										=hidden_sizes
        a__	:					List[str]										=downsampling_rates
        a__	:					List[str]										=num_attention_heads
        a__	:					Any										=is_training
        a__	:					int										=use_labels
        a__	:					Union[str, Any]										=hidden_act
        a__	:					List[Any]										=hidden_dropout_prob
        a__	:					Optional[int]										=attention_probs_dropout_prob
        a__	:					Optional[Any]										=initializer_range
        a__	:					List[Any]										=num_labels
        a__	:					Union[str, Any]										=scope







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



        '''simple docstring'''




        a__	:					str										=floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]							)

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

        a__	:					List[str]										=self.get_config()
        return config, pixel_values, labels







      def    _lowercase  (    self							)   ->						List[Any]:



        '''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		,      lowerCAmelCase__		,      lowerCAmelCase__		,      lowerCAmelCase__							)   ->						List[Any]:



        '''simple docstring'''




        a__	:					str										=SegformerModel(config=lowerCAmelCase__							)
        model.to(lowerCAmelCase__							)
        model.eval()
        a__	:					List[str]										=model(lowerCAmelCase__							)
        a__	:					List[Any]										=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		,      lowerCAmelCase__		,      lowerCAmelCase__		,      lowerCAmelCase__							)   ->						Any:



        '''simple docstring'''




        a__	:					Any										=self.num_labels
        a__	:					int										=SegformerForSemanticSegmentation(lowerCAmelCase__							)
        model.to(lowerCAmelCase__							)
        model.eval()
        a__	:					Tuple										=model(lowerCAmelCase__							)
        self.parent.assertEqual(
            result.logits.shape		,      (self.batch_size, self.num_labels, self.image_size // 4, self.image_size // 4)							)
        a__	:					Dict										=model(lowerCAmelCase__		,      labels=lowerCAmelCase__							)
        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		,      lowerCAmelCase__		,      lowerCAmelCase__		,      lowerCAmelCase__							)   ->						str:



        '''simple docstring'''




        a__	:					Union[str, Any]										=1
        a__	:					List[str]										=SegformerForSemanticSegmentation(config=lowerCAmelCase__							)
        model.to(lowerCAmelCase__							)
        model.eval()
        a__	:					Dict										=torch.randint(0		,      1		,      (self.batch_size, self.image_size, self.image_size)							).to(lowerCAmelCase__							)
        a__	:					List[Any]										=model(lowerCAmelCase__		,      labels=lowerCAmelCase__							)
        self.parent.assertGreater(result.loss		,      0.0							)







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



        '''simple docstring'''




        a__	:					int										=self.prepare_config_and_inputs()
        a__					,							a__					,							a__	:					List[str]										=config_and_inputs
        a__	:					Any										={"pixel_values": pixel_values}
        return config, inputs_dict







@require_torch
class 				__lowerCAmelCase       (    UpperCamelCase__						,				UpperCamelCase__						,				unittest.TestCase):
      _lowercase				:					int						=			(
          (
              SegformerModel,
              SegformerForSemanticSegmentation,
              SegformerForImageClassification,
          )
          if is_torch_available()
          else ()
      )
      _lowercase				:					str						=			(
          {
              """feature-extraction""": SegformerModel,
              """image-classification""": SegformerForImageClassification,
              """image-segmentation""": SegformerForSemanticSegmentation,
          }
          if is_torch_available()
          else {}
      )

      _lowercase				:					Dict						=			True
      _lowercase				:					List[str]						=			False
      _lowercase				:					List[str]						=			False
      _lowercase				:					Dict						=			False







      def    _lowercase  (    self							)   ->						Any:



        '''simple docstring'''




        a__	:					int										=SegformerModelTester(self							)
        a__	:					Optional[int]										=SegformerConfigTester(self		,      config_class=lowerCAmelCase__							)







      def    _lowercase  (    self							)   ->						str:



        '''simple docstring'''




        self.config_tester.run_common_tests()







      def    _lowercase  (    self							)   ->						int:



        '''simple docstring'''




        a__	:					str										=self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_model(*lowerCAmelCase__							)







      def    _lowercase  (    self							)   ->						int:



        '''simple docstring'''




        a__	:					str										=self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_for_binary_image_segmentation(*lowerCAmelCase__							)







      def    _lowercase  (    self							)   ->						List[Any]:



        '''simple docstring'''




        a__	:					Tuple										=self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_for_image_segmentation(*lowerCAmelCase__							)







      @unittest.skip("SegFormer does not use inputs_embeds"							)
      def    _lowercase  (    self							)   ->						int:



        '''simple docstring'''




        pass







      @unittest.skip("SegFormer does not have get_input_embeddings method and get_output_embeddings methods"							)
      def    _lowercase  (    self							)   ->						Union[str, Any]:



        '''simple docstring'''




        pass







      def    _lowercase  (    self							)   ->						List[str]:



        '''simple docstring'''




        a__					,							a__	:					Optional[Any]										=self.model_tester.prepare_config_and_inputs_for_common()

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

          a__	:					Any										=["pixel_values"]
          self.assertListEqual(arg_names[:1]		,      lowerCAmelCase__							)







      def    _lowercase  (    self							)   ->						Dict:



        '''simple docstring'''




        a__					,							a__	:					str										=self.model_tester.prepare_config_and_inputs_for_common()
        a__	:					Dict										=True

        for model_class in self.all_model_classes:
          a__	:					Tuple										=True
          a__	:					Tuple										=False
          a__	:					Dict										=True
          a__	:					Dict										=model_class(lowerCAmelCase__							)
          model.to(lowerCAmelCase__							)
          model.eval()
          with torch.no_grad():
            a__	:					List[Any]										=model(**self._prepare_for_class(lowerCAmelCase__		,      lowerCAmelCase__							)							)
          a__	:					Dict										=outputs.attentions

          a__	:					Optional[Any]										=sum(self.model_tester.depths							)
          self.assertEqual(len(lowerCAmelCase__							)		,      lowerCAmelCase__							)

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

          self.assertEqual(len(lowerCAmelCase__							)		,      lowerCAmelCase__							)

          # verify the first attentions (first block, first layer)
          a__	:					Union[str, Any]										=(self.model_tester.image_size // 4) ** 2
          a__	:					Tuple										=(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)
          a__	:					List[str]										=(self.model_tester.image_size // 3_2) ** 2
          a__	:					Dict										=(self.model_tester.image_size // (3_2 * 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]		,      )
          a__	:					List[Any]										=len(lowerCAmelCase__							)

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

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

          a__	:					List[str]										=outputs.attentions

          self.assertEqual(len(lowerCAmelCase__							)		,      lowerCAmelCase__							)
          # verify the first attentions (first block, first layer)
          a__	:					Any										=(self.model_tester.image_size // 4) ** 2
          a__	:					Union[str, Any]										=(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							)   ->						Union[str, Any]:



        '''simple docstring'''




        def check_hidden_states_output(lowerCAmelCase__		,      lowerCAmelCase__		,      lowerCAmelCase__							):
          a__	:					str										=model_class(lowerCAmelCase__							)
          model.to(lowerCAmelCase__							)
          model.eval()

          with torch.no_grad():
            a__	:					Any										=model(**self._prepare_for_class(lowerCAmelCase__		,      lowerCAmelCase__							)							)

          a__	:					Union[str, Any]										=outputs.hidden_states

          a__	:					Dict										=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,
              ]		,      )

        a__					,							a__	:					Optional[Any]										=self.model_tester.prepare_config_and_inputs_for_common()

        for model_class in self.all_model_classes:
          a__	:					List[Any]										=True
          check_hidden_states_output(lowerCAmelCase__		,      lowerCAmelCase__		,      lowerCAmelCase__							)

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

          check_hidden_states_output(lowerCAmelCase__		,      lowerCAmelCase__		,      lowerCAmelCase__							)







      def    _lowercase  (    self							)   ->						List[Any]:



        '''simple docstring'''




        if not self.model_tester.is_training:
          return

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

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

          a__	:					List[str]										=model_class(lowerCAmelCase__							)
          model.to(lowerCAmelCase__							)
          model.train()
          a__	:					List[str]										=self._prepare_for_class(lowerCAmelCase__		,      lowerCAmelCase__		,      return_labels=lowerCAmelCase__							)
          a__	:					str										=model(**lowerCAmelCase__							).loss
          loss.backward()







      @unittest.skip("Will be fixed soon by reducing the size of the model used for common tests."							)
      def    _lowercase  (    self							)   ->						Tuple:



        '''simple docstring'''




        pass







      @slow
      def    _lowercase  (    self							)   ->						List[str]:



        '''simple docstring'''




        for model_name in SEGFORMER_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
          a__	:					Optional[Any]										=SegformerModel.from_pretrained(lowerCAmelCase__							)
          self.assertIsNotNone(lowerCAmelCase__							)
def      _A       (		):





  """simple docstring"""



  a__	:					Union[str, Any]										=Image.open("./tests/fixtures/tests_samples/COCO/000000039769.png"						)
  return image







@require_torch
class 				__lowerCAmelCase       (    unittest.TestCase):







      @slow
      def    _lowercase  (    self							)   ->						Any:



        '''simple docstring'''




        a__	:					List[Any]										=SegformerImageProcessor(
            image_scale=(5_1_2, 5_1_2)		,      keep_ratio=lowerCAmelCase__		,      align=lowerCAmelCase__		,      do_random_crop=lowerCAmelCase__							)
        a__	:					List[str]										=SegformerForSemanticSegmentation.from_pretrained("nvidia/segformer-b0-finetuned-ade-512-512"							).to(
            lowerCAmelCase__							)

        a__	:					Optional[Any]										=prepare_img()
        a__	:					str										=image_processor(images=lowerCAmelCase__		,      return_tensors="pt"							)
        a__	:					Optional[Any]										=encoded_inputs.pixel_values.to(lowerCAmelCase__							)

        with torch.no_grad():
          a__	:					List[str]										=model(lowerCAmelCase__							)

        a__	:					Optional[int]										=torch.Size((1, model.config.num_labels, 1_2_8, 1_2_8)							)
        self.assertEqual(outputs.logits.shape		,      lowerCAmelCase__							)

        a__	:					int										=torch.tensor(
            [
                [[-4.63_10, -5.52_32, -6.23_56], [-5.19_21, -6.14_44, -6.59_96], [-5.44_24, -6.27_90, -6.75_74]],
                [[-12.13_91, -13.31_22, -13.95_54], [-12.87_32, -13.93_52, -14.35_63], [-12.94_38, -13.82_26, -14.25_13]],
                [[-12.51_34, -13.46_86, -14.49_15], [-12.86_69, -14.43_43, -14.77_58], [-13.25_23, -14.58_19, -15.06_94]],
            ]							).to(lowerCAmelCase__							)
        self.assertTrue(torch.allclose(outputs.logits[0, :3, :3, :3]		,      lowerCAmelCase__		,      atol=1E-4							)							)







      @slow
      def    _lowercase  (    self							)   ->						Optional[Any]:



        '''simple docstring'''




        a__	:					Optional[int]										=SegformerImageProcessor(
            image_scale=(5_1_2, 5_1_2)		,      keep_ratio=lowerCAmelCase__		,      align=lowerCAmelCase__		,      do_random_crop=lowerCAmelCase__							)
        a__	:					Dict										=SegformerForSemanticSegmentation.from_pretrained(
            "nvidia/segformer-b1-finetuned-cityscapes-1024-1024"							).to(lowerCAmelCase__							)

        a__	:					Dict										=prepare_img()
        a__	:					Dict										=image_processor(images=lowerCAmelCase__		,      return_tensors="pt"							)
        a__	:					Dict										=encoded_inputs.pixel_values.to(lowerCAmelCase__							)

        with torch.no_grad():
          a__	:					Dict										=model(lowerCAmelCase__							)

        a__	:					str										=torch.Size((1, model.config.num_labels, 1_2_8, 1_2_8)							)
        self.assertEqual(outputs.logits.shape		,      lowerCAmelCase__							)

        a__	:					Optional[int]										=torch.tensor(
            [
                [[-13.57_48, -13.91_11, -12.65_00], [-14.35_00, -15.36_83, -14.23_28], [-14.75_32, -16.04_24, -15.60_87]],
                [[-17.16_51, -15.87_25, -12.96_53], [-17.25_80, -17.37_18, -14.82_23], [-16.60_58, -16.87_83, -16.74_52]],
                [[-3.64_56, -3.02_09, -1.42_03], [-3.07_97, -3.19_59, -2.00_00], [-1.87_57, -1.92_17, -1.69_97]],
            ]							).to(lowerCAmelCase__							)
        self.assertTrue(torch.allclose(outputs.logits[0, :3, :3, :3]		,      lowerCAmelCase__		,      atol=1E-1							)							)







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



        '''simple docstring'''




        a__	:					List[str]										=SegformerImageProcessor(
            image_scale=(5_1_2, 5_1_2)		,      keep_ratio=lowerCAmelCase__		,      align=lowerCAmelCase__		,      do_random_crop=lowerCAmelCase__							)
        a__	:					Optional[Any]										=SegformerForSemanticSegmentation.from_pretrained("nvidia/segformer-b0-finetuned-ade-512-512"							).to(
            lowerCAmelCase__							)

        a__	:					Union[str, Any]										=prepare_img()
        a__	:					Optional[Any]										=image_processor(images=lowerCAmelCase__		,      return_tensors="pt"							)
        a__	:					List[str]										=encoded_inputs.pixel_values.to(lowerCAmelCase__							)

        with torch.no_grad():
          a__	:					Optional[Any]										=model(lowerCAmelCase__							)

        a__	:					int										=outputs.logits.detach().cpu()

        a__	:					Optional[Any]										=image_processor.post_process_semantic_segmentation(outputs=lowerCAmelCase__		,      target_sizes=[(5_0_0, 3_0_0)]							)
        a__	:					int										=torch.Size((5_0_0, 3_0_0)							)
        self.assertEqual(segmentation[0].shape		,      lowerCAmelCase__							)

        a__	:					Dict										=image_processor.post_process_semantic_segmentation(outputs=lowerCAmelCase__							)
        a__	:					int										=torch.Size((1_2_8, 1_2_8)							)
        self.assertEqual(segmentation[0].shape		,      lowerCAmelCase__							)













from diffusers.utils.testing_utils import require_onnxruntime







@require_onnxruntime
class 				__lowerCAmelCase       :
      pass






