





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






							def __init__(  self					,	__lowerCAmelCase					,	__lowerCAmelCase=2					,	__lowerCAmelCase=True					,	__lowerCAmelCase=False					,	__lowerCAmelCase=10					,	__lowerCAmelCase=3					,	__lowerCAmelCase=32 * 8					,	__lowerCAmelCase=32 * 8					,	__lowerCAmelCase=4					,	__lowerCAmelCase=64					,	):
														UpperCamelCase__			       =     parent
														UpperCamelCase__			       =     batch_size
														UpperCamelCase__			       =     is_training
														UpperCamelCase__			       =     use_auxiliary_loss
														UpperCamelCase__			       =     num_queries
														UpperCamelCase__			       =     num_channels
														UpperCamelCase__			       =     min_size
														UpperCamelCase__			       =     max_size
														UpperCamelCase__			       =     num_labels
														UpperCamelCase__			       =     hidden_dim
														UpperCamelCase__			       =     hidden_dim






							def 			_lowerCamelCase						(  self					):
														UpperCamelCase__			       =     floats_tensor([self.batch_size, self.num_channels, self.min_size, self.max_size]					).to(
														    _A					)

														UpperCamelCase__			       =     torch.ones([self.batch_size, self.min_size, self.max_size]					,	device=_A					)

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

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






							def 			_lowerCamelCase						(  self					):
														UpperCamelCase__			       =     MaskaFormerConfig(
														    hidden_size=self.hidden_dim					,	)
														UpperCamelCase__			       =     self.num_queries
														UpperCamelCase__			       =     self.num_labels

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

														UpperCamelCase__			       =     64
														UpperCamelCase__			       =     128
														UpperCamelCase__			       =     self.hidden_dim
														UpperCamelCase__			       =     self.hidden_dim
														UpperCamelCase__			       =     self.hidden_dim
														return config






							def 			_lowerCamelCase						(  self					):
														UpperCamelCase__			       =     self.prepare_config_and_inputs()
														UpperCamelCase__			       =     {'pixel_values': pixel_values, 'pixel_mask': pixel_mask}
														return config, inputs_dict






							def 			_lowerCamelCase						(  self					,	__lowerCAmelCase					,	__lowerCAmelCase					):
														UpperCamelCase__			       =     output.encoder_hidden_states
														UpperCamelCase__			       =     output.pixel_decoder_hidden_states
														UpperCamelCase__			       =     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					,	__lowerCAmelCase					,	__lowerCAmelCase					,	__lowerCAmelCase					,	__lowerCAmelCase=False					):
														with torch.no_grad():
																					UpperCamelCase__			       =     MaskaFormerModel(config=_A					)
																					model.to(_A					)
																					model.eval()

																					UpperCamelCase__			       =     model(pixel_values=_A					,	pixel_mask=_A					)
																					UpperCamelCase__			       =     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					,	__lowerCAmelCase					,	__lowerCAmelCase					,	__lowerCAmelCase					,	__lowerCAmelCase					,	__lowerCAmelCase					):
														UpperCamelCase__			       =     MaskaFormerForUniversalSegmentation(config=_A					)
														model.to(_A					)
														model.eval()

														def comm_check_on_output(__lowerCAmelCase					):
																					# 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__			       =     model(pixel_values=_A					,	pixel_mask=_A					)
																					UpperCamelCase__			       =     model(_A					)

																					comm_check_on_output(_A					)

																					UpperCamelCase__			       =     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   (       _a				,						_a				,						unittest.TestCase    ):
							snake_case       : Dict         =    (MaskaFormerModel, MaskaFormerForUniversalSegmentation) if is_torch_available() else ()
							snake_case       : Union[str, Any]         =    {"""feature-extraction""": MaskaFormerModel} if is_torch_available() else {}

							snake_case       : Union[str, Any]         =    False
							snake_case       : Dict         =    False
							snake_case       : Union[str, Any]         =    False
							snake_case       : Union[str, Any]         =    False






							def 			_lowerCamelCase						(  self					):
														UpperCamelCase__			       =     MaskaFormerModelTester(self					)
														UpperCamelCase__			       =     ConfigTester(self					,	config_class=_A					,	has_text_modality=_A					)






							def 			_lowerCamelCase						(  self					):
														self.config_tester.run_common_tests()






							def 			_lowerCamelCase						(  self					):
														UpperCamelCase__			       =     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					):
														UpperCamelCase__			       =     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					):
														pass






							@unittest.skip(reason="""Mask2Former does not have a get_input_embeddings method"""					)
							def 			_lowerCamelCase						(  self					):
														pass






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






							@unittest.skip(reason="""Mask2Former does not use token embeddings"""					)
							def 			_lowerCamelCase						(  self					):
														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					):
														pass






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






							def 			_lowerCamelCase						(  self					):
														UpperCamelCase__			       =     self.model_tester.prepare_config_and_inputs_for_common()

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

																					UpperCamelCase__			       =     ['pixel_values']
																					self.assertListEqual(arg_names[:1]					,	_A					)






							@slow
							def 			_lowerCamelCase						(  self					):
														for model_name in ["facebook/mask2former-swin-small-coco-instance"]:
																					UpperCamelCase__			       =     MaskaFormerModel.from_pretrained(_A					)
																					self.assertIsNotNone(_A					)






							def 			_lowerCamelCase						(  self					):
														UpperCamelCase__			       =     (self.model_tester.min_size,) * 2
														UpperCamelCase__			       =     {
														    '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__			       =     self.model_tester.get_config()

														UpperCamelCase__			       =     MaskaFormerForUniversalSegmentation(_A					).to(_A					)
														UpperCamelCase__			       =     model(**_A					)
														self.assertTrue(outputs.loss is not None					)






							def 			_lowerCamelCase						(  self					):
														UpperCamelCase__			       =     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					):
														UpperCamelCase__			       =     self.model_tester.prepare_config_and_inputs_for_common()

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






							def 			_lowerCamelCase						(  self					):
														if not self.model_tester.is_training:
																					return

														UpperCamelCase__			       =     self.all_model_classes[1]
														UpperCamelCase__			       =     self.model_tester.prepare_config_and_inputs()

														UpperCamelCase__			       =     model_class(_A					)
														model.to(_A					)
														model.train()

														UpperCamelCase__			       =     model(_A					,	mask_labels=_A					,	class_labels=_A					).loss
														loss.backward()







							def 			_lowerCamelCase						(  self					):
														UpperCamelCase__			       =     self.all_model_classes[1]
														UpperCamelCase__			       =     self.model_tester.prepare_config_and_inputs()
														UpperCamelCase__			       =     True
														UpperCamelCase__			       =     True

														UpperCamelCase__			       =     model_class(_A					).to(_A					)
														model.train()

														UpperCamelCase__			       =     model(_A					,	mask_labels=_A					,	class_labels=_A					)

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

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

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

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






UpperCamelCase__     =	1e-4





def       _UpperCamelCase       ():






							"""simple docstring"""





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

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






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






							@cached_property
							def 			_lowerCamelCase						(  self					):
														return MaskaFormerImageProcessor.from_pretrained(self.model_checkpoints					) if is_vision_available() else None






							def 			_lowerCamelCase						(  self					):
														UpperCamelCase__			       =     MaskaFormerModel.from_pretrained(self.model_checkpoints					).to(_A					)
														UpperCamelCase__			       =     self.default_image_processor
														UpperCamelCase__			       =     prepare_img()
														UpperCamelCase__			       =     image_processor(_A					,	return_tensors="""pt"""					).to(_A					)
														UpperCamelCase__			       =     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__			       =     model(**_A					)

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

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

														UpperCamelCase__			       =     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					):
														UpperCamelCase__			       =     MaskaFormerForUniversalSegmentation.from_pretrained(self.model_checkpoints					).to(_A					).eval()
														UpperCamelCase__			       =     self.default_image_processor
														UpperCamelCase__			       =     prepare_img()
														UpperCamelCase__			       =     image_processor(_A					,	return_tensors="""pt"""					).to(_A					)
														UpperCamelCase__			       =     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__			       =     model(**_A					)
														# masks_queries_logits
														UpperCamelCase__			       =     outputs.masks_queries_logits
														self.assertEqual(
														    masks_queries_logits.shape					,	(1, model.config.num_queries, inputs_shape[-2] // 4, inputs_shape[-1] // 4)					)
														UpperCamelCase__			       =     [
														    [-8.7839, -9.0056, -8.8121],
														    [-7.4104, -7.0313, -6.5401],
														    [-6.6105, -6.3427, -6.4675],
														]
														UpperCamelCase__			       =     torch.tensor(_A					).to(_A					)
														self.assertTrue(torch.allclose(masks_queries_logits[0, 0, :3, :3]					,	_A					,	atol=_A					)					)
														# class_queries_logits
														UpperCamelCase__			       =     outputs.class_queries_logits
														self.assertEqual(class_queries_logits.shape					,	(1, model.config.num_queries, model.config.num_labels + 1)					)
														UpperCamelCase__			       =     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					):
														UpperCamelCase__			       =     MaskaFormerForUniversalSegmentation.from_pretrained(self.model_checkpoints					).to(_A					).eval()
														UpperCamelCase__			       =     self.default_image_processor

														UpperCamelCase__			       =     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__			       =     inputs['pixel_values'].to(_A					)
														UpperCamelCase__			       =     [el.to(_A					) for el in inputs['mask_labels']]
														UpperCamelCase__			       =     [el.to(_A					) for el in inputs['class_labels']]

														with torch.no_grad():
																					UpperCamelCase__			       =     model(**_A					)

														self.assertTrue(outputs.loss is not None					)











import shutil
import tempfile
import unittest

import numpy as np
import pytest

from transformers import is_speech_available, is_vision_available
from transformers.testing_utils import require_torch


if is_vision_available():
					from transformers import TvltImageProcessor

if is_speech_available():
					from transformers import TvltFeatureExtractor

from transformers import TvltProcessor

@require_torch
class     __SCREAMING_SNAKE_CASE   (       unittest.TestCase    ):






							def 			_lowerCamelCase						(  self					):
														UpperCamelCase__			       =     """ZinengTang/tvlt-base"""
														UpperCamelCase__			       =     tempfile.mkdtemp()






							def 			_lowerCamelCase						(  self					,	**__lowerCAmelCase					):
														return TvltImageProcessor.from_pretrained(self.checkpoint					,	**__lowerCAmelCase					)






							def 			_lowerCamelCase						(  self					,	**__lowerCAmelCase					):
														return TvltFeatureExtractor.from_pretrained(self.checkpoint					,	**__lowerCAmelCase					)






							def 			_lowerCamelCase						(  self					):
														shutil.rmtree(self.tmpdirname					)






							def 			_lowerCamelCase						(  self					):
														UpperCamelCase__			       =     self.get_image_processor()
														UpperCamelCase__			       =     self.get_feature_extractor()

														UpperCamelCase__			       =     TvltProcessor(image_processor=__lowerCAmelCase					,	feature_extractor=__lowerCAmelCase					)
														processor.save_pretrained(self.tmpdirname					)
														UpperCamelCase__			       =     TvltProcessor.from_pretrained(self.tmpdirname					)

														self.assertIsInstance(processor.feature_extractor					,	__lowerCAmelCase					)
														self.assertIsInstance(processor.image_processor					,	__lowerCAmelCase					)






							def 			_lowerCamelCase						(  self					):
														UpperCamelCase__			       =     self.get_image_processor()
														UpperCamelCase__			       =     self.get_feature_extractor()

														UpperCamelCase__			       =     TvltProcessor(image_processor=__lowerCAmelCase					,	feature_extractor=__lowerCAmelCase					)

														UpperCamelCase__			       =     np.ones([12000]					)

														UpperCamelCase__			       =     feature_extractor(__lowerCAmelCase					,	return_tensors="""np"""					)
														UpperCamelCase__			       =     processor(audio=__lowerCAmelCase					,	return_tensors="""np"""					)

														for key in audio_dict.keys():
																					self.assertAlmostEqual(audio_dict[key].sum()					,	input_processor[key].sum()					,	delta=1E-2					)






							def 			_lowerCamelCase						(  self					):
														UpperCamelCase__			       =     self.get_image_processor()
														UpperCamelCase__			       =     self.get_feature_extractor()

														UpperCamelCase__			       =     TvltProcessor(image_processor=__lowerCAmelCase					,	feature_extractor=__lowerCAmelCase					)

														UpperCamelCase__			       =     np.ones([3, 224, 224]					)

														UpperCamelCase__			       =     image_processor(__lowerCAmelCase					,	return_tensors="""np"""					)
														UpperCamelCase__			       =     processor(images=__lowerCAmelCase					,	return_tensors="""np"""					)

														for key in image_dict.keys():
																					self.assertAlmostEqual(image_dict[key].sum()					,	input_processor[key].sum()					,	delta=1E-2					)






							def 			_lowerCamelCase						(  self					):
														UpperCamelCase__			       =     self.get_image_processor()
														UpperCamelCase__			       =     self.get_feature_extractor()

														UpperCamelCase__			       =     TvltProcessor(image_processor=__lowerCAmelCase					,	feature_extractor=__lowerCAmelCase					)

														UpperCamelCase__			       =     np.ones([12000]					)
														UpperCamelCase__			       =     np.ones([3, 224, 224]					)

														UpperCamelCase__			       =     processor(audio=__lowerCAmelCase					,	images=__lowerCAmelCase					)

														self.assertListEqual(list(inputs.keys()					)					,	["""audio_values""", """audio_mask""", """pixel_values""", """pixel_mask"""]					)

														# test if it raises when no input is passed
														with pytest.raises(__lowerCAmelCase					):
																					processor()







							def 			_lowerCamelCase						(  self					):
														UpperCamelCase__			       =     self.get_image_processor()
														UpperCamelCase__			       =     self.get_feature_extractor()

														UpperCamelCase__			       =     TvltProcessor(image_processor=__lowerCAmelCase					,	feature_extractor=__lowerCAmelCase					)

														self.assertListEqual(
														    processor.model_input_names					,	image_processor.model_input_names + feature_extractor.model_input_names					,	msg="""`processor` and `image_processor`+`feature_extractor` model input names do not match"""					,	)




