





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       __A     (						_A    ):


    def _snake_case       (		self							):
        lowerCamelCase	 =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       __A     :


    def __init__(		self   , UpperCAmelCase_   , UpperCAmelCase_=13   , UpperCAmelCase_=64   , UpperCAmelCase_=3   , UpperCAmelCase_=4   , UpperCAmelCase_=[2, 2, 2, 2]   , UpperCAmelCase_=[8, 4, 2, 1]   , UpperCAmelCase_=[16, 32, 64, 128]   , UpperCAmelCase_=[1, 4, 8, 16]   , UpperCAmelCase_=[1, 2, 4, 8]   , UpperCAmelCase_=True   , UpperCAmelCase_=True   , UpperCAmelCase_="gelu"   , UpperCAmelCase_=0.1   , UpperCAmelCase_=0.1   , UpperCAmelCase_=0.0_2   , UpperCAmelCase_=3   , UpperCAmelCase_=None   , ):
        lowerCamelCase	 =parent
        lowerCamelCase	 =batch_size
        lowerCamelCase	 =image_size
        lowerCamelCase	 =num_channels
        lowerCamelCase	 =num_encoder_blocks
        lowerCamelCase	 =sr_ratios
        lowerCamelCase	 =depths
        lowerCamelCase	 =hidden_sizes
        lowerCamelCase	 =downsampling_rates
        lowerCamelCase	 =num_attention_heads
        lowerCamelCase	 =is_training
        lowerCamelCase	 =use_labels
        lowerCamelCase	 =hidden_act
        lowerCamelCase	 =hidden_dropout_prob
        lowerCamelCase	 =attention_probs_dropout_prob
        lowerCamelCase	 =initializer_range
        lowerCamelCase	 =num_labels
        lowerCamelCase	 =scope


    def _snake_case       (		self							):
        lowerCamelCase	 =floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]							)

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

        lowerCamelCase	 =self.get_config()
        return config, pixel_values, labels


    def _snake_case       (		self							):
        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 _snake_case       (		self   , UpperCAmelCase_   , UpperCAmelCase_   , UpperCAmelCase_							):
        lowerCamelCase	 =SegformerModel(config=__lowerCamelCase							)
        model.to(__lowerCamelCase							)
        model.eval()
        lowerCamelCase	 =model(__lowerCamelCase							)
        lowerCamelCase	 =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 _snake_case       (		self   , UpperCAmelCase_   , UpperCAmelCase_   , UpperCAmelCase_							):
        lowerCamelCase	 =self.num_labels
        lowerCamelCase	 =SegformerForSemanticSegmentation(__lowerCamelCase							)
        model.to(__lowerCamelCase							)
        model.eval()
        lowerCamelCase	 =model(__lowerCamelCase							)
        self.parent.assertEqual(
            result.logits.shape   , (self.batch_size, self.num_labels, self.image_size // 4, self.image_size // 4)							)
        lowerCamelCase	 =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 _snake_case       (		self   , UpperCAmelCase_   , UpperCAmelCase_   , UpperCAmelCase_							):
        lowerCamelCase	 =1
        lowerCamelCase	 =SegformerForSemanticSegmentation(config=__lowerCamelCase							)
        model.to(__lowerCamelCase							)
        model.eval()
        lowerCamelCase	 =torch.randint(0   , 1   , (self.batch_size, self.image_size, self.image_size)							).to(__lowerCamelCase							)
        lowerCamelCase	 =model(__lowerCamelCase   , labels=__lowerCamelCase							)
        self.parent.assertGreater(result.loss   , 0.0							)


    def _snake_case       (		self							):
        lowerCamelCase	 =self.prepare_config_and_inputs()
        lowerCamelCase	 =config_and_inputs
        lowerCamelCase	 ={"pixel_values": pixel_values}
        return config, inputs_dict




@require_torch
class       __A     (						_A    ,					_A    ,					unittest.TestCase    ):
    __A      =				(
        (
            SegformerModel,
            SegformerForSemanticSegmentation,
            SegformerForImageClassification,
        )
        if is_torch_available()
        else ()
    )
    __A      =				(
        {
            """feature-extraction""": SegformerModel,
            """image-classification""": SegformerForImageClassification,
            """image-segmentation""": SegformerForSemanticSegmentation,
        }
        if is_torch_available()
        else {}
    )

    __A      =				True
    __A      =				False
    __A      =				False
    __A      =				False


    def _snake_case       (		self							):
        lowerCamelCase	 =SegformerModelTester(self							)
        lowerCamelCase	 =SegformerConfigTester(self   , config_class=__lowerCamelCase							)


    def _snake_case       (		self							):
        self.config_tester.run_common_tests()


    def _snake_case       (		self							):
        lowerCamelCase	 =self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_model(*__lowerCamelCase							)


    def _snake_case       (		self							):
        lowerCamelCase	 =self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_for_binary_image_segmentation(*__lowerCamelCase							)


    def _snake_case       (		self							):
        lowerCamelCase	 =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 _snake_case       (		self							):
        pass


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


    def _snake_case       (		self							):
        lowerCamelCase	 =self.model_tester.prepare_config_and_inputs_for_common()

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

            lowerCamelCase	 =["pixel_values"]
            self.assertListEqual(arg_names[:1]   , __lowerCamelCase							)


    def _snake_case       (		self							):
        lowerCamelCase	 =self.model_tester.prepare_config_and_inputs_for_common()
        lowerCamelCase	 =True

        for model_class in self.all_model_classes:
            lowerCamelCase	 =True
            lowerCamelCase	 =False
            lowerCamelCase	 =True
            lowerCamelCase	 =model_class(__lowerCamelCase							)
            model.to(__lowerCamelCase							)
            model.eval()
            with torch.no_grad():
                lowerCamelCase	 =model(**self._prepare_for_class(__lowerCamelCase   , __lowerCamelCase							)							)
            lowerCamelCase	 =outputs.attentions

            lowerCamelCase	 =sum(self.model_tester.depths							)
            self.assertEqual(len(__lowerCamelCase							)   , __lowerCamelCase							)

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

            self.assertEqual(len(__lowerCamelCase							)   , __lowerCamelCase							)

            # verify the first attentions (first block, first layer)
            lowerCamelCase	 =(self.model_tester.image_size // 4) ** 2
            lowerCamelCase	 =(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)
            lowerCamelCase	 =(self.model_tester.image_size // 32) ** 2
            lowerCamelCase	 =(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]   , )
            lowerCamelCase	 =len(__lowerCamelCase							)

            # Check attention is always last and order is fine
            lowerCamelCase	 =True
            lowerCamelCase	 =True
            lowerCamelCase	 =model_class(__lowerCamelCase							)
            model.to(__lowerCamelCase							)
            model.eval()
            with torch.no_grad():
                lowerCamelCase	 =model(**self._prepare_for_class(__lowerCamelCase   , __lowerCamelCase							)							)

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

            lowerCamelCase	 =outputs.attentions

            self.assertEqual(len(__lowerCamelCase							)   , __lowerCamelCase							)
            # verify the first attentions (first block, first layer)
            lowerCamelCase	 =(self.model_tester.image_size // 4) ** 2
            lowerCamelCase	 =(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 _snake_case       (		self							):
        def check_hidden_states_output(UpperCAmelCase_   , UpperCAmelCase_   , UpperCAmelCase_							):
            lowerCamelCase	 =model_class(__lowerCamelCase							)
            model.to(__lowerCamelCase							)
            model.eval()

            with torch.no_grad():
                lowerCamelCase	 =model(**self._prepare_for_class(__lowerCamelCase   , __lowerCamelCase							)							)

            lowerCamelCase	 =outputs.hidden_states

            lowerCamelCase	 =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,
                ]   , )

        lowerCamelCase	 =self.model_tester.prepare_config_and_inputs_for_common()

        for model_class in self.all_model_classes:
            lowerCamelCase	 =True
            check_hidden_states_output(__lowerCamelCase   , __lowerCamelCase   , __lowerCamelCase							)

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

            check_hidden_states_output(__lowerCamelCase   , __lowerCamelCase   , __lowerCamelCase							)


    def _snake_case       (		self							):
        if not self.model_tester.is_training:
            return

        lowerCamelCase	 =self.model_tester.prepare_config_and_inputs_for_common()
        lowerCamelCase	 =True

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

            lowerCamelCase	 =model_class(__lowerCamelCase							)
            model.to(__lowerCamelCase							)
            model.train()
            lowerCamelCase	 =self._prepare_for_class(__lowerCamelCase   , __lowerCamelCase   , return_labels=__lowerCamelCase							)
            lowerCamelCase	 =model(**__lowerCamelCase							).loss
            loss.backward()


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


    @slow
    def _snake_case       (		self							):
        for model_name in SEGFORMER_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
            lowerCamelCase	 =SegformerModel.from_pretrained(__lowerCamelCase							)
            self.assertIsNotNone(__lowerCamelCase							)





def      _lowercase					(						)						->  Optional[Any]:
    lowerCamelCase	 =Image.open("""./tests/fixtures/tests_samples/COCO/000000039769.png"""				)
    return image




@require_torch
class       __A     (						unittest.TestCase    ):


    @slow
    def _snake_case       (		self							):
        # only resize + normalize
        lowerCamelCase	 =SegformerImageProcessor(
            image_scale=(512, 512)   , keep_ratio=__lowerCamelCase   , align=__lowerCamelCase   , do_random_crop=__lowerCamelCase							)
        lowerCamelCase	 =SegformerForSemanticSegmentation.from_pretrained("""nvidia/segformer-b0-finetuned-ade-512-512"""							).to(
            __lowerCamelCase							)

        lowerCamelCase	 =prepare_img()
        lowerCamelCase	 =image_processor(images=__lowerCamelCase   , return_tensors="""pt"""							)
        lowerCamelCase	 =encoded_inputs.pixel_values.to(__lowerCamelCase							)

        with torch.no_grad():
            lowerCamelCase	 =model(__lowerCamelCase							)

        lowerCamelCase	 =torch.Size((1, model.config.num_labels, 128, 128)							)
        self.assertEqual(outputs.logits.shape   , __lowerCamelCase							)

        lowerCamelCase	 =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(__lowerCamelCase							)
        self.assertTrue(torch.allclose(outputs.logits[0, :3, :3, :3]   , __lowerCamelCase   , atol=1E-4							)							)


    @slow
    def _snake_case       (		self							):
        # only resize + normalize
        lowerCamelCase	 =SegformerImageProcessor(
            image_scale=(512, 512)   , keep_ratio=__lowerCamelCase   , align=__lowerCamelCase   , do_random_crop=__lowerCamelCase							)
        lowerCamelCase	 =SegformerForSemanticSegmentation.from_pretrained(
            """nvidia/segformer-b1-finetuned-cityscapes-1024-1024"""							).to(__lowerCamelCase							)

        lowerCamelCase	 =prepare_img()
        lowerCamelCase	 =image_processor(images=__lowerCamelCase   , return_tensors="""pt"""							)
        lowerCamelCase	 =encoded_inputs.pixel_values.to(__lowerCamelCase							)

        with torch.no_grad():
            lowerCamelCase	 =model(__lowerCamelCase							)

        lowerCamelCase	 =torch.Size((1, model.config.num_labels, 128, 128)							)
        self.assertEqual(outputs.logits.shape   , __lowerCamelCase							)

        lowerCamelCase	 =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(__lowerCamelCase							)
        self.assertTrue(torch.allclose(outputs.logits[0, :3, :3, :3]   , __lowerCamelCase   , atol=1E-1							)							)


    @slow
    def _snake_case       (		self							):
        # only resize + normalize
        lowerCamelCase	 =SegformerImageProcessor(
            image_scale=(512, 512)   , keep_ratio=__lowerCamelCase   , align=__lowerCamelCase   , do_random_crop=__lowerCamelCase							)
        lowerCamelCase	 =SegformerForSemanticSegmentation.from_pretrained("""nvidia/segformer-b0-finetuned-ade-512-512"""							).to(
            __lowerCamelCase							)

        lowerCamelCase	 =prepare_img()
        lowerCamelCase	 =image_processor(images=__lowerCamelCase   , return_tensors="""pt"""							)
        lowerCamelCase	 =encoded_inputs.pixel_values.to(__lowerCamelCase							)

        with torch.no_grad():
            lowerCamelCase	 =model(__lowerCamelCase							)

        lowerCamelCase	 =outputs.logits.detach().cpu()

        lowerCamelCase	 =image_processor.post_process_semantic_segmentation(outputs=__lowerCamelCase   , target_sizes=[(500, 300)]							)
        lowerCamelCase	 =torch.Size((500, 300)							)
        self.assertEqual(segmentation[0].shape   , __lowerCamelCase							)

        lowerCamelCase	 =image_processor.post_process_semantic_segmentation(outputs=__lowerCamelCase							)
        lowerCamelCase	 =torch.Size((128, 128)							)
        self.assertEqual(segmentation[0].shape   , __lowerCamelCase							)












from abc import ABC, abstractmethod
from argparse import ArgumentParser

class       __A     (						a    ):


				@staticmethod
				@abstractmethod
				def _snake_case       (		UpperCAmelCase_							):
								raise NotImplementedError()


				@abstractmethod
				def _snake_case       (		self							):
								raise NotImplementedError()





