

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     __snake_case    :
       """simple docstring"""







       def __init__(		self      , _UpperCamelCase      , _UpperCamelCase=2      , _UpperCamelCase=True      , _UpperCamelCase=False      , _UpperCamelCase=10      , _UpperCamelCase=3      , _UpperCamelCase=32 * 8      , _UpperCamelCase=32 * 8      , _UpperCamelCase=4      , _UpperCamelCase=64      , )      ->							Optional[int]:






              """simple docstring"""


              __snake_case       	=							parent
              __snake_case       	=							batch_size
              __snake_case       	=							is_training
              __snake_case       	=							use_auxiliary_loss
              __snake_case       	=							num_queries
              __snake_case       	=							num_channels
              __snake_case       	=							min_size
              __snake_case       	=							max_size
              __snake_case       	=							num_labels
              __snake_case       	=							hidden_dim
              __snake_case       	=							hidden_dim


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






              """simple docstring"""


              __snake_case       	=							floats_tensor([self.batch_size, self.num_channels, self.min_size, self.max_size]   ).to(
                  _UpperCamelCase   )

              __snake_case       	=							torch.ones([self.batch_size, self.min_size, self.max_size]      , device=_UpperCamelCase   )

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

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


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






              """simple docstring"""


              __snake_case       	=							MaskaFormerConfig(
                  hidden_size=self.hidden_dim      , )
              __snake_case       	=							self.num_queries
              __snake_case       	=							self.num_labels

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

              __snake_case       	=							64
              __snake_case       	=							1_28
              __snake_case       	=							self.hidden_dim
              __snake_case       	=							self.hidden_dim
              __snake_case       	=							self.hidden_dim
              return config


       def 			a		(		self   )      ->							Dict:






              """simple docstring"""


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


       def 			a		(		self      , _UpperCamelCase      , _UpperCamelCase   )      ->							str:






              """simple docstring"""


              __snake_case       	=							output.encoder_hidden_states
              __snake_case       	=							output.pixel_decoder_hidden_states
              __snake_case       	=							output.transformer_decoder_hidden_states

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


       def 			a		(		self      , _UpperCamelCase      , _UpperCamelCase      , _UpperCamelCase      , _UpperCamelCase=False   )      ->							Optional[Any]:






              """simple docstring"""


              with torch.no_grad():
                     __snake_case       	=							MaskaFormerModel(config=_UpperCamelCase   )
                     model.to(_UpperCamelCase   )
                     model.eval()

                     __snake_case       	=							model(pixel_values=_UpperCamelCase      , pixel_mask=_UpperCamelCase   )
                     __snake_case       	=							model(_UpperCamelCase      , output_hidden_states=_UpperCamelCase   )

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






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






              """simple docstring"""


              __snake_case       	=							MaskaFormerForUniversalSegmentation(config=_UpperCamelCase   )
              model.to(_UpperCamelCase   )
              model.eval()

              def comm_check_on_output(_UpperCamelCase   ):
                     # 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():
                     __snake_case       	=							model(pixel_values=_UpperCamelCase      , pixel_mask=_UpperCamelCase   )
                     __snake_case       	=							model(_UpperCamelCase   )

                     comm_check_on_output(_UpperCamelCase   )

                     __snake_case       	=							model(
                         pixel_values=_UpperCamelCase      , pixel_mask=_UpperCamelCase      , mask_labels=_UpperCamelCase      , class_labels=_UpperCamelCase   )

              comm_check_on_output(_UpperCamelCase   )

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






@require_torch
class     __snake_case    (				snake_case__			,	snake_case__			,	unittest.TestCase  ):
       """simple docstring"""







       __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 			a		(		self   )      ->							List[str]:






              """simple docstring"""


              __snake_case       	=							MaskaFormerModelTester(self   )
              __snake_case       	=							ConfigTester(self      , config_class=_UpperCamelCase      , has_text_modality=_UpperCamelCase   )


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






              """simple docstring"""


              self.config_tester.run_common_tests()


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






              """simple docstring"""


              __snake_case	,    __snake_case       	=							self.model_tester.prepare_config_and_inputs_for_common()
              self.model_tester.create_and_check_maskaformer_model(_UpperCamelCase      , **_UpperCamelCase      , output_hidden_states=_UpperCamelCase   )


       def 			a		(		self   )      ->							Dict:






              """simple docstring"""


              __snake_case       	=							self.model_tester.prepare_config_and_inputs()
              self.model_tester.create_and_check_maskaformer_instance_segmentation_head_model(*_UpperCamelCase   )


       @unittest.skip(reason="""Mask2Former does not use inputs_embeds"""   )
       def 			a		(		self   )      ->							Optional[int]:






              """simple docstring"""


              pass


       @unittest.skip(reason="""Mask2Former does not have a get_input_embeddings method"""   )
       def 			a		(		self   )      ->							Any:






              """simple docstring"""


              pass


       @unittest.skip(reason="""Mask2Former is not a generative model"""   )
       def 			a		(		self   )      ->							Any:






              """simple docstring"""


              pass


       @unittest.skip(reason="""Mask2Former does not use token embeddings"""   )
       def 			a		(		self   )      ->							Tuple:






              """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   )      ->							Union[str, Any]:






              """simple docstring"""


              pass


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






              """simple docstring"""


              pass


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






              """simple docstring"""


              __snake_case	,    __snake_case       	=							self.model_tester.prepare_config_and_inputs_for_common()

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

                     __snake_case       	=							["""pixel_values"""]
                     self.assertListEqual(arg_names[:1]      , _UpperCamelCase   )


       @slow
       def 			a		(		self   )      ->							Union[str, Any]:






              """simple docstring"""


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


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






              """simple docstring"""


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

              __snake_case       	=							MaskaFormerForUniversalSegmentation(_UpperCamelCase   ).to(_UpperCamelCase   )
              __snake_case       	=							model(**_UpperCamelCase   )
              self.assertTrue(outputs.loss is not None   )


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






              """simple docstring"""


              __snake_case	,    __snake_case       	=							self.model_tester.prepare_config_and_inputs_for_common()
              self.model_tester.create_and_check_maskaformer_model(_UpperCamelCase      , **_UpperCamelCase      , output_hidden_states=_UpperCamelCase   )


       def 			a		(		self   )      ->							str:






              """simple docstring"""


              __snake_case	,    __snake_case       	=							self.model_tester.prepare_config_and_inputs_for_common()

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


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






              """simple docstring"""


              if not self.model_tester.is_training:
                     return

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

              __snake_case       	=							model_class(_UpperCamelCase   )
              model.to(_UpperCamelCase   )
              model.train()

              __snake_case       	=							model(_UpperCamelCase      , mask_labels=_UpperCamelCase      , class_labels=_UpperCamelCase   ).loss
              loss.backward()






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






              """simple docstring"""


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

              __snake_case       	=							model_class(_UpperCamelCase   ).to(_UpperCamelCase   )
              model.train()

              __snake_case       	=							model(_UpperCamelCase      , mask_labels=_UpperCamelCase      , class_labels=_UpperCamelCase   )

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

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

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

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

              outputs.loss.backward(retain_graph=_UpperCamelCase   )

              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 						lowerCamelCase__      ( ):


       """simple docstring"""




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




@require_vision
@slow
class     __snake_case    (				unittest.TestCase  ):
       """simple docstring"""







       @cached_property
       def 			a		(		self   )      ->							Optional[Any]:






              """simple docstring"""


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


       @cached_property
       def 			a		(		self   )      ->							Any:






              """simple docstring"""


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


       def 			a		(		self   )      ->							int:






              """simple docstring"""


              __snake_case       	=							MaskaFormerModel.from_pretrained(self.model_checkpoints   ).to(_UpperCamelCase   )
              __snake_case       	=							self.default_image_processor
              __snake_case       	=							prepare_img()
              __snake_case       	=							image_processor(_UpperCamelCase      , return_tensors="""pt"""   ).to(_UpperCamelCase   )
              __snake_case       	=							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(_UpperCamelCase      , (1, 3, 3_84, 3_84)   )

              with torch.no_grad():
                     __snake_case       	=							model(**_UpperCamelCase   )

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

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

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


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






              """simple docstring"""


              __snake_case       	=							MaskaFormerForUniversalSegmentation.from_pretrained(self.model_checkpoints   ).to(_UpperCamelCase   ).eval()
              __snake_case       	=							self.default_image_processor
              __snake_case       	=							prepare_img()
              __snake_case       	=							image_processor(_UpperCamelCase      , return_tensors="""pt"""   ).to(_UpperCamelCase   )
              __snake_case       	=							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(_UpperCamelCase      , (1, 3, 3_84, 3_84)   )

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






       def 			a		(		self   )      ->							str:






              """simple docstring"""


              __snake_case       	=							MaskaFormerForUniversalSegmentation.from_pretrained(self.model_checkpoints   ).to(_UpperCamelCase   ).eval()
              __snake_case       	=							self.default_image_processor

              __snake_case       	=							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"""      , )

              __snake_case       	=							inputs["""pixel_values"""].to(_UpperCamelCase   )
              __snake_case       	=							[el.to(_UpperCamelCase   ) for el in inputs["""mask_labels"""]]
              __snake_case       	=							[el.to(_UpperCamelCase   ) for el in inputs["""class_labels"""]]

              with torch.no_grad():
                     __snake_case       	=							model(**_UpperCamelCase   )

              self.assertTrue(outputs.loss is not None   )




def 						lowerCamelCase__      ( __A    :int	,__A    :int		):


       """simple docstring"""




       if a < 0 or b < 0:
              raise ValueError("""the value of both inputs must be positive"""		)

       __snake_case       	=							str(bin(__A		)		)[2:]  # remove the leading "0b"
       __snake_case       	=							str(bin(__A		)		)[2:]  # remove the leading "0b"

       __snake_case       	=							max(len(__A		)	,len(__A		)		)

       return "0b" + "".join(
           str(int(char_a == """1""" and char_b == """1"""		)		)
           for char_a, char_b in zip(a_binary.zfill(__A		)	,b_binary.zfill(__A		)		)		)


if __name__ == "__main__":
    import doctest

    doctest.testmod()

