
'''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 							UpperCAmelCase_							:

 '''simple docstring'''







 def __init__(					self		:  List[Any]			,				snake_case__		:  int			,				snake_case__		:  Optional[Any]=2			,				snake_case__		:  Any=True			,				snake_case__		:  Tuple=False			,				snake_case__		:  Tuple=10			,				snake_case__		:  Optional[int]=3			,				snake_case__		:  Union[str, Any]=32 * 8			,				snake_case__		:  Optional[int]=32 * 8			,				snake_case__		:  Union[str, Any]=4			,				snake_case__		:  Dict=64			,				):



   '''simple docstring'''




   UpperCAmelCase__      :					List[Any]				       =							parent
   UpperCAmelCase__      :					List[str]				       =							batch_size
   UpperCAmelCase__      :					Dict				       =							is_training
   UpperCAmelCase__      :					Any				       =							use_auxiliary_loss
   UpperCAmelCase__      :					Any				       =							num_queries
   UpperCAmelCase__      :					Union[str, Any]				       =							num_channels
   UpperCAmelCase__      :					Optional[int]				       =							min_size
   UpperCAmelCase__      :					int				       =							max_size
   UpperCAmelCase__      :					Dict				       =							num_labels
   UpperCAmelCase__      :					Any				       =							hidden_dim
   UpperCAmelCase__      :					Tuple				       =							hidden_dim







 def 							UpperCamelCase  (					self		:  Tuple ):



   '''simple docstring'''




   UpperCAmelCase__      :					Tuple				       =							floats_tensor([self.batch_size, self.num_channels, self.min_size, self.max_size] ).to(
       snake_case__ )

   UpperCAmelCase__      :					int				       =							torch.ones([self.batch_size, self.min_size, self.max_size]			,				device=snake_case__ )

   UpperCAmelCase__      :					Any				       =							(
       torch.rand([self.batch_size, self.num_labels, self.min_size, self.max_size]			,				device=snake_case__ ) > 0.5
   ).float()
   UpperCAmelCase__      :					Dict				       =							(torch.rand((self.batch_size, self.num_labels)			,				device=snake_case__ ) > 0.5).long()

   UpperCAmelCase__      :					List[str]				       =							self.get_config()
   return config, pixel_values, pixel_mask, mask_labels, class_labels







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



   '''simple docstring'''




   UpperCAmelCase__      :					str				       =							MaskaFormerConfig(
       hidden_size=self.hidden_dim			,				)
   UpperCAmelCase__      :					Union[str, Any]				       =							self.num_queries
   UpperCAmelCase__      :					Any				       =							self.num_labels

   UpperCAmelCase__      :					Optional[Any]				       =							[1, 1, 1, 1]
   UpperCAmelCase__      :					Optional[Any]				       =							self.num_channels

   UpperCAmelCase__      :					str				       =							64
   UpperCAmelCase__      :					Optional[int]				       =							1_28
   UpperCAmelCase__      :					List[str]				       =							self.hidden_dim
   UpperCAmelCase__      :					Optional[Any]				       =							self.hidden_dim
   UpperCAmelCase__      :					str				       =							self.hidden_dim
   return config







 def 							UpperCamelCase  (					self		:  Optional[int] ):



   '''simple docstring'''




   UpperCAmelCase__       ,			UpperCAmelCase__       ,			UpperCAmelCase__       ,			UpperCAmelCase__       ,			UpperCAmelCase__      :					Optional[int]				       =							self.prepare_config_and_inputs()
   UpperCAmelCase__      :					Union[str, Any]				       =							{"pixel_values": pixel_values, "pixel_mask": pixel_mask}
   return config, inputs_dict







 def 							UpperCamelCase  (					self		:  int			,				snake_case__		:  Union[str, Any]			,				snake_case__		:  Dict ):



   '''simple docstring'''




   UpperCAmelCase__      :					Dict				       =							output.encoder_hidden_states
   UpperCAmelCase__      :					Optional[Any]				       =							output.pixel_decoder_hidden_states
   UpperCAmelCase__      :					Optional[Any]				       =							output.transformer_decoder_hidden_states

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







 def 							UpperCamelCase  (					self		:  Optional[int]			,				snake_case__		:  List[str]			,				snake_case__		:  List[str]			,				snake_case__		:  str			,				snake_case__		:  List[Any]=False ):



   '''simple docstring'''




   with torch.no_grad():
     UpperCAmelCase__      :					List[str]				       =							MaskaFormerModel(config=snake_case__ )
     model.to(snake_case__ )
     model.eval()

     UpperCAmelCase__      :					Dict				       =							model(pixel_values=snake_case__			,				pixel_mask=snake_case__ )
     UpperCAmelCase__      :					Dict				       =							model(snake_case__			,				output_hidden_states=snake_case__ )

   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(snake_case__			,				snake_case__ )







 def 							UpperCamelCase  (					self		:  Tuple			,				snake_case__		:  List[Any]			,				snake_case__		:  Any			,				snake_case__		:  Union[str, Any]			,				snake_case__		:  Any			,				snake_case__		:  Optional[int] ):



   '''simple docstring'''




   UpperCAmelCase__      :					Any				       =							MaskaFormerForUniversalSegmentation(config=snake_case__ )
   model.to(snake_case__ )
   model.eval()

   def comm_check_on_output(snake_case__		:  Union[str, Any] ):
     # 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__      :					Dict				       =							model(pixel_values=snake_case__			,				pixel_mask=snake_case__ )
     UpperCAmelCase__      :					Tuple				       =							model(snake_case__ )

     comm_check_on_output(snake_case__ )

     UpperCAmelCase__      :					str				       =							model(
         pixel_values=snake_case__			,				pixel_mask=snake_case__			,				mask_labels=snake_case__			,				class_labels=snake_case__ )

   comm_check_on_output(snake_case__ )

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






@require_torch
class 							UpperCAmelCase_							(    A					,				A					,				unittest.TestCase				):

 '''simple docstring'''
 lowercase_					:    Optional[Any]				=							(MaskaFormerModel, MaskaFormerForUniversalSegmentation) if is_torch_available() else ()
 lowercase_					:    Optional[Any]				=							{"feature-extraction": MaskaFormerModel} if is_torch_available() else {}

 lowercase_					:    List[str]				=							False
 lowercase_					:    Union[str, Any]				=							False
 lowercase_					:    int				=							False
 lowercase_					:    List[str]				=							False







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



   '''simple docstring'''




   UpperCAmelCase__      :					Any				       =							MaskaFormerModelTester(self )
   UpperCAmelCase__      :					List[str]				       =							ConfigTester(self			,				config_class=snake_case__			,				has_text_modality=snake_case__ )







 def 							UpperCamelCase  (					self		:  Dict ):



   '''simple docstring'''




   self.config_tester.run_common_tests()







 def 							UpperCamelCase  (					self		:  List[str] ):



   '''simple docstring'''




   UpperCAmelCase__       ,			UpperCAmelCase__      :					List[str]				       =							self.model_tester.prepare_config_and_inputs_for_common()
   self.model_tester.create_and_check_maskaformer_model(snake_case__			,				**snake_case__			,				output_hidden_states=snake_case__ )







 def 							UpperCamelCase  (					self		:  Dict ):



   '''simple docstring'''




   UpperCAmelCase__      :					Any				       =							self.model_tester.prepare_config_and_inputs()
   self.model_tester.create_and_check_maskaformer_instance_segmentation_head_model(*snake_case__ )







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



   '''simple docstring'''




   pass







 @unittest.skip(reason="Mask2Former does not have a get_input_embeddings method" )
 def 							UpperCamelCase  (					self		:  Optional[int] ):



   '''simple docstring'''




   pass







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



   '''simple docstring'''




   pass







 @unittest.skip(reason="Mask2Former does not use token embeddings" )
 def 							UpperCamelCase  (					self		:  Optional[int] ):



   '''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 							UpperCamelCase  (					self		:  Tuple ):



   '''simple docstring'''




   pass







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



   '''simple docstring'''




   pass







 def 							UpperCamelCase  (					self		:  Optional[int] ):



   '''simple docstring'''




   UpperCAmelCase__       ,			UpperCAmelCase__      :					Any				       =							self.model_tester.prepare_config_and_inputs_for_common()

   for model_class in self.all_model_classes:
     UpperCAmelCase__      :					Optional[Any]				       =							model_class(snake_case__ )
     UpperCAmelCase__      :					Optional[int]				       =							inspect.signature(model.forward )
     # signature.parameters is an OrderedDict => so arg_names order is deterministic
     UpperCAmelCase__      :					Dict				       =							[*signature.parameters.keys()]

     UpperCAmelCase__      :					List[str]				       =							["pixel_values"]
     self.assertListEqual(arg_names[:1]			,				snake_case__ )







 @slow
 def 							UpperCamelCase  (					self		:  Dict ):



   '''simple docstring'''




   for model_name in ["facebook/mask2former-swin-small-coco-instance"]:
     UpperCAmelCase__      :					Optional[int]				       =							MaskaFormerModel.from_pretrained(snake_case__ )
     self.assertIsNotNone(snake_case__ )







 def 							UpperCamelCase  (					self		:  List[str] ):



   '''simple docstring'''




   UpperCAmelCase__      :					Optional[int]				       =							(self.model_tester.min_size,) * 2
   UpperCAmelCase__      :					str				       =							{
       "pixel_values": torch.randn((2, 3, *size)			,				device=snake_case__ ),
       "mask_labels": torch.randn((2, 10, *size)			,				device=snake_case__ ),
       "class_labels": torch.zeros(2			,				10			,				device=snake_case__ ).long(),
   }
   UpperCAmelCase__      :					str				       =							self.model_tester.get_config()

   UpperCAmelCase__      :					Optional[Any]				       =							MaskaFormerForUniversalSegmentation(snake_case__ ).to(snake_case__ )
   UpperCAmelCase__      :					Optional[Any]				       =							model(**snake_case__ )
   self.assertTrue(outputs.loss is not None )







 def 							UpperCamelCase  (					self		:  Optional[int] ):



   '''simple docstring'''




   UpperCAmelCase__       ,			UpperCAmelCase__      :					List[str]				       =							self.model_tester.prepare_config_and_inputs_for_common()
   self.model_tester.create_and_check_maskaformer_model(snake_case__			,				**snake_case__			,				output_hidden_states=snake_case__ )







 def 							UpperCamelCase  (					self		:  Dict ):



   '''simple docstring'''




   UpperCAmelCase__       ,			UpperCAmelCase__      :					str				       =							self.model_tester.prepare_config_and_inputs_for_common()

   for model_class in self.all_model_classes:
     UpperCAmelCase__      :					int				       =							model_class(snake_case__ ).to(snake_case__ )
     UpperCAmelCase__      :					str				       =							model(**snake_case__			,				output_attentions=snake_case__ )
     self.assertTrue(outputs.attentions is not None )







 def 							UpperCamelCase  (					self		:  List[str] ):



   '''simple docstring'''




   if not self.model_tester.is_training:
     return

   UpperCAmelCase__      :					Any				       =							self.all_model_classes[1]
   UpperCAmelCase__       ,			UpperCAmelCase__       ,			UpperCAmelCase__       ,			UpperCAmelCase__       ,			UpperCAmelCase__      :					Dict				       =							self.model_tester.prepare_config_and_inputs()

   UpperCAmelCase__      :					Dict				       =							model_class(snake_case__ )
   model.to(snake_case__ )
   model.train()

   UpperCAmelCase__      :					Optional[int]				       =							model(snake_case__			,				mask_labels=snake_case__			,				class_labels=snake_case__ ).loss
   loss.backward()







 def 							UpperCamelCase  (					self		:  List[str] ):



   '''simple docstring'''




   UpperCAmelCase__      :					Optional[int]				       =							self.all_model_classes[1]
   UpperCAmelCase__       ,			UpperCAmelCase__       ,			UpperCAmelCase__       ,			UpperCAmelCase__       ,			UpperCAmelCase__      :					Any				       =							self.model_tester.prepare_config_and_inputs()
   UpperCAmelCase__      :					Any				       =							True
   UpperCAmelCase__      :					Optional[int]				       =							True

   UpperCAmelCase__      :					Dict				       =							model_class(snake_case__ ).to(snake_case__ )
   model.train()

   UpperCAmelCase__      :					Any				       =							model(snake_case__			,				mask_labels=snake_case__			,				class_labels=snake_case__ )

   UpperCAmelCase__      :					Dict				       =							outputs.encoder_hidden_states[0]
   encoder_hidden_states.retain_grad()

   UpperCAmelCase__      :					Any				       =							outputs.pixel_decoder_hidden_states[0]
   pixel_decoder_hidden_states.retain_grad()

   UpperCAmelCase__      :					Tuple				       =							outputs.transformer_decoder_hidden_states[0]
   transformer_decoder_hidden_states.retain_grad()

   UpperCAmelCase__      :					Dict				       =							outputs.attentions[0]
   attentions.retain_grad()

   outputs.loss.backward(retain_graph=snake_case__ )

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



SCREAMING_SNAKE_CASE     =       1E-4







def snake_case_     (					):
  UpperCAmelCase__      :					List[Any]				       =							Image.open("./tests/fixtures/tests_samples/COCO/000000039769.png"		)
  return image






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

 '''simple docstring'''







 @cached_property
 def 							UpperCamelCase  (					self		:  Tuple ):



   '''simple docstring'''




   return "facebook/mask2former-swin-small-coco-instance"







 @cached_property
 def 							UpperCamelCase  (					self		:  List[Any] ):



   '''simple docstring'''




   return MaskaFormerImageProcessor.from_pretrained(self.model_checkpoints ) if is_vision_available() else None







 def 							UpperCamelCase  (					self		:  Optional[Any] ):



   '''simple docstring'''




   UpperCAmelCase__      :					List[Any]				       =							MaskaFormerModel.from_pretrained(self.model_checkpoints ).to(snake_case__ )
   UpperCAmelCase__      :					Dict				       =							self.default_image_processor
   UpperCAmelCase__      :					int				       =							prepare_img()
   UpperCAmelCase__      :					int				       =							image_processor(snake_case__			,				return_tensors="pt" ).to(snake_case__ )
   UpperCAmelCase__      :					Optional[Any]				       =							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(snake_case__			,				(1, 3, 3_84, 3_84) )

   with torch.no_grad():
     UpperCAmelCase__      :					Any				       =							model(**snake_case__ )

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

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

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







 def 							UpperCamelCase  (					self		:  Tuple ):



   '''simple docstring'''




   UpperCAmelCase__      :					Tuple				       =							MaskaFormerForUniversalSegmentation.from_pretrained(self.model_checkpoints ).to(snake_case__ ).eval()
   UpperCAmelCase__      :					Tuple				       =							self.default_image_processor
   UpperCAmelCase__      :					Union[str, Any]				       =							prepare_img()
   UpperCAmelCase__      :					Optional[Any]				       =							image_processor(snake_case__			,				return_tensors="pt" ).to(snake_case__ )
   UpperCAmelCase__      :					Optional[int]				       =							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(snake_case__			,				(1, 3, 3_84, 3_84) )

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







 def 							UpperCamelCase  (					self		:  List[Any] ):



   '''simple docstring'''




   UpperCAmelCase__      :					List[str]				       =							MaskaFormerForUniversalSegmentation.from_pretrained(self.model_checkpoints ).to(snake_case__ ).eval()
   UpperCAmelCase__      :					List[Any]				       =							self.default_image_processor

   UpperCAmelCase__      :					int				       =							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"			,				)

   UpperCAmelCase__      :					str				       =							inputs["pixel_values"].to(snake_case__ )
   UpperCAmelCase__      :					Tuple				       =							[el.to(snake_case__ ) for el in inputs["mask_labels"]]
   UpperCAmelCase__      :					Union[str, Any]				       =							[el.to(snake_case__ ) for el in inputs["class_labels"]]

   with torch.no_grad():
     UpperCAmelCase__      :					str				       =							model(**snake_case__ )

   self.assertTrue(outputs.loss is not None )





'''simple docstring'''





from typing import TYPE_CHECKING

from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_sentencepiece_available


SCREAMING_SNAKE_CASE     =       {}

try:
      if not is_sentencepiece_available():
            raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
      pass
else:
      SCREAMING_SNAKE_CASE     =       ["""GPTSw3Tokenizer"""]


if TYPE_CHECKING:
      try:
            if not is_sentencepiece_available():
                  raise OptionalDependencyNotAvailable()
      except OptionalDependencyNotAvailable:
            pass
      else:
            from .tokenization_gpt_swa import GPTSwaTokenizer

else:
      import sys

      SCREAMING_SNAKE_CASE     =       _LazyModule(__name__, globals()["""__file__"""], _import_structure, module_spec=__spec__)



