


'''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 __UpperCamelCase   (			__lowerCamelCase		: List[str]				)      ->     Dict:







     '''simple docstring'''






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

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

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

     return config






def __UpperCamelCase   (			__lowerCamelCase		: Any				)      ->     List[Any]:







     '''simple docstring'''






     _a						=		[]

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







     '''simple docstring'''






     _a						=		dct.pop(__lowerCamelCase				)
     _a						=		val






def __UpperCamelCase   (			__lowerCamelCase		: Union[str, Any]      ,  __lowerCamelCase		: List[Any]				)      ->     List[str]:







     '''simple docstring'''






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






def __UpperCamelCase   (			__lowerCamelCase		: Optional[Any]				)      ->     List[str]:







     '''simple docstring'''






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






def __UpperCamelCase   (			__lowerCamelCase		: Any				)      ->     int:







     '''simple docstring'''






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

     return x






def __UpperCamelCase   (			__lowerCamelCase		: Optional[Any]				)      ->     List[str]:







     '''simple docstring'''






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






def __UpperCamelCase   (			__lowerCamelCase		: Any				)      ->     Dict:







     '''simple docstring'''






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






def __UpperCamelCase   (			__lowerCamelCase		: Dict      ,  __lowerCamelCase		: List[Any]      ,  __lowerCamelCase		: Dict				)      ->     Optional[int]:







     '''simple docstring'''






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

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

     _a						=		get_upernet_config(__lowerCamelCase				)
     _a						=		UperNetForSemanticSegmentation(__lowerCamelCase				)
     model.eval()

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

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

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

     model.load_state_dict(__lowerCamelCase				)

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

     _a						=		SegformerImageProcessor()
     _a						=		processor(__lowerCamelCase      ,  return_tensors="pt"				).pixel_values

     with torch.no_grad():
          _a						=		model(__lowerCamelCase				)
          _a						=		outputs.logits

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

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

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





'''simple docstring'''

import argparse
import os

import numpy as np
import tensorflow as tf
import torch

from transformers import BertModel






def __UpperCamelCase   (			__lowerCamelCase		: BertModel      ,  __lowerCamelCase		: str      ,  __lowerCamelCase		: str				)      ->     List[str]:







     '''simple docstring'''






     _a						=		("dense.weight", "attention.self.query", "attention.self.key", "attention.self.value")

     _a						=		(
         ("layer.", "layer_"),
         ("word_embeddings.weight", "word_embeddings"),
         ("position_embeddings.weight", "position_embeddings"),
         ("token_type_embeddings.weight", "token_type_embeddings"),
         (".", "/"),
         ("LayerNorm/weight", "LayerNorm/gamma"),
         ("LayerNorm/bias", "LayerNorm/beta"),
         ("weight", "kernel"),
     )

     if not os.path.isdir(__lowerCamelCase				):
          os.makedirs(__lowerCamelCase				)

     _a						=		model.state_dict()

     def to_tf_var_name(__lowerCamelCase		: str				):
          for patt, repl in iter(__lowerCamelCase				):
               _a						=		name.replace(__lowerCamelCase      ,  __lowerCamelCase				)
          return F"bert/{name}"

     def create_tf_var(__lowerCamelCase		: np.ndarray      ,  __lowerCamelCase		: str      ,  __lowerCamelCase		: tf.Session				):
          _a						=		tf.dtypes.as_dtype(tensor.dtype				)
          _a						=		tf.get_variable(dtype=__lowerCamelCase      ,  shape=tensor.shape      ,  name=__lowerCamelCase      ,  initializer=tf.zeros_initializer()				)
          session.run(tf.variables_initializer([tf_var]				)				)
          session.run(__lowerCamelCase				)
          return tf_var

     tf.reset_default_graph()
     with tf.Session() as session:
          for var_name in state_dict:
               _a						=		to_tf_var_name(__lowerCamelCase				)
               _a						=		state_dict[var_name].numpy()
               if any(x in var_name for x in tensors_to_transpose				):
                    _a						=		torch_tensor.T
               _a						=		create_tf_var(tensor=__lowerCamelCase      ,  name=__lowerCamelCase      ,  session=__lowerCamelCase				)
               tf.keras.backend.set_value(__lowerCamelCase      ,  __lowerCamelCase				)
               _a						=		session.run(__lowerCamelCase				)
               print(F"Successfully created {tf_name}: {np.allclose(__lowerCamelCase      ,  __lowerCamelCase				)}"				)

          _a						=		tf.train.Saver(tf.trainable_variables()				)
          saver.save(__lowerCamelCase      ,  os.path.join(__lowerCamelCase      ,  model_name.replace("-"      ,  "_"				) + ".ckpt"				)				)






def __UpperCamelCase   (			__lowerCamelCase		: str=None				)      ->     Optional[int]:







     '''simple docstring'''






     _a						=		argparse.ArgumentParser()
     parser.add_argument("--model_name"      ,  type=__lowerCamelCase      ,  required=__lowerCamelCase      ,  help="model name e.g. bert-base-uncased"				)
     parser.add_argument(
         "--cache_dir"      ,  type=__lowerCamelCase      ,  default=__lowerCamelCase      ,  required=__lowerCamelCase      ,  help="Directory containing pytorch model"				)
     parser.add_argument("--pytorch_model_path"      ,  type=__lowerCamelCase      ,  required=__lowerCamelCase      ,  help="/path/to/<pytorch-model-name>.bin"				)
     parser.add_argument("--tf_cache_dir"      ,  type=__lowerCamelCase      ,  required=__lowerCamelCase      ,  help="Directory in which to save tensorflow model"				)
     _a						=		parser.parse_args(__lowerCamelCase				)

     _a						=		BertModel.from_pretrained(
         pretrained_model_name_or_path=args.model_name      ,  state_dict=torch.load(args.pytorch_model_path				)      ,  cache_dir=args.cache_dir      ,  )

     convert_pytorch_checkpoint_to_tf(model=__lowerCamelCase      ,  ckpt_dir=args.tf_cache_dir      ,  model_name=args.model_name				)


if __name__ == "__main__":
   main()

