






import inspect
import unittest

from datasets import load_dataset
from packaging import version

from transformers import BeitConfig
from transformers.models.auto import get_values
from transformers.testing_utils import require_torch, require_torch_multi_gpu, require_vision, slow, torch_device
from transformers.utils import cached_property, 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 torch import nn

					from transformers import (
					    MODEL_MAPPING,
					    BeitForImageClassification,
					    BeitForMaskedImageModeling,
					    BeitForSemanticSegmentation,
					    BeitModel,
					)
					from transformers.models.beit.modeling_beit import BEIT_PRETRAINED_MODEL_ARCHIVE_LIST


if is_vision_available():
					import PIL
					from PIL import Image

					from transformers import BeitImageProcessor


class     __SCREAMING_SNAKE_CASE					:






						def __init__(	self     ,						_UpperCamelCase     ,						_UpperCamelCase=1_00     ,						_UpperCamelCase=13     ,						_UpperCamelCase=30     ,						_UpperCamelCase=2     ,						_UpperCamelCase=3     ,						_UpperCamelCase=True     ,						_UpperCamelCase=True     ,						_UpperCamelCase=32     ,						_UpperCamelCase=4     ,						_UpperCamelCase=4     ,						_UpperCamelCase=37     ,						_UpperCamelCase="gelu"     ,						_UpperCamelCase=0.1     ,						_UpperCamelCase=0.1     ,						_UpperCamelCase=10     ,						_UpperCamelCase=0.02     ,						_UpperCamelCase=3     ,						_UpperCamelCase=None     ,						_UpperCamelCase=[0, 1, 2, 3]     ,						):






												"""simple docstring"""
												lowerCAmelCase__    						=						parent
												lowerCAmelCase__    						=						1_00
												lowerCAmelCase__    						=						batch_size
												lowerCAmelCase__    						=						image_size
												lowerCAmelCase__    						=						patch_size
												lowerCAmelCase__    						=						num_channels
												lowerCAmelCase__    						=						is_training
												lowerCAmelCase__    						=						use_labels
												lowerCAmelCase__    						=						hidden_size
												lowerCAmelCase__    						=						num_hidden_layers
												lowerCAmelCase__    						=						num_attention_heads
												lowerCAmelCase__    						=						intermediate_size
												lowerCAmelCase__    						=						hidden_act
												lowerCAmelCase__    						=						hidden_dropout_prob
												lowerCAmelCase__    						=						attention_probs_dropout_prob
												lowerCAmelCase__    						=						type_sequence_label_size
												lowerCAmelCase__    						=						initializer_range
												lowerCAmelCase__    						=						scope
												lowerCAmelCase__    						=						out_indices
												lowerCAmelCase__    						=						num_labels

												# in BeiT, the seq length equals the number of patches + 1 (we add 1 for the [CLS] token)
												lowerCAmelCase__    						=						(image_size // patch_size) ** 2
												lowerCAmelCase__    						=						num_patches + 1






						def       UpperCamelCase__      (	self   ):






												"""simple docstring"""
												lowerCAmelCase__    						=						floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]   )

												lowerCAmelCase__    						=						None
												lowerCAmelCase__    						=						None
												if self.use_labels:
																		lowerCAmelCase__    						=						ids_tensor([self.batch_size]     ,						self.type_sequence_label_size   )
																		lowerCAmelCase__    						=						ids_tensor([self.batch_size, self.image_size, self.image_size]     ,						self.num_labels   )

												lowerCAmelCase__    						=						self.get_config()

												return config, pixel_values, labels, pixel_labels






						def       UpperCamelCase__      (	self   ):






												"""simple docstring"""
												return BeitConfig(
												    vocab_size=self.vocab_size     ,						image_size=self.image_size     ,						patch_size=self.patch_size     ,						num_channels=self.num_channels     ,						hidden_size=self.hidden_size     ,						num_hidden_layers=self.num_hidden_layers     ,						num_attention_heads=self.num_attention_heads     ,						intermediate_size=self.intermediate_size     ,						hidden_act=self.hidden_act     ,						hidden_dropout_prob=self.hidden_dropout_prob     ,						attention_probs_dropout_prob=self.attention_probs_dropout_prob     ,						is_decoder=_UpperCamelCase     ,						initializer_range=self.initializer_range     ,						out_indices=self.out_indices     ,						)






						def       UpperCamelCase__      (	self     ,						_UpperCamelCase     ,						_UpperCamelCase     ,						_UpperCamelCase     ,						_UpperCamelCase   ):






												"""simple docstring"""
												lowerCAmelCase__    						=						BeitModel(config=_UpperCamelCase   )
												model.to(_UpperCamelCase   )
												model.eval()
												lowerCAmelCase__    						=						model(_UpperCamelCase   )
												self.parent.assertEqual(result.last_hidden_state.shape     ,						(self.batch_size, self.seq_length, self.hidden_size)   )






						def       UpperCamelCase__      (	self     ,						_UpperCamelCase     ,						_UpperCamelCase     ,						_UpperCamelCase     ,						_UpperCamelCase   ):






												"""simple docstring"""
												lowerCAmelCase__    						=						BeitForMaskedImageModeling(config=_UpperCamelCase   )
												model.to(_UpperCamelCase   )
												model.eval()
												lowerCAmelCase__    						=						model(_UpperCamelCase   )
												self.parent.assertEqual(result.logits.shape     ,						(self.batch_size, self.seq_length - 1, self.vocab_size)   )






						def       UpperCamelCase__      (	self     ,						_UpperCamelCase     ,						_UpperCamelCase     ,						_UpperCamelCase     ,						_UpperCamelCase   ):






												"""simple docstring"""
												lowerCAmelCase__    						=						self.type_sequence_label_size
												lowerCAmelCase__    						=						BeitForImageClassification(_UpperCamelCase   )
												model.to(_UpperCamelCase   )
												model.eval()
												lowerCAmelCase__    						=						model(_UpperCamelCase     ,						labels=_UpperCamelCase   )
												self.parent.assertEqual(result.logits.shape     ,						(self.batch_size, self.type_sequence_label_size)   )

												# test greyscale images
												lowerCAmelCase__    						=						1
												lowerCAmelCase__    						=						BeitForImageClassification(_UpperCamelCase   )
												model.to(_UpperCamelCase   )
												model.eval()

												lowerCAmelCase__    						=						floats_tensor([self.batch_size, 1, self.image_size, self.image_size]   )
												lowerCAmelCase__    						=						model(_UpperCamelCase     ,						labels=_UpperCamelCase   )
												self.parent.assertEqual(result.logits.shape     ,						(self.batch_size, self.type_sequence_label_size)   )






						def       UpperCamelCase__      (	self     ,						_UpperCamelCase     ,						_UpperCamelCase     ,						_UpperCamelCase     ,						_UpperCamelCase   ):






												"""simple docstring"""
												lowerCAmelCase__    						=						self.num_labels
												lowerCAmelCase__    						=						BeitForSemanticSegmentation(_UpperCamelCase   )
												model.to(_UpperCamelCase   )
												model.eval()
												lowerCAmelCase__    						=						model(_UpperCamelCase   )
												self.parent.assertEqual(
												    result.logits.shape     ,						(self.batch_size, self.num_labels, self.image_size * 2, self.image_size * 2)   )
												lowerCAmelCase__    						=						model(_UpperCamelCase     ,						labels=_UpperCamelCase   )
												self.parent.assertEqual(
												    result.logits.shape     ,						(self.batch_size, self.num_labels, self.image_size * 2, self.image_size * 2)   )







						def       UpperCamelCase__      (	self   ):






												"""simple docstring"""
												lowerCAmelCase__    						=						self.prepare_config_and_inputs()
												lowerCAmelCase__						,       lowerCAmelCase__						,       lowerCAmelCase__						,       lowerCAmelCase__    						=						config_and_inputs
												lowerCAmelCase__    						=						{'pixel_values': pixel_values}
												return config, inputs_dict




@require_torch
class     __SCREAMING_SNAKE_CASE					(   __lowercase		,					__lowercase		,					unittest.TestCase):
						_SCREAMING_SNAKE_CASE	:       Union[str, Any]   =					(
						    (BeitModel, BeitForImageClassification, BeitForMaskedImageModeling, BeitForSemanticSegmentation)
						    if is_torch_available()
						    else ()
						)
						_SCREAMING_SNAKE_CASE	:       List[str]   =					(
						    {
						        '''feature-extraction''': BeitModel,
						        '''image-classification''': BeitForImageClassification,
						        '''image-segmentation''': BeitForSemanticSegmentation,
						    }
						    if is_torch_available()
						    else {}
						)

						_SCREAMING_SNAKE_CASE	:       Dict   =					False
						_SCREAMING_SNAKE_CASE	:       List[Any]   =					False
						_SCREAMING_SNAKE_CASE	:       List[Any]   =					False






						def       UpperCamelCase__      (	self   ):






												"""simple docstring"""
												lowerCAmelCase__    						=						BeitModelTester(self   )
												lowerCAmelCase__    						=						ConfigTester(self     ,						config_class=_UpperCamelCase     ,						has_text_modality=_UpperCamelCase     ,						hidden_size=37   )






						def       UpperCamelCase__      (	self   ):






												"""simple docstring"""
												self.config_tester.run_common_tests()






						@unittest.skip(reason='BEiT does not use inputs_embeds'   )
						def       UpperCamelCase__      (	self   ):






												"""simple docstring"""
												pass






						@require_torch_multi_gpu
						@unittest.skip(reason='BEiT has some layers using `add_module` which doesn\'t work well with `nn.DataParallel`'   )
						def       UpperCamelCase__      (	self   ):






												"""simple docstring"""
												pass






						def       UpperCamelCase__      (	self   ):






												"""simple docstring"""
												lowerCAmelCase__						,       lowerCAmelCase__    						=						self.model_tester.prepare_config_and_inputs_for_common()

												for model_class in self.all_model_classes:
																		lowerCAmelCase__    						=						model_class(_UpperCamelCase   )
																		self.assertIsInstance(model.get_input_embeddings()     ,						(nn.Module)   )
																		lowerCAmelCase__    						=						model.get_output_embeddings()
																		self.assertTrue(x is None or isinstance(_UpperCamelCase     ,						nn.Linear   )   )






						def       UpperCamelCase__      (	self   ):






												"""simple docstring"""
												lowerCAmelCase__						,       lowerCAmelCase__    						=						self.model_tester.prepare_config_and_inputs_for_common()

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

																		lowerCAmelCase__    						=						['pixel_values']
																		self.assertListEqual(arg_names[:1]     ,						_UpperCamelCase   )






						def       UpperCamelCase__      (	self   ):






												"""simple docstring"""
												lowerCAmelCase__    						=						self.model_tester.prepare_config_and_inputs()
												self.model_tester.create_and_check_model(*_UpperCamelCase   )






						def       UpperCamelCase__      (	self   ):






												"""simple docstring"""
												lowerCAmelCase__    						=						self.model_tester.prepare_config_and_inputs()
												self.model_tester.create_and_check_for_masked_lm(*_UpperCamelCase   )






						def       UpperCamelCase__      (	self   ):






												"""simple docstring"""
												lowerCAmelCase__    						=						self.model_tester.prepare_config_and_inputs()
												self.model_tester.create_and_check_for_image_classification(*_UpperCamelCase   )






						def       UpperCamelCase__      (	self   ):






												"""simple docstring"""
												lowerCAmelCase__    						=						self.model_tester.prepare_config_and_inputs()
												self.model_tester.create_and_check_for_semantic_segmentation(*_UpperCamelCase   )






						def       UpperCamelCase__      (	self   ):






												"""simple docstring"""
												if not self.model_tester.is_training:
																		return

												lowerCAmelCase__						,       lowerCAmelCase__    						=						self.model_tester.prepare_config_and_inputs_for_common()
												lowerCAmelCase__    						=						True

												for model_class in self.all_model_classes:
																		# we don't test BeitForMaskedImageModeling
																		if model_class in [*get_values(_UpperCamelCase   ), BeitForMaskedImageModeling]:
																								continue

																		lowerCAmelCase__    						=						model_class(_UpperCamelCase   )
																		model.to(_UpperCamelCase   )
																		model.train()
																		lowerCAmelCase__    						=						self._prepare_for_class(_UpperCamelCase     ,						_UpperCamelCase     ,						return_labels=_UpperCamelCase   )
																		lowerCAmelCase__    						=						model(**_UpperCamelCase   ).loss
																		loss.backward()






						def       UpperCamelCase__      (	self   ):






												"""simple docstring"""
												lowerCAmelCase__						,       lowerCAmelCase__    						=						self.model_tester.prepare_config_and_inputs_for_common()
												if not self.model_tester.is_training:
																		return

												lowerCAmelCase__    						=						False
												lowerCAmelCase__    						=						True

												for model_class in self.all_model_classes:
																		# we don't test BeitForMaskedImageModeling
																		if (
																		    model_class in [*get_values(_UpperCamelCase   ), BeitForMaskedImageModeling]
																		    or not model_class.supports_gradient_checkpointing
																		):
																								continue

																		lowerCAmelCase__    						=						model_class(_UpperCamelCase   )
																		model.gradient_checkpointing_enable()
																		model.to(_UpperCamelCase   )
																		model.train()
																		lowerCAmelCase__    						=						self._prepare_for_class(_UpperCamelCase     ,						_UpperCamelCase     ,						return_labels=_UpperCamelCase   )
																		lowerCAmelCase__    						=						model(**_UpperCamelCase   ).loss
																		loss.backward()






						def       UpperCamelCase__      (	self   ):






												"""simple docstring"""
												lowerCAmelCase__						,       lowerCAmelCase__    						=						self.model_tester.prepare_config_and_inputs_for_common()

												lowerCAmelCase__    						=						_config_zero_init(_UpperCamelCase   )
												for model_class in self.all_model_classes:
																		lowerCAmelCase__    						=						model_class(config=_UpperCamelCase   )
																		for name, param in model.named_parameters():
																								# we skip lambda parameters as these require special initial values
																								# determined by config.layer_scale_init_value
																								if "lambda" in name:
																														continue
																								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"     ,						)







						@slow
						def       UpperCamelCase__      (	self   ):






												"""simple docstring"""
												for model_name in BEIT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
																		lowerCAmelCase__    						=						BeitModel.from_pretrained(_UpperCamelCase   )
																		self.assertIsNotNone(_UpperCamelCase   )







def 	_UpperCamelCase   (    )  ->   str:

						"""simple docstring"""



						lowerCAmelCase__    						=						Image.open('./tests/fixtures/tests_samples/COCO/000000039769.png'		)
						return image


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






						@cached_property
						def       UpperCamelCase__      (	self   ):






												"""simple docstring"""
												return BeitImageProcessor.from_pretrained('microsoft/beit-base-patch16-224'   ) if is_vision_available() else None






						@slow
						def       UpperCamelCase__      (	self   ):






												"""simple docstring"""
												lowerCAmelCase__    						=						BeitForMaskedImageModeling.from_pretrained('microsoft/beit-base-patch16-224-pt22k'   ).to(_UpperCamelCase   )

												lowerCAmelCase__    						=						self.default_image_processor
												lowerCAmelCase__    						=						prepare_img()
												lowerCAmelCase__    						=						image_processor(images=_UpperCamelCase     ,						return_tensors='pt'   ).pixel_values.to(_UpperCamelCase   )

												# prepare bool_masked_pos
												lowerCAmelCase__    						=						torch.ones((1, 1_96)     ,						dtype=torch.bool   ).to(_UpperCamelCase   )

												# forward pass
												with torch.no_grad():
																		lowerCAmelCase__    						=						model(pixel_values=_UpperCamelCase     ,						bool_masked_pos=_UpperCamelCase   )
												lowerCAmelCase__    						=						outputs.logits

												# verify the logits
												lowerCAmelCase__    						=						torch.Size((1, 1_96, 81_92)   )
												self.assertEqual(logits.shape     ,						_UpperCamelCase   )

												lowerCAmelCase__    						=						torch.tensor(
												    [[-3.24_37, 0.50_72, -13.91_74], [-3.24_56, 0.49_48, -13.94_01], [-3.20_33, 0.51_21, -13.85_50]]   ).to(_UpperCamelCase   )

												self.assertTrue(torch.allclose(logits[bool_masked_pos][:3, :3]     ,						_UpperCamelCase     ,						atol=1E-2   )   )






						@slow
						def       UpperCamelCase__      (	self   ):






												"""simple docstring"""
												lowerCAmelCase__    						=						BeitForImageClassification.from_pretrained('microsoft/beit-base-patch16-224'   ).to(_UpperCamelCase   )

												lowerCAmelCase__    						=						self.default_image_processor
												lowerCAmelCase__    						=						prepare_img()
												lowerCAmelCase__    						=						image_processor(images=_UpperCamelCase     ,						return_tensors='pt'   ).to(_UpperCamelCase   )

												# forward pass
												with torch.no_grad():
																		lowerCAmelCase__    						=						model(**_UpperCamelCase   )
												lowerCAmelCase__    						=						outputs.logits

												# verify the logits
												lowerCAmelCase__    						=						torch.Size((1, 10_00)   )
												self.assertEqual(logits.shape     ,						_UpperCamelCase   )

												lowerCAmelCase__    						=						torch.tensor([-1.23_85, -1.09_87, -1.01_08]   ).to(_UpperCamelCase   )

												self.assertTrue(torch.allclose(logits[0, :3]     ,						_UpperCamelCase     ,						atol=1E-4   )   )

												lowerCAmelCase__    						=						2_81
												self.assertEqual(logits.argmax(-1   ).item()     ,						_UpperCamelCase   )






						@slow
						def       UpperCamelCase__      (	self   ):






												"""simple docstring"""
												lowerCAmelCase__    						=						BeitForImageClassification.from_pretrained('microsoft/beit-large-patch16-224-pt22k-ft22k'   ).to(
												    _UpperCamelCase   )

												lowerCAmelCase__    						=						self.default_image_processor
												lowerCAmelCase__    						=						prepare_img()
												lowerCAmelCase__    						=						image_processor(images=_UpperCamelCase     ,						return_tensors='pt'   ).to(_UpperCamelCase   )

												# forward pass
												with torch.no_grad():
																		lowerCAmelCase__    						=						model(**_UpperCamelCase   )
												lowerCAmelCase__    						=						outputs.logits

												# verify the logits
												lowerCAmelCase__    						=						torch.Size((1, 2_18_41)   )
												self.assertEqual(logits.shape     ,						_UpperCamelCase   )

												lowerCAmelCase__    						=						torch.tensor([1.68_81, -0.27_87, 0.59_01]   ).to(_UpperCamelCase   )

												self.assertTrue(torch.allclose(logits[0, :3]     ,						_UpperCamelCase     ,						atol=1E-4   )   )

												lowerCAmelCase__    						=						23_96
												self.assertEqual(logits.argmax(-1   ).item()     ,						_UpperCamelCase   )






						@slow
						def       UpperCamelCase__      (	self   ):






												"""simple docstring"""
												lowerCAmelCase__    						=						BeitForSemanticSegmentation.from_pretrained('microsoft/beit-base-finetuned-ade-640-640'   )
												lowerCAmelCase__    						=						model.to(_UpperCamelCase   )

												lowerCAmelCase__    						=						BeitImageProcessor(do_resize=_UpperCamelCase     ,						size=6_40     ,						do_center_crop=_UpperCamelCase   )

												lowerCAmelCase__    						=						load_dataset('hf-internal-testing/fixtures_ade20k'     ,						split='test'   )
												lowerCAmelCase__    						=						Image.open(ds[0]['file']   )
												lowerCAmelCase__    						=						image_processor(images=_UpperCamelCase     ,						return_tensors='pt'   ).to(_UpperCamelCase   )

												# forward pass
												with torch.no_grad():
																		lowerCAmelCase__    						=						model(**_UpperCamelCase   )
												lowerCAmelCase__    						=						outputs.logits

												# verify the logits
												lowerCAmelCase__    						=						torch.Size((1, 1_50, 1_60, 1_60)   )
												self.assertEqual(logits.shape     ,						_UpperCamelCase   )

												lowerCAmelCase__    						=						version.parse(PIL.__version__   ) < version.parse('9.0.0'   )

												if is_pillow_less_than_a:
																		lowerCAmelCase__    						=						torch.tensor(
																		    [
																		        [[-4.92_25, -2.39_54, -3.05_22], [-2.88_22, -1.00_46, -1.75_61], [-2.95_49, -1.32_28, -2.13_47]],
																		        [[-5.81_68, -3.41_29, -4.07_78], [-3.86_51, -2.22_14, -3.02_77], [-3.83_56, -2.46_43, -3.35_35]],
																		        [[-0.00_78, 3.99_52, 4.07_54], [2.98_56, 4.69_44, 5.00_35], [3.24_13, 4.78_13, 4.99_69]],
																		    ]     ,						device=_UpperCamelCase     ,						)
												else:
																		lowerCAmelCase__    						=						torch.tensor(
																		    [
																		        [[-4.89_60, -2.36_88, -3.03_55], [-2.84_78, -0.98_36, -1.74_18], [-2.94_49, -1.33_32, -2.14_56]],
																		        [[-5.80_81, -3.41_24, -4.10_06], [-3.85_61, -2.20_81, -3.03_23], [-3.83_65, -2.46_01, -3.36_69]],
																		        [[-0.03_09, 3.98_68, 4.05_40], [2.96_40, 4.68_77, 4.99_76], [3.20_81, 4.76_90, 4.99_42]],
																		    ]     ,						device=_UpperCamelCase     ,						)

												self.assertTrue(torch.allclose(logits[0, :3, :3, :3]     ,						_UpperCamelCase     ,						atol=1E-4   )   )







						@slow
						def       UpperCamelCase__      (	self   ):






												"""simple docstring"""
												lowerCAmelCase__    						=						BeitForSemanticSegmentation.from_pretrained('microsoft/beit-base-finetuned-ade-640-640'   )
												lowerCAmelCase__    						=						model.to(_UpperCamelCase   )

												lowerCAmelCase__    						=						BeitImageProcessor(do_resize=_UpperCamelCase     ,						size=6_40     ,						do_center_crop=_UpperCamelCase   )

												lowerCAmelCase__    						=						load_dataset('hf-internal-testing/fixtures_ade20k'     ,						split='test'   )
												lowerCAmelCase__    						=						Image.open(ds[0]['file']   )
												lowerCAmelCase__    						=						image_processor(images=_UpperCamelCase     ,						return_tensors='pt'   ).to(_UpperCamelCase   )

												# forward pass
												with torch.no_grad():
																		lowerCAmelCase__    						=						model(**_UpperCamelCase   )

												lowerCAmelCase__    						=						outputs.logits.detach().cpu()

												lowerCAmelCase__    						=						image_processor.post_process_semantic_segmentation(outputs=_UpperCamelCase     ,						target_sizes=[(5_00, 3_00)]   )
												lowerCAmelCase__    						=						torch.Size((5_00, 3_00)   )
												self.assertEqual(segmentation[0].shape     ,						_UpperCamelCase   )

												lowerCAmelCase__    						=						image_processor.post_process_semantic_segmentation(outputs=_UpperCamelCase   )
												lowerCAmelCase__    						=						torch.Size((1_60, 1_60)   )
												self.assertEqual(segmentation[0].shape     ,						_UpperCamelCase   )













from typing import Dict, Optional

import numpy as np

import datasets


__snake_case    :    str         =     """
IoU is the area of overlap between the predicted segmentation and the ground truth divided by the area of union
between the predicted segmentation and the ground truth. For binary (two classes) or multi-class segmentation,
the mean IoU of the image is calculated by taking the IoU of each class and averaging them.
"""

__snake_case    :    Tuple         =     """
Args:
    predictions (`List[ndarray]`):
        List of predicted segmentation maps, each of shape (height, width). Each segmentation map can be of a different size.
    references (`List[ndarray]`):
        List of ground truth segmentation maps, each of shape (height, width). Each segmentation map can be of a different size.
    num_labels (`int`):
        Number of classes (categories).
    ignore_index (`int`):
        Index that will be ignored during evaluation.
    nan_to_num (`int`, *optional*):
        If specified, NaN values will be replaced by the number defined by the user.
    label_map (`dict`, *optional*):
        If specified, dictionary mapping old label indices to new label indices.
    reduce_labels (`bool`, *optional*, defaults to `False`):
        Whether or not to reduce all label values of segmentation maps by 1. Usually used for datasets where 0 is used for background,
        and background itself is not included in all classes of a dataset (e.g. ADE20k). The background label will be replaced by 255.

Returns:
    `Dict[str, float | ndarray]` comprising various elements:
    - *mean_iou* (`float`):
        Mean Intersection-over-Union (IoU averaged over all categories).
    - *mean_accuracy* (`float`):
        Mean accuracy (averaged over all categories).
    - *overall_accuracy* (`float`):
        Overall accuracy on all images.
    - *per_category_accuracy* (`ndarray` of shape `(num_labels,)`):
        Per category accuracy.
    - *per_category_iou* (`ndarray` of shape `(num_labels,)`):
        Per category IoU.

Examples:

    >>> import numpy as np

    >>> mean_iou = datasets.load_metric(\"mean_iou\")

    >>> # suppose one has 3 different segmentation maps predicted
    >>> predicted_1 = np.array([[1, 2], [3, 4], [5, 255]])
    >>> actual_1 = np.array([[0, 3], [5, 4], [6, 255]])

    >>> predicted_2 = np.array([[2, 7], [9, 2], [3, 6]])
    >>> actual_2 = np.array([[1, 7], [9, 2], [3, 6]])

    >>> predicted_3 = np.array([[2, 2, 3], [8, 2, 4], [3, 255, 2]])
    >>> actual_3 = np.array([[1, 2, 2], [8, 2, 1], [3, 255, 1]])

    >>> predicted = [predicted_1, predicted_2, predicted_3]
    >>> ground_truth = [actual_1, actual_2, actual_3]

    >>> results = mean_iou.compute(predictions=predicted, references=ground_truth, num_labels=10, ignore_index=255, reduce_labels=False)
    >>> print(results) # doctest: +NORMALIZE_WHITESPACE
    {'mean_iou': 0.47750000000000004, 'mean_accuracy': 0.5916666666666666, 'overall_accuracy': 0.5263157894736842, 'per_category_iou': array([0.   , 0.   , 0.375, 0.4  , 0.5  , 0.   , 0.5  , 1.   , 1.   , 1.   ]), 'per_category_accuracy': array([0.        , 0.        , 0.75      , 0.66666667, 1.        , 0.        , 0.5       , 1.        , 1.        , 1.        ])}
"""

__snake_case    :    Any         =     """\
@software{MMSegmentation_Contributors_OpenMMLab_Semantic_Segmentation_2020,
author = {{MMSegmentation Contributors}},
license = {Apache-2.0},
month = {7},
title = {{OpenMMLab Semantic Segmentation Toolbox and Benchmark}},
url = {https://github.com/open-mmlab/mmsegmentation},
year = {2020}
}"""







def 	_UpperCamelCase   (    UpperCamelCase_       :					str    ,       UpperCamelCase_       :					Tuple    ,       UpperCamelCase_       :					Optional[Any]    ,       UpperCamelCase_       :					bool    ,       UpperCamelCase_       :					Optional[Dict[int, int]] = None    ,       UpperCamelCase_       :					bool = False    ,       )  ->   List[Any]:

						"""simple docstring"""



						if label_map is not None:
												for old_id, new_id in label_map.items():
																		lowerCAmelCase__    						=						new_id

    # turn into Numpy arrays
						lowerCAmelCase__    						=						np.array(UpperCamelCase_		)
						lowerCAmelCase__    						=						np.array(UpperCamelCase_		)

						if reduce_labels:
												lowerCAmelCase__    						=						255
												lowerCAmelCase__    						=						label - 1
												lowerCAmelCase__    						=						255

						lowerCAmelCase__    						=						label != ignore_index
						lowerCAmelCase__    						=						np.not_equal(UpperCamelCase_    ,       UpperCamelCase_		)
						lowerCAmelCase__    						=						pred_label[mask]
						lowerCAmelCase__    						=						np.array(UpperCamelCase_		)[mask]

						lowerCAmelCase__    						=						pred_label[pred_label == label]

						lowerCAmelCase__    						=						np.histogram(UpperCamelCase_    ,       bins=UpperCamelCase_    ,       range=(0, num_labels - 1)		)[0]
						lowerCAmelCase__    						=						np.histogram(UpperCamelCase_    ,       bins=UpperCamelCase_    ,       range=(0, num_labels - 1)		)[0]
						lowerCAmelCase__    						=						np.histogram(UpperCamelCase_    ,       bins=UpperCamelCase_    ,       range=(0, num_labels - 1)		)[0]

						lowerCAmelCase__    						=						area_pred_label + area_label - area_intersect

						return area_intersect, area_union, area_pred_label, area_label







def 	_UpperCamelCase   (    UpperCamelCase_       :					Dict    ,       UpperCamelCase_       :					str    ,       UpperCamelCase_       :					int    ,       UpperCamelCase_       :					bool    ,       UpperCamelCase_       :					Optional[Dict[int, int]] = None    ,       UpperCamelCase_       :					bool = False    ,       )  ->   Union[str, Any]:

						"""simple docstring"""



						lowerCAmelCase__    						=						np.zeros((num_labels,)    ,       dtype=np.floataa		)
						lowerCAmelCase__    						=						np.zeros((num_labels,)    ,       dtype=np.floataa		)
						lowerCAmelCase__    						=						np.zeros((num_labels,)    ,       dtype=np.floataa		)
						lowerCAmelCase__    						=						np.zeros((num_labels,)    ,       dtype=np.floataa		)
						for result, gt_seg_map in zip(UpperCamelCase_    ,       UpperCamelCase_		):
												lowerCAmelCase__						,       lowerCAmelCase__						,       lowerCAmelCase__						,       lowerCAmelCase__    						=						intersect_and_union(
												    UpperCamelCase_    ,       UpperCamelCase_    ,       UpperCamelCase_    ,       UpperCamelCase_    ,       UpperCamelCase_    ,       UpperCamelCase_		)
												total_area_intersect += area_intersect
												total_area_union += area_union
												total_area_pred_label += area_pred_label
												total_area_label += area_label
						return total_area_intersect, total_area_union, total_area_pred_label, total_area_label







def 	_UpperCamelCase   (    UpperCamelCase_       :					List[str]    ,       UpperCamelCase_       :					Optional[int]    ,       UpperCamelCase_       :					str    ,       UpperCamelCase_       :					bool    ,       UpperCamelCase_       :					Optional[int] = None    ,       UpperCamelCase_       :					Optional[Dict[int, int]] = None    ,       UpperCamelCase_       :					bool = False    ,       )  ->   List[str]:

						"""simple docstring"""



						lowerCAmelCase__						,       lowerCAmelCase__						,       lowerCAmelCase__						,       lowerCAmelCase__    						=						total_intersect_and_union(
						    UpperCamelCase_    ,       UpperCamelCase_    ,       UpperCamelCase_    ,       UpperCamelCase_    ,       UpperCamelCase_    ,       UpperCamelCase_		)

						# compute metrics
						lowerCAmelCase__    						=						{}

						lowerCAmelCase__    						=						total_area_intersect.sum() / total_area_label.sum()
						lowerCAmelCase__    						=						total_area_intersect / total_area_union
						lowerCAmelCase__    						=						total_area_intersect / total_area_label

						lowerCAmelCase__    						=						np.nanmean(UpperCamelCase_		)
						lowerCAmelCase__    						=						np.nanmean(UpperCamelCase_		)
						lowerCAmelCase__    						=						all_acc
						lowerCAmelCase__    						=						iou
						lowerCAmelCase__    						=						acc

						if nan_to_num is not None:
												lowerCAmelCase__    						=						{metric: np.nan_to_num(UpperCamelCase_    ,       nan=UpperCamelCase_		) for metric, metric_value in metrics.items()}

						return metrics


@datasets.utils.file_utils.add_start_docstrings(_DESCRIPTION		,					_KWARGS_DESCRIPTION)
class     __SCREAMING_SNAKE_CASE					(   datasets.Metric):






						def       UpperCamelCase__      (	self   ):






												"""simple docstring"""
												return datasets.MetricInfo(
												    description=_DESCRIPTION     ,						citation=_CITATION     ,						inputs_description=_KWARGS_DESCRIPTION     ,						features=datasets.Features(
												        # 1st Seq - height dim, 2nd - width dim
												        {
												            'predictions': datasets.Sequence(datasets.Sequence(datasets.Value('uint16'   )   )   ),
												            'references': datasets.Sequence(datasets.Sequence(datasets.Value('uint16'   )   )   ),
												        }   )     ,						reference_urls=[
												        'https://github.com/open-mmlab/mmsegmentation/blob/71c201b1813267d78764f306a297ca717827c4bf/mmseg/core/evaluation/metrics.py'
												    ]     ,						)







						def       UpperCamelCase__      (	self     ,						_UpperCamelCase     ,						_UpperCamelCase     ,						_UpperCamelCase     ,						_UpperCamelCase     ,						_UpperCamelCase = None     ,						_UpperCamelCase = None     ,						_UpperCamelCase = False     ,						):






												"""simple docstring"""
												lowerCAmelCase__    						=						mean_iou(
												    results=_UpperCamelCase     ,						gt_seg_maps=_UpperCamelCase     ,						num_labels=_UpperCamelCase     ,						ignore_index=_UpperCamelCase     ,						nan_to_num=_UpperCamelCase     ,						label_map=_UpperCamelCase     ,						reduce_labels=_UpperCamelCase     ,						)
												return iou_result





