

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   __UpperCAmelCase					(      __a :     List[str]      )     ->	Tuple:

					"""simple docstring"""

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

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

					_a      :       str   =						SwinConfig(
					    embed_dim=__a						,depths=__a						,num_heads=__a						,window_size=__a						,out_features=['''stage1''', '''stage2''', '''stage3''', '''stage4''']						,)
					_a      :       Tuple   =						UperNetConfig(
					    backbone_config=__a						,auxiliary_in_channels=__a						,num_labels=__a						,idalabel=__a						,labelaid=__a						,)

					return config



def   __UpperCAmelCase					(      __a :     int      )     ->	Optional[Any]:

					"""simple docstring"""

					_a      :       Union[str, 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   __UpperCAmelCase					(      __a :     int						,__a :     int						,__a :     Union[str, Any]      )     ->	Union[str, Any]:

					"""simple docstring"""

					_a      :       str   =						dct.pop(__a      )
					_a      :       Dict   =						val



def   __UpperCAmelCase					(      __a :     Dict						,__a :     str      )     ->	Dict:

					"""simple docstring"""

					_a      :       List[Any]   =						[int(backbone_config.embed_dim * 2**i      ) for i in range(len(backbone_config.depths      )      )]
					for i in range(len(backbone_config.depths      )      ):
										_a      :       int   =						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)
															_a      :       Optional[int]   =						state_dict.pop(F"""backbone.stages.{i}.blocks.{j}.attn.w_msa.qkv.weight"""      )
															_a      :       Dict   =						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
															_a      :       int   =						in_proj_weight[:dim, :]
															_a      :       Any   =						in_proj_bias[: dim]
															_a      :       Tuple   =						in_proj_weight[
															    dim : dim * 2, :
															]
															_a      :       Any   =						in_proj_bias[
															    dim : dim * 2
															]
															_a      :       List[Any]   =						in_proj_weight[
															    -dim :, :
															]
															_a      :       Any   =						in_proj_bias[-dim :]
															# fmt: on



def   __UpperCAmelCase					(      __a :     str      )     ->	str:

					"""simple docstring"""

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



def   __UpperCAmelCase					(      __a :     Tuple      )     ->	Optional[Any]:

					"""simple docstring"""

					_a ,   _a      :       Optional[int]   =						x.shape
					_a      :       List[str]   =						x.reshape(__a						,in_channel // 4						,4      )
					_a      :       Dict   =						x[:, :, [0, 2, 1, 3]].transpose(1						,2      ).reshape(__a						,__a      )

					return x



def   __UpperCAmelCase					(      __a :     List[Any]      )     ->	Dict:

					"""simple docstring"""

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



def   __UpperCAmelCase					(      __a :     List[Any]      )     ->	List[Any]:

					"""simple docstring"""

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



def   __UpperCAmelCase					(      __a :     Dict						,__a :     Optional[Any]						,__a :     Any      )     ->	str:

					"""simple docstring"""

					_a      :       Dict   =						{
					    '''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''',
					}
					_a      :       Any   =						model_name_to_url[model_name]
					_a      :       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      )

					_a      :       Union[str, Any]   =						get_upernet_config(__a      )
					_a      :       Optional[int]   =						UperNetForSemanticSegmentation(__a      )
					model.eval()

					# replace "bn" => "batch_norm"
					for key in state_dict.copy().keys():
										_a      :       Any   =						state_dict.pop(__a      )
										if "bn" in key:
															_a      :       Optional[int]   =						key.replace('''bn'''						,'''batch_norm'''      )
										_a      :       Any   =						val

					# rename keys
					_a      :       str   =						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:
																				_a      :       List[str]   =						reverse_correct_unfold_reduction_order(__a      )
															if "norm" in key:
																				_a      :       Optional[Any]   =						reverse_correct_unfold_norm_order(__a      )

					model.load_state_dict(__a      )

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

					_a      :       Union[str, Any]   =						SegformerImageProcessor()
					_a      :       List[Any]   =						processor(__a						,return_tensors='''pt'''      ).pixel_values

					with torch.no_grad():
										_a      :       Dict   =						model(__a      )
										_a      :       str   =						outputs.logits

					print(logits.shape      )
					print('''First values of logits:'''						,logits[0, 0, :3, :3]      )
					# assert values
					if model_name == "upernet-swin-tiny":
										_a      :       Any   =						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]]      )
					elif model_name == "upernet-swin-small":
										_a      :       Any   =						torch.tensor(
										    [[-7.19_21, -7.19_21, -6.95_32], [-7.19_21, -7.19_21, -6.95_32], [-7.09_08, -7.09_08, -6.85_34]]      )
					elif model_name == "upernet-swin-base":
										_a      :       int   =						torch.tensor(
										    [[-6.58_51, -6.58_51, -6.43_30], [-6.58_51, -6.58_51, -6.43_30], [-6.47_63, -6.47_63, -6.32_54]]      )
					elif model_name == "upernet-swin-large":
										_a      :       Any   =						torch.tensor(
										    [[-7.52_97, -7.52_97, -7.38_02], [-7.52_97, -7.52_97, -7.38_02], [-7.40_44, -7.40_44, -7.25_86]]      )
					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)






a__												=						'''Input must be a string of 8 numbers plus letter'''
a__												=						'''TRWAGMYFPDXBNJZSQVHLCKE'''



def   __UpperCAmelCase					(      __a :     str      )     ->	bool:

					"""simple docstring"""

					if not isinstance(__a						,__a      ):
										_a      :       List[str]   =						F"""Expected string as input, found {type(__a      ).__name__}"""
										raise TypeError(__a      )

					_a      :       List[Any]   =						spanish_id.replace('''-'''						,''''''      ).upper()
					if len(__a      ) != 9:
										raise ValueError(__a      )

					try:
										_a      :       Any   =						int(spanish_id_clean[0:8]      )
										_a      :       str   =						spanish_id_clean[8]
					except ValueError as ex:
										raise ValueError(__a      ) from ex

					if letter.isdigit():
										raise ValueError(__a      )

					return letter == LOOKUP_LETTERS[number % 23]


if __name__ == "__main__":
			import doctest

			doctest.testmod()



