





"""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      __SCREAMING_SNAKE_CASE    :







  '''simple docstring'''


  def __init__(		self   :       Union[str, Any]					,   __a   :       Any					,   __a   :       Tuple=2					,   __a   :       Optional[Any]=True					,   __a   :       Tuple=False					,   __a   :       Optional[Any]=10					,   __a   :       List[Any]=3					,   __a   :       str=32 * 8					,   __a   :       Optional[Any]=32 * 8					,   __a   :       List[str]=4					,   __a   :       Optional[Any]=64					,   )		->				Optional[Any]:
        _UpperCamelCase      :	List[Any]						=       parent
        _UpperCamelCase      :	str						=       batch_size
        _UpperCamelCase      :	Dict						=       is_training
        _UpperCamelCase      :	List[Any]						=       use_auxiliary_loss
        _UpperCamelCase      :	Optional[int]						=       num_queries
        _UpperCamelCase      :	List[str]						=       num_channels
        _UpperCamelCase      :	Tuple						=       min_size
        _UpperCamelCase      :	List[str]						=       max_size
        _UpperCamelCase      :	Tuple						=       num_labels
        _UpperCamelCase      :	Any						=       hidden_dim
        _UpperCamelCase      :	List[str]						=       hidden_dim

  def    __SCREAMING_SNAKE_CASE				(		self   :       Dict	)		->				Optional[int]:
        _UpperCamelCase      :	Union[str, Any]						=       floats_tensor([self.batch_size, self.num_channels, self.min_size, self.max_size]	).to(
            __a	)

        _UpperCamelCase      :	Any						=       torch.ones([self.batch_size, self.min_size, self.max_size]					,   device=__a	)

        _UpperCamelCase      :	Union[str, Any]						=       (
            torch.rand([self.batch_size, self.num_labels, self.min_size, self.max_size]					,   device=__a	) > 0.5
        ).float()
        _UpperCamelCase      :	Optional[Any]						=       (torch.rand((self.batch_size, self.num_labels)					,   device=__a	) > 0.5).long()

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

  def    __SCREAMING_SNAKE_CASE				(		self   :       Optional[Any]	)		->				str:
        _UpperCamelCase      :	Any						=       MaskaFormerConfig(
            hidden_size=self.hidden_dim					,   )
        _UpperCamelCase      :	int						=       self.num_queries
        _UpperCamelCase      :	Optional[Any]						=       self.num_labels

        _UpperCamelCase      :	Optional[Any]						=       [1, 1, 1, 1]
        _UpperCamelCase      :	List[str]						=       self.num_channels

        _UpperCamelCase      :	Optional[Any]						=       64
        _UpperCamelCase      :	str						=       128
        _UpperCamelCase      :	Union[str, Any]						=       self.hidden_dim
        _UpperCamelCase      :	int						=       self.hidden_dim
        _UpperCamelCase      :	Any						=       self.hidden_dim
        return config

  def    __SCREAMING_SNAKE_CASE				(		self   :       int	)		->				Dict:
        _UpperCamelCase,  _UpperCamelCase,  _UpperCamelCase,  _UpperCamelCase,  _UpperCamelCase      :	Optional[Any]						=       self.prepare_config_and_inputs()
        _UpperCamelCase      :	Dict						=       {"pixel_values": pixel_values, "pixel_mask": pixel_mask}
        return config, inputs_dict

  def    __SCREAMING_SNAKE_CASE				(		self   :       str					,   __a   :       Union[str, Any]					,   __a   :       str	)		->				Dict:
        _UpperCamelCase      :	int						=       output.encoder_hidden_states
        _UpperCamelCase      :	Dict						=       output.pixel_decoder_hidden_states
        _UpperCamelCase      :	Union[str, Any]						=       output.transformer_decoder_hidden_states

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

  def    __SCREAMING_SNAKE_CASE				(		self   :       Tuple					,   __a   :       List[Any]					,   __a   :       Tuple					,   __a   :       List[Any]					,   __a   :       Tuple=False	)		->				Dict:

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

              _UpperCamelCase      :	Optional[Any]						=       model(pixel_values=__a					,   pixel_mask=__a	)
              _UpperCamelCase      :	int						=       model(__a					,   output_hidden_states=__a	)

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


  def    __SCREAMING_SNAKE_CASE				(		self   :       List[str]					,   __a   :       int					,   __a   :       Union[str, Any]					,   __a   :       List[str]					,   __a   :       Dict					,   __a   :       Dict	)		->				List[Any]:
        _UpperCamelCase      :	Tuple						=       MaskaFormerForUniversalSegmentation(config=__a	)
        model.to(__a	)
        model.eval()

        def comm_check_on_output(__a   :       List[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      :	Optional[int]						=       model(pixel_values=__a					,   pixel_mask=__a	)
              _UpperCamelCase      :	str						=       model(__a	)

              comm_check_on_output(__a	)

              _UpperCamelCase      :	Tuple						=       model(
                  pixel_values=__a					,   pixel_mask=__a					,   mask_labels=__a					,   class_labels=__a	)

        comm_check_on_output(__a	)

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





@require_torch
class      __SCREAMING_SNAKE_CASE    (						__UpperCAmelCase		,  __UpperCAmelCase		,  unittest.TestCase   ):







  '''simple docstring'''


  SCREAMING_SNAKE_CASE__							:List[Any]	     =					(MaskaFormerModel, MaskaFormerForUniversalSegmentation) if is_torch_available() else ()
  SCREAMING_SNAKE_CASE__							:int	     =					{"feature-extraction": MaskaFormerModel} if is_torch_available() else {}

  SCREAMING_SNAKE_CASE__							:Optional[Any]	     =					False
  SCREAMING_SNAKE_CASE__							:Tuple	     =					False
  SCREAMING_SNAKE_CASE__							:Union[str, Any]	     =					False
  SCREAMING_SNAKE_CASE__							:Optional[int]	     =					False

  def    __SCREAMING_SNAKE_CASE				(		self   :       Optional[Any]	)		->				Tuple:
        _UpperCamelCase      :	Dict						=       MaskaFormerModelTester(self	)
        _UpperCamelCase      :	Any						=       ConfigTester(self					,   config_class=__a					,   has_text_modality=__a	)

  def    __SCREAMING_SNAKE_CASE				(		self   :       Optional[Any]	)		->				Any:
        self.config_tester.run_common_tests()

  def    __SCREAMING_SNAKE_CASE				(		self   :       Any	)		->				int:
        _UpperCamelCase,  _UpperCamelCase      :	Union[str, Any]						=       self.model_tester.prepare_config_and_inputs_for_common()
        self.model_tester.create_and_check_maskaformer_model(__a					,   **__a					,   output_hidden_states=__a	)

  def    __SCREAMING_SNAKE_CASE				(		self   :       Union[str, Any]	)		->				Optional[Any]:
        _UpperCamelCase      :	int						=       self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_maskaformer_instance_segmentation_head_model(*__a	)

  @unittest.skip(reason="Mask2Former does not use inputs_embeds"	)
  def    __SCREAMING_SNAKE_CASE				(		self   :       Dict	)		->				Any:
        pass

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

  @unittest.skip(reason="Mask2Former is not a generative model"	)
  def    __SCREAMING_SNAKE_CASE				(		self   :       Tuple	)		->				Optional[Any]:
        pass

  @unittest.skip(reason="Mask2Former does not use token embeddings"	)
  def    __SCREAMING_SNAKE_CASE				(		self   :       Dict	)		->				Optional[Any]:
        pass

  @require_torch_multi_gpu
  @unittest.skip(
      reason="Mask2Former has some layers using `add_module` which doesn\'t work well with `nn.DataParallel`"	)
  def    __SCREAMING_SNAKE_CASE				(		self   :       List[str]	)		->				Dict:
        pass

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

  def    __SCREAMING_SNAKE_CASE				(		self   :       Union[str, Any]	)		->				Dict:
        _UpperCamelCase,  _UpperCamelCase      :	Dict						=       self.model_tester.prepare_config_and_inputs_for_common()

        for model_class in self.all_model_classes:
              _UpperCamelCase      :	Any						=       model_class(__a	)
              _UpperCamelCase      :	Dict						=       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]					,   __a	)

  @slow
  def    __SCREAMING_SNAKE_CASE				(		self   :       Optional[int]	)		->				Union[str, Any]:

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

  def    __SCREAMING_SNAKE_CASE				(		self   :       Dict	)		->				List[Any]:
        _UpperCamelCase      :	Union[str, Any]						=       (self.model_tester.min_size,) * 2
        _UpperCamelCase      :	Dict						=       {
            "pixel_values": torch.randn((2, 3, *size)					,   device=__a	),
            "mask_labels": torch.randn((2, 10, *size)					,   device=__a	),
            "class_labels": torch.zeros(2					,   10					,   device=__a	).long(),
        }
        _UpperCamelCase      :	Optional[int]						=       self.model_tester.get_config()

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

  def    __SCREAMING_SNAKE_CASE				(		self   :       Any	)		->				int:
        _UpperCamelCase,  _UpperCamelCase      :	Any						=       self.model_tester.prepare_config_and_inputs_for_common()
        self.model_tester.create_and_check_maskaformer_model(__a					,   **__a					,   output_hidden_states=__a	)

  def    __SCREAMING_SNAKE_CASE				(		self   :       Tuple	)		->				List[Any]:
        _UpperCamelCase,  _UpperCamelCase      :	List[Any]						=       self.model_tester.prepare_config_and_inputs_for_common()

        for model_class in self.all_model_classes:
              _UpperCamelCase      :	Optional[int]						=       model_class(__a	).to(__a	)
              _UpperCamelCase      :	Optional[int]						=       model(**__a					,   output_attentions=__a	)
              self.assertTrue(outputs.attentions is not None	)

  def    __SCREAMING_SNAKE_CASE				(		self   :       Any	)		->				List[str]:

        if not self.model_tester.is_training:
              return

        _UpperCamelCase      :	str						=       self.all_model_classes[1]
        _UpperCamelCase,  _UpperCamelCase,  _UpperCamelCase,  _UpperCamelCase,  _UpperCamelCase      :	Union[str, Any]						=       self.model_tester.prepare_config_and_inputs()

        _UpperCamelCase      :	int						=       model_class(__a	)
        model.to(__a	)
        model.train()

        _UpperCamelCase      :	Tuple						=       model(__a					,   mask_labels=__a					,   class_labels=__a	).loss
        loss.backward()


  def    __SCREAMING_SNAKE_CASE				(		self   :       Dict	)		->				Any:
        _UpperCamelCase      :	int						=       self.all_model_classes[1]
        _UpperCamelCase,  _UpperCamelCase,  _UpperCamelCase,  _UpperCamelCase,  _UpperCamelCase      :	Union[str, Any]						=       self.model_tester.prepare_config_and_inputs()
        _UpperCamelCase      :	Union[str, Any]						=       True
        _UpperCamelCase      :	List[Any]						=       True

        _UpperCamelCase      :	List[Any]						=       model_class(__a	).to(__a	)
        model.train()

        _UpperCamelCase      :	str						=       model(__a					,   mask_labels=__a					,   class_labels=__a	)

        _UpperCamelCase      :	List[Any]						=       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      :	int						=       outputs.transformer_decoder_hidden_states[0]
        transformer_decoder_hidden_states.retain_grad()

        _UpperCamelCase      :	List[str]						=       outputs.attentions[0]
        attentions.retain_grad()

        outputs.loss.backward(retain_graph=__a	)

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





lowerCamelCase__				  =   1E-4




def 					lowercase__    (					)					->		Tuple:




      """simple docstring"""

      _UpperCamelCase      :	Optional[int]						=       Image.open("./tests/fixtures/tests_samples/COCO/000000039769.png"			)
      return image



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







  '''simple docstring'''


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

  @cached_property
  def    __SCREAMING_SNAKE_CASE				(		self   :       Any	)		->				Optional[int]:
        return MaskaFormerImageProcessor.from_pretrained(self.model_checkpoints	) if is_vision_available() else None

  def    __SCREAMING_SNAKE_CASE				(		self   :       Dict	)		->				Union[str, Any]:
        _UpperCamelCase      :	Any						=       MaskaFormerModel.from_pretrained(self.model_checkpoints	).to(__a	)
        _UpperCamelCase      :	Optional[Any]						=       self.default_image_processor
        _UpperCamelCase      :	int						=       prepare_img()
        _UpperCamelCase      :	str						=       image_processor(__a					,   return_tensors="pt"	).to(__a	)
        _UpperCamelCase      :	List[str]						=       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(__a					,   (1, 3, 384, 384)	)

        with torch.no_grad():
              _UpperCamelCase      :	Dict						=       model(**__a	)

        _UpperCamelCase      :	List[Any]						=       torch.tensor(
            [[-0.27_90, -1.07_17, -1.16_68], [-0.51_28, -0.31_28, -0.49_87], [-0.58_32, 0.19_71, -0.01_97]]	).to(__a	)
        self.assertTrue(
            torch.allclose(
                outputs.encoder_last_hidden_state[0, 0, :3, :3]					,   __a					,   atol=__a	)	)

        _UpperCamelCase      :	List[str]						=       torch.tensor(
            [[0.89_73, 1.18_47, 1.17_76], [1.19_34, 1.50_40, 1.51_28], [1.11_53, 1.44_86, 1.49_51]]	).to(__a	)
        self.assertTrue(
            torch.allclose(
                outputs.pixel_decoder_last_hidden_state[0, 0, :3, :3]					,   __a					,   atol=__a	)	)

        _UpperCamelCase      :	Tuple						=       torch.tensor(
            [[2.11_52, 1.70_00, -0.86_03], [1.58_08, 1.80_04, -0.93_53], [1.60_43, 1.74_95, -0.59_99]]	).to(__a	)
        self.assertTrue(
            torch.allclose(
                outputs.transformer_decoder_last_hidden_state[0, :3, :3]					,   __a					,   atol=__a	)	)

  def    __SCREAMING_SNAKE_CASE				(		self   :       List[Any]	)		->				List[str]:
        _UpperCamelCase      :	Optional[Any]						=       MaskaFormerForUniversalSegmentation.from_pretrained(self.model_checkpoints	).to(__a	).eval()
        _UpperCamelCase      :	List[str]						=       self.default_image_processor
        _UpperCamelCase      :	int						=       prepare_img()
        _UpperCamelCase      :	str						=       image_processor(__a					,   return_tensors="pt"	).to(__a	)
        _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(__a					,   (1, 3, 384, 384)	)

        with torch.no_grad():
              _UpperCamelCase      :	List[Any]						=       model(**__a	)
        # masks_queries_logits
        _UpperCamelCase      :	Union[str, Any]						=       outputs.masks_queries_logits
        self.assertEqual(
            masks_queries_logits.shape					,   (1, model.config.num_queries, inputs_shape[-2] // 4, inputs_shape[-1] // 4)	)
        _UpperCamelCase      :	str						=       [
            [-8.78_39, -9.00_56, -8.81_21],
            [-7.41_04, -7.03_13, -6.54_01],
            [-6.61_05, -6.34_27, -6.46_75],
        ]
        _UpperCamelCase      :	Tuple						=       torch.tensor(__a	).to(__a	)
        self.assertTrue(torch.allclose(masks_queries_logits[0, 0, :3, :3]					,   __a					,   atol=__a	)	)
        # class_queries_logits
        _UpperCamelCase      :	int						=       outputs.class_queries_logits
        self.assertEqual(class_queries_logits.shape					,   (1, model.config.num_queries, model.config.num_labels + 1)	)
        _UpperCamelCase      :	int						=       torch.tensor(
            [
                [1.83_24, -8.08_35, -4.19_22],
                [0.84_50, -9.00_50, -3.60_53],
                [0.30_45, -7.72_93, -3.02_75],
            ]	).to(__a	)
        self.assertTrue(torch.allclose(outputs.class_queries_logits[0, :3, :3]					,   __a					,   atol=__a	)	)


  def    __SCREAMING_SNAKE_CASE				(		self   :       Optional[int]	)		->				Any:
        _UpperCamelCase      :	List[Any]						=       MaskaFormerForUniversalSegmentation.from_pretrained(self.model_checkpoints	).to(__a	).eval()
        _UpperCamelCase      :	int						=       self.default_image_processor

        _UpperCamelCase      :	Tuple						=       image_processor(
            [np.zeros((3, 800, 1333)	), np.zeros((3, 800, 1333)	)]					,   segmentation_maps=[np.zeros((384, 384)	).astype(np.floataa	), np.zeros((384, 384)	).astype(np.floataa	)]					,   return_tensors="pt"					,   )

        _UpperCamelCase      :	Tuple						=       inputs["pixel_values"].to(__a	)
        _UpperCamelCase      :	int						=       [el.to(__a	) for el in inputs["mask_labels"]]
        _UpperCamelCase      :	Tuple						=       [el.to(__a	) for el in inputs["class_labels"]]

        with torch.no_grad():
              _UpperCamelCase      :	str						=       model(**__a	)

        self.assertTrue(outputs.loss is not None	)







"""simple docstring"""


class      __SCREAMING_SNAKE_CASE    :







  '''simple docstring'''


  def __init__(		self   :       Optional[Any]					,   __a   :       list[int]	)		->				None:
        _UpperCamelCase      :	Tuple						=       len(__a	)
        _UpperCamelCase      :	Dict						=       [0] * len_array

        if len_array > 0:
              _UpperCamelCase      :	Optional[Any]						=       array[0]

        for i in range(1					,   __a	):
              _UpperCamelCase      :	Tuple						=       self.prefix_sum[i - 1] + array[i]

  def    __SCREAMING_SNAKE_CASE				(		self   :       Dict					,   __a   :       int					,   __a   :       int	)		->				int:
        if start == 0:
              return self.prefix_sum[end]

        return self.prefix_sum[end] - self.prefix_sum[start - 1]


  def    __SCREAMING_SNAKE_CASE				(		self   :       Optional[Any]					,   __a   :       int	)		->				bool:
        _UpperCamelCase      :	int						=       {0}
        for sum_item in self.prefix_sum:
              if sum_item - target_sum in sums:
                    return True

              sums.add(__a	)

        return False


if __name__ == "__main__":
  import doctest

  doctest.testmod()
