






import inspect
import unittest

import numpy as np

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

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


if is_torch_available():
	import torch

	from transformers import MaskaFormerForUniversalSegmentation, MaskaFormerModel

	if is_vision_available():
		from transformers import MaskaFormerImageProcessor

if is_vision_available():
	from PIL import Image

class 		lowerCAmelCase     :

					'''simple docstring'''

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


												"""simple docstring"""

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



					def       lowerCAmelCase				(  self				:	Dict							)			->  Optional[int]:


												"""simple docstring"""

												__lowercase      :							int				=     floats_tensor([self.batch_size, self.num_channels, self.min_size, self.max_size]							).to(
												    __a							)

												__lowercase      :							List[str]				=     torch.ones([self.batch_size, self.min_size, self.max_size]		,     device=__a							)

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

												__lowercase      :							int				=     self.get_config()
												return config, pixel_values, pixel_mask, mask_labels, class_labels



					def       lowerCAmelCase				(  self				:	Optional[int]							)			->  Optional[Any]:


												"""simple docstring"""

												__lowercase      :							Optional[int]				=     MaskaFormerConfig(
												    hidden_size=self.hidden_dim		,     )
												__lowercase      :							List[Any]				=     self.num_queries
												__lowercase      :							Union[str, Any]				=     self.num_labels

												__lowercase      :							Dict				=     [1, 1, 1, 1]
												__lowercase      :							Optional[int]				=     self.num_channels

												__lowercase      :							List[str]				=     64
												__lowercase      :							List[str]				=     128
												__lowercase      :							List[str]				=     self.hidden_dim
												__lowercase      :							Optional[Any]				=     self.hidden_dim
												__lowercase      :							List[str]				=     self.hidden_dim
												return config



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


												"""simple docstring"""

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



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


												"""simple docstring"""

												__lowercase      :							List[Any]				=     output.encoder_hidden_states
												__lowercase      :							int				=     output.pixel_decoder_hidden_states
												__lowercase      :							Dict				=     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       lowerCAmelCase				(  self				:	Optional[Any]		,     __a				:	Any		,     __a				:	List[str]		,     __a				:	List[str]		,     __a				:	List[Any]=False							)			->  str:


												"""simple docstring"""

												with torch.no_grad():
																			__lowercase      :							Union[str, Any]				=     MaskaFormerModel(config=__a							)
																			model.to(__a							)
																			model.eval()

																			__lowercase      :							Optional[Any]				=     model(pixel_values=__a		,     pixel_mask=__a							)
																			__lowercase      :							Optional[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       lowerCAmelCase				(  self				:	Optional[int]		,     __a				:	List[str]		,     __a				:	List[str]		,     __a				:	Dict		,     __a				:	Tuple		,     __a				:	str							)			->  Tuple:


												"""simple docstring"""

												__lowercase      :							Tuple				=     MaskaFormerForUniversalSegmentation(config=__a							)
												model.to(__a							)
												model.eval()

												def comm_check_on_output(__a				:	Dict							):
																			# 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():
																			__lowercase      :							Tuple				=     model(pixel_values=__a		,     pixel_mask=__a							)
																			__lowercase      :							str				=     model(__a							)

																			comm_check_on_output(__a							)

																			__lowercase      :							int				=     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 		lowerCAmelCase     ( __a      ,       __a      ,       unittest.TestCase			):

					'''simple docstring'''

					_A						:			Tuple									=      (MaskaFormerModel, MaskaFormerForUniversalSegmentation) if is_torch_available() else ()
					_A						:			Union[str, Any]									=      {'''feature-extraction''': MaskaFormerModel} if is_torch_available() else {}

					_A						:			Optional[int]									=      False
					_A						:			Dict									=      False
					_A						:			List[str]									=      False
					_A						:			Dict									=      False



					def       lowerCAmelCase				(  self				:	List[Any]							)			->  Optional[Any]:


												"""simple docstring"""

												__lowercase      :							Union[str, Any]				=     MaskaFormerModelTester(self							)
												__lowercase      :							List[Any]				=     ConfigTester(self		,     config_class=__a		,     has_text_modality=__a							)



					def       lowerCAmelCase				(  self				:	Optional[Any]							)			->  List[Any]:


												"""simple docstring"""

												self.config_tester.run_common_tests()



					def       lowerCAmelCase				(  self				:	Dict							)			->  Any:


												"""simple docstring"""

												__lowercase			,			__lowercase      :							int				=     self.model_tester.prepare_config_and_inputs_for_common()
												self.model_tester.create_and_check_maskaformer_model(__a		,     **__a		,     output_hidden_states=__a							)



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


												"""simple docstring"""

												__lowercase      :							List[str]				=     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       lowerCAmelCase				(  self				:	Dict							)			->  Optional[int]:


												"""simple docstring"""

												pass



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


												"""simple docstring"""

												pass



					@unittest.skip(reason="""Mask2Former is not a generative model"""							)
					def       lowerCAmelCase				(  self				:	Tuple							)			->  List[str]:


												"""simple docstring"""

												pass



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


												"""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       lowerCAmelCase				(  self				:	Optional[Any]							)			->  Any:


												"""simple docstring"""

												pass



					@unittest.skip("""Will be fixed soon by reducing the size of the model used for common tests."""							)
					def       lowerCAmelCase				(  self				:	Optional[int]							)			->  Dict:


												"""simple docstring"""

												pass



					def       lowerCAmelCase				(  self				:	Optional[int]							)			->  str:


												"""simple docstring"""

												__lowercase			,			__lowercase      :							Optional[Any]				=     self.model_tester.prepare_config_and_inputs_for_common()

												for model_class in self.all_model_classes:
																			__lowercase      :							List[Any]				=     model_class(__a							)
																			__lowercase      :							Dict				=     inspect.signature(model.forward							)
																			# signature.parameters is an OrderedDict => so arg_names order is deterministic
																			__lowercase      :							str				=     [*signature.parameters.keys()]

																			__lowercase      :							Optional[int]				=     ["""pixel_values"""]
																			self.assertListEqual(arg_names[:1]		,     __a							)



					@slow
					def       lowerCAmelCase				(  self				:	Optional[int]							)			->  Optional[Any]:


												"""simple docstring"""

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



					def       lowerCAmelCase				(  self				:	Union[str, Any]							)			->  Optional[int]:


												"""simple docstring"""

												__lowercase      :							str				=     (self.model_tester.min_size,) * 2
												__lowercase      :							Tuple				=     {
												    """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(),
												}
												__lowercase      :							List[Any]				=     self.model_tester.get_config()

												__lowercase      :							str				=     MaskaFormerForUniversalSegmentation(__a							).to(__a							)
												__lowercase      :							Any				=     model(**__a							)
												self.assertTrue(outputs.loss is not None							)



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


												"""simple docstring"""

												__lowercase			,			__lowercase      :							Optional[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       lowerCAmelCase				(  self				:	Tuple							)			->  Union[str, Any]:


												"""simple docstring"""

												__lowercase			,			__lowercase      :							Optional[int]				=     self.model_tester.prepare_config_and_inputs_for_common()

												for model_class in self.all_model_classes:
																			__lowercase      :							Union[str, Any]				=     model_class(__a							).to(__a							)
																			__lowercase      :							Any				=     model(**__a		,     output_attentions=__a							)
																			self.assertTrue(outputs.attentions is not None							)



					def       lowerCAmelCase				(  self				:	str							)			->  Tuple:


												"""simple docstring"""

												if not self.model_tester.is_training:
																			return

												__lowercase      :							Optional[Any]				=     self.all_model_classes[1]
												__lowercase			,			__lowercase			,			__lowercase			,			__lowercase			,			__lowercase      :							Any				=     self.model_tester.prepare_config_and_inputs()

												__lowercase      :							Optional[int]				=     model_class(__a							)
												model.to(__a							)
												model.train()

												__lowercase      :							Any				=     model(__a		,     mask_labels=__a		,     class_labels=__a							).loss
												loss.backward()



					def       lowerCAmelCase				(  self				:	Optional[int]							)			->  str:


												"""simple docstring"""

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

												__lowercase      :							Optional[int]				=     model_class(__a							).to(__a							)
												model.train()

												__lowercase      :							Dict				=     model(__a		,     mask_labels=__a		,     class_labels=__a							)

												__lowercase      :							Union[str, Any]				=     outputs.encoder_hidden_states[0]
												encoder_hidden_states.retain_grad()

												__lowercase      :							int				=     outputs.pixel_decoder_hidden_states[0]
												pixel_decoder_hidden_states.retain_grad()

												__lowercase      :							Any				=     outputs.transformer_decoder_hidden_states[0]
												transformer_decoder_hidden_states.retain_grad()

												__lowercase      :							Optional[int]				=     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 : Optional[int]						   =       1E-4
def       snake_case_		(     ):
							__lowercase      :							List[str]				=     Image.open("""./tests/fixtures/tests_samples/COCO/000000039769.png"""      )
							return image



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

					'''simple docstring'''

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


												"""simple docstring"""

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



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


												"""simple docstring"""

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



					def       lowerCAmelCase				(  self				:	Optional[int]							)			->  str:


												"""simple docstring"""

												__lowercase      :							str				=     MaskaFormerModel.from_pretrained(self.model_checkpoints							).to(__a							)
												__lowercase      :							Tuple				=     self.default_image_processor
												__lowercase      :							int				=     prepare_img()
												__lowercase      :							Dict				=     image_processor(__a		,     return_tensors="""pt"""							).to(__a							)
												__lowercase      :							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(__a		,     (1, 3, 384, 384)							)

												with torch.no_grad():
																			__lowercase      :							Union[str, Any]				=     model(**__a							)

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

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

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



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


												"""simple docstring"""

												__lowercase      :							List[Any]				=     MaskaFormerForUniversalSegmentation.from_pretrained(self.model_checkpoints							).to(__a							).eval()
												__lowercase      :							Optional[Any]				=     self.default_image_processor
												__lowercase      :							Union[str, Any]				=     prepare_img()
												__lowercase      :							Any				=     image_processor(__a		,     return_tensors="""pt"""							).to(__a							)
												__lowercase      :							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():
																			__lowercase      :							Optional[Any]				=     model(**__a							)
												# masks_queries_logits
												__lowercase      :							Dict				=     outputs.masks_queries_logits
												self.assertEqual(
												    masks_queries_logits.shape		,     (1, model.config.num_queries, inputs_shape[-2] // 4, inputs_shape[-1] // 4)							)
												__lowercase      :							Dict				=     [
												    [-8.7839, -9.0056, -8.8121],
												    [-7.4104, -7.0313, -6.5401],
												    [-6.6105, -6.3427, -6.4675],
												]
												__lowercase      :							int				=     torch.tensor(__a							).to(__a							)
												self.assertTrue(torch.allclose(masks_queries_logits[0, 0, :3, :3]		,     __a		,     atol=__a							)							)
												# class_queries_logits
												__lowercase      :							Dict				=     outputs.class_queries_logits
												self.assertEqual(class_queries_logits.shape		,     (1, model.config.num_queries, model.config.num_labels + 1)							)
												__lowercase      :							List[str]				=     torch.tensor(
												    [
												        [1.8324, -8.0835, -4.1922],
												        [0.8450, -9.0050, -3.6053],
												        [0.3045, -7.7293, -3.0275],
												    ]							).to(__a							)
												self.assertTrue(torch.allclose(outputs.class_queries_logits[0, :3, :3]		,     __a		,     atol=__a							)							)



					def       lowerCAmelCase				(  self				:	List[str]							)			->  Optional[int]:


												"""simple docstring"""

												__lowercase      :							str				=     MaskaFormerForUniversalSegmentation.from_pretrained(self.model_checkpoints							).to(__a							).eval()
												__lowercase      :							Optional[int]				=     self.default_image_processor

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

												__lowercase      :							Optional[Any]				=     inputs["""pixel_values"""].to(__a							)
												__lowercase      :							Optional[int]				=     [el.to(__a							) for el in inputs["""mask_labels"""]]
												__lowercase      :							str				=     [el.to(__a							) for el in inputs["""class_labels"""]]

												with torch.no_grad():
																			__lowercase      :							Optional[Any]				=     model(**__a							)

												self.assertTrue(outputs.loss is not None							)







from .imports import is_tqdm_available


if is_tqdm_available():
	from tqdm.auto import tqdm as _tqdm

from ..state import PartialState
def       snake_case_		(     lowerCAmelCase_       :						bool = True       ,       *lowerCAmelCase_       :						int       ,       **lowerCAmelCase_       :						List[str]      ):
							if not is_tqdm_available():
														raise ImportError("""Accelerate's `tqdm` module requires `tqdm` to be installed. Please run `pip install tqdm`."""      )
							__lowercase      :							List[str]				=     False
							if main_process_only:
														__lowercase      :							Optional[int]				=     PartialState().local_process_index == 0
							return _tqdm(*lowerCAmelCase_       ,       **lowerCAmelCase_       ,       disable=lowerCAmelCase_      )