codingutf8 2022 meta platforms inc and 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 this class tracks a python dictionary by keeping track of which item is accessed args totrack dict the dictionary we wish to track this method returns a set difference between the keys in the tracked state dict and the one we have access so far this is an effective method to check if we have update all the keys returns liststr list of keys not yet updated proxy the call to the internal dictionary we will verify our results on an image of cute cats fake command line arguments needed by mask2formerdetectron implementation configfile str def setupcfgargs args load config from file and commandline arguments cfg getcfg adddeeplabconfigcfg addmaskformer2configcfg cfg mergefromfileargs configfile cfg freeze return cfg class originalmask2formerconfigtooursconverter def callself originalconfig object mask2formerconfig model originalconfig model repoid huggingfacelabelfiles if model semseghead numclasses 847 filename mask2formerade20kfullid2label json elif model semseghead numclasses 150 filename ade20kid2label json elif model semseghead numclasses 80 filename cocodetectionmmdetid2label json elif model semseghead numclasses 171 filename mask2formercocostuffid2label json elif model semseghead numclasses 133 filename cocopanopticid2label json elif model semseghead numclasses 19 filename cityscapesid2label json elif model semseghead numclasses 8 filename cityscapesinstanceid2label json elif model semseghead numclasses 65 filename mapillaryvistasid2label json id2label json loadopenhfhubdownloadrepoid filename repotypedataset r id2label intk v for k v in id2label items label2id label idx for idx label in id2label items if model swin embeddim 96 backboneconfig swinconfig frompretrained microsoftswintinypatch4window7224 outfeaturesstage1 stage2 stage3 stage4 elif model swin embeddim 128 backboneconfig swinconfig embeddim128 windowsize12 depths2 2 18 2 numheads4 8 16 32 outfeaturesstage1 stage2 stage3 stage4 elif model swin embeddim 192 backboneconfig swinconfig frompretrained microsoftswinlargepatch4window12384 outfeaturesstage1 stage2 stage3 stage4 else raise valueerrorfembed dim model swin embeddim not supported for swin backboneconfig droppathrate model swin droppathrate backboneconfig attentionprobsdropoutprob model swin attndroprate backboneconfig depths model swin depths config mask2formerconfig mask2formerconfig ignorevaluemodel semseghead ignorevalue numlabelsmodel semseghead numclasses numqueriesmodel maskformer numobjectqueries noobjectweightmodel maskformer noobjectweight classweightmodel maskformer classweight maskweightmodel maskformer maskweight diceweightmodel maskformer diceweight trainnumpointsmodel maskformer trainnumpoints oversampleratiomodel maskformer oversampleratio importancesampleratiomodel maskformer importancesampleratio initstd0 02 initxavierstd1 0 useauxiliarylossmodel maskformer deepsupervision featurestrides4 8 16 32 backboneconfigbackboneconfig id2labelid2label label2idlabel2id featuresizemodel semseghead convsdim maskfeaturesizemodel semseghead maskdim hiddendimmodel maskformer hiddendim encoderlayersmodel semseghead transformerenclayers encoderfeedforwarddim1024 decoderlayersmodel maskformer declayers numattentionheadsmodel maskformer nheads dropoutmodel maskformer dropout dimfeedforwardmodel maskformer dimfeedforward prenormmodel maskformer prenorm enforceinputprojmodel maskformer enforceinputproj commonstridemodel semseghead commonstride return config class originalmask2formerconfigtoimageprocessorconverter def callself originalconfig object mask2formerimageprocessor model originalconfig model modelinput originalconfig input return mask2formerimageprocessor imagemeantorch tensormodel pixelmean 255 tolist imagestdtorch tensormodel pixelstd 255 tolist sizemodelinput minsizetest maxsizemodelinput maxsizetest numlabelsmodel semseghead numclasses ignoreindexmodel semseghead ignorevalue sizedivisibility32 class originalmask2formercheckpointtooursconverter def initself originalmodel nn module config mask2formerconfig self originalmodel originalmodel self config config def popallself renamedkeys listtuplestr str dststatedict statedict srcstatedict statedict for srckey dstkey in renamedkeys dststatedictdstkey srcstatedict popsrckey def replacemaskformerswinbackbone self dststatedict statedict srcstatedict statedict config mask2formerconfig dstprefix str pixellevelmodule encoder srcprefix str backbone renamedkeys fsrcprefix patchembed proj weight fdstprefix model embeddings patchembeddings projection weight fsrcprefix patchembed proj bias fdstprefix model embeddings patchembeddings projection bias fsrcprefix patchembed norm weight fdstprefix model embeddings norm weight fsrcprefix patchembed norm bias fdstprefix model embeddings norm bias numlayers lenconfig backboneconfig depths for layeridx in rangenumlayers for blockidx in rangeconfig backboneconfig depthslayeridx renamedkeys extend src dst fsrcprefix layers layeridx blocks blockidx norm1 weight fdstprefix model encoder layers layeridx blocks blockidx layernormbefore weight fsrcprefix layers layeridx blocks blockidx norm1 bias fdstprefix model encoder layers layeridx blocks blockidx layernormbefore bias fsrcprefix layers layeridx blocks blockidx attn relativepositionbiastable fdstprefix model encoder layers layeridx blocks blockidx attention self relativepositionbiastable now we need to handle the attentions read in weights bias of input projection layer of crossattention srcattweight srcstatedictfsrcprefix layers layeridx blocks blockidx attn qkv weight srcattbias srcstatedictfsrcprefix layers layeridx blocks blockidx attn qkv bias size srcattweight shape0 offset size 3 dststatedict fdstprefix model encoder layers layeridx blocks blockidx attention self query weight srcattweight offset dststatedict fdstprefix model encoder layers layeridx blocks blockidx attention self query bias srcattbias offset dststatedict fdstprefix model encoder layers layeridx blocks blockidx attention self key weight srcattweightoffset offset 2 dststatedict fdstprefix model encoder layers layeridx blocks blockidx attention self key bias srcattbiasoffset offset 2 dststatedict fdstprefix model encoder layers layeridx blocks blockidx attention self value weight srcattweightoffset dststatedict fdstprefix model encoder layers layeridx blocks blockidx attention self value bias srcattbiasoffset let s pop them srcstatedict popfsrcprefix layers layeridx blocks blockidx attn qkv weight srcstatedict popfsrcprefix layers layeridx blocks blockidx attn qkv bias proj renamedkeys extend fsrcprefix layers layeridx blocks blockidx attn proj weight fdstprefix model encoder layers layeridx blocks blockidx attention output dense weight fsrcprefix layers layeridx blocks blockidx attn proj bias fdstprefix model encoder layers layeridx blocks blockidx attention output dense bias second norm renamedkeys extend fsrcprefix layers layeridx blocks blockidx norm2 weight fdstprefix model encoder layers layeridx blocks blockidx layernormafter weight fsrcprefix layers layeridx blocks blockidx norm2 bias fdstprefix model encoder layers layeridx blocks blockidx layernormafter bias mlp renamedkeys extend fsrcprefix layers layeridx blocks blockidx mlp fc1 weight fdstprefix model encoder layers layeridx blocks blockidx intermediate dense weight fsrcprefix layers layeridx blocks blockidx mlp fc1 bias fdstprefix model encoder layers layeridx blocks blockidx intermediate dense bias fsrcprefix layers layeridx blocks blockidx mlp fc2 weight fdstprefix model encoder layers layeridx blocks blockidx output dense weight fsrcprefix layers layeridx blocks blockidx mlp fc2 bias fdstprefix model encoder layers layeridx blocks blockidx output dense bias renamedkeys extend fsrcprefix layers layeridx blocks blockidx attn relativepositionindex fdstprefix model encoder layers layeridx blocks blockidx attention self relativepositionindex if layeridx numlayers 1 patch merging renamedkeys extend fsrcprefix layers layeridx downsample reduction weight fdstprefix model encoder layers layeridx downsample reduction weight fsrcprefix layers layeridx downsample norm weight fdstprefix model encoder layers layeridx downsample norm weight fsrcprefix layers layeridx downsample norm bias fdstprefix model encoder layers layeridx downsample norm bias hidden states norms renamedkeys extend fsrcprefix normlayeridx weight fdstprefix hiddenstatesnorms layeridx weight fsrcprefix normlayeridx bias fdstprefix hiddenstatesnorms layeridx bias self popallrenamedkeys dststatedict srcstatedict def replaceswinbackboneself dststatedict statedict srcstatedict statedict config mask2formerconfig dstprefix str pixellevelmodule encoder srcprefix str backbone renamedkeys fsrcprefix patchembed proj weight fdstprefix embeddings patchembeddings projection weight fsrcprefix patchembed proj bias fdstprefix embeddings patchembeddings projection bias fsrcprefix patchembed norm weight fdstprefix embeddings norm weight fsrcprefix patchembed norm bias fdstprefix embeddings norm bias for layeridx in rangelenconfig backboneconfig depths for blockidx in rangeconfig backboneconfig depthslayeridx renamedkeys extend src dst fsrcprefix layers layeridx blocks blockidx norm1 weight fdstprefix encoder layers layeridx blocks blockidx layernormbefore weight fsrcprefix layers layeridx blocks blockidx norm1 bias fdstprefix encoder layers layeridx blocks blockidx layernormbefore bias fsrcprefix layers layeridx blocks blockidx attn relativepositionbiastable fdstprefix encoder layers layeridx blocks blockidx attention self relativepositionbiastable now we need to handle the attentions read in weights bias of input projection layer of crossattention srcattweight srcstatedictfsrcprefix layers layeridx blocks blockidx attn qkv weight srcattbias srcstatedictfsrcprefix layers layeridx blocks blockidx attn qkv bias size srcattweight shape0 offset size 3 dststatedict fdstprefix encoder layers layeridx blocks blockidx attention self query weight srcattweight offset dststatedict fdstprefix encoder layers layeridx blocks blockidx attention self query bias srcattbias offset dststatedict fdstprefix encoder layers layeridx blocks blockidx attention self key weight srcattweightoffset offset 2 dststatedict fdstprefix encoder layers layeridx blocks blockidx attention self key bias srcattbiasoffset offset 2 dststatedict fdstprefix encoder layers layeridx blocks blockidx attention self value weight srcattweightoffset dststatedict fdstprefix encoder layers layeridx blocks blockidx attention self value bias srcattbiasoffset let s pop them srcstatedict popfsrcprefix layers layeridx blocks blockidx attn qkv weight srcstatedict popfsrcprefix layers layeridx blocks blockidx attn qkv bias proj renamedkeys extend fsrcprefix layers layeridx blocks blockidx attn proj weight fdstprefix encoder layers layeridx blocks blockidx attention output dense weight fsrcprefix layers layeridx blocks blockidx attn proj bias fdstprefix encoder layers layeridx blocks blockidx attention output dense bias second norm renamedkeys extend fsrcprefix layers layeridx blocks blockidx norm2 weight fdstprefix encoder layers layeridx blocks blockidx layernormafter weight fsrcprefix layers layeridx blocks blockidx norm2 bias fdstprefix encoder layers layeridx blocks blockidx layernormafter bias mlp renamedkeys extend fsrcprefix layers layeridx blocks blockidx mlp fc1 weight fdstprefix encoder layers layeridx blocks blockidx intermediate dense weight fsrcprefix layers layeridx blocks blockidx mlp fc1 bias fdstprefix encoder layers layeridx blocks blockidx intermediate dense bias fsrcprefix layers layeridx blocks blockidx mlp fc2 weight fdstprefix encoder layers layeridx blocks blockidx output dense weight fsrcprefix layers layeridx blocks blockidx mlp fc2 bias fdstprefix encoder layers layeridx blocks blockidx output dense bias renamedkeys extend fsrcprefix layers layeridx blocks blockidx attn relativepositionindex fdstprefix encoder layers layeridx blocks blockidx attention self relativepositionindex if layeridx 3 patch merging renamedkeys extend fsrcprefix layers layeridx downsample reduction weight fdstprefix encoder layers layeridx downsample reduction weight fsrcprefix layers layeridx downsample norm weight fdstprefix encoder layers layeridx downsample norm weight fsrcprefix layers layeridx downsample norm bias fdstprefix encoder layers layeridx downsample norm bias hidden states norms renamedkeys extend fsrcprefix normlayeridx weight fdstprefix hiddenstatesnorms stagelayeridx1 weight fsrcprefix normlayeridx bias fdstprefix hiddenstatesnorms stagelayeridx1 bias self popallrenamedkeys dststatedict srcstatedict backbone pixel decoder def replacepixelmoduleself dststatedict statedict srcstatedict statedict dstprefix str pixellevelmodule decoder srcprefix str semseghead pixeldecoder self replaceswinbackbonedststatedict srcstatedict self config def renamekeysforweightbiassrcprefix str dstprefix str return fsrcprefix weight fdstprefix weight fsrcprefix bias fdstprefix bias def renamekeysforselfattnsrcprefix str dstprefix str selfattnkeys selfattnkeys extend renamekeysforweightbiasfsrcprefix attentionweights fdstprefix attentionweights selfattnkeys extend renamekeysforweightbiasfsrcprefix outputproj fdstprefix outputproj selfattnkeys extend renamekeysforweightbiasfsrcprefix samplingoffsets fdstprefix samplingoffsets selfattnkeys extendrenamekeysforweightbiasfsrcprefix valueproj fdstprefix valueproj return selfattnkeys def renamekeysforencoderlayersrcprefix str dstprefix str encoderkeys encoderkeys extendrenamekeysforweightbiasfsrcprefix linear1 fdstprefix fc1 encoderkeys extendrenamekeysforweightbiasfsrcprefix linear2 fdstprefix fc2 encoderkeys extend renamekeysforweightbiasfsrcprefix norm1 fdstprefix selfattnlayernorm encoderkeys extendrenamekeysforweightbiasfsrcprefix norm2 fdstprefix finallayernorm encoderkeys extendrenamekeysforselfattnfsrcprefix selfattn fdstprefix selfattn return encoderkeys convolution layer for final features renamedkeys fsrcprefix adapter1 weight fdstprefix adapter1 0 weight fsrcprefix adapter1 norm weight fdstprefix adapter1 1 weight fsrcprefix adapter1 norm bias fdstprefix adapter1 1 bias renamedkeys extend fsrcprefix layer1 weight fdstprefix layer1 0 weight fsrcprefix layer1 norm weight fdstprefix layer1 1 weight fsrcprefix layer1 norm bias fdstprefix layer1 1 bias proj layers for i in range3 for j in range2 renamedkeys extend fsrcprefix inputproj i j weight fdstprefix inputprojections i j weight fsrcprefix inputproj i j bias fdstprefix inputprojections i j bias renamedkeys extendfsrcprefix transformer levelembed fdstprefix levelembed layers for layeridx in rangeself config encoderlayers renamedkeys extend renamekeysforencoderlayer fsrcprefix transformer encoder layers layeridx fdstprefix encoder layers layeridx proj renamedkeys extend fsrcprefix maskfeatures weight fdstprefix maskprojection weight fsrcprefix maskfeatures bias fdstprefix maskprojection bias self popallrenamedkeys dststatedict srcstatedict transformer decoder def renamekeysinmaskedattentiondecoderself dststatedict statedict srcstatedict statedict dstprefix str transformermodule decoder srcprefix str semseghead predictor renamekeys for i in rangeself config decoderlayers 1 renamekeys append fsrcprefix transformerselfattentionlayers i selfattn outproj weight fdstprefix layers i selfattn outproj weight renamekeys append fsrcprefix transformerselfattentionlayers i selfattn outproj bias fdstprefix layers i selfattn outproj bias renamekeys append fsrcprefix transformerselfattentionlayers i norm weight fdstprefix layers i selfattnlayernorm weight renamekeys append fsrcprefix transformerselfattentionlayers i norm bias fdstprefix layers i selfattnlayernorm bias renamekeys append fsrcprefix transformercrossattentionlayers i multiheadattn inprojweight fdstprefix layers i crossattn inprojweight renamekeys append fsrcprefix transformercrossattentionlayers i multiheadattn inprojbias fdstprefix layers i crossattn inprojbias renamekeys append fsrcprefix transformercrossattentionlayers i multiheadattn outproj weight fdstprefix layers i crossattn outproj weight renamekeys append fsrcprefix transformercrossattentionlayers i multiheadattn outproj bias fdstprefix layers i crossattn outproj bias renamekeys append fsrcprefix transformercrossattentionlayers i norm weight fdstprefix layers i crossattnlayernorm weight renamekeys append fsrcprefix transformercrossattentionlayers i norm bias fdstprefix layers i crossattnlayernorm bias renamekeys append fsrcprefix transformerffnlayers i linear1 weight fdstprefix layers i fc1 weight renamekeys append fsrcprefix transformerffnlayers i linear1 bias fdstprefix layers i fc1 bias renamekeys append fsrcprefix transformerffnlayers i linear2 weight fdstprefix layers i fc2 weight renamekeys append fsrcprefix transformerffnlayers i linear2 bias fdstprefix layers i fc2 bias renamekeys append fsrcprefix transformerffnlayers i norm weight fdstprefix layers i finallayernorm weight renamekeys append fsrcprefix transformerffnlayers i norm bias fdstprefix layers i finallayernorm bias return renamekeys def replacemaskedattentiondecoderself dststatedict statedict srcstatedict statedict dstprefix str transformermodule decoder srcprefix str semseghead predictor renamedkeys self renamekeysinmaskedattentiondecoderdststatedict srcstatedict add more renamedkeys extend fsrcprefix decodernorm weight fdstprefix layernorm weight fsrcprefix decodernorm bias fdstprefix layernorm bias mlplen 3 for i in rangemlplen renamedkeys extend fsrcprefix maskembed layers i weight fdstprefix maskpredictor maskembedder i 0 weight fsrcprefix maskembed layers i bias fdstprefix maskpredictor maskembedder i 0 bias self popallrenamedkeys dststatedict srcstatedict def replacekeysqkvtransformerdecoderself dststatedict statedict srcstatedict statedict dstprefix str transformermodule decoder layers srcprefix str semseghead predictor for i in rangeself config decoderlayers 1 read in weights bias of input projection layer of selfattention inprojweight srcstatedict pop fsrcprefix transformerselfattentionlayers i selfattn inprojweight inprojbias srcstatedict pop fsrcprefix transformerselfattentionlayers i selfattn inprojbias next add query keys and values in that order to the state dict dststatedictfdstprefix i selfattn qproj weight inprojweight 256 dststatedictfdstprefix i selfattn qproj bias inprojbias 256 dststatedictfdstprefix i selfattn kproj weight inprojweight256 512 dststatedictfdstprefix i selfattn kproj bias inprojbias256 512 dststatedictfdstprefix i selfattn vproj weight inprojweight256 dststatedictfdstprefix i selfattn vproj bias inprojbias256 def replacetransformermoduleself dststatedict statedict srcstatedict statedict dstprefix str transformermodule srcprefix str semseghead predictor self replacemaskedattentiondecoderdststatedict srcstatedict renamedkeys fsrcprefix queryembed weight fdstprefix queriesembedder weight fsrcprefix queryfeat weight fdstprefix queriesfeatures weight fsrcprefix levelembed weight fdstprefix levelembed weight self popallrenamedkeys dststatedict srcstatedict self replacekeysqkvtransformerdecoderdststatedict srcstatedict def replaceuniversalsegmentationmoduleself dststatedict statedict srcstatedict statedict dstprefix str srcprefix str semseghead predictor renamedkeys fsrcprefix classembed weight fdstprefixclasspredictor weight fsrcprefix classembed bias fdstprefixclasspredictor bias logger infofreplacing keys pformatrenamedkeys self popallrenamedkeys dststatedict srcstatedict def convertself mask2former mask2formermodel mask2formermodel dststatedict trackedstatedictmask2former statedict srcstatedict self originalmodel statedict self replacepixelmoduledststatedict srcstatedict self replacetransformermoduledststatedict srcstatedict logger infofmissed keys are pformatdststatedict diff logger infofnot copied keys are pformatsrcstatedict keys logger info done statedict key dststatedictkey for key in dststatedict totrack keys mask2former loadstatedictstatedict return mask2former def convertuniversalsegmentation self mask2former mask2formerforuniversalsegmentation mask2formerforuniversalsegmentation dststatedict trackedstatedictmask2former statedict srcstatedict self originalmodel statedict self replaceuniversalsegmentationmoduledststatedict srcstatedict statedict key dststatedictkey for key in dststatedict totrack keys mask2former loadstatedictstatedict return mask2former staticmethod def usingdirscheckpointsdir path configdir path iteratortupleobject path path checkpoints listpath checkpointsdir glob pkl for checkpoint in checkpoints logger infof converting checkpoint stem find associated config file datasetname e g coco datasetname checkpoint parents2 stem if datasetname ade datasetname datasetname replaceade ade20k task type e g instancesegmentation segmentationtask checkpoint parents1 stem config file corresponding to checkpoint configfilename fcheckpoint parents0 stem yaml config path configdir datasetname segmentationtask swin configfilename yield config checkpoint def test originalmodel ourmodel mask2formerforuniversalsegmentation imageprocessor mask2formerimageprocessor tolerance float with torch nograd originalmodel originalmodel eval ourmodel ourmodel eval im prepareimg x imageprocessorimagesim returntensorsptpixelvalues originalmodelbackbonefeatures originalmodel backbonex clone ourmodeloutput mask2formermodeloutput ourmodel modelx clone outputhiddenstatestrue test backbone for originalmodelfeature ourmodelfeature in zip originalmodelbackbonefeatures values ourmodeloutput encoderhiddenstates assert torch allclose originalmodelfeature ourmodelfeature atoltolerance the backbone features are not the same test pixel decoder maskfeatures multiscalefeatures originalmodel semseghead pixeldecoder forwardfeatures originalmodelbackbonefeatures for originalmodelfeature ourmodelfeature in zip multiscalefeatures ourmodeloutput pixeldecoderhiddenstates assert torch allclose originalmodelfeature ourmodelfeature atoltolerance the pixel decoder feature are not the same let s test the full model trcomplete t compose t resize384 384 t totensor y trcompleteim 255 0 totorch int float modify original mask2former code to return mask and class logits originalclasslogits originalmasklogits originalmodelimage y clone squeeze0 ourmodelout mask2formerforuniversalsegmentationoutput ourmodelx clone ourmasklogits ourmodelout masksquerieslogits ourclasslogits ourmodelout classquerieslogits assert originalmasklogits shape ourmasklogits shape output masks shapes are not matching assert originalclasslogits shape ourclasslogits shape output class logits shapes are not matching assert torch allclose originalclasslogits ourclasslogits atoltolerance the class logits are not the same assert torch allclose originalmasklogits ourmasklogits atoltolerance the predicted masks are not the same logger info test passed def getmodelnamecheckpointfile path modelnameraw is something like maskformer2swinsmallbs1650ep modelnameraw str checkpointfile parents0 stem segmentationtasktype must be one of the following instancesegmentation panopticsegmentation semanticsegmentation segmentationtaskname str checkpointfile parents1 stem if segmentationtaskname not in instancesegmentation panopticsegmentation semanticsegmentation raise valueerror fsegmentationtaskname must be wrong since acceptable values are instancesegmentation panopticsegmentation semanticsegmentation dataset name must be one of the following coco ade cityscapes mapillaryvistas datasetname str checkpointfile parents2 stem if datasetname not in coco ade cityscapes mapillaryvistas raise valueerror fdatasetname must be wrong since we didn t find coco or ade or cityscapes or mapillaryvistas in it backbone swin backbonetypes tiny small basein21k base large backbonetype listfilterlambda x x in modelnameraw backbonetypes0 replace modelname fmask2formerbackbonebackbonetypedatasetnamesegmentationtaskname split 0 return modelname if name main parser argumentparser descriptioncommand line to convert the original mask2formers with swin backbone to our implementations parser addargument checkpointsdir typepath help a directory containing the model s checkpoints the directory has to have the following structure dirnamedatasetnamesegmentationtasknameconfigname pkl parser addargument configsdir typepath help a directory containing the model s configs see detectron2 doc the directory has to have the following structure dirnamedatasetnamesegmentationtasknameconfigname yaml parser addargument mask2formerdir requiredtrue typepath help a path to mask2former s original implementation directory you can download from here https github comfacebookresearchmask2former args parser parseargs checkpointsdir path args checkpointsdir configdir path args configsdir mask2formerdir path args mask2formerdir append the path to the parents to mask2former dir sys path appendstrmask2formerdir parent import original mask2former config and model from original source code repo from mask2former mask2former config import addmaskformer2config from mask2former mask2former maskformermodel import maskformer as originalmask2former for configfile checkpointfile in originalmask2formercheckpointtooursconverter usingdirs checkpointsdir configdir modelname getmodelnamecheckpointfile imageprocessor originalmask2formerconfigtoimageprocessorconverter setupcfgargsconfigfileconfigfile imageprocessor size height 384 width 384 originalconfig setupcfgargsconfigfileconfigfile mask2formerkwargs originalmask2former fromconfigoriginalconfig originalmodel originalmask2formermask2formerkwargs eval detectioncheckpointeroriginalmodel loadstrcheckpointfile config mask2formerconfig originalmask2formerconfigtooursconverteroriginalconfig mask2former mask2formermodelconfigconfig eval converter originalmask2formercheckpointtooursconverteroriginalmodel config mask2former converter convertmask2former mask2formerforsegmentation mask2formerforuniversalsegmentationconfigconfig eval mask2formerforsegmentation model mask2former mask2formerforsegmentation converter convertuniversalsegmentationmask2formerforsegmentation tolerance 3e1 hightolerancemodels mask2formerswinbasein21kcocoinstance mask2formerswinbasecocoinstance mask2formerswinsmallcityscapessemantic if modelname in hightolerancemodels tolerance 3e1 logger infof testing modelname testoriginalmodel mask2formerforsegmentation imageprocessor tolerance logger infof pushing modelname to hub imageprocessor pushtohubmodelname mask2formerforsegmentation pushtohubmodelname coding utf 8 2022 meta platforms inc and 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 this class tracks a python dictionary by keeping track of which item is accessed args to_track dict the dictionary we wish to track this method returns a set difference between the keys in the tracked state dict and the one we have access so far this is an effective method to check if we have update all the keys returns list str list of keys not yet updated proxy the call to the internal dictionary we will verify our results on an image of cute cats fake command line arguments needed by mask2former detectron implementation load config from file and command line arguments src dst now we need to handle the attentions read in weights bias of input projection layer of cross attention let s pop them proj second norm mlp patch merging hidden states norms src dst now we need to handle the attentions read in weights bias of input projection layer of cross attention let s pop them proj second norm mlp patch merging hidden states norms backbone pixel decoder convolution layer for final features proj layers layers proj transformer decoder add more read in weights bias of input projection layer of self attention next add query keys and values in that order to the state dict find associated config file dataset_name e g coco task type e g instance segmentation config file corresponding to checkpoint test backbone test pixel decoder let s test the full model modify original mask2former code to return mask and class logits model_name_raw is something like maskformer2_swin_small_bs16_50ep segmentation_task_type must be one of the following instance segmentation panoptic segmentation semantic segmentation dataset name must be one of the following coco ade cityscapes mapillary vistas append the path to the parents to mask2former dir import original mask2former config and model from original source code repo
import json
import sys
from argparse import ArgumentParser
from dataclasses import dataclass
from pathlib import Path
from pprint import pformat
from typing import Any, Dict, Iterator, List, Set, Tuple

import requests
import torch
import torchvision.transforms as T
from detectron2.checkpoint import DetectionCheckpointer
from detectron2.config import get_cfg
from detectron2.projects.deeplab import add_deeplab_config
from huggingface_hub import hf_hub_download
from PIL import Image
from torch import Tensor, nn

from transformers import (
    Mask2FormerConfig,
    Mask2FormerForUniversalSegmentation,
    Mask2FormerImageProcessor,
    Mask2FormerModel,
    SwinConfig,
)
from transformers.models.mask2former.modeling_mask2former import (
    Mask2FormerForUniversalSegmentationOutput,
    Mask2FormerModelOutput,
)
from transformers.utils import logging


StateDict = Dict[str, Tensor]

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

torch.manual_seed(0)


class TrackedStateDict:
    def __init__(self, to_track: Dict):
        
        self.to_track = to_track
        self._seen: Set[str] = set()

    def __getitem__(self, key: str) -> Any:
        return self.to_track[key]

    def __setitem__(self, key: str, item: Any):
        self._seen.add(key)
        self.to_track[key] = item

    def diff(self) -> List[str]:
        
        return set(self.to_track.keys()) - self._seen

    def copy(self) -> Dict:
        
        return self.to_track.copy()



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


@dataclass
class Args:
    

    config_file: str


def setup_cfg(args: Args):
    
    cfg = get_cfg()
    add_deeplab_config(cfg)
    add_maskformer2_config(cfg)
    cfg.merge_from_file(args.config_file)
    cfg.freeze()
    return cfg


class OriginalMask2FormerConfigToOursConverter:
    def __call__(self, original_config: object) -> Mask2FormerConfig:
        model = original_config.MODEL

        repo_id = "huggingface/label-files"
        if model.SEM_SEG_HEAD.NUM_CLASSES == 847:
            filename = "mask2former-ade20k-full-id2label.json"
        elif model.SEM_SEG_HEAD.NUM_CLASSES == 150:
            filename = "ade20k-id2label.json"
        elif model.SEM_SEG_HEAD.NUM_CLASSES == 80:
            filename = "coco-detection-mmdet-id2label.json"
        elif model.SEM_SEG_HEAD.NUM_CLASSES == 171:
            filename = "mask2former-coco-stuff-id2label.json"
        elif model.SEM_SEG_HEAD.NUM_CLASSES == 133:
            filename = "coco-panoptic-id2label.json"
        elif model.SEM_SEG_HEAD.NUM_CLASSES == 19:
            filename = "cityscapes-id2label.json"
        elif model.SEM_SEG_HEAD.NUM_CLASSES == 8:
            filename = "cityscapes-instance-id2label.json"
        elif model.SEM_SEG_HEAD.NUM_CLASSES == 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()}
        label2id = {label: idx for idx, label in id2label.items()}

        if model.SWIN.EMBED_DIM == 96:
            backbone_config = SwinConfig.from_pretrained(
                "microsoft/swin-tiny-patch4-window7-224", out_features=["stage1", "stage2", "stage3", "stage4"]
            )
        elif model.SWIN.EMBED_DIM == 128:
            backbone_config = SwinConfig(
                embed_dim=128,
                window_size=12,
                depths=(2, 2, 18, 2),
                num_heads=(4, 8, 16, 32),
                out_features=["stage1", "stage2", "stage3", "stage4"],
            )

        elif model.SWIN.EMBED_DIM == 192:
            backbone_config = SwinConfig.from_pretrained(
                "microsoft/swin-large-patch4-window12-384", out_features=["stage1", "stage2", "stage3", "stage4"]
            )
        else:
            raise ValueError(f"embed dim {model.SWIN.EMBED_DIM} not supported for Swin!")

        backbone_config.drop_path_rate = model.SWIN.DROP_PATH_RATE
        backbone_config.attention_probs_dropout_prob = model.SWIN.ATTN_DROP_RATE
        backbone_config.depths = model.SWIN.DEPTHS

        config: Mask2FormerConfig = Mask2FormerConfig(
            ignore_value=model.SEM_SEG_HEAD.IGNORE_VALUE,
            num_labels=model.SEM_SEG_HEAD.NUM_CLASSES,
            num_queries=model.MASK_FORMER.NUM_OBJECT_QUERIES,
            no_object_weight=model.MASK_FORMER.NO_OBJECT_WEIGHT,
            class_weight=model.MASK_FORMER.CLASS_WEIGHT,
            mask_weight=model.MASK_FORMER.MASK_WEIGHT,
            dice_weight=model.MASK_FORMER.DICE_WEIGHT,
            train_num_points=model.MASK_FORMER.TRAIN_NUM_POINTS,
            oversample_ratio=model.MASK_FORMER.OVERSAMPLE_RATIO,
            importance_sample_ratio=model.MASK_FORMER.IMPORTANCE_SAMPLE_RATIO,
            init_std=0.02,
            init_xavier_std=1.0,
            use_auxiliary_loss=model.MASK_FORMER.DEEP_SUPERVISION,
            feature_strides=[4, 8, 16, 32],
            backbone_config=backbone_config,
            id2label=id2label,
            label2id=label2id,
            feature_size=model.SEM_SEG_HEAD.CONVS_DIM,
            mask_feature_size=model.SEM_SEG_HEAD.MASK_DIM,
            hidden_dim=model.MASK_FORMER.HIDDEN_DIM,
            encoder_layers=model.SEM_SEG_HEAD.TRANSFORMER_ENC_LAYERS,
            encoder_feedforward_dim=1024,
            decoder_layers=model.MASK_FORMER.DEC_LAYERS,
            num_attention_heads=model.MASK_FORMER.NHEADS,
            dropout=model.MASK_FORMER.DROPOUT,
            dim_feedforward=model.MASK_FORMER.DIM_FEEDFORWARD,
            pre_norm=model.MASK_FORMER.PRE_NORM,
            enforce_input_proj=model.MASK_FORMER.ENFORCE_INPUT_PROJ,
            common_stride=model.SEM_SEG_HEAD.COMMON_STRIDE,
        )
        return config


class OriginalMask2FormerConfigToImageProcessorConverter:
    def __call__(self, original_config: object) -> Mask2FormerImageProcessor:
        model = original_config.MODEL
        model_input = original_config.INPUT

        return Mask2FormerImageProcessor(
            image_mean=(torch.tensor(model.PIXEL_MEAN) / 255).tolist(),
            image_std=(torch.tensor(model.PIXEL_STD) / 255).tolist(),
            size=model_input.MIN_SIZE_TEST,
            max_size=model_input.MAX_SIZE_TEST,
            num_labels=model.SEM_SEG_HEAD.NUM_CLASSES,
            ignore_index=model.SEM_SEG_HEAD.IGNORE_VALUE,
            size_divisibility=32,
        )


class OriginalMask2FormerCheckpointToOursConverter:
    def __init__(self, original_model: nn.Module, config: Mask2FormerConfig):
        self.original_model = original_model
        self.config = config

    def pop_all(self, renamed_keys: List[Tuple[str, str]], dst_state_dict: StateDict, src_state_dict: StateDict):
        for src_key, dst_key in renamed_keys:
            dst_state_dict[dst_key] = src_state_dict.pop(src_key)

    def replace_maskformer_swin_backbone(
        self, dst_state_dict: StateDict, src_state_dict: StateDict, config: Mask2FormerConfig
    ):
        dst_prefix: str = "pixel_level_module.encoder"
        src_prefix: str = "backbone"

        renamed_keys = [
            (
                f"{src_prefix}.patch_embed.proj.weight",
                f"{dst_prefix}.model.embeddings.patch_embeddings.projection.weight",
            ),
            (f"{src_prefix}.patch_embed.proj.bias", f"{dst_prefix}.model.embeddings.patch_embeddings.projection.bias"),
            (f"{src_prefix}.patch_embed.norm.weight", f"{dst_prefix}.model.embeddings.norm.weight"),
            (f"{src_prefix}.patch_embed.norm.bias", f"{dst_prefix}.model.embeddings.norm.bias"),
        ]
        num_layers = len(config.backbone_config.depths)
        for layer_idx in range(num_layers):
            for block_idx in range(config.backbone_config.depths[layer_idx]):
                renamed_keys.extend(
                    [  
                        (
                            f"{src_prefix}.layers.{layer_idx}.blocks.{block_idx}.norm1.weight",
                            f"{dst_prefix}.model.encoder.layers.{layer_idx}.blocks.{block_idx}.layernorm_before.weight",
                        ),
                        (
                            f"{src_prefix}.layers.{layer_idx}.blocks.{block_idx}.norm1.bias",
                            f"{dst_prefix}.model.encoder.layers.{layer_idx}.blocks.{block_idx}.layernorm_before.bias",
                        ),
                        (
                            f"{src_prefix}.layers.{layer_idx}.blocks.{block_idx}.attn.relative_position_bias_table",
                            f"{dst_prefix}.model.encoder.layers.{layer_idx}.blocks.{block_idx}.attention.self.relative_position_bias_table",
                        ),
                    ]
                )
                
                

                src_att_weight = src_state_dict[f"{src_prefix}.layers.{layer_idx}.blocks.{block_idx}.attn.qkv.weight"]
                src_att_bias = src_state_dict[f"{src_prefix}.layers.{layer_idx}.blocks.{block_idx}.attn.qkv.bias"]

                size = src_att_weight.shape[0]
                offset = size // 3
                dst_state_dict[
                    f"{dst_prefix}.model.encoder.layers.{layer_idx}.blocks.{block_idx}.attention.self.query.weight"
                ] = src_att_weight[:offset, :]
                dst_state_dict[
                    f"{dst_prefix}.model.encoder.layers.{layer_idx}.blocks.{block_idx}.attention.self.query.bias"
                ] = src_att_bias[:offset]

                dst_state_dict[
                    f"{dst_prefix}.model.encoder.layers.{layer_idx}.blocks.{block_idx}.attention.self.key.weight"
                ] = src_att_weight[offset : offset * 2, :]
                dst_state_dict[
                    f"{dst_prefix}.model.encoder.layers.{layer_idx}.blocks.{block_idx}.attention.self.key.bias"
                ] = src_att_bias[offset : offset * 2]

                dst_state_dict[
                    f"{dst_prefix}.model.encoder.layers.{layer_idx}.blocks.{block_idx}.attention.self.value.weight"
                ] = src_att_weight[-offset:, :]
                dst_state_dict[
                    f"{dst_prefix}.model.encoder.layers.{layer_idx}.blocks.{block_idx}.attention.self.value.bias"
                ] = src_att_bias[-offset:]

                
                src_state_dict.pop(f"{src_prefix}.layers.{layer_idx}.blocks.{block_idx}.attn.qkv.weight")
                src_state_dict.pop(f"{src_prefix}.layers.{layer_idx}.blocks.{block_idx}.attn.qkv.bias")
                
                renamed_keys.extend(
                    [
                        (
                            f"{src_prefix}.layers.{layer_idx}.blocks.{block_idx}.attn.proj.weight",
                            f"{dst_prefix}.model.encoder.layers.{layer_idx}.blocks.{block_idx}.attention.output.dense.weight",
                        ),
                        (
                            f"{src_prefix}.layers.{layer_idx}.blocks.{block_idx}.attn.proj.bias",
                            f"{dst_prefix}.model.encoder.layers.{layer_idx}.blocks.{block_idx}.attention.output.dense.bias",
                        ),
                    ]
                )

                
                renamed_keys.extend(
                    [
                        (
                            f"{src_prefix}.layers.{layer_idx}.blocks.{block_idx}.norm2.weight",
                            f"{dst_prefix}.model.encoder.layers.{layer_idx}.blocks.{block_idx}.layernorm_after.weight",
                        ),
                        (
                            f"{src_prefix}.layers.{layer_idx}.blocks.{block_idx}.norm2.bias",
                            f"{dst_prefix}.model.encoder.layers.{layer_idx}.blocks.{block_idx}.layernorm_after.bias",
                        ),
                    ]
                )

                
                renamed_keys.extend(
                    [
                        (
                            f"{src_prefix}.layers.{layer_idx}.blocks.{block_idx}.mlp.fc1.weight",
                            f"{dst_prefix}.model.encoder.layers.{layer_idx}.blocks.{block_idx}.intermediate.dense.weight",
                        ),
                        (
                            f"{src_prefix}.layers.{layer_idx}.blocks.{block_idx}.mlp.fc1.bias",
                            f"{dst_prefix}.model.encoder.layers.{layer_idx}.blocks.{block_idx}.intermediate.dense.bias",
                        ),
                        (
                            f"{src_prefix}.layers.{layer_idx}.blocks.{block_idx}.mlp.fc2.weight",
                            f"{dst_prefix}.model.encoder.layers.{layer_idx}.blocks.{block_idx}.output.dense.weight",
                        ),
                        (
                            f"{src_prefix}.layers.{layer_idx}.blocks.{block_idx}.mlp.fc2.bias",
                            f"{dst_prefix}.model.encoder.layers.{layer_idx}.blocks.{block_idx}.output.dense.bias",
                        ),
                    ]
                )

                renamed_keys.extend(
                    [
                        (
                            f"{src_prefix}.layers.{layer_idx}.blocks.{block_idx}.attn.relative_position_index",
                            f"{dst_prefix}.model.encoder.layers.{layer_idx}.blocks.{block_idx}.attention.self.relative_position_index",
                        )
                    ]
                )

            if layer_idx < num_layers - 1:
                
                renamed_keys.extend(
                    [
                        (
                            f"{src_prefix}.layers.{layer_idx}.downsample.reduction.weight",
                            f"{dst_prefix}.model.encoder.layers.{layer_idx}.downsample.reduction.weight",
                        ),
                        (
                            f"{src_prefix}.layers.{layer_idx}.downsample.norm.weight",
                            f"{dst_prefix}.model.encoder.layers.{layer_idx}.downsample.norm.weight",
                        ),
                        (
                            f"{src_prefix}.layers.{layer_idx}.downsample.norm.bias",
                            f"{dst_prefix}.model.encoder.layers.{layer_idx}.downsample.norm.bias",
                        ),
                    ]
                )

            
            renamed_keys.extend(
                [
                    (
                        f"{src_prefix}.norm{layer_idx}.weight",
                        f"{dst_prefix}.hidden_states_norms.{layer_idx}.weight",
                    ),
                    (
                        f"{src_prefix}.norm{layer_idx}.bias",
                        f"{dst_prefix}.hidden_states_norms.{layer_idx}.bias",
                    ),
                ]
            )
        self.pop_all(renamed_keys, dst_state_dict, src_state_dict)

    def replace_swin_backbone(self, dst_state_dict: StateDict, src_state_dict: StateDict, config: Mask2FormerConfig):
        dst_prefix: str = "pixel_level_module.encoder"
        src_prefix: str = "backbone"

        renamed_keys = [
            (
                f"{src_prefix}.patch_embed.proj.weight",
                f"{dst_prefix}.embeddings.patch_embeddings.projection.weight",
            ),
            (f"{src_prefix}.patch_embed.proj.bias", f"{dst_prefix}.embeddings.patch_embeddings.projection.bias"),
            (f"{src_prefix}.patch_embed.norm.weight", f"{dst_prefix}.embeddings.norm.weight"),
            (f"{src_prefix}.patch_embed.norm.bias", f"{dst_prefix}.embeddings.norm.bias"),
        ]

        for layer_idx in range(len(config.backbone_config.depths)):
            for block_idx in range(config.backbone_config.depths[layer_idx]):
                renamed_keys.extend(
                    [  
                        (
                            f"{src_prefix}.layers.{layer_idx}.blocks.{block_idx}.norm1.weight",
                            f"{dst_prefix}.encoder.layers.{layer_idx}.blocks.{block_idx}.layernorm_before.weight",
                        ),
                        (
                            f"{src_prefix}.layers.{layer_idx}.blocks.{block_idx}.norm1.bias",
                            f"{dst_prefix}.encoder.layers.{layer_idx}.blocks.{block_idx}.layernorm_before.bias",
                        ),
                        (
                            f"{src_prefix}.layers.{layer_idx}.blocks.{block_idx}.attn.relative_position_bias_table",
                            f"{dst_prefix}.encoder.layers.{layer_idx}.blocks.{block_idx}.attention.self.relative_position_bias_table",
                        ),
                    ]
                )
                
                

                src_att_weight = src_state_dict[f"{src_prefix}.layers.{layer_idx}.blocks.{block_idx}.attn.qkv.weight"]
                src_att_bias = src_state_dict[f"{src_prefix}.layers.{layer_idx}.blocks.{block_idx}.attn.qkv.bias"]

                size = src_att_weight.shape[0]
                offset = size // 3
                dst_state_dict[
                    f"{dst_prefix}.encoder.layers.{layer_idx}.blocks.{block_idx}.attention.self.query.weight"
                ] = src_att_weight[:offset, :]
                dst_state_dict[
                    f"{dst_prefix}.encoder.layers.{layer_idx}.blocks.{block_idx}.attention.self.query.bias"
                ] = src_att_bias[:offset]

                dst_state_dict[
                    f"{dst_prefix}.encoder.layers.{layer_idx}.blocks.{block_idx}.attention.self.key.weight"
                ] = src_att_weight[offset : offset * 2, :]
                dst_state_dict[
                    f"{dst_prefix}.encoder.layers.{layer_idx}.blocks.{block_idx}.attention.self.key.bias"
                ] = src_att_bias[offset : offset * 2]

                dst_state_dict[
                    f"{dst_prefix}.encoder.layers.{layer_idx}.blocks.{block_idx}.attention.self.value.weight"
                ] = src_att_weight[-offset:, :]
                dst_state_dict[
                    f"{dst_prefix}.encoder.layers.{layer_idx}.blocks.{block_idx}.attention.self.value.bias"
                ] = src_att_bias[-offset:]

                
                src_state_dict.pop(f"{src_prefix}.layers.{layer_idx}.blocks.{block_idx}.attn.qkv.weight")
                src_state_dict.pop(f"{src_prefix}.layers.{layer_idx}.blocks.{block_idx}.attn.qkv.bias")
                
                renamed_keys.extend(
                    [
                        (
                            f"{src_prefix}.layers.{layer_idx}.blocks.{block_idx}.attn.proj.weight",
                            f"{dst_prefix}.encoder.layers.{layer_idx}.blocks.{block_idx}.attention.output.dense.weight",
                        ),
                        (
                            f"{src_prefix}.layers.{layer_idx}.blocks.{block_idx}.attn.proj.bias",
                            f"{dst_prefix}.encoder.layers.{layer_idx}.blocks.{block_idx}.attention.output.dense.bias",
                        ),
                    ]
                )

                
                renamed_keys.extend(
                    [
                        (
                            f"{src_prefix}.layers.{layer_idx}.blocks.{block_idx}.norm2.weight",
                            f"{dst_prefix}.encoder.layers.{layer_idx}.blocks.{block_idx}.layernorm_after.weight",
                        ),
                        (
                            f"{src_prefix}.layers.{layer_idx}.blocks.{block_idx}.norm2.bias",
                            f"{dst_prefix}.encoder.layers.{layer_idx}.blocks.{block_idx}.layernorm_after.bias",
                        ),
                    ]
                )

                
                renamed_keys.extend(
                    [
                        (
                            f"{src_prefix}.layers.{layer_idx}.blocks.{block_idx}.mlp.fc1.weight",
                            f"{dst_prefix}.encoder.layers.{layer_idx}.blocks.{block_idx}.intermediate.dense.weight",
                        ),
                        (
                            f"{src_prefix}.layers.{layer_idx}.blocks.{block_idx}.mlp.fc1.bias",
                            f"{dst_prefix}.encoder.layers.{layer_idx}.blocks.{block_idx}.intermediate.dense.bias",
                        ),
                        (
                            f"{src_prefix}.layers.{layer_idx}.blocks.{block_idx}.mlp.fc2.weight",
                            f"{dst_prefix}.encoder.layers.{layer_idx}.blocks.{block_idx}.output.dense.weight",
                        ),
                        (
                            f"{src_prefix}.layers.{layer_idx}.blocks.{block_idx}.mlp.fc2.bias",
                            f"{dst_prefix}.encoder.layers.{layer_idx}.blocks.{block_idx}.output.dense.bias",
                        ),
                    ]
                )

                renamed_keys.extend(
                    [
                        (
                            f"{src_prefix}.layers.{layer_idx}.blocks.{block_idx}.attn.relative_position_index",
                            f"{dst_prefix}.encoder.layers.{layer_idx}.blocks.{block_idx}.attention.self.relative_position_index",
                        )
                    ]
                )

            if layer_idx < 3:
                
                renamed_keys.extend(
                    [
                        (
                            f"{src_prefix}.layers.{layer_idx}.downsample.reduction.weight",
                            f"{dst_prefix}.encoder.layers.{layer_idx}.downsample.reduction.weight",
                        ),
                        (
                            f"{src_prefix}.layers.{layer_idx}.downsample.norm.weight",
                            f"{dst_prefix}.encoder.layers.{layer_idx}.downsample.norm.weight",
                        ),
                        (
                            f"{src_prefix}.layers.{layer_idx}.downsample.norm.bias",
                            f"{dst_prefix}.encoder.layers.{layer_idx}.downsample.norm.bias",
                        ),
                    ]
                )

            
            renamed_keys.extend(
                [
                    (
                        f"{src_prefix}.norm{layer_idx}.weight",
                        f"{dst_prefix}.hidden_states_norms.stage{layer_idx+1}.weight",
                    ),
                    (
                        f"{src_prefix}.norm{layer_idx}.bias",
                        f"{dst_prefix}.hidden_states_norms.stage{layer_idx+1}.bias",
                    ),
                ]
            )
        self.pop_all(renamed_keys, dst_state_dict, src_state_dict)

    
    def replace_pixel_module(self, dst_state_dict: StateDict, src_state_dict: StateDict):
        dst_prefix: str = "pixel_level_module.decoder"
        src_prefix: str = "sem_seg_head.pixel_decoder"

        self.replace_swin_backbone(dst_state_dict, src_state_dict, self.config)

        def rename_keys_for_weight_bias(src_prefix: str, dst_prefix: str):
            return [
                (f"{src_prefix}.weight", f"{dst_prefix}.weight"),
                (f"{src_prefix}.bias", f"{dst_prefix}.bias"),
            ]

        def rename_keys_for_self_attn(src_prefix: str, dst_prefix: str):
            self_attn_keys = []
            self_attn_keys.extend(
                rename_keys_for_weight_bias(f"{src_prefix}.attention_weights", f"{dst_prefix}.attention_weights")
            )
            self_attn_keys.extend(
                rename_keys_for_weight_bias(f"{src_prefix}.output_proj", f"{dst_prefix}.output_proj")
            )
            self_attn_keys.extend(
                rename_keys_for_weight_bias(f"{src_prefix}.sampling_offsets", f"{dst_prefix}.sampling_offsets")
            )
            self_attn_keys.extend(rename_keys_for_weight_bias(f"{src_prefix}.value_proj", f"{dst_prefix}.value_proj"))

            return self_attn_keys

        def rename_keys_for_encoder_layer(src_prefix: str, dst_prefix: str):
            encoder_keys = []
            encoder_keys.extend(rename_keys_for_weight_bias(f"{src_prefix}.linear1", f"{dst_prefix}.fc1"))
            encoder_keys.extend(rename_keys_for_weight_bias(f"{src_prefix}.linear2", f"{dst_prefix}.fc2"))
            encoder_keys.extend(
                rename_keys_for_weight_bias(f"{src_prefix}.norm1", f"{dst_prefix}.self_attn_layer_norm")
            )
            encoder_keys.extend(rename_keys_for_weight_bias(f"{src_prefix}.norm2", f"{dst_prefix}.final_layer_norm"))
            encoder_keys.extend(rename_keys_for_self_attn(f"{src_prefix}.self_attn", f"{dst_prefix}.self_attn"))

            return encoder_keys

        
        renamed_keys = [
            (f"{src_prefix}.adapter_1.weight", f"{dst_prefix}.adapter_1.0.weight"),
            (f"{src_prefix}.adapter_1.norm.weight", f"{dst_prefix}.adapter_1.1.weight"),
            (f"{src_prefix}.adapter_1.norm.bias", f"{dst_prefix}.adapter_1.1.bias"),
        ]

        renamed_keys.extend(
            [
                (f"{src_prefix}.layer_1.weight", f"{dst_prefix}.layer_1.0.weight"),
                (f"{src_prefix}.layer_1.norm.weight", f"{dst_prefix}.layer_1.1.weight"),
                (f"{src_prefix}.layer_1.norm.bias", f"{dst_prefix}.layer_1.1.bias"),
            ]
        )

        
        for i in range(3):
            for j in range(2):
                renamed_keys.extend(
                    [
                        (f"{src_prefix}.input_proj.{i}.{j}.weight", f"{dst_prefix}.input_projections.{i}.{j}.weight"),
                        (f"{src_prefix}.input_proj.{i}.{j}.bias", f"{dst_prefix}.input_projections.{i}.{j}.bias"),
                    ]
                )

        renamed_keys.extend([(f"{src_prefix}.transformer.level_embed", f"{dst_prefix}.level_embed")])

        
        for layer_idx in range(self.config.encoder_layers):
            renamed_keys.extend(
                rename_keys_for_encoder_layer(
                    f"{src_prefix}.transformer.encoder.layers.{layer_idx}", f"{dst_prefix}.encoder.layers.{layer_idx}"
                )
            )

        
        renamed_keys.extend(
            [
                (f"{src_prefix}.mask_features.weight", f"{dst_prefix}.mask_projection.weight"),
                (f"{src_prefix}.mask_features.bias", f"{dst_prefix}.mask_projection.bias"),
            ]
        )
        self.pop_all(renamed_keys, dst_state_dict, src_state_dict)

    
    def rename_keys_in_masked_attention_decoder(self, dst_state_dict: StateDict, src_state_dict: StateDict):
        dst_prefix: str = "transformer_module.decoder"
        src_prefix: str = "sem_seg_head.predictor"

        rename_keys = []
        for i in range(self.config.decoder_layers - 1):
            rename_keys.append(
                (
                    f"{src_prefix}.transformer_self_attention_layers.{i}.self_attn.out_proj.weight",
                    f"{dst_prefix}.layers.{i}.self_attn.out_proj.weight",
                )
            )
            rename_keys.append(
                (
                    f"{src_prefix}.transformer_self_attention_layers.{i}.self_attn.out_proj.bias",
                    f"{dst_prefix}.layers.{i}.self_attn.out_proj.bias",
                )
            )

            rename_keys.append(
                (
                    f"{src_prefix}.transformer_self_attention_layers.{i}.norm.weight",
                    f"{dst_prefix}.layers.{i}.self_attn_layer_norm.weight",
                )
            )
            rename_keys.append(
                (
                    f"{src_prefix}.transformer_self_attention_layers.{i}.norm.bias",
                    f"{dst_prefix}.layers.{i}.self_attn_layer_norm.bias",
                )
            )

            rename_keys.append(
                (
                    f"{src_prefix}.transformer_cross_attention_layers.{i}.multihead_attn.in_proj_weight",
                    f"{dst_prefix}.layers.{i}.cross_attn.in_proj_weight",
                )
            )
            rename_keys.append(
                (
                    f"{src_prefix}.transformer_cross_attention_layers.{i}.multihead_attn.in_proj_bias",
                    f"{dst_prefix}.layers.{i}.cross_attn.in_proj_bias",
                )
            )
            rename_keys.append(
                (
                    f"{src_prefix}.transformer_cross_attention_layers.{i}.multihead_attn.out_proj.weight",
                    f"{dst_prefix}.layers.{i}.cross_attn.out_proj.weight",
                )
            )
            rename_keys.append(
                (
                    f"{src_prefix}.transformer_cross_attention_layers.{i}.multihead_attn.out_proj.bias",
                    f"{dst_prefix}.layers.{i}.cross_attn.out_proj.bias",
                )
            )

            rename_keys.append(
                (
                    f"{src_prefix}.transformer_cross_attention_layers.{i}.norm.weight",
                    f"{dst_prefix}.layers.{i}.cross_attn_layer_norm.weight",
                )
            )
            rename_keys.append(
                (
                    f"{src_prefix}.transformer_cross_attention_layers.{i}.norm.bias",
                    f"{dst_prefix}.layers.{i}.cross_attn_layer_norm.bias",
                )
            )

            rename_keys.append(
                (f"{src_prefix}.transformer_ffn_layers.{i}.linear1.weight", f"{dst_prefix}.layers.{i}.fc1.weight")
            )
            rename_keys.append(
                (f"{src_prefix}.transformer_ffn_layers.{i}.linear1.bias", f"{dst_prefix}.layers.{i}.fc1.bias")
            )
            rename_keys.append(
                (f"{src_prefix}.transformer_ffn_layers.{i}.linear2.weight", f"{dst_prefix}.layers.{i}.fc2.weight")
            )
            rename_keys.append(
                (f"{src_prefix}.transformer_ffn_layers.{i}.linear2.bias", f"{dst_prefix}.layers.{i}.fc2.bias")
            )
            rename_keys.append(
                (
                    f"{src_prefix}.transformer_ffn_layers.{i}.norm.weight",
                    f"{dst_prefix}.layers.{i}.final_layer_norm.weight",
                )
            )
            rename_keys.append(
                (
                    f"{src_prefix}.transformer_ffn_layers.{i}.norm.bias",
                    f"{dst_prefix}.layers.{i}.final_layer_norm.bias",
                )
            )

        return rename_keys

    def replace_masked_attention_decoder(self, dst_state_dict: StateDict, src_state_dict: StateDict):
        dst_prefix: str = "transformer_module.decoder"
        src_prefix: str = "sem_seg_head.predictor"

        renamed_keys = self.rename_keys_in_masked_attention_decoder(dst_state_dict, src_state_dict)

        
        renamed_keys.extend(
            [
                (f"{src_prefix}.decoder_norm.weight", f"{dst_prefix}.layernorm.weight"),
                (f"{src_prefix}.decoder_norm.bias", f"{dst_prefix}.layernorm.bias"),
            ]
        )

        mlp_len = 3
        for i in range(mlp_len):
            renamed_keys.extend(
                [
                    (
                        f"{src_prefix}.mask_embed.layers.{i}.weight",
                        f"{dst_prefix}.mask_predictor.mask_embedder.{i}.0.weight",
                    ),
                    (
                        f"{src_prefix}.mask_embed.layers.{i}.bias",
                        f"{dst_prefix}.mask_predictor.mask_embedder.{i}.0.bias",
                    ),
                ]
            )

        self.pop_all(renamed_keys, dst_state_dict, src_state_dict)

    def replace_keys_qkv_transformer_decoder(self, dst_state_dict: StateDict, src_state_dict: StateDict):
        dst_prefix: str = "transformer_module.decoder.layers"
        src_prefix: str = "sem_seg_head.predictor"
        for i in range(self.config.decoder_layers - 1):
            
            in_proj_weight = src_state_dict.pop(
                f"{src_prefix}.transformer_self_attention_layers.{i}.self_attn.in_proj_weight"
            )
            in_proj_bias = src_state_dict.pop(
                f"{src_prefix}.transformer_self_attention_layers.{i}.self_attn.in_proj_bias"
            )
            
            dst_state_dict[f"{dst_prefix}.{i}.self_attn.q_proj.weight"] = in_proj_weight[:256, :]
            dst_state_dict[f"{dst_prefix}.{i}.self_attn.q_proj.bias"] = in_proj_bias[:256]
            dst_state_dict[f"{dst_prefix}.{i}.self_attn.k_proj.weight"] = in_proj_weight[256:512, :]
            dst_state_dict[f"{dst_prefix}.{i}.self_attn.k_proj.bias"] = in_proj_bias[256:512]
            dst_state_dict[f"{dst_prefix}.{i}.self_attn.v_proj.weight"] = in_proj_weight[-256:, :]
            dst_state_dict[f"{dst_prefix}.{i}.self_attn.v_proj.bias"] = in_proj_bias[-256:]

    def replace_transformer_module(self, dst_state_dict: StateDict, src_state_dict: StateDict):
        dst_prefix: str = "transformer_module"
        src_prefix: str = "sem_seg_head.predictor"

        self.replace_masked_attention_decoder(dst_state_dict, src_state_dict)

        renamed_keys = [
            (f"{src_prefix}.query_embed.weight", f"{dst_prefix}.queries_embedder.weight"),
            (f"{src_prefix}.query_feat.weight", f"{dst_prefix}.queries_features.weight"),
            (f"{src_prefix}.level_embed.weight", f"{dst_prefix}.level_embed.weight"),
        ]

        self.pop_all(renamed_keys, dst_state_dict, src_state_dict)
        self.replace_keys_qkv_transformer_decoder(dst_state_dict, src_state_dict)

    def replace_universal_segmentation_module(self, dst_state_dict: StateDict, src_state_dict: StateDict):
        dst_prefix: str = ""
        src_prefix: str = "sem_seg_head.predictor"

        renamed_keys = [
            (f"{src_prefix}.class_embed.weight", f"{dst_prefix}class_predictor.weight"),
            (f"{src_prefix}.class_embed.bias", f"{dst_prefix}class_predictor.bias"),
        ]

        logger.info(f"Replacing keys {pformat(renamed_keys)}")
        self.pop_all(renamed_keys, dst_state_dict, src_state_dict)

    def convert(self, mask2former: Mask2FormerModel) -> Mask2FormerModel:
        dst_state_dict = TrackedStateDict(mask2former.state_dict())
        src_state_dict = self.original_model.state_dict()

        self.replace_pixel_module(dst_state_dict, src_state_dict)
        self.replace_transformer_module(dst_state_dict, src_state_dict)

        logger.info(f"Missed keys are {pformat(dst_state_dict.diff())}")
        logger.info(f"Not copied keys are {pformat(src_state_dict.keys())}")
        logger.info("🙌 Done")

        state_dict = {key: dst_state_dict[key] for key in dst_state_dict.to_track.keys()}
        mask2former.load_state_dict(state_dict)
        return mask2former

    def convert_universal_segmentation(
        self, mask2former: Mask2FormerForUniversalSegmentation
    ) -> Mask2FormerForUniversalSegmentation:
        dst_state_dict = TrackedStateDict(mask2former.state_dict())
        src_state_dict = self.original_model.state_dict()

        self.replace_universal_segmentation_module(dst_state_dict, src_state_dict)

        state_dict = {key: dst_state_dict[key] for key in dst_state_dict.to_track.keys()}
        mask2former.load_state_dict(state_dict)

        return mask2former

    @staticmethod
    def using_dirs(checkpoints_dir: Path, config_dir: Path) -> Iterator[Tuple[object, Path, Path]]:
        checkpoints: List[Path] = checkpoints_dir.glob("**/*.pkl")

        for checkpoint in checkpoints:
            logger.info(f"💪 Converting {checkpoint.stem}")
            

            
            dataset_name = checkpoint.parents[2].stem
            if dataset_name == "ade":
                dataset_name = dataset_name.replace("ade", "ade20k")

            
            segmentation_task = checkpoint.parents[1].stem

            
            config_file_name = f"{checkpoint.parents[0].stem}.yaml"

            config: Path = config_dir / dataset_name / segmentation_task / "swin" / config_file_name
            yield config, checkpoint


def test(
    original_model,
    our_model: Mask2FormerForUniversalSegmentation,
    image_processor: Mask2FormerImageProcessor,
    tolerance: float,
):
    with torch.no_grad():
        original_model = original_model.eval()
        our_model = our_model.eval()

        im = prepare_img()
        x = image_processor(images=im, return_tensors="pt")["pixel_values"]

        original_model_backbone_features = original_model.backbone(x.clone())
        our_model_output: Mask2FormerModelOutput = our_model.model(x.clone(), output_hidden_states=True)

        
        for original_model_feature, our_model_feature in zip(
            original_model_backbone_features.values(), our_model_output.encoder_hidden_states
        ):
            assert torch.allclose(
                original_model_feature, our_model_feature, atol=tolerance
            ), "The backbone features are not the same."

        
        mask_features, _, multi_scale_features = original_model.sem_seg_head.pixel_decoder.forward_features(
            original_model_backbone_features
        )

        for original_model_feature, our_model_feature in zip(
            multi_scale_features, our_model_output.pixel_decoder_hidden_states
        ):
            assert torch.allclose(
                original_model_feature, our_model_feature, atol=tolerance
            ), "The pixel decoder feature are not the same"

        
        tr_complete = T.Compose(
            [T.Resize((384, 384)), T.ToTensor()],
        )
        y = (tr_complete(im) * 255.0).to(torch.int).float()

        
        original_class_logits, original_mask_logits = original_model([{"image": y.clone().squeeze(0)}])

        our_model_out: Mask2FormerForUniversalSegmentationOutput = our_model(x.clone())
        our_mask_logits = our_model_out.masks_queries_logits
        our_class_logits = our_model_out.class_queries_logits

        assert original_mask_logits.shape == our_mask_logits.shape, "Output masks shapes are not matching."
        assert original_class_logits.shape == our_class_logits.shape, "Output class logits shapes are not matching."
        assert torch.allclose(
            original_class_logits, our_class_logits, atol=tolerance
        ), "The class logits are not the same."
        assert torch.allclose(
            original_mask_logits, our_mask_logits, atol=tolerance
        ), "The predicted masks are not the same."

        logger.info("✅ Test passed!")


def get_model_name(checkpoint_file: Path):
    
    model_name_raw: str = checkpoint_file.parents[0].stem

    
    segmentation_task_name: str = checkpoint_file.parents[1].stem
    if segmentation_task_name not in ["instance-segmentation", "panoptic-segmentation", "semantic-segmentation"]:
        raise ValueError(
            f"{segmentation_task_name} must be wrong since acceptable values are: instance-segmentation,"
            " panoptic-segmentation, semantic-segmentation."
        )

    
    dataset_name: str = checkpoint_file.parents[2].stem
    if dataset_name not in ["coco", "ade", "cityscapes", "mapillary-vistas"]:
        raise ValueError(
            f"{dataset_name} must be wrong since we didn't find 'coco' or 'ade' or 'cityscapes' or 'mapillary-vistas'"
            " in it "
        )

    backbone = "swin"
    backbone_types = ["tiny", "small", "base_IN21k", "base", "large"]
    backbone_type = list(filter(lambda x: x in model_name_raw, backbone_types))[0].replace("_", "-")

    model_name = f"mask2former-{backbone}-{backbone_type}-{dataset_name}-{segmentation_task_name.split('-')[0]}"

    return model_name


if __name__ == "__main__":
    parser = ArgumentParser(
        description="Command line to convert the original mask2formers (with swin backbone) to our implementations."
    )

    parser.add_argument(
        "--checkpoints_dir",
        type=Path,
        help=(
            "A directory containing the model's checkpoints. The directory has to have the following structure:"
            " <DIR_NAME>/<DATASET_NAME>/<SEGMENTATION_TASK_NAME>/<CONFIG_NAME>.pkl"
        ),
    )
    parser.add_argument(
        "--configs_dir",
        type=Path,
        help=(
            "A directory containing the model's configs, see detectron2 doc. The directory has to have the following"
            " structure: <DIR_NAME>/<DATASET_NAME>/<SEGMENTATION_TASK_NAME>/<CONFIG_NAME>.yaml"
        ),
    )
    parser.add_argument(
        "--mask2former_dir",
        required=True,
        type=Path,
        help=(
            "A path to Mask2Former's original implementation directory. You can download from here:"
            " https://github.com/facebookresearch/Mask2Former"
        ),
    )

    args = parser.parse_args()

    checkpoints_dir: Path = args.checkpoints_dir
    config_dir: Path = args.configs_dir
    mask2former_dir: Path = args.mask2former_dir
    
    sys.path.append(str(mask2former_dir.parent))
    
    from Mask2Former.mask2former.config import add_maskformer2_config
    from Mask2Former.mask2former.maskformer_model import MaskFormer as OriginalMask2Former

    for config_file, checkpoint_file in OriginalMask2FormerCheckpointToOursConverter.using_dirs(
        checkpoints_dir, config_dir
    ):
        model_name = get_model_name(checkpoint_file)
        image_processor = OriginalMask2FormerConfigToImageProcessorConverter()(
            setup_cfg(Args(config_file=config_file))
        )
        image_processor.size = {"height": 384, "width": 384}

        original_config = setup_cfg(Args(config_file=config_file))
        mask2former_kwargs = OriginalMask2Former.from_config(original_config)
        original_model = OriginalMask2Former(**mask2former_kwargs).eval()

        DetectionCheckpointer(original_model).load(str(checkpoint_file))

        config: Mask2FormerConfig = OriginalMask2FormerConfigToOursConverter()(original_config)
        mask2former = Mask2FormerModel(config=config).eval()

        converter = OriginalMask2FormerCheckpointToOursConverter(original_model, config)
        mask2former = converter.convert(mask2former)

        mask2former_for_segmentation = Mask2FormerForUniversalSegmentation(config=config).eval()
        mask2former_for_segmentation.model = mask2former

        mask2former_for_segmentation = converter.convert_universal_segmentation(mask2former_for_segmentation)

        tolerance = 3e-1
        high_tolerance_models = [
            "mask2former-swin-base-IN21k-coco-instance",
            "mask2former-swin-base-coco-instance",
            "mask2former-swin-small-cityscapes-semantic",
        ]

        if model_name in high_tolerance_models:
            tolerance = 3e-1

        logger.info(f"🪄 Testing {model_name}...")
        test(original_model, mask2former_for_segmentation, image_processor, tolerance)
        logger.info(f"🪄 Pushing {model_name} to hub...")

        image_processor.push_to_hub(model_name)
        mask2former_for_segmentation.push_to_hub(model_name)