
import argparse
import json
from collections import OrderedDict
from pathlib import Path

import requests
import torch
from huggingface_hub import hf_hub_download
from PIL import Image

from transformers import (
    SegformerConfig,
    SegformerForImageClassification,
    SegformerForSemanticSegmentation,
    SegformerImageProcessor,
)
from transformers.utils import logging


logging.set_verbosity_info()
lowerCamelCase_						     =					logging.get_logger(__name__)






def 						__magic_name__			(   __a      :	int	,		__a      :	Optional[int]=False						):



     '''simple docstring'''




     UpperCamelCase__				=    OrderedDict()
     for key, value in state_dict.items():
          if encoder_only and not key.startswith("""head"""						):
               UpperCamelCase__				=    """segformer.encoder.""" + key
          if key.startswith("""backbone"""						):
               UpperCamelCase__				=    key.replace("""backbone"""	,		"""segformer.encoder"""						)
          if "patch_embed" in key:
               # replace for example patch_embed1 by patch_embeddings.0
               UpperCamelCase__				=    key[key.find("""patch_embed"""						) + len("""patch_embed"""						)]
               UpperCamelCase__				=    key.replace(f"patch_embed{idx}"	,		f"patch_embeddings.{int(__a						)-1}"						)
          if "norm" in key:
               UpperCamelCase__				=    key.replace("""norm"""	,		"""layer_norm"""						)
          if "segformer.encoder.layer_norm" in key:
               # replace for example layer_norm1 by layer_norm.0
               UpperCamelCase__				=    key[key.find("""segformer.encoder.layer_norm"""						) + len("""segformer.encoder.layer_norm"""						)]
               UpperCamelCase__				=    key.replace(f"layer_norm{idx}"	,		f"layer_norm.{int(__a						)-1}"						)
          if "layer_norm1" in key:
               UpperCamelCase__				=    key.replace("""layer_norm1"""	,		"""layer_norm_1"""						)
          if "layer_norm2" in key:
               UpperCamelCase__				=    key.replace("""layer_norm2"""	,		"""layer_norm_2"""						)
          if "block" in key:
               # replace for example block1 by block.0
               UpperCamelCase__				=    key[key.find("""block"""						) + len("""block"""						)]
               UpperCamelCase__				=    key.replace(f"block{idx}"	,		f"block.{int(__a						)-1}"						)
          if "attn.q" in key:
               UpperCamelCase__				=    key.replace("""attn.q"""	,		"""attention.self.query"""						)
          if "attn.proj" in key:
               UpperCamelCase__				=    key.replace("""attn.proj"""	,		"""attention.output.dense"""						)
          if "attn" in key:
               UpperCamelCase__				=    key.replace("""attn"""	,		"""attention.self"""						)
          if "fc1" in key:
               UpperCamelCase__				=    key.replace("""fc1"""	,		"""dense1"""						)
          if "fc2" in key:
               UpperCamelCase__				=    key.replace("""fc2"""	,		"""dense2"""						)
          if "linear_pred" in key:
               UpperCamelCase__				=    key.replace("""linear_pred"""	,		"""classifier"""						)
          if "linear_fuse" in key:
               UpperCamelCase__				=    key.replace("""linear_fuse.conv"""	,		"""linear_fuse"""						)
               UpperCamelCase__				=    key.replace("""linear_fuse.bn"""	,		"""batch_norm"""						)
          if "linear_c" in key:
               # replace for example linear_c4 by linear_c.3
               UpperCamelCase__				=    key[key.find("""linear_c"""						) + len("""linear_c"""						)]
               UpperCamelCase__				=    key.replace(f"linear_c{idx}"	,		f"linear_c.{int(__a						)-1}"						)
          if key.startswith("""head"""						):
               UpperCamelCase__				=    key.replace("""head"""	,		"""classifier"""						)
          UpperCamelCase__				=    value

     return new_state_dict






def 						__magic_name__			(   __a      :	List[Any]	,		__a      :	Optional[int]						):



     '''simple docstring'''




     for i in range(config.num_encoder_blocks						):
          for j in range(config.depths[i]						):
               # read in weights + bias of keys and values (which is a single matrix in the original implementation)
               UpperCamelCase__				=    state_dict.pop(f"segformer.encoder.block.{i}.{j}.attention.self.kv.weight"						)
               UpperCamelCase__				=    state_dict.pop(f"segformer.encoder.block.{i}.{j}.attention.self.kv.bias"						)
               # next, add keys and values (in that order) to the state dict
               UpperCamelCase__				=    kv_weight[
                   : config.hidden_sizes[i], :
               ]
               UpperCamelCase__				=    kv_bias[: config.hidden_sizes[i]]
               UpperCamelCase__				=    kv_weight[
                   config.hidden_sizes[i] :, :
               ]
               UpperCamelCase__				=    kv_bias[
                   config.hidden_sizes[i] :
               ]






def 						__magic_name__			(   ):



     '''simple docstring'''




     UpperCamelCase__				=    """http://images.cocodataset.org/val2017/000000039769.jpg"""
     UpperCamelCase__				=    Image.open(requests.get(__a	,		stream=__a						).raw						)

     return image






@torch.no_grad()
def 						__magic_name__			(   __a      :	Union[str, Any]	,		__a      :	Optional[int]	,		__a      :	Any						):



     '''simple docstring'''




     UpperCamelCase__				=    SegformerConfig()
     UpperCamelCase__				=    False

     # set attributes based on model_name
     UpperCamelCase__				=    """huggingface/label-files"""
     if "segformer" in model_name:
          UpperCamelCase__				=    model_name[len("""segformer."""						) : len("""segformer."""						) + 2]
          if "ade" in model_name:
               UpperCamelCase__				=    150
               UpperCamelCase__				=    """ade20k-id2label.json"""
               UpperCamelCase__				=    (1, 150, 128, 128)
          elif "city" in model_name:
               UpperCamelCase__				=    19
               UpperCamelCase__				=    """cityscapes-id2label.json"""
               UpperCamelCase__				=    (1, 19, 128, 128)
          else:
               raise ValueError(f"Model {model_name} not supported"						)
     elif "mit" in model_name:
          UpperCamelCase__				=    True
          UpperCamelCase__				=    model_name[4:6]
          UpperCamelCase__				=    1_000
          UpperCamelCase__				=    """imagenet-1k-id2label.json"""
          UpperCamelCase__				=    (1, 1_000)
     else:
          raise ValueError(f"Model {model_name} not supported"						)

     # set config attributes
     UpperCamelCase__				=    json.load(open(hf_hub_download(__a	,		__a	,		repo_type="""dataset"""						)	,		"""r"""						)						)
     UpperCamelCase__				=    {int(__a						): v for k, v in idalabel.items()}
     UpperCamelCase__				=    idalabel
     UpperCamelCase__				=    {v: k for k, v in idalabel.items()}
     if size == "b0":
          pass
     elif size == "b1":
          UpperCamelCase__				=    [64, 128, 320, 512]
          UpperCamelCase__				=    256
     elif size == "b2":
          UpperCamelCase__				=    [64, 128, 320, 512]
          UpperCamelCase__				=    768
          UpperCamelCase__				=    [3, 4, 6, 3]
     elif size == "b3":
          UpperCamelCase__				=    [64, 128, 320, 512]
          UpperCamelCase__				=    768
          UpperCamelCase__				=    [3, 4, 18, 3]
     elif size == "b4":
          UpperCamelCase__				=    [64, 128, 320, 512]
          UpperCamelCase__				=    768
          UpperCamelCase__				=    [3, 8, 27, 3]
     elif size == "b5":
          UpperCamelCase__				=    [64, 128, 320, 512]
          UpperCamelCase__				=    768
          UpperCamelCase__				=    [3, 6, 40, 3]
     else:
          raise ValueError(f"Size {size} not supported"						)

     # load image processor (only resize + normalize)
     UpperCamelCase__				=    SegformerImageProcessor(
         image_scale=(512, 512)	,		keep_ratio=__a	,		align=__a	,		do_random_crop=__a						)

     # prepare image
     UpperCamelCase__				=    prepare_img()
     UpperCamelCase__				=    image_processor(images=__a	,		return_tensors="""pt"""						).pixel_values

     logger.info(f"Converting model {model_name}..."						)

     # load original state dict
     if encoder_only:
          UpperCamelCase__				=    torch.load(__a	,		map_location=torch.device("""cpu"""						)						)
     else:
          UpperCamelCase__				=    torch.load(__a	,		map_location=torch.device("""cpu"""						)						)["""state_dict"""]

     # rename keys
     UpperCamelCase__				=    rename_keys(__a	,		encoder_only=__a						)
     if not encoder_only:
          del state_dict["decode_head.conv_seg.weight"]
          del state_dict["decode_head.conv_seg.bias"]

     # key and value matrices need special treatment
     read_in_k_v(__a	,		__a						)

     # create HuggingFace model and load state dict
     if encoder_only:
          UpperCamelCase__				=    False
          UpperCamelCase__				=    SegformerForImageClassification(__a						)
     else:
          UpperCamelCase__				=    SegformerForSemanticSegmentation(__a						)
     model.load_state_dict(__a						)
     model.eval()

     # forward pass
     UpperCamelCase__				=    model(__a						)
     UpperCamelCase__				=    outputs.logits

     # set expected_slice based on model name
     # ADE20k checkpoints
     if model_name == "segformer.b0.512x512.ade.160k":
          UpperCamelCase__				=    torch.tensor(
              [
                  [[-4.6_310, -5.5_232, -6.2_356], [-5.1_921, -6.1_444, -6.5_996], [-5.4_424, -6.2_790, -6.7_574]],
                  [[-12.1_391, -13.3_122, -13.9_554], [-12.8_732, -13.9_352, -14.3_563], [-12.9_438, -13.8_226, -14.2_513]],
                  [[-12.5_134, -13.4_686, -14.4_915], [-12.8_669, -14.4_343, -14.7_758], [-13.2_523, -14.5_819, -15.0_694]],
              ]						)
     elif model_name == "segformer.b1.512x512.ade.160k":
          UpperCamelCase__				=    torch.tensor(
              [
                  [[-7.5_820, -8.7_231, -8.3_215], [-8.0_600, -10.3_529, -10.0_304], [-7.5_208, -9.4_103, -9.6_239]],
                  [[-12.6_918, -13.8_994, -13.7_137], [-13.3_196, -15.7_523, -15.4_789], [-12.9_343, -14.8_757, -14.9_689]],
                  [[-11.1_911, -11.9_421, -11.3_243], [-11.3_342, -13.6_839, -13.3_581], [-10.3_909, -12.1_832, -12.4_858]],
              ]						)
     elif model_name == "segformer.b2.512x512.ade.160k":
          UpperCamelCase__				=    torch.tensor(
              [
                  [[-11.8_173, -14.3_850, -16.3_128], [-14.5_648, -16.5_804, -18.6_568], [-14.7_223, -15.7_387, -18.4_218]],
                  [[-15.7_290, -17.9_171, -19.4_423], [-18.3_105, -19.9_448, -21.4_661], [-17.9_296, -18.6_497, -20.7_910]],
                  [[-15.0_783, -17.0_336, -18.2_789], [-16.8_771, -18.6_870, -20.1_612], [-16.2_454, -17.1_426, -19.5_055]],
              ]						)
     elif model_name == "segformer.b3.512x512.ade.160k":
          UpperCamelCase__				=    torch.tensor(
              [
                  [[-9.0_878, -10.2_081, -10.1_891], [-9.3_144, -10.7_941, -10.9_843], [-9.2_294, -10.3_855, -10.5_704]],
                  [[-12.2_316, -13.9_068, -13.6_102], [-12.9_161, -14.3_702, -14.3_235], [-12.5_233, -13.7_174, -13.7_932]],
                  [[-14.6_275, -15.2_490, -14.9_727], [-14.3_400, -15.9_687, -16.2_827], [-14.1_484, -15.4_033, -15.8_937]],
              ]						)
     elif model_name == "segformer.b4.512x512.ade.160k":
          UpperCamelCase__				=    torch.tensor(
              [
                  [[-12.3_144, -13.2_447, -14.0_802], [-13.3_614, -14.5_816, -15.6_117], [-13.3_340, -14.4_433, -16.2_219]],
                  [[-19.2_781, -20.4_128, -20.7_506], [-20.6_153, -21.6_566, -22.0_998], [-19.9_800, -21.0_430, -22.1_494]],
                  [[-18.8_739, -19.7_804, -21.1_834], [-20.1_233, -21.6_765, -23.2_944], [-20.0_315, -21.2_641, -23.6_944]],
              ]						)
     elif model_name == "segformer.b5.640x640.ade.160k":
          UpperCamelCase__				=    torch.tensor(
              [
                  [[-9.5_524, -12.0_835, -11.7_348], [-10.5_229, -13.6_446, -14.5_662], [-9.5_842, -12.8_851, -13.9_414]],
                  [[-15.3_432, -17.5_323, -17.0_818], [-16.3_330, -18.9_255, -19.2_101], [-15.1_340, -17.7_848, -18.3_971]],
                  [[-12.6_072, -14.9_486, -14.6_631], [-13.7_629, -17.0_907, -17.7_745], [-12.7_899, -16.1_695, -17.1_671]],
              ]						)
     # Cityscapes checkpoints
     elif model_name == "segformer.b0.1024x1024.city.160k":
          UpperCamelCase__				=    torch.tensor(
              [
                  [[-11.9_295, -13.4_057, -14.8_106], [-13.3_431, -14.8_179, -15.3_781], [-14.2_836, -15.5_942, -16.1_588]],
                  [[-11.4_906, -12.8_067, -13.6_564], [-13.1_189, -14.0_500, -14.1_543], [-13.8_748, -14.5_136, -14.8_789]],
                  [[0.5_374, 0.1_067, -0.4_742], [0.1_141, -0.2_255, -0.7_099], [-0.3_000, -0.5_924, -1.3_105]],
              ]						)
     elif model_name == "segformer.b0.512x1024.city.160k":
          UpperCamelCase__				=    torch.tensor(
              [
                  [[-7.8_217, -9.8_767, -10.1_717], [-9.4_438, -10.9_058, -11.4_047], [-9.7_939, -12.3_495, -12.1_079]],
                  [[-7.1_514, -9.5_336, -10.0_860], [-9.7_776, -11.6_822, -11.8_439], [-10.1_411, -12.7_655, -12.8_972]],
                  [[0.3_021, 0.0_805, -0.2_310], [-0.0_328, -0.1_605, -0.2_714], [-0.1_408, -0.5_477, -0.6_976]],
              ]						)
     elif model_name == "segformer.b0.640x1280.city.160k":
          UpperCamelCase__				=    torch.tensor(
              [
                  [
                      [-1.1_3_7_2E0_1, -1.2_7_8_7E0_1, -1.3_4_7_7E0_1],
                      [-1.2_5_3_6E0_1, -1.4_1_9_4E0_1, -1.4_4_0_9E0_1],
                      [-1.3_2_1_7E0_1, -1.4_8_8_8E0_1, -1.5_3_2_7E0_1],
                  ],
                  [
                      [-1.4_7_9_1E0_1, -1.7_1_2_2E0_1, -1.8_2_7_7E0_1],
                      [-1.7_1_6_3E0_1, -1.9_1_9_2E0_1, -1.9_5_3_3E0_1],
                      [-1.7_8_9_7E0_1, -1.9_9_9_1E0_1, -2.0_3_1_5E0_1],
                  ],
                  [
                      [7.6_7_2_3E-0_1, 4.1_9_2_1E-0_1, -7.7_8_7_8E-0_2],
                      [4.7_7_7_2E-0_1, 9.5_5_5_7E-0_3, -2.8_0_8_2E-0_1],
                      [3.6_0_3_2E-0_1, -2.4_8_2_6E-0_1, -5.1_1_6_8E-0_1],
                  ],
              ]						)
     elif model_name == "segformer.b0.768x768.city.160k":
          UpperCamelCase__				=    torch.tensor(
              [
                  [[-9.4_959, -11.3_087, -11.7_479], [-11.0_025, -12.6_540, -12.3_319], [-11.4_064, -13.0_487, -12.9_905]],
                  [[-9.8_905, -11.3_084, -12.0_854], [-11.1_726, -12.7_698, -12.9_583], [-11.5_985, -13.3_278, -14.1_774]],
                  [[0.2_213, 0.0_192, -0.2_466], [-0.1_731, -0.4_213, -0.4_874], [-0.3_126, -0.6_541, -1.1_389]],
              ]						)
     elif model_name == "segformer.b1.1024x1024.city.160k":
          UpperCamelCase__				=    torch.tensor(
              [
                  [[-13.5_748, -13.9_111, -12.6_500], [-14.3_500, -15.3_683, -14.2_328], [-14.7_532, -16.0_424, -15.6_087]],
                  [[-17.1_651, -15.8_725, -12.9_653], [-17.2_580, -17.3_718, -14.8_223], [-16.6_058, -16.8_783, -16.7_452]],
                  [[-3.6_456, -3.0_209, -1.4_203], [-3.0_797, -3.1_959, -2.0_000], [-1.8_757, -1.9_217, -1.6_997]],
              ]						)
     elif model_name == "segformer.b2.1024x1024.city.160k":
          UpperCamelCase__				=    torch.tensor(
              [
                  [[-16.0_976, -16.4_856, -17.3_962], [-16.6_234, -19.0_342, -19.7_685], [-16.0_900, -18.0_661, -19.1_180]],
                  [[-18.4_750, -18.8_488, -19.5_074], [-19.4_030, -22.1_570, -22.5_977], [-19.1_191, -20.8_486, -22.3_783]],
                  [[-4.5_178, -5.5_037, -6.5_109], [-5.0_884, -7.2_174, -8.0_334], [-4.4_156, -5.8_117, -7.2_970]],
              ]						)
     elif model_name == "segformer.b3.1024x1024.city.160k":
          UpperCamelCase__				=    torch.tensor(
              [
                  [[-14.2_081, -14.4_732, -14.1_977], [-14.5_867, -16.4_423, -16.6_356], [-13.4_441, -14.9_685, -16.8_696]],
                  [[-14.4_576, -14.7_073, -15.0_451], [-15.0_816, -17.6_237, -17.9_873], [-14.4_213, -16.0_199, -18.5_992]],
                  [[-4.7_349, -4.9_588, -5.0_966], [-4.3_210, -6.9_325, -7.2_591], [-3.4_312, -4.7_484, -7.1_917]],
              ]						)
     elif model_name == "segformer.b4.1024x1024.city.160k":
          UpperCamelCase__				=    torch.tensor(
              [
                  [[-11.7_737, -11.9_526, -11.3_273], [-13.6_692, -14.4_574, -13.8_878], [-13.8_937, -14.6_924, -15.9_345]],
                  [[-14.6_706, -14.5_330, -14.1_306], [-16.1_502, -16.8_180, -16.4_269], [-16.8_338, -17.8_939, -20.1_746]],
                  [[1.0_491, 0.8_289, 1.0_310], [1.1_044, 0.5_219, 0.8_055], [1.0_899, 0.6_926, 0.5_590]],
              ]						)
     elif model_name == "segformer.b5.1024x1024.city.160k":
          UpperCamelCase__				=    torch.tensor(
              [
                  [[-12.5_641, -13.4_777, -13.0_684], [-13.9_587, -15.8_983, -16.6_557], [-13.3_109, -15.7_350, -16.3_141]],
                  [[-14.7_074, -15.4_352, -14.5_944], [-16.6_353, -18.1_663, -18.6_120], [-15.1_702, -18.0_329, -18.1_547]],
                  [[-1.7_990, -2.0_951, -1.7_784], [-2.6_397, -3.8_245, -3.9_686], [-1.5_264, -2.8_126, -2.9_316]],
              ]						)
     else:
          UpperCamelCase__				=    logits.argmax(-1						).item()
          print("""Predicted class:"""	,		model.config.idalabel[predicted_class_idx]						)

     # verify logits
     if not encoder_only:
          assert logits.shape == expected_shape
          assert torch.allclose(logits[0, :3, :3, :3]	,		__a	,		atol=1E-2						)

     # finally, save model and image processor
     logger.info(f"Saving PyTorch model and image processor to {pytorch_dump_folder_path}..."						)
     Path(__a						).mkdir(exist_ok=__a						)
     model.save_pretrained(__a						)
     image_processor.save_pretrained(__a						)


if __name__ == "__main__":
    lowerCamelCase_						     =					argparse.ArgumentParser()

    parser.add_argument(
        '''--model_name''',
        default='''segformer.b0.512x512.ade.160k''',
        type=str,
        help='''Name of the model you\'d like to convert.''',
    )
    parser.add_argument(
        '''--checkpoint_path''', default=None, type=str, help='''Path to the original PyTorch checkpoint (.pth file).'''
    )
    parser.add_argument(
        '''--pytorch_dump_folder_path''', default=None, type=str, help='''Path to the folder to output PyTorch model.'''
    )
    lowerCamelCase_						     =					parser.parse_args()
    convert_segformer_checkpoint(args.model_name, args.checkpoint_path, args.pytorch_dump_folder_path)







import argparse
from collections import defaultdict

import yaml


lowerCamelCase_						     =					'''docs/source/en/_toctree.yml'''






def 						__magic_name__			(   __a      :	Union[str, Any]						):



     '''simple docstring'''




     UpperCamelCase__				=    defaultdict(__a						)
     for doc in model_doc:
          counts[doc["local"]] += 1
     UpperCamelCase__				=    [key for key, value in counts.items() if value > 1]

     UpperCamelCase__				=    []
     for duplicate_key in duplicates:
          UpperCamelCase__				=    list({doc["""title"""] for doc in model_doc if doc["""local"""] == duplicate_key}						)
          if len(__a						) > 1:
               raise ValueError(
                   f"{duplicate_key} is present several times in the documentation table of content at "
                   """`docs/source/en/_toctree.yml` with different *Title* values. Choose one of those and remove the """
                   """others."""						)
          # Only add this once
          new_doc.append({"""local""": duplicate_key, """title""": titles[0]}						)

     # Add none duplicate-keys
     new_doc.extend([doc for doc in model_doc if counts[doc["""local"""]] == 1]						)

     # Sort
     return sorted(__a	,		key=lambda __a						: s["title"].lower()						)






def 						__magic_name__			(   __a      :	int=False						):



     '''simple docstring'''




     with open(__a	,		encoding="""utf-8"""						) as f:
          UpperCamelCase__				=    yaml.safe_load(f.read()						)

     # Get to the API doc
     UpperCamelCase__				=    0
     while content[api_idx]["title"] != "API":
          api_idx += 1
     UpperCamelCase__				=    content[api_idx]["""sections"""]

     # Then to the model doc
     UpperCamelCase__				=    0
     while api_doc[model_idx]["title"] != "Models":
          model_idx += 1

     UpperCamelCase__				=    api_doc[model_idx]["""sections"""]

     UpperCamelCase__				=    [(idx, section) for idx, section in enumerate(__a						) if """sections""" in section]
     UpperCamelCase__				=    False
     for idx, modality_doc in modalities_docs:
          UpperCamelCase__				=    modality_doc["""sections"""]
          UpperCamelCase__				=    clean_model_doc_toc(__a						)

          if old_modality_doc != new_modality_doc:
               UpperCamelCase__				=    True
               if overwrite:
                    UpperCamelCase__				=    new_modality_doc

     if diff:
          if overwrite:
               UpperCamelCase__				=    model_doc
               UpperCamelCase__				=    api_doc
               with open(__a	,		"""w"""	,		encoding="""utf-8"""						) as f:
                    f.write(yaml.dump(__a	,		allow_unicode=__a						)						)
          else:
               raise ValueError(
                   """The model doc part of the table of content is not properly sorted, run `make style` to fix this."""						)


if __name__ == "__main__":
    lowerCamelCase_						     =					argparse.ArgumentParser()
    parser.add_argument('''--fix_and_overwrite''', action='''store_true''', help='''Whether to fix inconsistencies.''')
    lowerCamelCase_						     =					parser.parse_args()

    check_model_doc(args.fix_and_overwrite)





