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__       (			a__ )      ->Union[str, Any]:




  '''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(a__    ,     a__    ,     repo_type="dataset" )    ,     "r" ) )
  _UpperCamelCase           =    {int(a__ ): v for k, v in idalabel.items()}
  _UpperCamelCase           =    {v: k for k, v in idalabel.items()}

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

  return config


def 				lowerCAmelCase__       (			a__ )      ->str:




  '''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 				lowerCAmelCase__       (			a__    ,     a__    ,     a__ )      ->int:




  '''simple docstring'''

  _UpperCamelCase           =    dct.pop(a__ )
  _UpperCamelCase           =    val


def 				lowerCAmelCase__       (			a__    ,     a__ )      ->Dict:




  '''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 				lowerCAmelCase__       (			a__ )      ->Optional[Any]:




  '''simple docstring'''

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


def 				lowerCAmelCase__       (			a__ )      ->Union[str, Any]:




  '''simple docstring'''

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

  return x


def 				lowerCAmelCase__       (			a__ )      ->Optional[int]:




  '''simple docstring'''

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


def 				lowerCAmelCase__       (			a__ )      ->Optional[Any]:




  '''simple docstring'''

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


def 				lowerCAmelCase__       (			a__    ,     a__    ,     a__ )      ->List[Any]:




  '''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(a__    ,     map_location="cpu"    ,     file_name=a__ )[
      "state_dict"
  ]

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

  _UpperCamelCase           =    get_upernet_config(a__ )
  _UpperCamelCase           =    UperNetForSemanticSegmentation(a__ )
  model.eval()

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

  # rename keys
  _UpperCamelCase           =    create_rename_keys(a__ )
  for src, dest in rename_keys:
    rename_key(a__    ,     a__    ,     a__ )
  read_in_q_k_v(a__    ,     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(a__ )
      if "norm" in key:
        _UpperCamelCase           =    reverse_correct_unfold_norm_order(a__ )

  model.load_state_dict(a__ )

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

  _UpperCamelCase           =    SegformerImageProcessor()
  _UpperCamelCase           =    processor(a__    ,     return_tensors="pt" ).pixel_values

  with torch.no_grad():
    _UpperCamelCase           =    model(a__ )
    _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.5958, -7.5958, -7.4302], [-7.5958, -7.5958, -7.4302], [-7.4797, -7.4797, -7.3068]] )
  elif model_name == "upernet-swin-small":
    _UpperCamelCase           =    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":
    _UpperCamelCase           =    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":
    _UpperCamelCase           =    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]    ,     a__    ,     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(a__ )
    print(f'Saving processor to {pytorch_dump_folder_path}' )
    processor.save_pretrained(a__ )

  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)





def 				lowerCAmelCase__       (			a__ )      ->int:




  '''simple docstring'''

  if not isinstance(a__    ,     a__ ) or number < 0:
    raise ValueError("Input must be a non-negative integer" )

  _UpperCamelCase           =    0
  while number:
    # This way we arrive at next set bit (next 1) instead of looping
    # through each bit and checking for 1s hence the
    # loop won't run 32 times it will only run the number of `1` times
    number &= number - 1
    count += 1
  return count


if __name__ == "__main__":
   import doctest

   doctest.testmod()




