





'''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     __magic_name__							(				__UpperCAmelCase   )    ->       str:

   '''simple docstring'''







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

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

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

   return config

def     __magic_name__							(				__UpperCAmelCase   )    ->       Optional[Any]:

   '''simple docstring'''







   __SCREAMING_SNAKE_CASE   				=							[]

   # 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     __magic_name__							(				__UpperCAmelCase				,						__UpperCAmelCase				,						__UpperCAmelCase   )    ->       Tuple:

   '''simple docstring'''







   __SCREAMING_SNAKE_CASE   				=							dct.pop(__UpperCAmelCase   )
   __SCREAMING_SNAKE_CASE   				=							val

def     __magic_name__							(				__UpperCAmelCase				,						__UpperCAmelCase   )    ->       Optional[Any]:

   '''simple docstring'''







   __SCREAMING_SNAKE_CASE   				=							[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   				=							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   				=							state_dict.pop(f"""backbone.stages.{i}.blocks.{j}.attn.w_msa.qkv.weight"""   )
         __SCREAMING_SNAKE_CASE   				=							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   				=							in_proj_weight[:dim, :]
         __SCREAMING_SNAKE_CASE   				=							in_proj_bias[: dim]
         __SCREAMING_SNAKE_CASE   				=							in_proj_weight[
             dim : dim * 2, :
         ]
         __SCREAMING_SNAKE_CASE   				=							in_proj_bias[
             dim : dim * 2
         ]
         __SCREAMING_SNAKE_CASE   				=							in_proj_weight[
             -dim :, :
         ]
         __SCREAMING_SNAKE_CASE   				=							in_proj_bias[-dim :]
         # fmt: on

def     __magic_name__							(				__UpperCAmelCase   )    ->       Union[str, Any]:

   '''simple docstring'''







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

def     __magic_name__							(				__UpperCAmelCase   )    ->       Tuple:

   '''simple docstring'''







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

   return x

def     __magic_name__							(				__UpperCAmelCase   )    ->       Any:

   '''simple docstring'''







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

def     __magic_name__							(				__UpperCAmelCase   )    ->       Dict:

   '''simple docstring'''







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

def     __magic_name__							(				__UpperCAmelCase				,						__UpperCAmelCase				,						__UpperCAmelCase   )    ->       int:

   '''simple docstring'''







   __SCREAMING_SNAKE_CASE   				=							{
       """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   				=							model_name_to_url[model_name]
   __SCREAMING_SNAKE_CASE   				=							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   )

   __SCREAMING_SNAKE_CASE   				=							get_upernet_config(__UpperCAmelCase   )
   __SCREAMING_SNAKE_CASE   				=							UperNetForSemanticSegmentation(__UpperCAmelCase   )
   model.eval()

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

   # rename keys
   __SCREAMING_SNAKE_CASE   				=							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:
            __SCREAMING_SNAKE_CASE   				=							reverse_correct_unfold_reduction_order(__UpperCAmelCase   )
         if "norm" in key:
            __SCREAMING_SNAKE_CASE   				=							reverse_correct_unfold_norm_order(__UpperCAmelCase   )

   model.load_state_dict(__UpperCAmelCase   )

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

   __SCREAMING_SNAKE_CASE   				=							SegformerImageProcessor()
   __SCREAMING_SNAKE_CASE   				=							processor(__UpperCAmelCase				,						return_tensors="""pt"""   ).pixel_values

   with torch.no_grad():
      __SCREAMING_SNAKE_CASE   				=							model(__UpperCAmelCase   )
      __SCREAMING_SNAKE_CASE   				=							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   				=							torch.tensor(
          [[-7.5_9_5_8, -7.5_9_5_8, -7.4_3_0_2], [-7.5_9_5_8, -7.5_9_5_8, -7.4_3_0_2], [-7.4_7_9_7, -7.4_7_9_7, -7.3_0_6_8]]   )
   elif model_name == "upernet-swin-small":
      __SCREAMING_SNAKE_CASE   				=							torch.tensor(
          [[-7.1_9_2_1, -7.1_9_2_1, -6.9_5_3_2], [-7.1_9_2_1, -7.1_9_2_1, -6.9_5_3_2], [-7.0_9_0_8, -7.0_9_0_8, -6.8_5_3_4]]   )
   elif model_name == "upernet-swin-base":
      __SCREAMING_SNAKE_CASE   				=							torch.tensor(
          [[-6.5_8_5_1, -6.5_8_5_1, -6.4_3_3_0], [-6.5_8_5_1, -6.5_8_5_1, -6.4_3_3_0], [-6.4_7_6_3, -6.4_7_6_3, -6.3_2_5_4]]   )
   elif model_name == "upernet-swin-large":
      __SCREAMING_SNAKE_CASE   				=							torch.tensor(
          [[-7.5_2_9_7, -7.5_2_9_7, -7.3_8_0_2], [-7.5_2_9_7, -7.5_2_9_7, -7.3_8_0_2], [-7.4_0_4_4, -7.4_0_4_4, -7.2_5_8_6]]   )
   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__":
    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)








"""simple docstring"""





def     __SCREAMING_SNAKE_CASE    ( A_					):
	lowerCAmelCase__      :		Any    							=      int(A_					)
	if decimal in (0, 1):  # Exit cases for the recursion
		return str(A_					)
	lowerCAmelCase__							,lowerCAmelCase__      :		Any    							=      divmod(A_		,	2					)
	return binary_recursive(A_					) + str(A_					)

def     __SCREAMING_SNAKE_CASE    ( A_					):
	lowerCAmelCase__      :		Union[str, Any]    							=      str(A_					).strip()
	if not number:
		raise ValueError('''No input value was provided'''					)
	lowerCAmelCase__      :		Any    							=      '''-''' if number.startswith('''-'''					) else ''''''
	lowerCAmelCase__      :		Dict    							=      number.lstrip('''-'''					)
	if not number.isnumeric():
		raise ValueError('''Input value is not an integer'''					)
	return f'{negative}0b{binary_recursive(int(A_					)					)}'


if __name__ == "__main__":
							from doctest import testmod

							testmod()




