


"""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   __lowerCAmelCase						(						lowercase  :  Union[str, Any]			)     ->						Union[str, Any]:


				"""simple docstring"""





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

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

				snake_case   :				Tuple				 =   SwinConfig(
				    embed_dim=UpperCamelCase__						,							depths=UpperCamelCase__						,							num_heads=UpperCamelCase__						,							window_size=UpperCamelCase__						,							out_features=["stage1", "stage2", "stage3", "stage4"]						,							)
				snake_case   :				Dict				 =   UperNetConfig(
				    backbone_config=UpperCamelCase__						,							auxiliary_in_channels=UpperCamelCase__						,							num_labels=UpperCamelCase__						,							idalabel=UpperCamelCase__						,							labelaid=UpperCamelCase__						,							)

				return config







def   __lowerCAmelCase						(						lowercase  :  Dict			)     ->						Any:


				"""simple docstring"""





				snake_case   :				Optional[int]				 =   []

				# 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						(						lowercase  :  List[Any]						,							lowercase  :  Union[str, Any]						,							lowercase  :  Any			)     ->						List[str]:


				"""simple docstring"""





				snake_case   :				List[str]				 =   dct.pop(UpperCamelCase__			)
				snake_case   :				List[Any]				 =   val







def   __lowerCAmelCase						(						lowercase  :  Union[str, Any]						,							lowercase  :  Optional[int]			)     ->						Union[str, Any]:


				"""simple docstring"""





				snake_case   :				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			)			):
								snake_case   :				str				 =   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)
												snake_case   :				Union[str, Any]				 =   state_dict.pop(F'backbone.stages.{i}.blocks.{j}.attn.w_msa.qkv.weight'			)
												snake_case   :				Optional[Any]				 =   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
												snake_case   :				Any				 =   in_proj_weight[:dim, :]
												snake_case   :				Optional[Any]				 =   in_proj_bias[: dim]
												snake_case   :				Optional[Any]				 =   in_proj_weight[
												    dim : dim * 2, :
												]
												snake_case   :				str				 =   in_proj_bias[
												    dim : dim * 2
												]
												snake_case   :				List[Any]				 =   in_proj_weight[
												    -dim :, :
												]
												snake_case   :				Union[str, Any]				 =   in_proj_bias[-dim :]
												# fmt: on







def   __lowerCAmelCase						(						lowercase  :  int			)     ->						Dict:


				"""simple docstring"""





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







def   __lowerCAmelCase						(						lowercase  :  List[str]			)     ->						Optional[Any]:


				"""simple docstring"""





				snake_case   :				List[str]				 =   x.shape
				snake_case   :				Any				 =   x.reshape(UpperCamelCase__						,							in_channel // 4						,							4			)
				snake_case   :				Tuple				 =   x[:, :, [0, 2, 1, 3]].transpose(1						,							2			).reshape(UpperCamelCase__						,							UpperCamelCase__			)

				return x







def   __lowerCAmelCase						(						lowercase  :  Union[str, Any]			)     ->						Optional[Any]:


				"""simple docstring"""





				snake_case   :				Union[str, Any]				 =   x.shape[0]
				snake_case   :				Tuple				 =   x.reshape(4						,							in_channel // 4			)
				snake_case   :				str				 =   x[[0, 2, 1, 3], :].transpose(0						,							1			).reshape(UpperCamelCase__			)
				return x







def   __lowerCAmelCase						(						lowercase  :  List[Any]			)     ->						int:


				"""simple docstring"""





				snake_case   :				str				 =   x.shape[0]
				snake_case   :				List[Any]				 =   x.reshape(in_channel // 4						,							4			)
				snake_case   :				Optional[Any]				 =   x[:, [0, 2, 1, 3]].transpose(0						,							1			).reshape(UpperCamelCase__			)
				return x







def   __lowerCAmelCase						(						lowercase  :  Optional[Any]						,							lowercase  :  Union[str, Any]						,							lowercase  :  Union[str, Any]			)     ->						Union[str, Any]:


				"""simple docstring"""





				snake_case   :				Optional[int]				 =   {
				    '''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''',
				}
				snake_case   :				int				 =   model_name_to_url[model_name]
				snake_case   :				Optional[Any]				 =   torch.hub.load_state_dict_from_url(UpperCamelCase__						,							map_location="cpu"						,							file_name=UpperCamelCase__			)[
				    '''state_dict'''
				]

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

				snake_case   :				Optional[int]				 =   get_upernet_config(UpperCamelCase__			)
				snake_case   :				List[str]				 =   UperNetForSemanticSegmentation(UpperCamelCase__			)
				model.eval()

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

				# rename keys
				snake_case   :				Optional[int]				 =   create_rename_keys(UpperCamelCase__			)
				for src, dest in rename_keys:
								rename_key(UpperCamelCase__						,							UpperCamelCase__						,							UpperCamelCase__			)
				read_in_q_k_v(UpperCamelCase__						,							config.backbone_config			)

				# fix downsample parameters
				for key, value in state_dict.items():
								if "downsample" in key:
												if "reduction" in key:
																snake_case   :				Optional[Any]				 =   reverse_correct_unfold_reduction_order(UpperCamelCase__			)
												if "norm" in key:
																snake_case   :				Tuple				 =   reverse_correct_unfold_norm_order(UpperCamelCase__			)

				model.load_state_dict(UpperCamelCase__			)

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

				snake_case   :				Optional[int]				 =   SegformerImageProcessor()
				snake_case   :				List[Any]				 =   processor(UpperCamelCase__						,							return_tensors="pt"			).pixel_values

				with torch.no_grad():
								snake_case   :				Union[str, Any]				 =   model(UpperCamelCase__			)
								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":
								snake_case   :				Optional[int]				 =   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":
								snake_case   :				str				 =   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":
								snake_case   :				Union[str, Any]				 =   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":
								snake_case   :				List[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]						,							UpperCamelCase__						,							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(UpperCamelCase__			)
								print(F'Saving processor to {pytorch_dump_folder_path}'			)
								processor.save_pretrained(UpperCamelCase__			)

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

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









"""simple docstring"""



def   __lowerCAmelCase						(						lowercase  :  int			)     ->						int:


    """simple docstring"""





    if not isinstance(lowercase						,							lowercase			):
        raise ValueError("Input must be an integer"			)
    if input_num <= 0:
        raise ValueError("Input must be positive"			)
    return sum(
        divisor for divisor in range(1						,							input_num // 2 + 1			) if input_num % divisor == 0			)


if __name__ == "__main__":
  import doctest

  doctest.testmod()





