






"""simple docstring"""







import inspect
import unittest

import numpy as np

from tests.test_modeling_common import floats_tensor
from transformers import MaskaFormerConfig, is_torch_available, is_vision_available
from transformers.testing_utils import require_torch, require_torch_multi_gpu, require_vision, slow, torch_device
from transformers.utils import cached_property

from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
  import torch

  from transformers import MaskaFormerForUniversalSegmentation, MaskaFormerModel

  if is_vision_available():
    from transformers import MaskaFormerImageProcessor

if is_vision_available():
  from PIL import Image







class    lowerCamelCase__   :






    def __init__(					self      ,A      ,A=2      ,A=True      ,A=False      ,A=10      ,A=3      ,A=32 * 8      ,A=32 * 8      ,A=4      ,A=64      ,):
          UpperCAmelCase							     =			parent
          UpperCAmelCase							     =			batch_size
          UpperCAmelCase							     =			is_training
          UpperCAmelCase							     =			use_auxiliary_loss
          UpperCAmelCase							     =			num_queries
          UpperCAmelCase							     =			num_channels
          UpperCAmelCase							     =			min_size
          UpperCAmelCase							     =			max_size
          UpperCAmelCase							     =			num_labels
          UpperCAmelCase							     =			hidden_dim
          UpperCAmelCase							     =			hidden_dim






    def 							_UpperCamelCase				(					self   ):
          UpperCAmelCase							     =			floats_tensor([self.batch_size, self.num_channels, self.min_size, self.max_size]   ).to(
              __lowerCAmelCase   )

          UpperCAmelCase							     =			torch.ones([self.batch_size, self.min_size, self.max_size]      ,device=__lowerCAmelCase   )

          UpperCAmelCase							     =			(
              torch.rand([self.batch_size, self.num_labels, self.min_size, self.max_size]      ,device=__lowerCAmelCase   ) > 0.5
          ).float()
          UpperCAmelCase							     =			(torch.rand((self.batch_size, self.num_labels)      ,device=__lowerCAmelCase   ) > 0.5).long()

          UpperCAmelCase							     =			self.get_config()
          return config, pixel_values, pixel_mask, mask_labels, class_labels






    def 							_UpperCamelCase				(					self   ):
          UpperCAmelCase							     =			MaskaFormerConfig(
              hidden_size=self.hidden_dim      ,)
          UpperCAmelCase							     =			self.num_queries
          UpperCAmelCase							     =			self.num_labels

          UpperCAmelCase							     =			[1, 1, 1, 1]
          UpperCAmelCase							     =			self.num_channels

          UpperCAmelCase							     =			64
          UpperCAmelCase							     =			128
          UpperCAmelCase							     =			self.hidden_dim
          UpperCAmelCase							     =			self.hidden_dim
          UpperCAmelCase							     =			self.hidden_dim
          return config






    def 							_UpperCamelCase				(					self   ):
          UpperCAmelCase     ,   UpperCAmelCase     ,   UpperCAmelCase     ,   UpperCAmelCase     ,   UpperCAmelCase							     =			self.prepare_config_and_inputs()
          UpperCAmelCase							     =			{"""pixel_values""": pixel_values, """pixel_mask""": pixel_mask}
          return config, inputs_dict






    def 							_UpperCamelCase				(					self      ,A      ,A   ):
          UpperCAmelCase							     =			output.encoder_hidden_states
          UpperCAmelCase							     =			output.pixel_decoder_hidden_states
          UpperCAmelCase							     =			output.transformer_decoder_hidden_states

          self.parent.assertTrue(len(__lowerCAmelCase   )      ,len(config.backbone_config.depths   )   )
          self.parent.assertTrue(len(__lowerCAmelCase   )      ,len(config.backbone_config.depths   )   )
          self.parent.assertTrue(len(__lowerCAmelCase   )      ,config.decoder_layers   )






    def 							_UpperCamelCase				(					self      ,A      ,A      ,A      ,A=False   ):


          with torch.no_grad():
                UpperCAmelCase							     =			MaskaFormerModel(config=__lowerCAmelCase   )
                model.to(__lowerCAmelCase   )
                model.eval()

                UpperCAmelCase							     =			model(pixel_values=__lowerCAmelCase      ,pixel_mask=__lowerCAmelCase   )
                UpperCAmelCase							     =			model(__lowerCAmelCase      ,output_hidden_states=__lowerCAmelCase   )

          self.parent.assertEqual(
              output.transformer_decoder_last_hidden_state.shape      ,(self.batch_size, self.num_queries, self.hidden_dim)      ,)
          # let's ensure the other two hidden state exists
          self.parent.assertTrue(output.pixel_decoder_last_hidden_state is not None   )
          self.parent.assertTrue(output.encoder_last_hidden_state is not None   )

          if output_hidden_states:
                self.check_output_hidden_state(__lowerCAmelCase      ,__lowerCAmelCase   )






    def 							_UpperCamelCase				(					self      ,A      ,A      ,A      ,A      ,A   ):
          UpperCAmelCase							     =			MaskaFormerForUniversalSegmentation(config=__lowerCAmelCase   )
          model.to(__lowerCAmelCase   )
          model.eval()

          def comm_check_on_output(A   ):
                # let's still check that all the required stuff is there
                self.parent.assertTrue(result.transformer_decoder_last_hidden_state is not None   )
                self.parent.assertTrue(result.pixel_decoder_last_hidden_state is not None   )
                self.parent.assertTrue(result.encoder_last_hidden_state is not None   )
                # okay, now we need to check the logits shape
                # due to the encoder compression, masks have a //4 spatial size
                self.parent.assertEqual(
                    result.masks_queries_logits.shape      ,(self.batch_size, self.num_queries, self.min_size // 4, self.max_size // 4)      ,)
                # + 1 for null class
                self.parent.assertEqual(
                    result.class_queries_logits.shape      ,(self.batch_size, self.num_queries, self.num_labels + 1)   )

          with torch.no_grad():
                UpperCAmelCase							     =			model(pixel_values=__lowerCAmelCase      ,pixel_mask=__lowerCAmelCase   )
                UpperCAmelCase							     =			model(__lowerCAmelCase   )

                comm_check_on_output(__lowerCAmelCase   )

                UpperCAmelCase							     =			model(
                    pixel_values=__lowerCAmelCase      ,pixel_mask=__lowerCAmelCase      ,mask_labels=__lowerCAmelCase      ,class_labels=__lowerCAmelCase   )

          comm_check_on_output(__lowerCAmelCase   )

          self.parent.assertTrue(result.loss is not None   )
          self.parent.assertEqual(result.loss.shape      ,torch.Size([1]   )   )







@require_torch
class    lowerCamelCase__   (			lowerCAmelCase__       ,	lowerCAmelCase__       ,	unittest.TestCase					):
    SCREAMING_SNAKE_CASE						=     (MaskaFormerModel, MaskaFormerForUniversalSegmentation) if is_torch_available() else ()
    SCREAMING_SNAKE_CASE						=     {'''feature-extraction''': MaskaFormerModel} if is_torch_available() else {}

    SCREAMING_SNAKE_CASE						=     False
    SCREAMING_SNAKE_CASE						=     False
    SCREAMING_SNAKE_CASE						=     False
    SCREAMING_SNAKE_CASE						=     False






    def 							_UpperCamelCase				(					self   ):
          UpperCAmelCase							     =			MaskaFormerModelTester(self   )
          UpperCAmelCase							     =			ConfigTester(self      ,config_class=__lowerCAmelCase      ,has_text_modality=__lowerCAmelCase   )






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






    def 							_UpperCamelCase				(					self   ):
          UpperCAmelCase     ,   UpperCAmelCase							     =			self.model_tester.prepare_config_and_inputs_for_common()
          self.model_tester.create_and_check_maskaformer_model(__lowerCAmelCase      ,**__lowerCAmelCase      ,output_hidden_states=__lowerCAmelCase   )






    def 							_UpperCamelCase				(					self   ):
          UpperCAmelCase							     =			self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_maskaformer_instance_segmentation_head_model(*__lowerCAmelCase   )






    @unittest.skip(reason="""Mask2Former does not use inputs_embeds"""   )
    def 							_UpperCamelCase				(					self   ):
          pass






    @unittest.skip(reason="""Mask2Former does not have a get_input_embeddings method"""   )
    def 							_UpperCamelCase				(					self   ):
          pass






    @unittest.skip(reason="""Mask2Former is not a generative model"""   )
    def 							_UpperCamelCase				(					self   ):
          pass






    @unittest.skip(reason="""Mask2Former does not use token embeddings"""   )
    def 							_UpperCamelCase				(					self   ):
          pass






    @require_torch_multi_gpu
    @unittest.skip(
        reason="""Mask2Former has some layers using `add_module` which doesn't work well with `nn.DataParallel`"""   )
    def 							_UpperCamelCase				(					self   ):
          pass






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






    def 							_UpperCamelCase				(					self   ):
          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 							_UpperCamelCase				(					self   ):


          for model_name in ["facebook/mask2former-swin-small-coco-instance"]:
                UpperCAmelCase							     =			MaskaFormerModel.from_pretrained(__lowerCAmelCase   )
                self.assertIsNotNone(__lowerCAmelCase   )






    def 							_UpperCamelCase				(					self   ):
          UpperCAmelCase							     =			(self.model_tester.min_size,) * 2
          UpperCAmelCase							     =			{
              """pixel_values""": torch.randn((2, 3, *size)      ,device=__lowerCAmelCase   ),
              """mask_labels""": torch.randn((2, 10, *size)      ,device=__lowerCAmelCase   ),
              """class_labels""": torch.zeros(2      ,10      ,device=__lowerCAmelCase   ).long(),
          }
          UpperCAmelCase							     =			self.model_tester.get_config()

          UpperCAmelCase							     =			MaskaFormerForUniversalSegmentation(__lowerCAmelCase   ).to(__lowerCAmelCase   )
          UpperCAmelCase							     =			model(**__lowerCAmelCase   )
          self.assertTrue(outputs.loss is not None   )






    def 							_UpperCamelCase				(					self   ):
          UpperCAmelCase     ,   UpperCAmelCase							     =			self.model_tester.prepare_config_and_inputs_for_common()
          self.model_tester.create_and_check_maskaformer_model(__lowerCAmelCase      ,**__lowerCAmelCase      ,output_hidden_states=__lowerCAmelCase   )






    def 							_UpperCamelCase				(					self   ):
          UpperCAmelCase     ,   UpperCAmelCase							     =			self.model_tester.prepare_config_and_inputs_for_common()

          for model_class in self.all_model_classes:
                UpperCAmelCase							     =			model_class(__lowerCAmelCase   ).to(__lowerCAmelCase   )
                UpperCAmelCase							     =			model(**__lowerCAmelCase      ,output_attentions=__lowerCAmelCase   )
                self.assertTrue(outputs.attentions is not None   )






    def 							_UpperCamelCase				(					self   ):


          if not self.model_tester.is_training:
                return

          UpperCAmelCase							     =			self.all_model_classes[1]
          UpperCAmelCase     ,   UpperCAmelCase     ,   UpperCAmelCase     ,   UpperCAmelCase     ,   UpperCAmelCase							     =			self.model_tester.prepare_config_and_inputs()

          UpperCAmelCase							     =			model_class(__lowerCAmelCase   )
          model.to(__lowerCAmelCase   )
          model.train()

          UpperCAmelCase							     =			model(__lowerCAmelCase      ,mask_labels=__lowerCAmelCase      ,class_labels=__lowerCAmelCase   ).loss
          loss.backward()






    def 							_UpperCamelCase				(					self   ):
          UpperCAmelCase							     =			self.all_model_classes[1]
          UpperCAmelCase     ,   UpperCAmelCase     ,   UpperCAmelCase     ,   UpperCAmelCase     ,   UpperCAmelCase							     =			self.model_tester.prepare_config_and_inputs()
          UpperCAmelCase							     =			True
          UpperCAmelCase							     =			True

          UpperCAmelCase							     =			model_class(__lowerCAmelCase   ).to(__lowerCAmelCase   )
          model.train()

          UpperCAmelCase							     =			model(__lowerCAmelCase      ,mask_labels=__lowerCAmelCase      ,class_labels=__lowerCAmelCase   )

          UpperCAmelCase							     =			outputs.encoder_hidden_states[0]
          encoder_hidden_states.retain_grad()

          UpperCAmelCase							     =			outputs.pixel_decoder_hidden_states[0]
          pixel_decoder_hidden_states.retain_grad()

          UpperCAmelCase							     =			outputs.transformer_decoder_hidden_states[0]
          transformer_decoder_hidden_states.retain_grad()

          UpperCAmelCase							     =			outputs.attentions[0]
          attentions.retain_grad()

          outputs.loss.backward(retain_graph=__lowerCAmelCase   )

          self.assertIsNotNone(encoder_hidden_states.grad   )
          self.assertIsNotNone(pixel_decoder_hidden_states.grad   )
          self.assertIsNotNone(transformer_decoder_hidden_states.grad   )
          self.assertIsNotNone(attentions.grad   )







_UpperCamelCase	     =    1E-4

def _a					(						):







      """simple docstring"""






      UpperCAmelCase							     =			Image.open("""./tests/fixtures/tests_samples/COCO/000000039769.png"""				)
      return image







@require_vision
@slow
class    lowerCamelCase__   (			unittest.TestCase					):






    @cached_property
    def 							_UpperCamelCase				(					self   ):
          return "facebook/mask2former-swin-small-coco-instance"






    @cached_property
    def 							_UpperCamelCase				(					self   ):
          return MaskaFormerImageProcessor.from_pretrained(self.model_checkpoints   ) if is_vision_available() else None






    def 							_UpperCamelCase				(					self   ):
          UpperCAmelCase							     =			MaskaFormerModel.from_pretrained(self.model_checkpoints   ).to(__lowerCAmelCase   )
          UpperCAmelCase							     =			self.default_image_processor
          UpperCAmelCase							     =			prepare_img()
          UpperCAmelCase							     =			image_processor(__lowerCAmelCase      ,return_tensors="""pt"""   ).to(__lowerCAmelCase   )
          UpperCAmelCase							     =			inputs["""pixel_values"""].shape
          # check size is divisible by 32
          self.assertTrue((inputs_shape[-1] % 32) == 0 and (inputs_shape[-2] % 32) == 0   )
          # check size
          self.assertEqual(__lowerCAmelCase      ,(1, 3, 384, 384)   )

          with torch.no_grad():
                UpperCAmelCase							     =			model(**__lowerCAmelCase   )

          UpperCAmelCase							     =			torch.tensor(
              [[-0.2790, -1.0717, -1.1668], [-0.5128, -0.3128, -0.4987], [-0.5832, 0.1971, -0.0197]]   ).to(__lowerCAmelCase   )
          self.assertTrue(
              torch.allclose(
                  outputs.encoder_last_hidden_state[0, 0, :3, :3]      ,__lowerCAmelCase      ,atol=__lowerCAmelCase   )   )

          UpperCAmelCase							     =			torch.tensor(
              [[0.8973, 1.1847, 1.1776], [1.1934, 1.5040, 1.5128], [1.1153, 1.4486, 1.4951]]   ).to(__lowerCAmelCase   )
          self.assertTrue(
              torch.allclose(
                  outputs.pixel_decoder_last_hidden_state[0, 0, :3, :3]      ,__lowerCAmelCase      ,atol=__lowerCAmelCase   )   )

          UpperCAmelCase							     =			torch.tensor(
              [[2.1152, 1.7000, -0.8603], [1.5808, 1.8004, -0.9353], [1.6043, 1.7495, -0.5999]]   ).to(__lowerCAmelCase   )
          self.assertTrue(
              torch.allclose(
                  outputs.transformer_decoder_last_hidden_state[0, :3, :3]      ,__lowerCAmelCase      ,atol=__lowerCAmelCase   )   )






    def 							_UpperCamelCase				(					self   ):
          UpperCAmelCase							     =			MaskaFormerForUniversalSegmentation.from_pretrained(self.model_checkpoints   ).to(__lowerCAmelCase   ).eval()
          UpperCAmelCase							     =			self.default_image_processor
          UpperCAmelCase							     =			prepare_img()
          UpperCAmelCase							     =			image_processor(__lowerCAmelCase      ,return_tensors="""pt"""   ).to(__lowerCAmelCase   )
          UpperCAmelCase							     =			inputs["""pixel_values"""].shape
          # check size is divisible by 32
          self.assertTrue((inputs_shape[-1] % 32) == 0 and (inputs_shape[-2] % 32) == 0   )
          # check size
          self.assertEqual(__lowerCAmelCase      ,(1, 3, 384, 384)   )

          with torch.no_grad():
                UpperCAmelCase							     =			model(**__lowerCAmelCase   )
          # masks_queries_logits
          UpperCAmelCase							     =			outputs.masks_queries_logits
          self.assertEqual(
              masks_queries_logits.shape      ,(1, model.config.num_queries, inputs_shape[-2] // 4, inputs_shape[-1] // 4)   )
          UpperCAmelCase							     =			[
              [-8.7839, -9.0056, -8.8121],
              [-7.4104, -7.0313, -6.5401],
              [-6.6105, -6.3427, -6.4675],
          ]
          UpperCAmelCase							     =			torch.tensor(__lowerCAmelCase   ).to(__lowerCAmelCase   )
          self.assertTrue(torch.allclose(masks_queries_logits[0, 0, :3, :3]      ,__lowerCAmelCase      ,atol=__lowerCAmelCase   )   )
          # class_queries_logits
          UpperCAmelCase							     =			outputs.class_queries_logits
          self.assertEqual(class_queries_logits.shape      ,(1, model.config.num_queries, model.config.num_labels + 1)   )
          UpperCAmelCase							     =			torch.tensor(
              [
                  [1.8324, -8.0835, -4.1922],
                  [0.8450, -9.0050, -3.6053],
                  [0.3045, -7.7293, -3.0275],
              ]   ).to(__lowerCAmelCase   )
          self.assertTrue(torch.allclose(outputs.class_queries_logits[0, :3, :3]      ,__lowerCAmelCase      ,atol=__lowerCAmelCase   )   )






    def 							_UpperCamelCase				(					self   ):
          UpperCAmelCase							     =			MaskaFormerForUniversalSegmentation.from_pretrained(self.model_checkpoints   ).to(__lowerCAmelCase   ).eval()
          UpperCAmelCase							     =			self.default_image_processor

          UpperCAmelCase							     =			image_processor(
              [np.zeros((3, 800, 1_333)   ), np.zeros((3, 800, 1_333)   )]      ,segmentation_maps=[np.zeros((384, 384)   ).astype(np.floataa   ), np.zeros((384, 384)   ).astype(np.floataa   )]      ,return_tensors="""pt"""      ,)

          UpperCAmelCase							     =			inputs["""pixel_values"""].to(__lowerCAmelCase   )
          UpperCAmelCase							     =			[el.to(__lowerCAmelCase   ) for el in inputs["""mask_labels"""]]
          UpperCAmelCase							     =			[el.to(__lowerCAmelCase   ) for el in inputs["""class_labels"""]]

          with torch.no_grad():
                UpperCAmelCase							     =			model(**__lowerCAmelCase   )

          self.assertTrue(outputs.loss is not None   )




'''simple docstring'''




import warnings
from typing import List, Optional, Union

from ...processing_utils import ProcessorMixin
from ...tokenization_utils_base import BatchEncoding, PaddingStrategy, PreTokenizedInput, TextInput, TruncationStrategy
from ...utils import TensorType






class 							a__(   lowerCAmelCase__      ):


			'''simple docstring'''






			UpperCAmelCase_     :	str						=			['''image_processor''', '''tokenizer''']
			UpperCAmelCase_     :	str						=			'''ViltImageProcessor'''
			UpperCAmelCase_     :	Any						=			('''BertTokenizer''', '''BertTokenizerFast''')



			def __init__(  self			,       __lowerCAmelCase=None			,       __lowerCAmelCase=None			,       **__lowerCAmelCase):




					"""simple docstring"""


					lowerCAmelCase     =    None
					if "feature_extractor" in kwargs:
							warnings.warn(
							    """The `feature_extractor` argument is deprecated and will be removed in v5, use `image_processor`"""
							    """ instead."""			,       __lowerCAmelCase			,       )
							lowerCAmelCase     =    kwargs.pop("""feature_extractor""")

					lowerCAmelCase     =    image_processor if image_processor is not None else feature_extractor
					if image_processor is None:
							raise ValueError("""You need to specify an `image_processor`.""")
					if tokenizer is None:
							raise ValueError("""You need to specify a `tokenizer`.""")

					super().__init__(__lowerCAmelCase			,       __lowerCAmelCase)
					lowerCAmelCase     =    self.image_processor



			def __call__(  self			,       __lowerCAmelCase			,       __lowerCAmelCase = None			,       __lowerCAmelCase = True			,       __lowerCAmelCase = False			,       __lowerCAmelCase = None			,       __lowerCAmelCase = None			,       __lowerCAmelCase = 0			,       __lowerCAmelCase = None			,       __lowerCAmelCase = None			,       __lowerCAmelCase = None			,       __lowerCAmelCase = False			,       __lowerCAmelCase = False			,       __lowerCAmelCase = False			,       __lowerCAmelCase = False			,       __lowerCAmelCase = True			,       __lowerCAmelCase = None			,       **__lowerCAmelCase			,       ):




					"""simple docstring"""


					lowerCAmelCase     =    self.tokenizer(
					    text=__lowerCAmelCase			,       add_special_tokens=__lowerCAmelCase			,       padding=__lowerCAmelCase			,       truncation=__lowerCAmelCase			,       max_length=__lowerCAmelCase			,       stride=__lowerCAmelCase			,       pad_to_multiple_of=__lowerCAmelCase			,       return_token_type_ids=__lowerCAmelCase			,       return_attention_mask=__lowerCAmelCase			,       return_overflowing_tokens=__lowerCAmelCase			,       return_special_tokens_mask=__lowerCAmelCase			,       return_offsets_mapping=__lowerCAmelCase			,       return_length=__lowerCAmelCase			,       verbose=__lowerCAmelCase			,       return_tensors=__lowerCAmelCase			,       **__lowerCAmelCase			,       )
					# add pixel_values + pixel_mask
					lowerCAmelCase     =    self.image_processor(__lowerCAmelCase			,       return_tensors=__lowerCAmelCase)
					encoding.update(__lowerCAmelCase)

					return encoding



			def 				a_			(  self			,       *__lowerCAmelCase			,       **__lowerCAmelCase):




					"""simple docstring"""


					return self.tokenizer.batch_decode(*__lowerCAmelCase			,       **__lowerCAmelCase)



			def 				a_			(  self			,       *__lowerCAmelCase			,       **__lowerCAmelCase):




					"""simple docstring"""


					return self.tokenizer.decode(*__lowerCAmelCase			,       **__lowerCAmelCase)



			@property
			def 				a_			(  self):




					"""simple docstring"""


					lowerCAmelCase     =    self.tokenizer.model_input_names
					lowerCAmelCase     =    self.image_processor.model_input_names
					return list(dict.fromkeys(tokenizer_input_names + image_processor_input_names))



			@property
			def 				a_			(  self):




					"""simple docstring"""


					warnings.warn(
					    """`feature_extractor_class` is deprecated and will be removed in v5. Use `image_processor_class` instead."""			,       __lowerCAmelCase			,       )
					return self.image_processor_class







			@property
			def 				a_			(  self):




					"""simple docstring"""


					warnings.warn(
					    """`feature_extractor` is deprecated and will be removed in v5. Use `image_processor` instead."""			,       __lowerCAmelCase			,       )
					return self.image_processor




