






'''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__):
 lowerCamelCase__  =     384
 lowerCamelCase__  =     7
 if "tiny" in model_name:
  lowerCamelCase__  =     96
  lowerCamelCase__  =     (2, 2, 6, 2)
  lowerCamelCase__  =     (3, 6, 12, 24)
 elif "small" in model_name:
  lowerCamelCase__  =     96
  lowerCamelCase__  =     (2, 2, 18, 2)
  lowerCamelCase__  =     (3, 6, 12, 24)
 elif "base" in model_name:
  lowerCamelCase__  =     128
  lowerCamelCase__  =     (2, 2, 18, 2)
  lowerCamelCase__  =     (4, 8, 16, 32)
  lowerCamelCase__  =     12
  lowerCamelCase__  =     512
 elif "large" in model_name:
  lowerCamelCase__  =     192
  lowerCamelCase__  =     (2, 2, 18, 2)
  lowerCamelCase__  =     (6, 12, 24, 48)
  lowerCamelCase__  =     12
  lowerCamelCase__  =     768

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

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

 return config






def    lowerCamelCase_     (    lowercase__):
 lowerCamelCase__  =     []

 # 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__				,     lowercase__				,     lowercase__):
 lowerCamelCase__  =     dct.pop(lowercase__)
 lowerCamelCase__  =     val






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






def    lowerCamelCase_     (    lowercase__):
 lowerCamelCase__  ,	lowerCamelCase__  =     x.shape
 lowerCamelCase__  =     x.reshape(lowercase__				,     4				,     in_channel // 4)
 lowerCamelCase__  =     x[:, [0, 2, 1, 3], :].transpose(1				,     2).reshape(lowercase__				,     lowercase__)
 return x






def    lowerCamelCase_     (    lowercase__):
 lowerCamelCase__  ,	lowerCamelCase__  =     x.shape
 lowerCamelCase__  =     x.reshape(lowercase__				,     in_channel // 4				,     4)
 lowerCamelCase__  =     x[:, :, [0, 2, 1, 3]].transpose(1				,     2).reshape(lowercase__				,     lowercase__)

 return x






def    lowerCamelCase_     (    lowercase__):
 lowerCamelCase__  =     x.shape[0]
 lowerCamelCase__  =     x.reshape(4				,     in_channel // 4)
 lowerCamelCase__  =     x[[0, 2, 1, 3], :].transpose(0				,     1).reshape(lowercase__)
 return x






def    lowerCamelCase_     (    lowercase__):
 lowerCamelCase__  =     x.shape[0]
 lowerCamelCase__  =     x.reshape(in_channel // 4				,     4)
 lowerCamelCase__  =     x[:, [0, 2, 1, 3]].transpose(0				,     1).reshape(lowercase__)
 return x






def    lowerCamelCase_     (    lowercase__				,     lowercase__				,     lowercase__):
 lowerCamelCase__  =     {
     "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",
 }
 lowerCamelCase__  =     model_name_to_url[model_name]
 lowerCamelCase__  =     torch.hub.load_state_dict_from_url(lowercase__				,     map_location="cpu"				,     file_name=lowercase__)[
     "state_dict"
 ]

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

 lowerCamelCase__  =     get_upernet_config(lowercase__)
 lowerCamelCase__  =     UperNetForSemanticSegmentation(lowercase__)
 model.eval()

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

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

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

 model.load_state_dict(lowercase__)

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

 lowerCamelCase__  =     SegformerImageProcessor()
 lowerCamelCase__  =     processor(lowercase__				,     return_tensors="pt").pixel_values

 with torch.no_grad():
  lowerCamelCase__  =     model(lowercase__)
  lowerCamelCase__  =     outputs.logits

 print(logits.shape)
 print("First values of logits:"				,     logits[0, 0, :3, :3])
 # assert values
 if model_name == "upernet-swin-tiny":
  lowerCamelCase__  =     torch.tensor(
      [[-7.5_958, -7.5_958, -7.4_302], [-7.5_958, -7.5_958, -7.4_302], [-7.4_797, -7.4_797, -7.3_068]])
 elif model_name == "upernet-swin-small":
  lowerCamelCase__  =     torch.tensor(
      [[-7.1_921, -7.1_921, -6.9_532], [-7.1_921, -7.1_921, -6.9_532], [-7.0_908, -7.0_908, -6.8_534]])
 elif model_name == "upernet-swin-base":
  lowerCamelCase__  =     torch.tensor(
      [[-6.5_851, -6.5_851, -6.4_330], [-6.5_851, -6.5_851, -6.4_330], [-6.4_763, -6.4_763, -6.3_254]])
 elif model_name == "upernet-swin-large":
  lowerCamelCase__  =     torch.tensor(
      [[-7.5_297, -7.5_297, -7.3_802], [-7.5_297, -7.5_297, -7.3_802], [-7.4_044, -7.4_044, -7.2_586]])
 print("Logits:"				,     outputs.logits[0, 0, :3, :3])
 assert torch.allclose(outputs.logits[0, 0, :3, :3]				,     lowercase__				,     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(lowercase__)
  print(F'''Saving processor to {pytorch_dump_folder_path}''')
  processor.save_pretrained(lowercase__)

 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							:							str  							=		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							:							Optional[int]  							=		parser.parse_args()
 convert_upernet_checkpoint(args.model_name, args.pytorch_dump_folder_path, args.push_to_hub)









'''simple docstring'''





from collections import OrderedDict
from typing import Mapping

from packaging import version

from ...configuration_utils import PretrainedConfig
from ...onnx import OnnxConfig
from ...utils import logging


__A							:							Optional[Any]  							=		logging.get_logger(__name__)

__A							:							int  							=		{
    """google/vit-base-patch16-224""": """https://huggingface.co/vit-base-patch16-224/resolve/main/config.json""",
    # See all ViT models at https://huggingface.co/models?filter=vit
}
class 				lowercase       (							_lowerCamelCase      ):


    '''simple docstring'''

    lowerCAmelCase__      				=							"vit"

    def __init__(					self     :     Optional[int]				,     __lowerCamelCase     :     Tuple=768				,     __lowerCamelCase     :     Tuple=12				,     __lowerCamelCase     :     Tuple=12				,     __lowerCamelCase     :     int=3072				,     __lowerCamelCase     :     Optional[Any]="gelu"				,     __lowerCamelCase     :     Optional[int]=0.0				,     __lowerCamelCase     :     Tuple=0.0				,     __lowerCamelCase     :     Optional[Any]=0.0_2				,     __lowerCamelCase     :     int=1E-12				,     __lowerCamelCase     :     Tuple=224				,     __lowerCamelCase     :     int=16				,     __lowerCamelCase     :     Optional[int]=3				,     __lowerCamelCase     :     List[str]=True				,     __lowerCamelCase     :     str=16				,     **__lowerCamelCase     :     Optional[int]				,     )							->							str:

     '''simple docstring'''




     super().__init__(**__lowerCamelCase     )

     lowerCamelCase__  =     hidden_size
     lowerCamelCase__  =     num_hidden_layers
     lowerCamelCase__  =     num_attention_heads
     lowerCamelCase__  =     intermediate_size
     lowerCamelCase__  =     hidden_act
     lowerCamelCase__  =     hidden_dropout_prob
     lowerCamelCase__  =     attention_probs_dropout_prob
     lowerCamelCase__  =     initializer_range
     lowerCamelCase__  =     layer_norm_eps
     lowerCamelCase__  =     image_size
     lowerCamelCase__  =     patch_size
     lowerCamelCase__  =     num_channels
     lowerCamelCase__  =     qkv_bias
     lowerCamelCase__  =     encoder_stride
class 				lowercase       (							_lowerCamelCase      ):


    '''simple docstring'''

    lowerCAmelCase__      				=							version.parse("1.11"      )

    @property
    def    a__				(					self     :     Optional[int]     )							->							Mapping[str, Mapping[int, str]]:

     '''simple docstring'''




     return OrderedDict(
         [
             ("pixel_values", {0: "batch", 1: "num_channels", 2: "height", 3: "width"}),
         ]     )




    @property
    def    a__				(					self     :     Optional[int]     )							->							float:

     '''simple docstring'''




     return 1E-4

