

from scipy.stats import spearmanr

import datasets


a_						     =    '\nThe Spearman rank-order correlation coefficient is a measure of the\nrelationship between two datasets. Like other correlation coefficients,\nthis one varies between -1 and +1 with 0 implying no correlation.\nPositive correlations imply that as data in dataset x increases, so\ndoes data in dataset y. Negative correlations imply that as x increases,\ny decreases. Correlations of -1 or +1 imply an exact monotonic relationship.\n\nUnlike the Pearson correlation, the Spearman correlation does not\nassume that both datasets are normally distributed.\n\nThe p-value roughly indicates the probability of an uncorrelated system\nproducing datasets that have a Spearman correlation at least as extreme\nas the one computed from these datasets. The p-values are not entirely\nreliable but are probably reasonable for datasets larger than 500 or so.\n'

a_						     =    '\nArgs:\n    predictions (`List[float]`): Predicted labels, as returned by a model.\n    references (`List[float]`): Ground truth labels.\n    return_pvalue (`bool`): If `True`, returns the p-value. If `False`, returns\n            only the spearmanr score. Defaults to `False`.\nReturns:\n    spearmanr (`float`): Spearman correlation coefficient.\n    p-value (`float`): p-value. **Note**: is only returned if `return_pvalue=True` is input.\nExamples:\n    Example 1:\n        >>> spearmanr_metric = datasets.load_metric(\"spearmanr\")\n        >>> results = spearmanr_metric.compute(references=[1, 2, 3, 4, 5], predictions=[10, 9, 2.5, 6, 4])\n        >>> print(results)\n        {\'spearmanr\': -0.7}\n\n    Example 2:\n        >>> spearmanr_metric = datasets.load_metric(\"spearmanr\")\n        >>> results = spearmanr_metric.compute(references=[1, 2, 3, 4, 5],\n        ...                                     predictions=[10, 9, 2.5, 6, 4],\n        ...                                     return_pvalue=True)\n        >>> print(results[\'spearmanr\'])\n        -0.7\n        >>> print(round(results[\'spearmanr_pvalue\'], 2))\n        0.19\n'

a_						     =    r'\\n@book{kokoska2000crc,\n  title={CRC standard probability and statistics tables and formulae},\n  author={Kokoska, Stephen and Zwillinger, Daniel},\n  year={2000},\n  publisher={Crc Press}\n}\n@article{2020SciPy-NMeth,\n  author  = {Virtanen, Pauli and Gommers, Ralf and Oliphant, Travis E. and\n            Haberland, Matt and Reddy, Tyler and Cournapeau, David and\n            Burovski, Evgeni and Peterson, Pearu and Weckesser, Warren and\n            Bright, Jonathan and {van der Walt}, St{\'e}fan J. and\n            Brett, Matthew and Wilson, Joshua and Millman, K. Jarrod and\n            Mayorov, Nikolay and Nelson, Andrew R. J. and Jones, Eric and\n            Kern, Robert and Larson, Eric and Carey, C J and\n            Polat, {\.I}lhan and Feng, Yu and Moore, Eric W. and\n            {VanderPlas}, Jake and Laxalde, Denis and Perktold, Josef and\n            Cimrman, Robert and Henriksen, Ian and Quintero, E. A. and\n            Harris, Charles R. and Archibald, Anne M. and\n            Ribeiro, Ant{\^o}nio H. and Pedregosa, Fabian and\n            {van Mulbregt}, Paul and {SciPy 1.0 Contributors}},\n  title   = {{{SciPy} 1.0: Fundamental Algorithms for Scientific\n            Computing in Python}},\n  journal = {Nature Methods},\n  year    = {2020},\n  volume  = {17},\n  pages   = {261--272},\n  adsurl  = {https://rdcu.be/b08Wh},\n  doi     = {10.1038/s41592-019-0686-2},\n}\n'






@datasets.utils.file_utils.add_start_docstrings(_DESCRIPTION							,	_KWARGS_DESCRIPTION   )
class _UpperCamelCase							(							datasets.Metric   ):
   '''simple docstring'''







   def 							__UpperCamelCase (			self							: Union[str, Any]	)							->       Tuple:

    """simple docstring"""

    return datasets.MetricInfo(
        description=_DESCRIPTION      ,      citation=_CITATION      ,      inputs_description=_KWARGS_DESCRIPTION      ,      features=datasets.Features(
            {
                "predictions": datasets.Value("float"	),
                "references": datasets.Value("float"	),
            }	)      ,      reference_urls=["https://docs.scipy.org/doc/scipy/reference/generated/scipy.stats.spearmanr.html"]      ,      )






   def 							__UpperCamelCase (			self							: int      ,      a							: str      ,      a							: Any      ,      a							: str=False	)							->       Optional[int]:

    """simple docstring"""

    SCREAMING_SNAKE_CASE   :      Optional[int]				     =       spearmanr(a      ,      a	)
    if return_pvalue:
     return {"spearmanr": results[0], "spearmanr_pvalue": results[1]}
    else:
     return {"spearmanr": results[0]}


import argparse
import json

import requests
import torch
from huggingface_hub import hf_hub_download
from PIL import Image

from transformers import SegformerImageProcessor, SwinConfig, UperNetConfig, UperNetForSemanticSegmentation

def  lowerCamelCase__ (					_a):
	SCREAMING_SNAKE_CASE   :      Tuple				     =       384
	SCREAMING_SNAKE_CASE   :      Dict				     =       7
	if "tiny" in model_name:
		SCREAMING_SNAKE_CASE   :      List[Any]				     =       96
		SCREAMING_SNAKE_CASE   :      Union[str, Any]				     =       (2, 2, 6, 2)
		SCREAMING_SNAKE_CASE   :      Dict				     =       (3, 6, 12, 24)
	elif "small" in model_name:
		SCREAMING_SNAKE_CASE   :      Union[str, Any]				     =       96
		SCREAMING_SNAKE_CASE   :      Dict				     =       (2, 2, 18, 2)
		SCREAMING_SNAKE_CASE   :      Dict				     =       (3, 6, 12, 24)
	elif "base" in model_name:
		SCREAMING_SNAKE_CASE   :      Optional[int]				     =       128
		SCREAMING_SNAKE_CASE   :      Tuple				     =       (2, 2, 18, 2)
		SCREAMING_SNAKE_CASE   :      str				     =       (4, 8, 16, 32)
		SCREAMING_SNAKE_CASE   :      Any				     =       12
		SCREAMING_SNAKE_CASE   :      int				     =       512
	elif "large" in model_name:
		SCREAMING_SNAKE_CASE   :      Any				     =       192
		SCREAMING_SNAKE_CASE   :      List[Any]				     =       (2, 2, 18, 2)
		SCREAMING_SNAKE_CASE   :      Optional[Any]				     =       (6, 12, 24, 48)
		SCREAMING_SNAKE_CASE   :      Any				     =       12
		SCREAMING_SNAKE_CASE   :      Any				     =       768

	# set label information
	SCREAMING_SNAKE_CASE   :      Optional[int]				     =       150
	SCREAMING_SNAKE_CASE   :      int				     =       "huggingface/label-files"
	SCREAMING_SNAKE_CASE   :      Dict				     =       "ade20k-id2label.json"
	SCREAMING_SNAKE_CASE   :      List[Any]				     =       json.load(open(hf_hub_download(_a							,			_a							,			repo_type="dataset")							,			"r"))
	SCREAMING_SNAKE_CASE   :      List[Any]				     =       {int(_a): v for k, v in idalabel.items()}
	SCREAMING_SNAKE_CASE   :      int				     =       {v: k for k, v in idalabel.items()}

	SCREAMING_SNAKE_CASE   :      Union[str, Any]				     =       SwinConfig(
	    embed_dim=_a							,			depths=_a							,			num_heads=_a							,			window_size=_a							,			out_features=["stage1", "stage2", "stage3", "stage4"]							,			)
	SCREAMING_SNAKE_CASE   :      List[str]				     =       UperNetConfig(
	    backbone_config=_a							,			auxiliary_in_channels=_a							,			num_labels=_a							,			idalabel=_a							,			labelaid=_a							,			)

	return config

def  lowerCamelCase__ (					_a):
	SCREAMING_SNAKE_CASE   :      List[Any]				     =       []

	# fmt: off
	# stem
	rename_keys.append(("backbone.patch_embed.projection.weight", "backbone.embeddings.patch_embeddings.projection.weight"))
	rename_keys.append(("backbone.patch_embed.projection.bias", "backbone.embeddings.patch_embeddings.projection.bias"))
	rename_keys.append(("backbone.patch_embed.norm.weight", "backbone.embeddings.norm.weight"))
	rename_keys.append(("backbone.patch_embed.norm.bias", "backbone.embeddings.norm.bias"))
	# stages
	for i in range(len(config.backbone_config.depths)):
		for j in range(config.backbone_config.depths[i]):
			rename_keys.append((f"backbone.stages.{i}.blocks.{j}.norm1.weight", f"backbone.encoder.layers.{i}.blocks.{j}.layernorm_before.weight"))
			rename_keys.append((f"backbone.stages.{i}.blocks.{j}.norm1.bias", f"backbone.encoder.layers.{i}.blocks.{j}.layernorm_before.bias"))
			rename_keys.append((f"backbone.stages.{i}.blocks.{j}.attn.w_msa.relative_position_bias_table", f"backbone.encoder.layers.{i}.blocks.{j}.attention.self.relative_position_bias_table"))
			rename_keys.append((f"backbone.stages.{i}.blocks.{j}.attn.w_msa.relative_position_index", f"backbone.encoder.layers.{i}.blocks.{j}.attention.self.relative_position_index"))
			rename_keys.append((f"backbone.stages.{i}.blocks.{j}.attn.w_msa.proj.weight", f"backbone.encoder.layers.{i}.blocks.{j}.attention.output.dense.weight"))
			rename_keys.append((f"backbone.stages.{i}.blocks.{j}.attn.w_msa.proj.bias", f"backbone.encoder.layers.{i}.blocks.{j}.attention.output.dense.bias"))
			rename_keys.append((f"backbone.stages.{i}.blocks.{j}.norm2.weight", f"backbone.encoder.layers.{i}.blocks.{j}.layernorm_after.weight"))
			rename_keys.append((f"backbone.stages.{i}.blocks.{j}.norm2.bias", f"backbone.encoder.layers.{i}.blocks.{j}.layernorm_after.bias"))
			rename_keys.append((f"backbone.stages.{i}.blocks.{j}.ffn.layers.0.0.weight", f"backbone.encoder.layers.{i}.blocks.{j}.intermediate.dense.weight"))
			rename_keys.append((f"backbone.stages.{i}.blocks.{j}.ffn.layers.0.0.bias", f"backbone.encoder.layers.{i}.blocks.{j}.intermediate.dense.bias"))
			rename_keys.append((f"backbone.stages.{i}.blocks.{j}.ffn.layers.1.weight", f"backbone.encoder.layers.{i}.blocks.{j}.output.dense.weight"))
			rename_keys.append((f"backbone.stages.{i}.blocks.{j}.ffn.layers.1.bias", f"backbone.encoder.layers.{i}.blocks.{j}.output.dense.bias"))

		if i < 3:
			rename_keys.append((f"backbone.stages.{i}.downsample.reduction.weight", f"backbone.encoder.layers.{i}.downsample.reduction.weight"))
			rename_keys.append((f"backbone.stages.{i}.downsample.norm.weight", f"backbone.encoder.layers.{i}.downsample.norm.weight"))
			rename_keys.append((f"backbone.stages.{i}.downsample.norm.bias", f"backbone.encoder.layers.{i}.downsample.norm.bias"))
		rename_keys.append((f"backbone.norm{i}.weight", f"backbone.hidden_states_norms.stage{i+1}.weight"))
		rename_keys.append((f"backbone.norm{i}.bias", f"backbone.hidden_states_norms.stage{i+1}.bias"))

	# decode head
	rename_keys.extend(
	    [
	        ("decode_head.conv_seg.weight", "decode_head.classifier.weight"),
	        ("decode_head.conv_seg.bias", "decode_head.classifier.bias"),
	        ("auxiliary_head.conv_seg.weight", "auxiliary_head.classifier.weight"),
	        ("auxiliary_head.conv_seg.bias", "auxiliary_head.classifier.bias"),
	    ])
	# fmt: on

	return rename_keys

def  lowerCamelCase__ (					_a							,			_a							,			_a):
	SCREAMING_SNAKE_CASE   :      Any				     =       dct.pop(_a)
	SCREAMING_SNAKE_CASE   :      Optional[Any]				     =       val

def  lowerCamelCase__ (					_a							,			_a):
	SCREAMING_SNAKE_CASE   :      Optional[int]				     =       [int(backbone_config.embed_dim * 2**i) for i in range(len(backbone_config.depths))]
	for i in range(len(backbone_config.depths)):
		SCREAMING_SNAKE_CASE   :      Tuple				     =       num_features[i]
		for j in range(backbone_config.depths[i]):
			# fmt: off
			# read in weights + bias of input projection layer (in original implementation, this is a single matrix + bias)
			SCREAMING_SNAKE_CASE   :      int				     =       state_dict.pop(f"backbone.stages.{i}.blocks.{j}.attn.w_msa.qkv.weight")
			SCREAMING_SNAKE_CASE   :      List[str]				     =       state_dict.pop(f"backbone.stages.{i}.blocks.{j}.attn.w_msa.qkv.bias")
			# next, add query, keys and values (in that order) to the state dict
			SCREAMING_SNAKE_CASE   :      Optional[int]				     =       in_proj_weight[:dim, :]
			SCREAMING_SNAKE_CASE   :      int				     =       in_proj_bias[: dim]
			SCREAMING_SNAKE_CASE   :      Optional[Any]				     =       in_proj_weight[
			    dim : dim * 2, :
			]
			SCREAMING_SNAKE_CASE   :      str				     =       in_proj_bias[
			    dim : dim * 2
			]
			SCREAMING_SNAKE_CASE   :      List[str]				     =       in_proj_weight[
			    -dim :, :
			]
			SCREAMING_SNAKE_CASE   :      Union[str, Any]				     =       in_proj_bias[-dim :]
			# fmt: on

def  lowerCamelCase__ (					_a):
	SCREAMING_SNAKE_CASE     ,SCREAMING_SNAKE_CASE   :      List[Any]				     =       x.shape
	SCREAMING_SNAKE_CASE   :      List[str]				     =       x.reshape(_a							,			4							,			in_channel // 4)
	SCREAMING_SNAKE_CASE   :      str				     =       x[:, [0, 2, 1, 3], :].transpose(1							,			2).reshape(_a							,			_a)
	return x

def  lowerCamelCase__ (					_a):
	SCREAMING_SNAKE_CASE     ,SCREAMING_SNAKE_CASE   :      Any				     =       x.shape
	SCREAMING_SNAKE_CASE   :      Tuple				     =       x.reshape(_a							,			in_channel // 4							,			4)
	SCREAMING_SNAKE_CASE   :      Optional[int]				     =       x[:, :, [0, 2, 1, 3]].transpose(1							,			2).reshape(_a							,			_a)

	return x

def  lowerCamelCase__ (					_a):
	SCREAMING_SNAKE_CASE   :      str				     =       x.shape[0]
	SCREAMING_SNAKE_CASE   :      int				     =       x.reshape(4							,			in_channel // 4)
	SCREAMING_SNAKE_CASE   :      int				     =       x[[0, 2, 1, 3], :].transpose(0							,			1).reshape(_a)
	return x

def  lowerCamelCase__ (					_a):
	SCREAMING_SNAKE_CASE   :      Any				     =       x.shape[0]
	SCREAMING_SNAKE_CASE   :      Dict				     =       x.reshape(in_channel // 4							,			4)
	SCREAMING_SNAKE_CASE   :      Dict				     =       x[:, [0, 2, 1, 3]].transpose(0							,			1).reshape(_a)
	return x

def  lowerCamelCase__ (					_a							,			_a							,			_a):
	SCREAMING_SNAKE_CASE   :      Any				     =       {
	    "upernet-swin-tiny": "https://download.openmmlab.com/mmsegmentation/v0.5/swin/upernet_swin_tiny_patch4_window7_512x512_160k_ade20k_pretrain_224x224_1K/upernet_swin_tiny_patch4_window7_512x512_160k_ade20k_pretrain_224x224_1K_20210531_112542-e380ad3e.pth",
	    "upernet-swin-small": "https://download.openmmlab.com/mmsegmentation/v0.5/swin/upernet_swin_small_patch4_window7_512x512_160k_ade20k_pretrain_224x224_1K/upernet_swin_small_patch4_window7_512x512_160k_ade20k_pretrain_224x224_1K_20210526_192015-ee2fff1c.pth",
	    "upernet-swin-base": "https://download.openmmlab.com/mmsegmentation/v0.5/swin/upernet_swin_base_patch4_window12_512x512_160k_ade20k_pretrain_384x384_22K/upernet_swin_base_patch4_window12_512x512_160k_ade20k_pretrain_384x384_22K_20210531_125459-429057bf.pth",
	    "upernet-swin-large": "https://download.openmmlab.com/mmsegmentation/v0.5/swin/upernet_swin_large_patch4_window12_512x512_pretrain_384x384_22K_160k_ade20k/upernet_swin_large_patch4_window12_512x512_pretrain_384x384_22K_160k_ade20k_20220318_091743-9ba68901.pth",
	}
	SCREAMING_SNAKE_CASE   :      Optional[int]				     =       model_name_to_url[model_name]
	SCREAMING_SNAKE_CASE   :      Union[str, Any]				     =       torch.hub.load_state_dict_from_url(_a							,			map_location="cpu"							,			file_name=_a)[
	    "state_dict"
	]

	for name, param in state_dict.items():
		print(_a							,			param.shape)

	SCREAMING_SNAKE_CASE   :      Union[str, Any]				     =       get_upernet_config(_a)
	SCREAMING_SNAKE_CASE   :      Dict				     =       UperNetForSemanticSegmentation(_a)
	model.eval()

	# replace "bn" => "batch_norm"
	for key in state_dict.copy().keys():
		SCREAMING_SNAKE_CASE   :      Union[str, Any]				     =       state_dict.pop(_a)
		if "bn" in key:
			SCREAMING_SNAKE_CASE   :      Any				     =       key.replace("bn"							,			"batch_norm")
		SCREAMING_SNAKE_CASE   :      List[Any]				     =       val

	# rename keys
	SCREAMING_SNAKE_CASE   :      Optional[Any]				     =       create_rename_keys(_a)
	for src, dest in rename_keys:
		rename_key(_a							,			_a							,			_a)
	read_in_q_k_v(_a							,			config.backbone_config)

	# fix downsample parameters
	for key, value in state_dict.items():
		if "downsample" in key:
			if "reduction" in key:
				SCREAMING_SNAKE_CASE   :      Union[str, Any]				     =       reverse_correct_unfold_reduction_order(_a)
			if "norm" in key:
				SCREAMING_SNAKE_CASE   :      str				     =       reverse_correct_unfold_norm_order(_a)

	model.load_state_dict(_a)

	# verify on image
	SCREAMING_SNAKE_CASE   :      Optional[int]				     =       "https://huggingface.co/datasets/hf-internal-testing/fixtures_ade20k/resolve/main/ADE_val_00000001.jpg"
	SCREAMING_SNAKE_CASE   :      List[Any]				     =       Image.open(requests.get(_a							,			stream=_a).raw).convert("RGB")

	SCREAMING_SNAKE_CASE   :      Tuple				     =       SegformerImageProcessor()
	SCREAMING_SNAKE_CASE   :      Dict				     =       processor(_a							,			return_tensors="pt").pixel_values

	with torch.no_grad():
		SCREAMING_SNAKE_CASE   :      str				     =       model(_a)
		SCREAMING_SNAKE_CASE   :      List[Any]				     =       outputs.logits

	print(logits.shape)
	print("First values of logits:"							,			logits[0, 0, :3, :3])
	# assert values
	if model_name == "upernet-swin-tiny":
		SCREAMING_SNAKE_CASE   :      Optional[Any]				     =       torch.tensor(
		    [[-7.5958, -7.5958, -7.4302], [-7.5958, -7.5958, -7.4302], [-7.4797, -7.4797, -7.3068]])
	elif model_name == "upernet-swin-small":
		SCREAMING_SNAKE_CASE   :      Tuple				     =       torch.tensor(
		    [[-7.1921, -7.1921, -6.9532], [-7.1921, -7.1921, -6.9532], [-7.0908, -7.0908, -6.8534]])
	elif model_name == "upernet-swin-base":
		SCREAMING_SNAKE_CASE   :      str				     =       torch.tensor(
		    [[-6.5851, -6.5851, -6.4330], [-6.5851, -6.5851, -6.4330], [-6.4763, -6.4763, -6.3254]])
	elif model_name == "upernet-swin-large":
		SCREAMING_SNAKE_CASE   :      Any				     =       torch.tensor(
		    [[-7.5297, -7.5297, -7.3802], [-7.5297, -7.5297, -7.3802], [-7.4044, -7.4044, -7.2586]])
	print("Logits:"							,			outputs.logits[0, 0, :3, :3])
	assert torch.allclose(outputs.logits[0, 0, :3, :3]							,			_a							,			atol=1E-4)
	print("Looks ok!")

	if pytorch_dump_folder_path is not None:
		print(f"Saving model {model_name} to {pytorch_dump_folder_path}")
		model.save_pretrained(_a)
		print(f"Saving processor to {pytorch_dump_folder_path}")
		processor.save_pretrained(_a)

	if push_to_hub:
		print(f"Pushing model and processor for {model_name} to hub")
		model.push_to_hub(f"openmmlab/{model_name}")
		processor.push_to_hub(f"openmmlab/{model_name}")


if __name__ == "__main__":
			a_						     =    argparse.ArgumentParser()
			# Required parameters
			parser.add_argument(
			    '--model_name',
			    default='upernet-swin-tiny',
			    type=str,
			    choices=[F'''upernet-swin-{size}''' for size in ['tiny', 'small', 'base', 'large']],
			    help='Name of the Swin + UperNet model you\'d like to convert.',
			)
			parser.add_argument(
			    '--pytorch_dump_folder_path', default=None, type=str, help='Path to the output PyTorch model directory.'
			)
			parser.add_argument(
			    '--push_to_hub', action='store_true', help='Whether or not to push the converted model to the 🤗 hub.'
			)

			a_						     =    parser.parse_args()
			convert_upernet_checkpoint(args.model_name, args.pytorch_dump_folder_path, args.push_to_hub)