





"""simple docstring"""





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       a__					(					_SCREAMING_SNAKE_CASE    ):




		"""simple docstring"""


		UpperCamelCase   =	384
		UpperCamelCase   =	7
		if "tiny" in model_name:
				UpperCamelCase   =	96
				UpperCamelCase   =	(2, 2, 6, 2)
				UpperCamelCase   =	(3, 6, 12, 24)
		elif "small" in model_name:
				UpperCamelCase   =	96
				UpperCamelCase   =	(2, 2, 18, 2)
				UpperCamelCase   =	(3, 6, 12, 24)
		elif "base" in model_name:
				UpperCamelCase   =	128
				UpperCamelCase   =	(2, 2, 18, 2)
				UpperCamelCase   =	(4, 8, 16, 32)
				UpperCamelCase   =	12
				UpperCamelCase   =	512
		elif "large" in model_name:
				UpperCamelCase   =	192
				UpperCamelCase   =	(2, 2, 18, 2)
				UpperCamelCase   =	(6, 12, 24, 48)
				UpperCamelCase   =	12
				UpperCamelCase   =	768

		# set label information
		UpperCamelCase   =	150
		UpperCamelCase   =	"huggingface/label-files"
		UpperCamelCase   =	"ade20k-id2label.json"
		UpperCamelCase   =	json.load(open(hf_hub_download(_SCREAMING_SNAKE_CASE  ,	_SCREAMING_SNAKE_CASE  ,	repo_type="dataset"    )  ,	"r"    )    )
		UpperCamelCase   =	{int(_SCREAMING_SNAKE_CASE    ): v for k, v in idalabel.items()}
		UpperCamelCase   =	{v: k for k, v in idalabel.items()}

		UpperCamelCase   =	SwinConfig(
		    embed_dim=_SCREAMING_SNAKE_CASE  ,	depths=_SCREAMING_SNAKE_CASE  ,	num_heads=_SCREAMING_SNAKE_CASE  ,	window_size=_SCREAMING_SNAKE_CASE  ,	out_features=["stage1", "stage2", "stage3", "stage4"]  ,	)
		UpperCamelCase   =	UperNetConfig(
		    backbone_config=_SCREAMING_SNAKE_CASE  ,	auxiliary_in_channels=_SCREAMING_SNAKE_CASE  ,	num_labels=_SCREAMING_SNAKE_CASE  ,	idalabel=_SCREAMING_SNAKE_CASE  ,	labelaid=_SCREAMING_SNAKE_CASE  ,	)

		return config


def       a__					(					_SCREAMING_SNAKE_CASE    ):




		"""simple docstring"""


		UpperCamelCase   =	[]

		# 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       a__					(					_SCREAMING_SNAKE_CASE  ,	_SCREAMING_SNAKE_CASE  ,	_SCREAMING_SNAKE_CASE    ):




		"""simple docstring"""


		UpperCamelCase   =	dct.pop(_SCREAMING_SNAKE_CASE    )
		UpperCamelCase   =	val


def       a__					(					_SCREAMING_SNAKE_CASE  ,	_SCREAMING_SNAKE_CASE    ):




		"""simple docstring"""


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


def       a__					(					_SCREAMING_SNAKE_CASE    ):




		"""simple docstring"""


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


def       a__					(					_SCREAMING_SNAKE_CASE    ):




		"""simple docstring"""


		UpperCamelCase   ,				UpperCamelCase   =	x.shape
		UpperCamelCase   =	x.reshape(_SCREAMING_SNAKE_CASE  ,	in_channel // 4  ,	4    )
		UpperCamelCase   =	x[:, :, [0, 2, 1, 3]].transpose(1  ,	2    ).reshape(_SCREAMING_SNAKE_CASE  ,	_SCREAMING_SNAKE_CASE    )

		return x


def       a__					(					_SCREAMING_SNAKE_CASE    ):




		"""simple docstring"""


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


def       a__					(					_SCREAMING_SNAKE_CASE    ):




		"""simple docstring"""


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


def       a__					(					_SCREAMING_SNAKE_CASE  ,	_SCREAMING_SNAKE_CASE  ,	_SCREAMING_SNAKE_CASE    ):




		"""simple docstring"""


		UpperCamelCase   =	{
		    "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",
		}
		UpperCamelCase   =	model_name_to_url[model_name]
		UpperCamelCase   =	torch.hub.load_state_dict_from_url(_SCREAMING_SNAKE_CASE  ,	map_location="cpu"  ,	file_name=_SCREAMING_SNAKE_CASE    )[
		    "state_dict"
		]

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

		UpperCamelCase   =	get_upernet_config(_SCREAMING_SNAKE_CASE    )
		UpperCamelCase   =	UperNetForSemanticSegmentation(_SCREAMING_SNAKE_CASE    )
		model.eval()

		# replace "bn" => "batch_norm"
		for key in state_dict.copy().keys():
				UpperCamelCase   =	state_dict.pop(_SCREAMING_SNAKE_CASE    )
				if "bn" in key:
						UpperCamelCase   =	key.replace("bn"  ,	"batch_norm"    )
				UpperCamelCase   =	val

		# rename keys
		UpperCamelCase   =	create_rename_keys(_SCREAMING_SNAKE_CASE    )
		for src, dest in rename_keys:
				rename_key(_SCREAMING_SNAKE_CASE  ,	_SCREAMING_SNAKE_CASE  ,	_SCREAMING_SNAKE_CASE    )
		read_in_q_k_v(_SCREAMING_SNAKE_CASE  ,	config.backbone_config    )

		# fix downsample parameters
		for key, value in state_dict.items():
				if "downsample" in key:
						if "reduction" in key:
								UpperCamelCase   =	reverse_correct_unfold_reduction_order(_SCREAMING_SNAKE_CASE    )
						if "norm" in key:
								UpperCamelCase   =	reverse_correct_unfold_norm_order(_SCREAMING_SNAKE_CASE    )

		model.load_state_dict(_SCREAMING_SNAKE_CASE    )

		# verify on image
		UpperCamelCase   =	"https://huggingface.co/datasets/hf-internal-testing/fixtures_ade20k/resolve/main/ADE_val_00000001.jpg"
		UpperCamelCase   =	Image.open(requests.get(_SCREAMING_SNAKE_CASE  ,	stream=_SCREAMING_SNAKE_CASE    ).raw    ).convert("RGB"    )

		UpperCamelCase   =	SegformerImageProcessor()
		UpperCamelCase   =	processor(_SCREAMING_SNAKE_CASE  ,	return_tensors="pt"    ).pixel_values

		with torch.no_grad():
				UpperCamelCase   =	model(_SCREAMING_SNAKE_CASE    )
				UpperCamelCase   =	outputs.logits

		print(logits.shape    )
		print("First values of logits:"  ,	logits[0, 0, :3, :3]    )
		# assert values
		if model_name == "upernet-swin-tiny":
				UpperCamelCase   =	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":
				UpperCamelCase   =	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":
				UpperCamelCase   =	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":
				UpperCamelCase   =	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]  ,	_SCREAMING_SNAKE_CASE  ,	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(_SCREAMING_SNAKE_CASE    )
				print(F"Saving processor to {pytorch_dump_folder_path}"    )
				processor.save_pretrained(_SCREAMING_SNAKE_CASE    )

		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__":
						lowerCAmelCase__     			=			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.'''
						)

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













"""simple docstring"""





from typing import List

from .keymap import KEYMAP, get_character


def       a__					(					_SCREAMING_SNAKE_CASE    ):




		"""simple docstring"""


		def decorator(_SCREAMING_SNAKE_CASE    ):
				UpperCamelCase   =	getattr(_SCREAMING_SNAKE_CASE  ,	"handle_key"  ,	[]    )
				handle += [key]
				setattr(_SCREAMING_SNAKE_CASE  ,	"handle_key"  ,	_SCREAMING_SNAKE_CASE    )
				return func

		return decorator


def       a__					(					*_SCREAMING_SNAKE_CASE    ):




		"""simple docstring"""


		def decorator(_SCREAMING_SNAKE_CASE    ):
				UpperCamelCase   =	getattr(_SCREAMING_SNAKE_CASE  ,	"handle_key"  ,	[]    )
				handle += keys
				setattr(_SCREAMING_SNAKE_CASE  ,	"handle_key"  ,	_SCREAMING_SNAKE_CASE    )
				return func

		return decorator





class        _lowerCamelCase      (       _lowercase ):


						def __new__(cls   ,							__a   ,							__a   ,							__a		) ->   Any:
								UpperCamelCase   =	super().__new__(cls   ,							__a   ,							__a   ,							__a		)
								if not hasattr(__a   ,							"key_handler"		):
										setattr(__a   ,							"key_handler"   ,							{}		)
								setattr(__a   ,							"handle_input"   ,							KeyHandler.handle_input		)

								for value in attrs.values():
										UpperCamelCase   =	getattr(__a   ,							"handle_key"   ,							[]		)
										for key in handled_keys:
												UpperCamelCase   =	value
								return new_cls




						@staticmethod
						def     snake_case_   (cls		) ->   Optional[Any]:
								UpperCamelCase   =	get_character()
								if char != KEYMAP["undefined"]:
										UpperCamelCase   =	ord(__a		)
								UpperCamelCase   =	cls.key_handler.get(__a		)
								if handler:
										UpperCamelCase   =	char
										return handler(cls		)
								else:
										return None



def       a__					(					cls    ):




		"""simple docstring"""


		return KeyHandler(cls.__name__  ,	cls.__bases__  ,	cls.__dict__.copy()    )






