
'''simple docstring'''
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    UpperCAmelCase_				:





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





 def        snake_case__		(    self):
     snake_case_ :    Dict     			=						floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size])

     snake_case_ :    str     			=						None
     if self.use_labels:
         snake_case_ :    Any     			=						ids_tensor([self.batch_size] ,       self.type_sequence_label_size)

     snake_case_ :    int     			=						self.get_config()

     return config, pixel_values, labels





 def        snake_case__		(    self):
     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        snake_case__		(    self):
     return UperNetConfig(
         backbone_config=self.get_backbone_config() ,       hidden_size=5_12 ,       pool_scales=[1, 2, 3, 6] ,       use_auxiliary_head=lowercase_ ,       auxiliary_loss_weight=0.4 ,       auxiliary_in_channels=40 ,       auxiliary_channels=2_56 ,       auxiliary_num_convs=1 ,       auxiliary_concat_input=lowercase_ ,       loss_ignore_index=2_55 ,       num_labels=self.num_labels ,       )





 def        snake_case__		(    self ,       lowercase_ ,       lowercase_ ,       lowercase_):
     snake_case_ :    Any     			=						UperNetForSemanticSegmentation(config=lowercase_)
     model.to(lowercase_)
     model.eval()
     snake_case_ :    Optional[int]     			=						model(lowercase_)
     self.parent.assertEqual(
         result.logits.shape ,       (self.batch_size, self.num_labels, self.image_size, self.image_size))





 def        snake_case__		(    self):
     snake_case_ :    Tuple     			=						self.prepare_config_and_inputs()
     (
         (
         snake_case_
     )			, (
         snake_case_
     )			, (
         snake_case_
     )			, 
     ) :    str     			=						config_and_inputs
     snake_case_ :    str     			=						{"pixel_values": pixel_values}
     return config, inputs_dict






@require_torch
class    UpperCAmelCase_				(	snake_case__    ,						snake_case__    ,						unittest.TestCase					):
 UpperCAmelCase_		    =						(UperNetForSemanticSegmentation,) if is_torch_available() else ()
 UpperCAmelCase_		    =						{"""image-segmentation""": UperNetForSemanticSegmentation} if is_torch_available() else {}
 UpperCAmelCase_		    =						False
 UpperCAmelCase_		    =						False
 UpperCAmelCase_		    =						False
 UpperCAmelCase_		    =						False
 UpperCAmelCase_		    =						False
 UpperCAmelCase_		    =						False





 def        snake_case__		(    self):
     snake_case_ :    List[Any]     			=						UperNetModelTester(self)
     snake_case_ :    Dict     			=						ConfigTester(self ,       config_class=lowercase_ ,       has_text_modality=lowercase_ ,       hidden_size=37)





 def        snake_case__		(    self):
     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        snake_case__		(    self):
     return





 def        snake_case__		(    self):
     snake_case_			, snake_case_ :    List[Any]     			=						self.model_tester.prepare_config_and_inputs_for_common()

     for model_class in self.all_model_classes:
         snake_case_ :    int     			=						model_class(lowercase_)
         snake_case_ :    Union[str, Any]     			=						inspect.signature(model.forward)
         # signature.parameters is an OrderedDict => so arg_names order is deterministic
         snake_case_ :    Union[str, Any]     			=						[*signature.parameters.keys()]

         snake_case_ :    Union[str, Any]     			=						["pixel_values"]
         self.assertListEqual(arg_names[:1] ,       lowercase_)





 def        snake_case__		(    self):
     snake_case_ :    List[str]     			=						self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_for_semantic_segmentation(*lowercase_)





 @unittest.skip(reason="UperNet does not use inputs_embeds")
 def        snake_case__		(    self):
     pass





 @unittest.skip(reason="UperNet does not support input and output embeddings")
 def        snake_case__		(    self):
     pass





 @unittest.skip(reason="UperNet does not have a base model")
 def        snake_case__		(    self):
     pass





 @unittest.skip(reason="UperNet does not have a base model")
 def        snake_case__		(    self):
     pass





 @require_torch_multi_gpu
 @unittest.skip(reason="UperNet has some layers using `add_module` which doesn't work well with `nn.DataParallel`")
 def        snake_case__		(    self):
     pass





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





 def        snake_case__		(    self):
     def check_hidden_states_output(lowercase_ ,       lowercase_ ,       lowercase_):
         snake_case_ :    Optional[Any]     			=						model_class(lowercase_)
         model.to(lowercase_)
         model.eval()

         with torch.no_grad():
             snake_case_ :    Optional[int]     			=						model(**self._prepare_for_class(lowercase_ ,       lowercase_))

         snake_case_ :    List[Any]     			=						outputs.encoder_hidden_states if config.is_encoder_decoder else outputs.hidden_states

         snake_case_ :    Optional[int]     			=						self.model_tester.num_stages
         self.assertEqual(len(lowercase_) ,       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] ,       )

     snake_case_			, snake_case_ :    List[str]     			=						self.model_tester.prepare_config_and_inputs_for_common()

     for model_class in self.all_model_classes:
         snake_case_ :    Dict     			=						True
         check_hidden_states_output(lowercase_ ,       lowercase_ ,       lowercase_)

         # check that output_hidden_states also work using config
         del inputs_dict["output_hidden_states"]
         snake_case_ :    Optional[int]     			=						True

         check_hidden_states_output(lowercase_ ,       lowercase_ ,       lowercase_)





 def        snake_case__		(    self):
     snake_case_			, snake_case_ :    Any     			=						self.model_tester.prepare_config_and_inputs_for_common()

     snake_case_ :    Tuple     			=						_config_zero_init(lowercase_)
     snake_case_ :    Optional[Any]     			=						_config_zero_init(configs_no_init.backbone_config)
     for model_class in self.all_model_classes:
         snake_case_ :    str     			=						model_class(config=lowercase_)
         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        snake_case__		(    self):
     pass





 @slow
 def        snake_case__		(    self):
     for model_name in UPERNET_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
         snake_case_ :    List[str]     			=						UperNetForSemanticSegmentation.from_pretrained(lowercase_)
         self.assertIsNotNone(lowercase_)


def 						UpperCamelCase_     (	):
    """simple docstring"""






    snake_case_ :    int     			=						hf_hub_download(
        repo_id="hf-internal-testing/fixtures_ade20k",		repo_type="dataset",		filename="ADE_val_00000001.jpg"       )
    snake_case_ :    str     			=						Image.open(__SCREAMING_SNAKE_CASE       ).convert("RGB"       )
    return image






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





 def        snake_case__		(    self):
     snake_case_ :    Tuple     			=						AutoImageProcessor.from_pretrained("openmmlab/upernet-swin-tiny")
     snake_case_ :    Optional[int]     			=						UperNetForSemanticSegmentation.from_pretrained("openmmlab/upernet-swin-tiny").to(lowercase_)

     snake_case_ :    List[Any]     			=						prepare_img()
     snake_case_ :    Optional[int]     			=						processor(images=lowercase_ ,       return_tensors="pt").to(lowercase_)

     with torch.no_grad():
         snake_case_ :    Dict     			=						model(**lowercase_)

     snake_case_ :    List[Any]     			=						torch.Size((1, model.config.num_labels, 5_12, 5_12))
     self.assertEqual(outputs.logits.shape ,       lowercase_)

     snake_case_ :    str     			=						torch.tensor(
         [[-7.5_958, -7.5_958, -7.4_302], [-7.5_958, -7.5_958, -7.4_302], [-7.4_797, -7.4_797, -7.3_068]]).to(lowercase_)
     self.assertTrue(torch.allclose(outputs.logits[0, 0, :3, :3] ,       lowercase_ ,       atol=1E-4))





 def        snake_case__		(    self):
     snake_case_ :    int     			=						AutoImageProcessor.from_pretrained("openmmlab/upernet-convnext-tiny")
     snake_case_ :    Tuple     			=						UperNetForSemanticSegmentation.from_pretrained("openmmlab/upernet-convnext-tiny").to(lowercase_)

     snake_case_ :    Optional[Any]     			=						prepare_img()
     snake_case_ :    Optional[int]     			=						processor(images=lowercase_ ,       return_tensors="pt").to(lowercase_)

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

     snake_case_ :    Any     			=						torch.Size((1, model.config.num_labels, 5_12, 5_12))
     self.assertEqual(outputs.logits.shape ,       lowercase_)

     snake_case_ :    List[str]     			=						torch.tensor(
         [[-8.8_110, -8.8_110, -8.6_521], [-8.8_110, -8.8_110, -8.6_521], [-8.7_746, -8.7_746, -8.6_130]]).to(lowercase_)
     self.assertTrue(torch.allclose(outputs.logits[0, 0, :3, :3] ,       lowercase_ ,       atol=1E-4))






'''simple docstring'''
import contextlib
from multiprocessing import Pool, RLock

from tqdm.auto import tqdm

from ..utils import experimental, logging


a_											=						logging.get_logger(__name__)






class    UpperCAmelCase_				:
 UpperCAmelCase_		    =						None


@experimental
def 						UpperCamelCase_     (	__SCREAMING_SNAKE_CASE,		__SCREAMING_SNAKE_CASE,		__SCREAMING_SNAKE_CASE,		__SCREAMING_SNAKE_CASE,		__SCREAMING_SNAKE_CASE,		__SCREAMING_SNAKE_CASE,		__SCREAMING_SNAKE_CASE       ):
    """simple docstring"""






    if ParallelBackendConfig.backend_name is None:
        return _map_with_multiprocessing_pool(
            __SCREAMING_SNAKE_CASE,		__SCREAMING_SNAKE_CASE,		__SCREAMING_SNAKE_CASE,		__SCREAMING_SNAKE_CASE,		__SCREAMING_SNAKE_CASE,		__SCREAMING_SNAKE_CASE,		__SCREAMING_SNAKE_CASE       )

    return _map_with_joblib(__SCREAMING_SNAKE_CASE,		__SCREAMING_SNAKE_CASE,		__SCREAMING_SNAKE_CASE,		__SCREAMING_SNAKE_CASE,		__SCREAMING_SNAKE_CASE,		__SCREAMING_SNAKE_CASE,		__SCREAMING_SNAKE_CASE       )

def 						UpperCamelCase_     (	__SCREAMING_SNAKE_CASE,		__SCREAMING_SNAKE_CASE,		__SCREAMING_SNAKE_CASE,		__SCREAMING_SNAKE_CASE,		__SCREAMING_SNAKE_CASE,		__SCREAMING_SNAKE_CASE,		__SCREAMING_SNAKE_CASE       ):
    """simple docstring"""






    snake_case_ :    str     			=						num_proc if num_proc <= len(__SCREAMING_SNAKE_CASE       ) else len(__SCREAMING_SNAKE_CASE       )
    snake_case_ :    Optional[int]     			=						[]  # We organize the splits ourselve (contiguous splits)
    for index in range(__SCREAMING_SNAKE_CASE       ):
        snake_case_ :    List[str]     			=						len(__SCREAMING_SNAKE_CASE       ) // num_proc
        snake_case_ :    Optional[Any]     			=						len(__SCREAMING_SNAKE_CASE       ) % num_proc
        snake_case_ :    Optional[Any]     			=						div * index + min(__SCREAMING_SNAKE_CASE,		__SCREAMING_SNAKE_CASE       )
        snake_case_ :    Union[str, Any]     			=						start + div + (1 if index < mod else 0)
        split_kwds.append((function, iterable[start:end], types, index, disable_tqdm, desc)       )

    if len(__SCREAMING_SNAKE_CASE       ) != sum(len(i[1]       ) for i in split_kwds       ):
        raise ValueError(
            f'Error dividing inputs iterable among processes. '
            f'Total number of objects {len(__SCREAMING_SNAKE_CASE       )}, '
            f'length: {sum(len(i[1]       ) for i in split_kwds       )}'       )

    logger.info(
        f'Spawning {num_proc} processes for {len(__SCREAMING_SNAKE_CASE       )} objects in slices of {[len(i[1]       ) for i in split_kwds]}'       )
    snake_case_			, snake_case_ :    Tuple     			=						None, None
    if not disable_tqdm:
        snake_case_			, snake_case_ :    int     			=						(RLock(),), tqdm.set_lock
    with Pool(__SCREAMING_SNAKE_CASE,		initargs=__SCREAMING_SNAKE_CASE,		initializer=__SCREAMING_SNAKE_CASE       ) as pool:
        snake_case_ :    Union[str, Any]     			=						pool.map(__SCREAMING_SNAKE_CASE,		__SCREAMING_SNAKE_CASE       )
    logger.info(f'Finished {num_proc} processes'       )
    snake_case_ :    str     			=						[obj for proc_res in mapped for obj in proc_res]
    logger.info(f'Unpacked {len(__SCREAMING_SNAKE_CASE       )} objects'       )

    return mapped

def 						UpperCamelCase_     (	__SCREAMING_SNAKE_CASE,		__SCREAMING_SNAKE_CASE,		__SCREAMING_SNAKE_CASE,		__SCREAMING_SNAKE_CASE,		__SCREAMING_SNAKE_CASE,		__SCREAMING_SNAKE_CASE,		__SCREAMING_SNAKE_CASE       ):
    """simple docstring"""






    import joblib

    with joblib.parallel_backend(ParallelBackendConfig.backend_name,		n_jobs=__SCREAMING_SNAKE_CASE       ):
        return joblib.Parallel()(
            joblib.delayed(__SCREAMING_SNAKE_CASE       )((function, obj, types, None, True, None)       ) for obj in iterable       )


@experimental
@contextlib.contextmanager
def 						UpperCamelCase_     (	__SCREAMING_SNAKE_CASE       ):
    """simple docstring"""






    snake_case_ :    Union[str, Any]     			=						backend_name

    if backend_name == "spark":
        from joblibspark import register_spark

        register_spark()

        # TODO: call create_cache_and_write_probe if "download" in steps
        # TODO: raise NotImplementedError when Dataset.map etc is called

    try:
        yield
    finally:
        snake_case_ :    Any     			=						None




