import inspect
import unittest

from huggingface_hub import hf_hub_download

from transformers import ConvNextConfig, UperNetConfig
from transformers.testing_utils import require_torch, require_torch_multi_gpu, require_vision, slow, torch_device
from transformers.utils import is_torch_available, is_vision_available

from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, _config_zero_init, floats_tensor, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
     import torch

     from transformers import UperNetForSemanticSegmentation
     from transformers.models.upernet.modeling_upernet import UPERNET_PRETRAINED_MODEL_ARCHIVE_LIST


if is_vision_available():
     from PIL import Image

     from transformers import AutoImageProcessor
class    A_ :

  def __init__(		self :      int     ,							__SCREAMING_SNAKE_CASE :      List[str]     ,							__SCREAMING_SNAKE_CASE :      int=13     ,							__SCREAMING_SNAKE_CASE :      Optional[Any]=32     ,							__SCREAMING_SNAKE_CASE :      Optional[int]=3     ,							__SCREAMING_SNAKE_CASE :      List[Any]=4     ,							__SCREAMING_SNAKE_CASE :      Dict=[10, 20, 30, 40]     ,							__SCREAMING_SNAKE_CASE :      List[str]=[2, 2, 3, 2]     ,							__SCREAMING_SNAKE_CASE :      Union[str, Any]=True     ,							__SCREAMING_SNAKE_CASE :      Optional[Any]=True     ,							__SCREAMING_SNAKE_CASE :      str=37     ,							__SCREAMING_SNAKE_CASE :      str="gelu"     ,							__SCREAMING_SNAKE_CASE :      Dict=10     ,							__SCREAMING_SNAKE_CASE :      Union[str, Any]=0.02     ,							__SCREAMING_SNAKE_CASE :      List[str]=["stage2", "stage3", "stage4"]     ,							__SCREAMING_SNAKE_CASE :      Union[str, Any]=3     ,							__SCREAMING_SNAKE_CASE :      Any=None     ,							):
   __a       =  parent
   __a       =  batch_size
   __a       =  image_size
   __a       =  num_channels
   __a       =  num_stages
   __a       =  hidden_sizes
   __a       =  depths
   __a       =  is_training
   __a       =  use_labels
   __a       =  intermediate_size
   __a       =  hidden_act
   __a       =  type_sequence_label_size
   __a       =  initializer_range
   __a       =  out_features
   __a       =  num_labels
   __a       =  scope
   __a       =  num_stages

  def 		_UpperCAmelCase   (		self :      Dict       ):
   __a       =  floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]       )

   __a       =  None
   if self.use_labels:
    __a       =  ids_tensor([self.batch_size]     ,							self.type_sequence_label_size       )

   __a       =  self.get_config()

   return config, pixel_values, labels

  def 		_UpperCAmelCase   (		self :      int       ):
   return ConvNextConfig(
       num_channels=self.num_channels     ,							num_stages=self.num_stages     ,							hidden_sizes=self.hidden_sizes     ,							depths=self.depths     ,							is_training=self.is_training     ,							intermediate_size=self.intermediate_size     ,							hidden_act=self.hidden_act     ,							out_features=self.out_features     ,							)

  def 		_UpperCAmelCase   (		self :      str       ):
   return UperNetConfig(
       backbone_config=self.get_backbone_config()     ,							hidden_size=5_12     ,							pool_scales=[1, 2, 3, 6]     ,							use_auxiliary_head=__lowerCamelCase     ,							auxiliary_loss_weight=0.4     ,							auxiliary_in_channels=40     ,							auxiliary_channels=2_56     ,							auxiliary_num_convs=1     ,							auxiliary_concat_input=__lowerCamelCase     ,							loss_ignore_index=2_55     ,							num_labels=self.num_labels     ,							)

  def 		_UpperCAmelCase   (		self :      Union[str, Any]     ,							__SCREAMING_SNAKE_CASE :      int     ,							__SCREAMING_SNAKE_CASE :      Union[str, Any]     ,							__SCREAMING_SNAKE_CASE :      int       ):
   __a       =  UperNetForSemanticSegmentation(config=__lowerCamelCase       )
   model.to(__lowerCamelCase       )
   model.eval()
   __a       =  model(__lowerCamelCase       )
   self.parent.assertEqual(
       result.logits.shape     ,							(self.batch_size, self.num_labels, self.image_size, self.image_size)       )

  def 		_UpperCAmelCase   (		self :      str       ):
   __a       =  self.prepare_config_and_inputs()
   (
       __a     
   )       =  config_and_inputs
   __a       =  {"pixel_values": pixel_values}
   return config, inputs_dict
@require_torch
class    A_ (       lowercase__			,       lowercase__			,       unittest.TestCase ):
  _SCREAMING_SNAKE_CASE    				=			(UperNetForSemanticSegmentation,) if is_torch_available() else ()
  _SCREAMING_SNAKE_CASE    				=			{"""image-segmentation""": UperNetForSemanticSegmentation} if is_torch_available() else {}
  _SCREAMING_SNAKE_CASE    				=			False
  _SCREAMING_SNAKE_CASE    				=			False
  _SCREAMING_SNAKE_CASE    				=			False
  _SCREAMING_SNAKE_CASE    				=			False
  _SCREAMING_SNAKE_CASE    				=			False
  _SCREAMING_SNAKE_CASE    				=			False

  def 		_UpperCAmelCase   (		self :      Union[str, Any]       ):
   __a       =  UperNetModelTester(self       )
   __a       =  ConfigTester(self     ,							config_class=__lowerCamelCase     ,							has_text_modality=__lowerCamelCase     ,							hidden_size=37       )

  def 		_UpperCAmelCase   (		self :      Tuple       ):
   self.create_and_test_config_common_properties()
   self.config_tester.create_and_test_config_to_json_string()
   self.config_tester.create_and_test_config_to_json_file()
   self.config_tester.create_and_test_config_from_and_save_pretrained()
   self.config_tester.create_and_test_config_with_num_labels()
   self.config_tester.check_config_can_be_init_without_params()
   self.config_tester.check_config_arguments_init()

  def 		_UpperCAmelCase   (		self :      Optional[Any]       ):
   return

  def 		_UpperCAmelCase   (		self :      Union[str, Any]       ):
   __a       =  self.model_tester.prepare_config_and_inputs_for_common()

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

    __a       =  ["pixel_values"]
    self.assertListEqual(arg_names[:1]     ,							__lowerCamelCase       )

  def 		_UpperCAmelCase   (		self :      List[str]       ):
   __a       =  self.model_tester.prepare_config_and_inputs()
   self.model_tester.create_and_check_for_semantic_segmentation(*__lowerCamelCase       )

  @unittest.skip(reason="UperNet does not use inputs_embeds"       )
  def 		_UpperCAmelCase   (		self :      List[str]       ):
   pass

  @unittest.skip(reason="UperNet does not support input and output embeddings"       )
  def 		_UpperCAmelCase   (		self :      Optional[Any]       ):
   pass

  @unittest.skip(reason="UperNet does not have a base model"       )
  def 		_UpperCAmelCase   (		self :      Optional[int]       ):
   pass

  @unittest.skip(reason="UperNet does not have a base model"       )
  def 		_UpperCAmelCase   (		self :      Optional[Any]       ):
   pass

  @require_torch_multi_gpu
  @unittest.skip(reason="UperNet has some layers using `add_module` which doesn't work well with `nn.DataParallel`"       )
  def 		_UpperCAmelCase   (		self :      Union[str, Any]       ):
   pass

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

  def 		_UpperCAmelCase   (		self :      Optional[int]       ):



   def check_hidden_states_output(__SCREAMING_SNAKE_CASE :      Optional[Any]     ,							__SCREAMING_SNAKE_CASE :      List[str]     ,							__SCREAMING_SNAKE_CASE :      List[Any]       ):
    __a       =  model_class(__lowerCamelCase       )
    model.to(__lowerCamelCase       )
    model.eval()

    with torch.no_grad():
     __a       =  model(**self._prepare_for_class(__lowerCamelCase     ,							__lowerCamelCase       )       )

    __a       =  outputs.encoder_hidden_states if config.is_encoder_decoder else outputs.hidden_states

    __a       =  self.model_tester.num_stages
    self.assertEqual(len(__lowerCamelCase       )     ,							expected_num_stages + 1       )

    # ConvNext's feature maps are of shape (batch_size, num_channels, height, width)
    self.assertListEqual(
        list(hidden_states[0].shape[-2:]       )     ,							[self.model_tester.image_size // 4, self.model_tester.image_size // 4]     ,							)

   __a       =  self.model_tester.prepare_config_and_inputs_for_common()

   for model_class in self.all_model_classes:
    __a       =  True
    check_hidden_states_output(__lowerCamelCase     ,							__lowerCamelCase     ,							__lowerCamelCase       )

    # check that output_hidden_states also work using config
    del inputs_dict["output_hidden_states"]
    __a       =  True

    check_hidden_states_output(__lowerCamelCase     ,							__lowerCamelCase     ,							__lowerCamelCase       )

  def 		_UpperCAmelCase   (		self :      Tuple       ):
   __a       =  self.model_tester.prepare_config_and_inputs_for_common()

   __a       =  _config_zero_init(__lowerCamelCase       )
   __a       =  _config_zero_init(configs_no_init.backbone_config       )
   for model_class in self.all_model_classes:
    __a       =  model_class(config=__lowerCamelCase       )
    for name, param in model.named_parameters():
     if param.requires_grad:
      self.assertIn(
          ((param.data.mean() * 1E9).round() / 1E9).item()     ,							[0.0, 1.0]     ,							msg=f"""Parameter {name} of model {model_class} seems not properly initialized"""     ,							)

  @unittest.skip(reason="UperNet does not have tied weights"       )
  def 		_UpperCAmelCase   (		self :      List[str]       ):
   pass

  @slow
  def 		_UpperCAmelCase   (		self :      Union[str, Any]       ):



   for model_name in UPERNET_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
    __a       =  UperNetForSemanticSegmentation.from_pretrained(__lowerCamelCase       )
    self.assertIsNotNone(__lowerCamelCase       )



def   __A		(			):



 """simple docstring"""


 __a       =  hf_hub_download(
     repo_id="hf-internal-testing/fixtures_ade20k"							,				repo_type="dataset"							,				filename="ADE_val_00000001.jpg"					)
 __a       =  Image.open(_A					).convert("RGB"					)
 return image






@require_torch
@require_vision
@slow
class    A_ (       unittest.TestCase ):

  def 		_UpperCAmelCase   (		self :      str       ):
   __a       =  AutoImageProcessor.from_pretrained("openmmlab/upernet-swin-tiny"       )
   __a       =  UperNetForSemanticSegmentation.from_pretrained("openmmlab/upernet-swin-tiny"       ).to(__lowerCamelCase       )

   __a       =  prepare_img()
   __a       =  processor(images=__lowerCamelCase     ,							return_tensors="pt"       ).to(__lowerCamelCase       )

   with torch.no_grad():
    __a       =  model(**__lowerCamelCase       )

   __a       =  torch.Size((1, model.config.num_labels, 5_12, 5_12)       )
   self.assertEqual(outputs.logits.shape     ,							__lowerCamelCase       )

   __a       =  torch.tensor(
       [[-7.59_58, -7.59_58, -7.43_02], [-7.59_58, -7.59_58, -7.43_02], [-7.47_97, -7.47_97, -7.30_68]]       ).to(__lowerCamelCase       )
   self.assertTrue(torch.allclose(outputs.logits[0, 0, :3, :3]     ,							__lowerCamelCase     ,							atol=1E-4       )       )

  def 		_UpperCAmelCase   (		self :      List[str]       ):
   __a       =  AutoImageProcessor.from_pretrained("openmmlab/upernet-convnext-tiny"       )
   __a       =  UperNetForSemanticSegmentation.from_pretrained("openmmlab/upernet-convnext-tiny"       ).to(__lowerCamelCase       )

   __a       =  prepare_img()
   __a       =  processor(images=__lowerCamelCase     ,							return_tensors="pt"       ).to(__lowerCamelCase       )

   with torch.no_grad():
    __a       =  model(**__lowerCamelCase       )

   __a       =  torch.Size((1, model.config.num_labels, 5_12, 5_12)       )
   self.assertEqual(outputs.logits.shape     ,							__lowerCamelCase       )

   __a       =  torch.tensor(
       [[-8.81_10, -8.81_10, -8.65_21], [-8.81_10, -8.81_10, -8.65_21], [-8.77_46, -8.77_46, -8.61_30]]       ).to(__lowerCamelCase       )
   self.assertTrue(torch.allclose(outputs.logits[0, 0, :3, :3]     ,							__lowerCamelCase     ,							atol=1E-4       )       )



import pytest

import datasets.config
from datasets.utils.info_utils import is_small_dataset



@pytest.mark.parametrize("dataset_size"							,				[None, 400 * 2**20, 600 * 2**20]					)
@pytest.mark.parametrize("input_in_memory_max_size"							,				["default", 0, 100 * 2**20, 900 * 2**20]					)
def   __A		(			_A							,				_A							,				_A					):



 """simple docstring"""


 if input_in_memory_max_size != "default":
  monkeypatch.setattr(datasets.config							,				"IN_MEMORY_MAX_SIZE"							,				_A					)
 __a       =  datasets.config.IN_MEMORY_MAX_SIZE
 if input_in_memory_max_size == "default":
  assert in_memory_max_size == 0
 else:
  assert in_memory_max_size == input_in_memory_max_size
 if dataset_size and in_memory_max_size:
  __a       =  dataset_size < in_memory_max_size
 else:
  __a       =  False
 __a       =  is_small_dataset(_A					)
 assert result == expected


