




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     SCREAMING_SNAKE_CASE__		:
     def __init__(						self	:		int       ,						__lowerCamelCase	:		Union[str, Any]       ,						__lowerCamelCase	:		Dict=2       ,						__lowerCamelCase	:		Union[str, Any]=True       ,						__lowerCamelCase	:		int=False       ,						__lowerCamelCase	:		Tuple=10       ,						__lowerCamelCase	:		Optional[int]=3       ,						__lowerCamelCase	:		Dict=32 * 8       ,						__lowerCamelCase	:		Tuple=32 * 8       ,						__lowerCamelCase	:		Dict=4       ,						__lowerCamelCase	:		Any=64       ,						):



       """simple docstring"""
       lowerCAmelCase__										=      parent
       lowerCAmelCase__										=      batch_size
       lowerCAmelCase__										=      is_training
       lowerCAmelCase__										=      use_auxiliary_loss
       lowerCAmelCase__										=      num_queries
       lowerCAmelCase__										=      num_channels
       lowerCAmelCase__										=      min_size
       lowerCAmelCase__										=      max_size
       lowerCAmelCase__										=      num_labels
       lowerCAmelCase__										=      hidden_dim
       lowerCAmelCase__										=      hidden_dim
     def        A__      (						self	:		List[Any]						):



       """simple docstring"""
       lowerCAmelCase__										=      floats_tensor([self.batch_size, self.num_channels, self.min_size, self.max_size]						).to(
           __lowerCamelCase						)

       lowerCAmelCase__										=      torch.ones([self.batch_size, self.min_size, self.max_size]       ,						device=__lowerCamelCase						)

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

       lowerCAmelCase__										=      self.get_config()
       return config, pixel_values, pixel_mask, mask_labels, class_labels
     def        A__      (						self	:		Optional[Any]						):



       """simple docstring"""
       lowerCAmelCase__										=      MaskaFormerConfig(
           hidden_size=self.hidden_dim       ,						)
       lowerCAmelCase__										=      self.num_queries
       lowerCAmelCase__										=      self.num_labels

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

       lowerCAmelCase__										=      64
       lowerCAmelCase__										=      1_28
       lowerCAmelCase__										=      self.hidden_dim
       lowerCAmelCase__										=      self.hidden_dim
       lowerCAmelCase__										=      self.hidden_dim
       return config
     def        A__      (						self	:		Optional[Any]						):



       """simple docstring"""
       lowerCAmelCase__				,      lowerCAmelCase__				,      lowerCAmelCase__				,      lowerCAmelCase__				,      lowerCAmelCase__										=      self.prepare_config_and_inputs()
       lowerCAmelCase__										=      {'''pixel_values''': pixel_values, '''pixel_mask''': pixel_mask}
       return config, inputs_dict
     def        A__      (						self	:		Any       ,						__lowerCamelCase	:		Dict       ,						__lowerCamelCase	:		List[str]						):



       """simple docstring"""
       lowerCAmelCase__										=      output.encoder_hidden_states
       lowerCAmelCase__										=      output.pixel_decoder_hidden_states
       lowerCAmelCase__										=      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        A__      (						self	:		Optional[int]       ,						__lowerCamelCase	:		Optional[Any]       ,						__lowerCamelCase	:		List[Any]       ,						__lowerCamelCase	:		Tuple       ,						__lowerCamelCase	:		Dict=False						):



       """simple docstring"""
       with torch.no_grad():
         lowerCAmelCase__										=      MaskaFormerModel(config=__lowerCamelCase						)
         model.to(__lowerCamelCase						)
         model.eval()

         lowerCAmelCase__										=      model(pixel_values=__lowerCamelCase       ,						pixel_mask=__lowerCamelCase						)
         lowerCAmelCase__										=      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        A__      (						self	:		int       ,						__lowerCamelCase	:		List[str]       ,						__lowerCamelCase	:		Any       ,						__lowerCamelCase	:		Optional[int]       ,						__lowerCamelCase	:		Optional[Any]       ,						__lowerCamelCase	:		Union[str, Any]						):



       """simple docstring"""
       lowerCAmelCase__										=      MaskaFormerForUniversalSegmentation(config=__lowerCamelCase						)
       model.to(__lowerCamelCase						)
       model.eval()

       def comm_check_on_output(__lowerCamelCase	:		List[str]						):
         # 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():
         lowerCAmelCase__										=      model(pixel_values=__lowerCamelCase       ,						pixel_mask=__lowerCamelCase						)
         lowerCAmelCase__										=      model(__lowerCamelCase						)

         comm_check_on_output(__lowerCamelCase						)

         lowerCAmelCase__										=      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     SCREAMING_SNAKE_CASE__		(_a    ,					_a    ,					unittest.TestCase			):
     lowercase_ :      Optional[int]				= (MaskaFormerModel, MaskaFormerForUniversalSegmentation) if is_torch_available() else ()
     lowercase_ :      Optional[Any]				= {"feature-extraction": MaskaFormerModel} if is_torch_available() else {}

     lowercase_ :      Dict				= False
     lowercase_ :      Union[str, Any]				= False
     lowercase_ :      Optional[Any]				= False
     lowercase_ :      List[str]				= False
     def        A__      (						self	:		Dict						):



       """simple docstring"""
       lowerCAmelCase__										=      MaskaFormerModelTester(self						)
       lowerCAmelCase__										=      ConfigTester(self       ,						config_class=__lowerCamelCase       ,						has_text_modality=__lowerCamelCase						)
     def        A__      (						self	:		Optional[Any]						):



       """simple docstring"""
       self.config_tester.run_common_tests()
     def        A__      (						self	:		Optional[Any]						):



       """simple docstring"""
       lowerCAmelCase__				,      lowerCAmelCase__										=      self.model_tester.prepare_config_and_inputs_for_common()
       self.model_tester.create_and_check_maskaformer_model(__lowerCamelCase       ,						**__lowerCamelCase       ,						output_hidden_states=__lowerCamelCase						)
     def        A__      (						self	:		Tuple						):



       """simple docstring"""
       lowerCAmelCase__										=      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        A__      (						self	:		List[Any]						):



       """simple docstring"""
       pass
     @unittest.skip(reason='''Mask2Former does not have a get_input_embeddings method'''						)
     def        A__      (						self	:		Tuple						):



       """simple docstring"""
       pass
     @unittest.skip(reason='''Mask2Former is not a generative model'''						)
     def        A__      (						self	:		Optional[int]						):



       """simple docstring"""
       pass
     @unittest.skip(reason='''Mask2Former does not use token embeddings'''						)
     def        A__      (						self	:		Union[str, Any]						):



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



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



       """simple docstring"""
       pass
     def        A__      (						self	:		int						):



       """simple docstring"""
       lowerCAmelCase__				,      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						)
     @slow
     def        A__      (						self	:		List[Any]						):



       """simple docstring"""
       for model_name in ["facebook/mask2former-swin-small-coco-instance"]:
         lowerCAmelCase__										=      MaskaFormerModel.from_pretrained(__lowerCamelCase						)
         self.assertIsNotNone(__lowerCamelCase						)
     def        A__      (						self	:		Any						):



       """simple docstring"""
       lowerCAmelCase__										=      (self.model_tester.min_size,) * 2
       lowerCAmelCase__										=      {
           '''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(),
       }
       lowerCAmelCase__										=      self.model_tester.get_config()

       lowerCAmelCase__										=      MaskaFormerForUniversalSegmentation(__lowerCamelCase						).to(__lowerCamelCase						)
       lowerCAmelCase__										=      model(**__lowerCamelCase						)
       self.assertTrue(outputs.loss is not None						)
     def        A__      (						self	:		Optional[int]						):



       """simple docstring"""
       lowerCAmelCase__				,      lowerCAmelCase__										=      self.model_tester.prepare_config_and_inputs_for_common()
       self.model_tester.create_and_check_maskaformer_model(__lowerCamelCase       ,						**__lowerCamelCase       ,						output_hidden_states=__lowerCamelCase						)
     def        A__      (						self	:		str						):



       """simple docstring"""
       lowerCAmelCase__				,      lowerCAmelCase__										=      self.model_tester.prepare_config_and_inputs_for_common()

       for model_class in self.all_model_classes:
         lowerCAmelCase__										=      model_class(__lowerCamelCase						).to(__lowerCamelCase						)
         lowerCAmelCase__										=      model(**__lowerCamelCase       ,						output_attentions=__lowerCamelCase						)
         self.assertTrue(outputs.attentions is not None						)
     def        A__      (						self	:		str						):



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

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

       lowerCAmelCase__										=      model_class(__lowerCamelCase						)
       model.to(__lowerCamelCase						)
       model.train()

       lowerCAmelCase__										=      model(__lowerCamelCase       ,						mask_labels=__lowerCamelCase       ,						class_labels=__lowerCamelCase						).loss
       loss.backward()



     def        A__      (						self	:		List[Any]						):



       """simple docstring"""
       lowerCAmelCase__										=      self.all_model_classes[1]
       lowerCAmelCase__				,      lowerCAmelCase__				,      lowerCAmelCase__				,      lowerCAmelCase__				,      lowerCAmelCase__										=      self.model_tester.prepare_config_and_inputs()
       lowerCAmelCase__										=      True
       lowerCAmelCase__										=      True

       lowerCAmelCase__										=      model_class(__lowerCamelCase						).to(__lowerCamelCase						)
       model.train()

       lowerCAmelCase__										=      model(__lowerCamelCase       ,						mask_labels=__lowerCamelCase       ,						class_labels=__lowerCamelCase						)

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

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

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

       lowerCAmelCase__										=      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						)



__magic_name__  :  Union[str, Any]					=       1E-4
def 		a_							( ):
  lowerCAmelCase__										=      Image.open('''./tests/fixtures/tests_samples/COCO/000000039769.png'''				)
  return image





@require_vision
@slow
class     SCREAMING_SNAKE_CASE__		(unittest.TestCase			):
     @cached_property
     def        A__      (						self	:		Optional[Any]						):



       """simple docstring"""
       return "facebook/mask2former-swin-small-coco-instance"
     @cached_property
     def        A__      (						self	:		Tuple						):



       """simple docstring"""
       return MaskaFormerImageProcessor.from_pretrained(self.model_checkpoints						) if is_vision_available() else None
     def        A__      (						self	:		List[Any]						):



       """simple docstring"""
       lowerCAmelCase__										=      MaskaFormerModel.from_pretrained(self.model_checkpoints						).to(__lowerCamelCase						)
       lowerCAmelCase__										=      self.default_image_processor
       lowerCAmelCase__										=      prepare_img()
       lowerCAmelCase__										=      image_processor(__lowerCamelCase       ,						return_tensors='''pt'''						).to(__lowerCamelCase						)
       lowerCAmelCase__										=      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, 3_84, 3_84)						)

       with torch.no_grad():
         lowerCAmelCase__										=      model(**__lowerCamelCase						)

       lowerCAmelCase__										=      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						)						)

       lowerCAmelCase__										=      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						)						)

       lowerCAmelCase__										=      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        A__      (						self	:		Tuple						):



       """simple docstring"""
       lowerCAmelCase__										=      MaskaFormerForUniversalSegmentation.from_pretrained(self.model_checkpoints						).to(__lowerCamelCase						).eval()
       lowerCAmelCase__										=      self.default_image_processor
       lowerCAmelCase__										=      prepare_img()
       lowerCAmelCase__										=      image_processor(__lowerCamelCase       ,						return_tensors='''pt'''						).to(__lowerCamelCase						)
       lowerCAmelCase__										=      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, 3_84, 3_84)						)

       with torch.no_grad():
         lowerCAmelCase__										=      model(**__lowerCamelCase						)
       # masks_queries_logits
       lowerCAmelCase__										=      outputs.masks_queries_logits
       self.assertEqual(
           masks_queries_logits.shape       ,						(1, model.config.num_queries, inputs_shape[-2] // 4, inputs_shape[-1] // 4)						)
       lowerCAmelCase__										=      [
           [-8.7839, -9.0056, -8.8121],
           [-7.4104, -7.0313, -6.5401],
           [-6.6105, -6.3427, -6.4675],
       ]
       lowerCAmelCase__										=      torch.tensor(__lowerCamelCase						).to(__lowerCamelCase						)
       self.assertTrue(torch.allclose(masks_queries_logits[0, 0, :3, :3]       ,						__lowerCamelCase       ,						atol=__lowerCamelCase						)						)
       # class_queries_logits
       lowerCAmelCase__										=      outputs.class_queries_logits
       self.assertEqual(class_queries_logits.shape       ,						(1, model.config.num_queries, model.config.num_labels + 1)						)
       lowerCAmelCase__										=      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        A__      (						self	:		List[Any]						):



       """simple docstring"""
       lowerCAmelCase__										=      MaskaFormerForUniversalSegmentation.from_pretrained(self.model_checkpoints						).to(__lowerCamelCase						).eval()
       lowerCAmelCase__										=      self.default_image_processor

       lowerCAmelCase__										=      image_processor(
           [np.zeros((3, 8_00, 13_33)						), np.zeros((3, 8_00, 13_33)						)]       ,						segmentation_maps=[np.zeros((3_84, 3_84)						).astype(np.floataa						), np.zeros((3_84, 3_84)						).astype(np.floataa						)]       ,						return_tensors='''pt'''       ,						)

       lowerCAmelCase__										=      inputs['''pixel_values'''].to(__lowerCamelCase						)
       lowerCAmelCase__										=      [el.to(__lowerCamelCase						) for el in inputs['''mask_labels''']]
       lowerCAmelCase__										=      [el.to(__lowerCamelCase						) for el in inputs['''class_labels''']]

       with torch.no_grad():
         lowerCAmelCase__										=      model(**__lowerCamelCase						)

       self.assertTrue(outputs.loss is not None						)







import logging

from transformers.configuration_utils import PretrainedConfig


__magic_name__  :  Tuple					=       logging.getLogger(__name__)





class     SCREAMING_SNAKE_CASE__		(_a			):
     lowercase_ :      Tuple				= "masked_bert"
     def __init__(						self	:		List[Any]       ,						__lowerCamelCase	:		Optional[Any]=3_05_22       ,						__lowerCamelCase	:		Dict=7_68       ,						__lowerCamelCase	:		Tuple=12       ,						__lowerCamelCase	:		int=12       ,						__lowerCamelCase	:		Tuple=30_72       ,						__lowerCamelCase	:		List[Any]="gelu"       ,						__lowerCamelCase	:		Union[str, Any]=0.1       ,						__lowerCamelCase	:		Dict=0.1       ,						__lowerCamelCase	:		Optional[int]=5_12       ,						__lowerCamelCase	:		List[Any]=2       ,						__lowerCamelCase	:		Union[str, Any]=0.02       ,						__lowerCamelCase	:		Optional[int]=1e-12       ,						__lowerCamelCase	:		Optional[int]=0       ,						__lowerCamelCase	:		List[str]="topK"       ,						__lowerCamelCase	:		Union[str, Any]="constant"       ,						__lowerCamelCase	:		Union[str, Any]=0.0       ,						**__lowerCamelCase	:		Tuple       ,						):



       """simple docstring"""
       super().__init__(pad_token_id=__lowerCamelCase       ,						**__lowerCamelCase						)

       lowerCAmelCase__										=      vocab_size
       lowerCAmelCase__										=      hidden_size
       lowerCAmelCase__										=      num_hidden_layers
       lowerCAmelCase__										=      num_attention_heads
       lowerCAmelCase__										=      hidden_act
       lowerCAmelCase__										=      intermediate_size
       lowerCAmelCase__										=      hidden_dropout_prob
       lowerCAmelCase__										=      attention_probs_dropout_prob
       lowerCAmelCase__										=      max_position_embeddings
       lowerCAmelCase__										=      type_vocab_size
       lowerCAmelCase__										=      initializer_range
       lowerCAmelCase__										=      layer_norm_eps
       lowerCAmelCase__										=      pruning_method
       lowerCAmelCase__										=      mask_init
       lowerCAmelCase__										=      mask_scale

