codingutf8 2022 the huggingface inc team licensed under the apache license version 2 0 the license you may not use this file except in compliance with the license you may obtain a copy of the license at http www apache orglicenseslicense2 0 unless required by applicable law or agreed to in writing software distributed under the license is distributed on an as is basis without warranties or conditions of any kind either express or implied see the license for the specific language governing permissions and limitations under the license convert maskformer checkpoints with swin backbone from the original repository url https github comfacebookresearchmaskformer import argparse import json import pickle from pathlib import path import requests import torch from huggingfacehub import hfhubdownload from pil import image from transformers import maskformerconfig maskformerforinstancesegmentation maskformerimageprocessor swinconfig from transformers utils import logging logging setverbosityinfo logger logging getloggername def getmaskformerconfigmodelname str backboneconfig swinconfig frompretrained microsoftswintinypatch4window7224 outfeaturesstage1 stage2 stage3 stage4 config maskformerconfigbackboneconfigbackboneconfig repoid huggingfacelabelfiles if ade20kfull in modelname this should be ok config numlabels 847 filename maskformerade20kfullid2label json elif ade in modelname this should be ok config numlabels 150 filename ade20kid2label json elif cocostuff in modelname this should be ok config numlabels 171 filename maskformercocostuffid2label json elif coco in modelname todo config numlabels 133 filename cocopanopticid2label json elif cityscapes in modelname this should be ok config numlabels 19 filename cityscapesid2label json elif vistas in modelname this should be ok config numlabels 65 filename mapillaryvistasid2label json id2label json loadopenhfhubdownloadrepoid filename repotypedataset r id2label intk v for k v in id2label items return config def createrenamekeysconfig renamekeys stem fmt off renamekeys appendbackbone patchembed proj weight model pixellevelmodule encoder model embeddings patchembeddings projection weight renamekeys appendbackbone patchembed proj bias model pixellevelmodule encoder model embeddings patchembeddings projection bias renamekeys appendbackbone patchembed norm weight model pixellevelmodule encoder model embeddings norm weight renamekeys appendbackbone patchembed norm bias model pixellevelmodule encoder model embeddings norm bias stages for i in rangelenconfig backboneconfig depths for j in rangeconfig backboneconfig depthsi renamekeys appendfbackbone layers i blocks j norm1 weight fmodel pixellevelmodule encoder model encoder layers i blocks j layernormbefore weight renamekeys appendfbackbone layers i blocks j norm1 bias fmodel pixellevelmodule encoder model encoder layers i blocks j layernormbefore bias renamekeys appendfbackbone layers i blocks j attn relativepositionbiastable fmodel pixellevelmodule encoder model encoder layers i blocks j attention self relativepositionbiastable renamekeys appendfbackbone layers i blocks j attn relativepositionindex fmodel pixellevelmodule encoder model encoder layers i blocks j attention self relativepositionindex renamekeys appendfbackbone layers i blocks j attn proj weight fmodel pixellevelmodule encoder model encoder layers i blocks j attention output dense weight renamekeys appendfbackbone layers i blocks j attn proj bias fmodel pixellevelmodule encoder model encoder layers i blocks j attention output dense bias renamekeys appendfbackbone layers i blocks j norm2 weight fmodel pixellevelmodule encoder model encoder layers i blocks j layernormafter weight renamekeys appendfbackbone layers i blocks j norm2 bias fmodel pixellevelmodule encoder model encoder layers i blocks j layernormafter bias renamekeys appendfbackbone layers i blocks j mlp fc1 weight fmodel pixellevelmodule encoder model encoder layers i blocks j intermediate dense weight renamekeys appendfbackbone layers i blocks j mlp fc1 bias fmodel pixellevelmodule encoder model encoder layers i blocks j intermediate dense bias renamekeys appendfbackbone layers i blocks j mlp fc2 weight fmodel pixellevelmodule encoder model encoder layers i blocks j output dense weight renamekeys appendfbackbone layers i blocks j mlp fc2 bias fmodel pixellevelmodule encoder model encoder layers i blocks j output dense bias if i 3 renamekeys appendfbackbone layers i downsample reduction weight fmodel pixellevelmodule encoder model encoder layers i downsample reduction weight renamekeys appendfbackbone layers i downsample norm weight fmodel pixellevelmodule encoder model encoder layers i downsample norm weight renamekeys appendfbackbone layers i downsample norm bias fmodel pixellevelmodule encoder model encoder layers i downsample norm bias renamekeys appendfbackbone normi weight fmodel pixellevelmodule encoder hiddenstatesnorms i weight renamekeys appendfbackbone normi bias fmodel pixellevelmodule encoder hiddenstatesnorms i bias fpn renamekeys appendsemseghead layer4 weight model pixellevelmodule decoder fpn stem 0 weight renamekeys appendsemseghead layer4 norm weight model pixellevelmodule decoder fpn stem 1 weight renamekeys appendsemseghead layer4 norm bias model pixellevelmodule decoder fpn stem 1 bias for sourceindex targetindex in ziprange3 0 1 range0 3 renamekeys appendfsemseghead adaptersourceindex weight fmodel pixellevelmodule decoder fpn layers targetindex proj 0 weight renamekeys appendfsemseghead adaptersourceindex norm weight fmodel pixellevelmodule decoder fpn layers targetindex proj 1 weight renamekeys appendfsemseghead adaptersourceindex norm bias fmodel pixellevelmodule decoder fpn layers targetindex proj 1 bias renamekeys appendfsemseghead layersourceindex weight fmodel pixellevelmodule decoder fpn layers targetindex block 0 weight renamekeys appendfsemseghead layersourceindex norm weight fmodel pixellevelmodule decoder fpn layers targetindex block 1 weight renamekeys appendfsemseghead layersourceindex norm bias fmodel pixellevelmodule decoder fpn layers targetindex block 1 bias renamekeys appendsemseghead maskfeatures weight model pixellevelmodule decoder maskprojection weight renamekeys appendsemseghead maskfeatures bias model pixellevelmodule decoder maskprojection bias transformer decoder for idx in rangeconfig decoderconfig decoderlayers selfattention out projection renamekeys appendfsemseghead predictor transformer decoder layers idx selfattn outproj weight fmodel transformermodule decoder layers idx selfattn outproj weight renamekeys appendfsemseghead predictor transformer decoder layers idx selfattn outproj bias fmodel transformermodule decoder layers idx selfattn outproj bias crossattention out projection renamekeys appendfsemseghead predictor transformer decoder layers idx multiheadattn outproj weight fmodel transformermodule decoder layers idx encoderattn outproj weight renamekeys appendfsemseghead predictor transformer decoder layers idx multiheadattn outproj bias fmodel transformermodule decoder layers idx encoderattn outproj bias mlp 1 renamekeys appendfsemseghead predictor transformer decoder layers idx linear1 weight fmodel transformermodule decoder layers idx fc1 weight renamekeys appendfsemseghead predictor transformer decoder layers idx linear1 bias fmodel transformermodule decoder layers idx fc1 bias mlp 2 renamekeys appendfsemseghead predictor transformer decoder layers idx linear2 weight fmodel transformermodule decoder layers idx fc2 weight renamekeys appendfsemseghead predictor transformer decoder layers idx linear2 bias fmodel transformermodule decoder layers idx fc2 bias layernorm 1 selfattention layernorm renamekeys appendfsemseghead predictor transformer decoder layers idx norm1 weight fmodel transformermodule decoder layers idx selfattnlayernorm weight renamekeys appendfsemseghead predictor transformer decoder layers idx norm1 bias fmodel transformermodule decoder layers idx selfattnlayernorm bias layernorm 2 crossattention layernorm renamekeys appendfsemseghead predictor transformer decoder layers idx norm2 weight fmodel transformermodule decoder layers idx encoderattnlayernorm weight renamekeys appendfsemseghead predictor transformer decoder layers idx norm2 bias fmodel transformermodule decoder layers idx encoderattnlayernorm bias layernorm 3 final layernorm renamekeys appendfsemseghead predictor transformer decoder layers idx norm3 weight fmodel transformermodule decoder layers idx finallayernorm weight renamekeys appendfsemseghead predictor transformer decoder layers idx norm3 bias fmodel transformermodule decoder layers idx finallayernorm bias renamekeys appendsemseghead predictor transformer decoder norm weight model transformermodule decoder layernorm weight renamekeys appendsemseghead predictor transformer decoder norm bias model transformermodule decoder layernorm bias heads on top renamekeys appendsemseghead predictor queryembed weight model transformermodule queriesembedder weight renamekeys appendsemseghead predictor inputproj weight model transformermodule inputprojection weight renamekeys appendsemseghead predictor inputproj bias model transformermodule inputprojection bias renamekeys appendsemseghead predictor classembed weight classpredictor weight renamekeys appendsemseghead predictor classembed bias classpredictor bias for i in range3 renamekeys appendfsemseghead predictor maskembed layers i weight fmaskembedder i 0 weight renamekeys appendfsemseghead predictor maskembed layers i bias fmaskembedder i 0 bias fmt on return renamekeys def renamekeydct old new val dct popold dctnew val we split up the matrix of each encoder layer into queries keys and values def readinswinqkvstatedict backboneconfig numfeatures intbackboneconfig embeddim 2i for i in rangelenbackboneconfig depths for i in rangelenbackboneconfig depths dim numfeaturesi for j in rangebackboneconfig depthsi fmt off read in weights bias of input projection layer in original implementation this is a single matrix bias inprojweight statedict popfbackbone layers i blocks j attn qkv weight inprojbias statedict popfbackbone layers i blocks j attn qkv bias next add query keys and values in that order to the state dict statedictfmodel pixellevelmodule encoder model encoder layers i blocks j attention self query weight inprojweight dim statedictfmodel pixellevelmodule encoder model encoder layers i blocks j attention self query bias inprojbias dim statedictfmodel pixellevelmodule encoder model encoder layers i blocks j attention self key weight inprojweight dim dim 2 statedictfmodel pixellevelmodule encoder model encoder layers i blocks j attention self key bias inprojbias dim dim 2 statedictfmodel pixellevelmodule encoder model encoder layers i blocks j attention self value weight inprojweight dim statedictfmodel pixellevelmodule encoder model encoder layers i blocks j attention self value bias inprojbiasdim fmt on we split up the matrix of each encoder layer into queries keys and values def readindecoderqkvstatedict config fmt off hiddensize config decoderconfig hiddensize for idx in rangeconfig decoderconfig decoderlayers read in weights bias of selfattention input projection layer in the original implementation this is a single matrix bias inprojweight statedict popfsemseghead predictor transformer decoder layers idx selfattn inprojweight inprojbias statedict popfsemseghead predictor transformer decoder layers idx selfattn inprojbias next add query keys and values in that order to the state dict statedictfmodel transformermodule decoder layers idx selfattn qproj weight inprojweight hiddensize statedictfmodel transformermodule decoder layers idx selfattn qproj bias inprojbias config hiddensize statedictfmodel transformermodule decoder layers idx selfattn kproj weight inprojweighthiddensize hiddensize 2 statedictfmodel transformermodule decoder layers idx selfattn kproj bias inprojbiashiddensize hiddensize 2 statedictfmodel transformermodule decoder layers idx selfattn vproj weight inprojweighthiddensize statedictfmodel transformermodule decoder layers idx selfattn vproj bias inprojbiashiddensize read in weights bias of crossattention input projection layer in the original implementation this is a single matrix bias inprojweight statedict popfsemseghead predictor transformer decoder layers idx multiheadattn inprojweight inprojbias statedict popfsemseghead predictor transformer decoder layers idx multiheadattn inprojbias next add query keys and values in that order to the state dict statedictfmodel transformermodule decoder layers idx encoderattn qproj weight inprojweight hiddensize statedictfmodel transformermodule decoder layers idx encoderattn qproj bias inprojbias config hiddensize statedictfmodel transformermodule decoder layers idx encoderattn kproj weight inprojweighthiddensize hiddensize 2 statedictfmodel transformermodule decoder layers idx encoderattn kproj bias inprojbiashiddensize hiddensize 2 statedictfmodel transformermodule decoder layers idx encoderattn vproj weight inprojweighthiddensize statedictfmodel transformermodule decoder layers idx encoderattn vproj bias inprojbiashiddensize fmt on we will verify our results on an image of cute cats def prepareimg torch tensor url http images cocodataset orgval2017000000039769 jpg im image openrequests geturl streamtrue raw return im torch nograd def convertmaskformercheckpoint modelname str checkpointpath str pytorchdumpfolderpath str pushtohub bool false config getmaskformerconfigmodelname load original statedict with opencheckpointpath rb as f data pickle loadf statedict datamodel for name param in statedict items printname param shape rename keys renamekeys createrenamekeysconfig for src dest in renamekeys renamekeystatedict src dest readinswinqkvstatedict config backboneconfig readindecoderqkvstatedict config update to torch tensors for key value in statedict items statedictkey torch fromnumpyvalue load model model maskformerforinstancesegmentationconfig model eval for name param in model namedparameters printname param shape missingkeys unexpectedkeys model loadstatedictstatedict strictfalse assert missingkeys model pixellevelmodule encoder model layernorm weight model pixellevelmodule encoder model layernorm bias assert lenunexpectedkeys 0 funexpected keys unexpectedkeys verify results image prepareimg if vistas in modelname ignoreindex 65 elif cityscapes in modelname ignoreindex 65535 else ignoreindex 255 reducelabels true if ade in modelname else false imageprocessor maskformerimageprocessorignoreindexignoreindex reducelabelsreducelabels inputs imageprocessorimage returntensorspt outputs modelinputs printlogits outputs classquerieslogits0 3 3 if modelname maskformerswintinyade expectedlogits torch tensor 3 6353 4 4770 2 6065 0 5081 4 2394 3 5343 2 1909 5 0353 1 9323 assert torch allcloseoutputs classquerieslogits0 3 3 expectedlogits atol1e4 printlooks ok if pytorchdumpfolderpath is not none printfsaving model and image processor to pytorchdumpfolderpath pathpytorchdumpfolderpath mkdirexistoktrue model savepretrainedpytorchdumpfolderpath imageprocessor savepretrainedpytorchdumpfolderpath if pushtohub printpushing model and image processor to the hub model pushtohubfnielsrmodelname imageprocessor pushtohubfnielsrmodelname if name main parser argparse argumentparser required parameters parser addargument modelname defaultmaskformerswintinyade typestr helpname of the maskformer model you d like to convert parser addargument checkpointpath defaultusersnielsroggedocumentsmaskformercheckpointsmaskformerswintinyade20kmodel pkl typestr helppath to the original state dict pth file parser addargument pytorchdumpfolderpath defaultnone typestr helppath to the output pytorch model directory parser addargument pushtohub actionstoretrue helpwhether or not to push the converted model to the hub args parser parseargs convertmaskformercheckpoint args modelname args checkpointpath args pytorchdumpfolderpath args pushtohub coding utf 8 2022 the huggingface inc team licensed under the apache license version 2 0 the license you may not use this file except in compliance with the license you may obtain a copy of the license at http www apache org licenses license 2 0 unless required by applicable law or agreed to in writing software distributed under the license is distributed on an as is basis without warranties or conditions of any kind either express or implied see the license for the specific language governing permissions and limitations under the license convert maskformer checkpoints with swin backbone from the original repository url https github com facebookresearch maskformer this should be ok this should be ok this should be ok todo this should be ok this should be ok stem fmt off stages fpn transformer decoder self attention out projection cross attention out projection mlp 1 mlp 2 layernorm 1 self attention layernorm layernorm 2 cross attention layernorm layernorm 3 final layernorm heads on top fmt on we split up the matrix of each encoder layer into queries keys and values fmt off read in weights bias of input projection layer in original implementation this is a single matrix bias next add query keys and values in that order to the state dict fmt on we split up the matrix of each encoder layer into queries keys and values fmt off read in weights bias of self attention input projection layer in the original implementation this is a single matrix bias next add query keys and values in that order to the state dict read in weights bias of cross attention input projection layer in the original implementation this is a single matrix bias next add query keys and values in that order to the state dict fmt on we will verify our results on an image of cute cats copy paste tweak model s weights to our maskformer structure load original state_dict for name param in state_dict items print name param shape rename keys update to torch tensors load model verify results required parameters
import argparse
import json
import pickle
from pathlib import Path

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

from transformers import MaskFormerConfig, MaskFormerForInstanceSegmentation, MaskFormerImageProcessor, SwinConfig
from transformers.utils import logging


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


def get_maskformer_config(model_name: str):
    backbone_config = SwinConfig.from_pretrained(
        "microsoft/swin-tiny-patch4-window7-224", out_features=["stage1", "stage2", "stage3", "stage4"]
    )
    config = MaskFormerConfig(backbone_config=backbone_config)

    repo_id = "huggingface/label-files"
    if "ade20k-full" in model_name:
        
        config.num_labels = 847
        filename = "maskformer-ade20k-full-id2label.json"
    elif "ade" in model_name:
        
        config.num_labels = 150
        filename = "ade20k-id2label.json"
    elif "coco-stuff" in model_name:
        
        config.num_labels = 171
        filename = "maskformer-coco-stuff-id2label.json"
    elif "coco" in model_name:
        
        config.num_labels = 133
        filename = "coco-panoptic-id2label.json"
    elif "cityscapes" in model_name:
        
        config.num_labels = 19
        filename = "cityscapes-id2label.json"
    elif "vistas" in model_name:
        
        config.num_labels = 65
        filename = "mapillary-vistas-id2label.json"

    id2label = json.load(open(hf_hub_download(repo_id, filename, repo_type="dataset"), "r"))
    id2label = {int(k): v for k, v in id2label.items()}

    return config


def create_rename_keys(config):
    rename_keys = []
    
    
    rename_keys.append(("backbone.patch_embed.proj.weight", "model.pixel_level_module.encoder.model.embeddings.patch_embeddings.projection.weight"))
    rename_keys.append(("backbone.patch_embed.proj.bias", "model.pixel_level_module.encoder.model.embeddings.patch_embeddings.projection.bias"))
    rename_keys.append(("backbone.patch_embed.norm.weight", "model.pixel_level_module.encoder.model.embeddings.norm.weight"))
    rename_keys.append(("backbone.patch_embed.norm.bias", "model.pixel_level_module.encoder.model.embeddings.norm.bias"))
    
    for i in range(len(config.backbone_config.depths)):
        for j in range(config.backbone_config.depths[i]):
            rename_keys.append((f"backbone.layers.{i}.blocks.{j}.norm1.weight", f"model.pixel_level_module.encoder.model.encoder.layers.{i}.blocks.{j}.layernorm_before.weight"))
            rename_keys.append((f"backbone.layers.{i}.blocks.{j}.norm1.bias", f"model.pixel_level_module.encoder.model.encoder.layers.{i}.blocks.{j}.layernorm_before.bias"))
            rename_keys.append((f"backbone.layers.{i}.blocks.{j}.attn.relative_position_bias_table", f"model.pixel_level_module.encoder.model.encoder.layers.{i}.blocks.{j}.attention.self.relative_position_bias_table"))
            rename_keys.append((f"backbone.layers.{i}.blocks.{j}.attn.relative_position_index", f"model.pixel_level_module.encoder.model.encoder.layers.{i}.blocks.{j}.attention.self.relative_position_index"))
            rename_keys.append((f"backbone.layers.{i}.blocks.{j}.attn.proj.weight", f"model.pixel_level_module.encoder.model.encoder.layers.{i}.blocks.{j}.attention.output.dense.weight"))
            rename_keys.append((f"backbone.layers.{i}.blocks.{j}.attn.proj.bias", f"model.pixel_level_module.encoder.model.encoder.layers.{i}.blocks.{j}.attention.output.dense.bias"))
            rename_keys.append((f"backbone.layers.{i}.blocks.{j}.norm2.weight", f"model.pixel_level_module.encoder.model.encoder.layers.{i}.blocks.{j}.layernorm_after.weight"))
            rename_keys.append((f"backbone.layers.{i}.blocks.{j}.norm2.bias", f"model.pixel_level_module.encoder.model.encoder.layers.{i}.blocks.{j}.layernorm_after.bias"))
            rename_keys.append((f"backbone.layers.{i}.blocks.{j}.mlp.fc1.weight", f"model.pixel_level_module.encoder.model.encoder.layers.{i}.blocks.{j}.intermediate.dense.weight"))
            rename_keys.append((f"backbone.layers.{i}.blocks.{j}.mlp.fc1.bias", f"model.pixel_level_module.encoder.model.encoder.layers.{i}.blocks.{j}.intermediate.dense.bias"))
            rename_keys.append((f"backbone.layers.{i}.blocks.{j}.mlp.fc2.weight", f"model.pixel_level_module.encoder.model.encoder.layers.{i}.blocks.{j}.output.dense.weight"))
            rename_keys.append((f"backbone.layers.{i}.blocks.{j}.mlp.fc2.bias", f"model.pixel_level_module.encoder.model.encoder.layers.{i}.blocks.{j}.output.dense.bias"))

        if i < 3:
            rename_keys.append((f"backbone.layers.{i}.downsample.reduction.weight", f"model.pixel_level_module.encoder.model.encoder.layers.{i}.downsample.reduction.weight"))
            rename_keys.append((f"backbone.layers.{i}.downsample.norm.weight", f"model.pixel_level_module.encoder.model.encoder.layers.{i}.downsample.norm.weight"))
            rename_keys.append((f"backbone.layers.{i}.downsample.norm.bias", f"model.pixel_level_module.encoder.model.encoder.layers.{i}.downsample.norm.bias"))
        rename_keys.append((f"backbone.norm{i}.weight", f"model.pixel_level_module.encoder.hidden_states_norms.{i}.weight"))
        rename_keys.append((f"backbone.norm{i}.bias", f"model.pixel_level_module.encoder.hidden_states_norms.{i}.bias"))

    
    rename_keys.append(("sem_seg_head.layer_4.weight", "model.pixel_level_module.decoder.fpn.stem.0.weight"))
    rename_keys.append(("sem_seg_head.layer_4.norm.weight", "model.pixel_level_module.decoder.fpn.stem.1.weight"))
    rename_keys.append(("sem_seg_head.layer_4.norm.bias", "model.pixel_level_module.decoder.fpn.stem.1.bias"))
    for source_index, target_index in zip(range(3, 0, -1), range(0, 3)):
        rename_keys.append((f"sem_seg_head.adapter_{source_index}.weight", f"model.pixel_level_module.decoder.fpn.layers.{target_index}.proj.0.weight"))
        rename_keys.append((f"sem_seg_head.adapter_{source_index}.norm.weight", f"model.pixel_level_module.decoder.fpn.layers.{target_index}.proj.1.weight"))
        rename_keys.append((f"sem_seg_head.adapter_{source_index}.norm.bias", f"model.pixel_level_module.decoder.fpn.layers.{target_index}.proj.1.bias"))
        rename_keys.append((f"sem_seg_head.layer_{source_index}.weight", f"model.pixel_level_module.decoder.fpn.layers.{target_index}.block.0.weight"))
        rename_keys.append((f"sem_seg_head.layer_{source_index}.norm.weight", f"model.pixel_level_module.decoder.fpn.layers.{target_index}.block.1.weight"))
        rename_keys.append((f"sem_seg_head.layer_{source_index}.norm.bias", f"model.pixel_level_module.decoder.fpn.layers.{target_index}.block.1.bias"))
    rename_keys.append(("sem_seg_head.mask_features.weight", "model.pixel_level_module.decoder.mask_projection.weight"))
    rename_keys.append(("sem_seg_head.mask_features.bias", "model.pixel_level_module.decoder.mask_projection.bias"))

    
    for idx in range(config.decoder_config.decoder_layers):
        
        rename_keys.append((f"sem_seg_head.predictor.transformer.decoder.layers.{idx}.self_attn.out_proj.weight", f"model.transformer_module.decoder.layers.{idx}.self_attn.out_proj.weight"))
        rename_keys.append((f"sem_seg_head.predictor.transformer.decoder.layers.{idx}.self_attn.out_proj.bias", f"model.transformer_module.decoder.layers.{idx}.self_attn.out_proj.bias"))
        
        rename_keys.append((f"sem_seg_head.predictor.transformer.decoder.layers.{idx}.multihead_attn.out_proj.weight", f"model.transformer_module.decoder.layers.{idx}.encoder_attn.out_proj.weight"))
        rename_keys.append((f"sem_seg_head.predictor.transformer.decoder.layers.{idx}.multihead_attn.out_proj.bias", f"model.transformer_module.decoder.layers.{idx}.encoder_attn.out_proj.bias"))
        
        rename_keys.append((f"sem_seg_head.predictor.transformer.decoder.layers.{idx}.linear1.weight", f"model.transformer_module.decoder.layers.{idx}.fc1.weight"))
        rename_keys.append((f"sem_seg_head.predictor.transformer.decoder.layers.{idx}.linear1.bias", f"model.transformer_module.decoder.layers.{idx}.fc1.bias"))
        
        rename_keys.append((f"sem_seg_head.predictor.transformer.decoder.layers.{idx}.linear2.weight", f"model.transformer_module.decoder.layers.{idx}.fc2.weight"))
        rename_keys.append((f"sem_seg_head.predictor.transformer.decoder.layers.{idx}.linear2.bias", f"model.transformer_module.decoder.layers.{idx}.fc2.bias"))
        
        rename_keys.append((f"sem_seg_head.predictor.transformer.decoder.layers.{idx}.norm1.weight", f"model.transformer_module.decoder.layers.{idx}.self_attn_layer_norm.weight"))
        rename_keys.append((f"sem_seg_head.predictor.transformer.decoder.layers.{idx}.norm1.bias", f"model.transformer_module.decoder.layers.{idx}.self_attn_layer_norm.bias"))
        
        rename_keys.append((f"sem_seg_head.predictor.transformer.decoder.layers.{idx}.norm2.weight", f"model.transformer_module.decoder.layers.{idx}.encoder_attn_layer_norm.weight"))
        rename_keys.append((f"sem_seg_head.predictor.transformer.decoder.layers.{idx}.norm2.bias", f"model.transformer_module.decoder.layers.{idx}.encoder_attn_layer_norm.bias"))
        
        rename_keys.append((f"sem_seg_head.predictor.transformer.decoder.layers.{idx}.norm3.weight", f"model.transformer_module.decoder.layers.{idx}.final_layer_norm.weight"))
        rename_keys.append((f"sem_seg_head.predictor.transformer.decoder.layers.{idx}.norm3.bias", f"model.transformer_module.decoder.layers.{idx}.final_layer_norm.bias"))

    rename_keys.append(("sem_seg_head.predictor.transformer.decoder.norm.weight", "model.transformer_module.decoder.layernorm.weight"))
    rename_keys.append(("sem_seg_head.predictor.transformer.decoder.norm.bias", "model.transformer_module.decoder.layernorm.bias"))

    
    rename_keys.append(("sem_seg_head.predictor.query_embed.weight", "model.transformer_module.queries_embedder.weight"))

    rename_keys.append(("sem_seg_head.predictor.input_proj.weight", "model.transformer_module.input_projection.weight"))
    rename_keys.append(("sem_seg_head.predictor.input_proj.bias", "model.transformer_module.input_projection.bias"))

    rename_keys.append(("sem_seg_head.predictor.class_embed.weight", "class_predictor.weight"))
    rename_keys.append(("sem_seg_head.predictor.class_embed.bias", "class_predictor.bias"))

    for i in range(3):
        rename_keys.append((f"sem_seg_head.predictor.mask_embed.layers.{i}.weight", f"mask_embedder.{i}.0.weight"))
        rename_keys.append((f"sem_seg_head.predictor.mask_embed.layers.{i}.bias", f"mask_embedder.{i}.0.bias"))
    

    return rename_keys


def rename_key(dct, old, new):
    val = dct.pop(old)
    dct[new] = val



def read_in_swin_q_k_v(state_dict, backbone_config):
    num_features = [int(backbone_config.embed_dim * 2**i) for i in range(len(backbone_config.depths))]
    for i in range(len(backbone_config.depths)):
        dim = num_features[i]
        for j in range(backbone_config.depths[i]):
            
            
            in_proj_weight = state_dict.pop(f"backbone.layers.{i}.blocks.{j}.attn.qkv.weight")
            in_proj_bias = state_dict.pop(f"backbone.layers.{i}.blocks.{j}.attn.qkv.bias")
            
            state_dict[f"model.pixel_level_module.encoder.model.encoder.layers.{i}.blocks.{j}.attention.self.query.weight"] = in_proj_weight[:dim, :]
            state_dict[f"model.pixel_level_module.encoder.model.encoder.layers.{i}.blocks.{j}.attention.self.query.bias"] = in_proj_bias[: dim]
            state_dict[f"model.pixel_level_module.encoder.model.encoder.layers.{i}.blocks.{j}.attention.self.key.weight"] = in_proj_weight[
                dim : dim * 2, :
            ]
            state_dict[f"model.pixel_level_module.encoder.model.encoder.layers.{i}.blocks.{j}.attention.self.key.bias"] = in_proj_bias[
                dim : dim * 2
            ]
            state_dict[f"model.pixel_level_module.encoder.model.encoder.layers.{i}.blocks.{j}.attention.self.value.weight"] = in_proj_weight[
                -dim :, :
            ]
            state_dict[f"model.pixel_level_module.encoder.model.encoder.layers.{i}.blocks.{j}.attention.self.value.bias"] = in_proj_bias[-dim :]
            



def read_in_decoder_q_k_v(state_dict, config):
    
    hidden_size = config.decoder_config.hidden_size
    for idx in range(config.decoder_config.decoder_layers):
        
        in_proj_weight = state_dict.pop(f"sem_seg_head.predictor.transformer.decoder.layers.{idx}.self_attn.in_proj_weight")
        in_proj_bias = state_dict.pop(f"sem_seg_head.predictor.transformer.decoder.layers.{idx}.self_attn.in_proj_bias")
        
        state_dict[f"model.transformer_module.decoder.layers.{idx}.self_attn.q_proj.weight"] = in_proj_weight[: hidden_size, :]
        state_dict[f"model.transformer_module.decoder.layers.{idx}.self_attn.q_proj.bias"] = in_proj_bias[:config.hidden_size]
        state_dict[f"model.transformer_module.decoder.layers.{idx}.self_attn.k_proj.weight"] = in_proj_weight[hidden_size : hidden_size * 2, :]
        state_dict[f"model.transformer_module.decoder.layers.{idx}.self_attn.k_proj.bias"] = in_proj_bias[hidden_size : hidden_size * 2]
        state_dict[f"model.transformer_module.decoder.layers.{idx}.self_attn.v_proj.weight"] = in_proj_weight[-hidden_size :, :]
        state_dict[f"model.transformer_module.decoder.layers.{idx}.self_attn.v_proj.bias"] = in_proj_bias[-hidden_size :]
        
        in_proj_weight = state_dict.pop(f"sem_seg_head.predictor.transformer.decoder.layers.{idx}.multihead_attn.in_proj_weight")
        in_proj_bias = state_dict.pop(f"sem_seg_head.predictor.transformer.decoder.layers.{idx}.multihead_attn.in_proj_bias")
        
        state_dict[f"model.transformer_module.decoder.layers.{idx}.encoder_attn.q_proj.weight"] = in_proj_weight[: hidden_size, :]
        state_dict[f"model.transformer_module.decoder.layers.{idx}.encoder_attn.q_proj.bias"] = in_proj_bias[:config.hidden_size]
        state_dict[f"model.transformer_module.decoder.layers.{idx}.encoder_attn.k_proj.weight"] = in_proj_weight[hidden_size : hidden_size * 2, :]
        state_dict[f"model.transformer_module.decoder.layers.{idx}.encoder_attn.k_proj.bias"] = in_proj_bias[hidden_size : hidden_size * 2]
        state_dict[f"model.transformer_module.decoder.layers.{idx}.encoder_attn.v_proj.weight"] = in_proj_weight[-hidden_size :, :]
        state_dict[f"model.transformer_module.decoder.layers.{idx}.encoder_attn.v_proj.bias"] = in_proj_bias[-hidden_size :]
    



def prepare_img() -> torch.Tensor:
    url = "http://images.cocodataset.org/val2017/000000039769.jpg"
    im = Image.open(requests.get(url, stream=True).raw)
    return im


@torch.no_grad()
def convert_maskformer_checkpoint(
    model_name: str, checkpoint_path: str, pytorch_dump_folder_path: str, push_to_hub: bool = False
):
    
    config = get_maskformer_config(model_name)

    
    with open(checkpoint_path, "rb") as f:
        data = pickle.load(f)
    state_dict = data["model"]

    
    

    
    rename_keys = create_rename_keys(config)
    for src, dest in rename_keys:
        rename_key(state_dict, src, dest)
    read_in_swin_q_k_v(state_dict, config.backbone_config)
    read_in_decoder_q_k_v(state_dict, config)

    
    for key, value in state_dict.items():
        state_dict[key] = torch.from_numpy(value)

    
    model = MaskFormerForInstanceSegmentation(config)
    model.eval()

    for name, param in model.named_parameters():
        print(name, param.shape)

    missing_keys, unexpected_keys = model.load_state_dict(state_dict, strict=False)
    assert missing_keys == [
        "model.pixel_level_module.encoder.model.layernorm.weight",
        "model.pixel_level_module.encoder.model.layernorm.bias",
    ]
    assert len(unexpected_keys) == 0, f"Unexpected keys: {unexpected_keys}"

    
    image = prepare_img()
    if "vistas" in model_name:
        ignore_index = 65
    elif "cityscapes" in model_name:
        ignore_index = 65535
    else:
        ignore_index = 255
    reduce_labels = True if "ade" in model_name else False
    image_processor = MaskFormerImageProcessor(ignore_index=ignore_index, reduce_labels=reduce_labels)

    inputs = image_processor(image, return_tensors="pt")

    outputs = model(**inputs)

    print("Logits:", outputs.class_queries_logits[0, :3, :3])

    if model_name == "maskformer-swin-tiny-ade":
        expected_logits = torch.tensor(
            [[3.6353, -4.4770, -2.6065], [0.5081, -4.2394, -3.5343], [2.1909, -5.0353, -1.9323]]
        )
    assert torch.allclose(outputs.class_queries_logits[0, :3, :3], expected_logits, atol=1e-4)
    print("Looks ok!")

    if pytorch_dump_folder_path is not None:
        print(f"Saving model and image processor to {pytorch_dump_folder_path}")
        Path(pytorch_dump_folder_path).mkdir(exist_ok=True)
        model.save_pretrained(pytorch_dump_folder_path)
        image_processor.save_pretrained(pytorch_dump_folder_path)

    if push_to_hub:
        print("Pushing model and image processor to the hub...")
        model.push_to_hub(f"nielsr/{model_name}")
        image_processor.push_to_hub(f"nielsr/{model_name}")


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    
    parser.add_argument(
        "--model_name",
        default="maskformer-swin-tiny-ade",
        type=str,
        help=("Name of the MaskFormer model you'd like to convert",),
    )
    parser.add_argument(
        "--checkpoint_path",
        default="/Users/nielsrogge/Documents/MaskFormer_checkpoints/MaskFormer-Swin-tiny-ADE20k/model.pkl",
        type=str,
        help="Path to the original state dict (.pth file).",
    )
    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."
    )

    args = parser.parse_args()
    convert_maskformer_checkpoint(
        args.model_name, args.checkpoint_path, args.pytorch_dump_folder_path, args.push_to_hub
    )