codingutf8 2022 meta platforms inc s 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 pytorch maskformer model import math from dataclasses import dataclass from numbers import number from typing import dict list optional tuple import numpy as np import torch from torch import tensor nn from import autobackbone from activations import act2fn from modelingattnmaskutils import prepare4dattentionmask from modelingoutputs import basemodeloutputwithcrossattentions from modelingutils import pretrainedmodel from utils import modeloutput addstartdocstrings addstartdocstringstomodelforward isscipyavailable logging replacereturndocstrings requiresbackends from detr import detrconfig from configurationmaskformer import maskformerconfig from configurationmaskformerswin import maskformerswinconfig if isscipyavailable from scipy optimize import linearsumassignment logger logging getloggername configfordoc maskformerconfig checkpointfordoc facebookmaskformerswinbaseade maskformerpretrainedmodelarchivelist facebookmaskformerswinbaseade see all maskformer models at https huggingface comodels filtermaskformer dataclass copied from transformers models detr modelingdetr detrdecoderoutput class detrdecoderoutputbasemodeloutputwithcrossattentions intermediatehiddenstates optionaltorch floattensor none dataclass class maskformerpixellevelmoduleoutputmodeloutput encoderlasthiddenstate optionaltorch floattensor none decoderlasthiddenstate optionaltorch floattensor none encoderhiddenstates optionaltupletorch floattensor none decoderhiddenstates optionaltupletorch floattensor none dataclass class maskformerpixeldecoderoutputmodeloutput lasthiddenstate torch floattensor none hiddenstates optionaltupletorch floattensor none attentions optionaltupletorch floattensor none dataclass class maskformermodeloutputmodeloutput encoderlasthiddenstate optionaltorch floattensor none pixeldecoderlasthiddenstate optionaltorch floattensor none transformerdecoderlasthiddenstate optionaltorch floattensor none encoderhiddenstates optionaltupletorch floattensor none pixeldecoderhiddenstates optionaltupletorch floattensor none transformerdecoderhiddenstates optionaltupletorch floattensor none hiddenstates optionaltupletorch floattensor none attentions optionaltupletorch floattensor none dataclass class maskformerforinstancesegmentationoutputmodeloutput loss optionaltorch floattensor none classquerieslogits torch floattensor none masksquerieslogits torch floattensor none auxiliarylogits torch floattensor none encoderlasthiddenstate optionaltorch floattensor none pixeldecoderlasthiddenstate optionaltorch floattensor none transformerdecoderlasthiddenstate optionaltorch floattensor none encoderhiddenstates optionaltupletorch floattensor none pixeldecoderhiddenstates optionaltupletorch floattensor none transformerdecoderhiddenstates optionaltupletorch floattensor none hiddenstates optionaltupletorch floattensor none attentions optionaltupletorch floattensor none def upsamplelikepixelvalues tensor like tensor mode str bilinear tensor height width like shape upsampled nn functional interpolatepixelvalues sizeheight width modemode aligncornersfalse return upsampled refactored from original implementation def dicelossinputs tensor labels tensor nummasks int tensor r compute the dice loss similar to generalized iou for masks as follows mathcalltextdicex y 1 frac2 x cap y x cup y 1 in practice since labels is a binary mask only 0s and 1s dice can be computed as follow mathcalltextdicex y 1 frac2 x y x y 1 args inputs torch tensor a tensor representing a mask labels torch tensor a tensor with the same shape as inputs stores the binary classification labels for each element in inputs 0 for the negative class and 1 for the positive class nummasks int the number of masks present in the current batch used for normalization returns torch tensor the computed loss refactored from original implementation criterion nn bcewithlogitslossreductionnone probs inputs sigmoid crossentropyloss criterioninputs labels pt probs labels 1 probs 1 labels loss crossentropyloss 1 pt gamma if alpha 0 alphat alpha labels 1 alpha 1 labels loss alphat loss loss loss mean1 sum nummasks return loss refactored from original implementation def pairwisedicelossinputs tensor labels tensor tensor inputs inputs sigmoid flatten1 numerator 2 torch matmulinputs labels t using broadcasting to get a numqueries numclasses matrix denominator inputs sum1 none labels sum1none loss 1 numerator 1 denominator 1 return loss refactored from original implementation def pairwisesigmoidfocallossinputs tensor labels tensor alpha float 0 25 gamma float 2 0 tensor r a pair wise version of the focal loss see sigmoidfocalloss for usage args inputs torch tensor a tensor representing a mask labels torch tensor a tensor with the same shape as inputs stores the binary classification labels for each element in inputs 0 for the negative class and 1 for the positive class alpha float optional defaults to 0 25 weighting factor in range 0 1 to balance positive vs negative examples gamma float optional defaults to 2 0 exponent of the modulating factor 1 pt to balance easy vs hard examples returns torch tensor the computed loss between each pairs copied from transformers models detr modelingdetr detrattention multiheaded attention from attention is all you need paper here we add position embeddings to the queries and keys as explained in the detr paper input shape batch x time x channel positionembeddings kwargs poppositionebmeddings none keyvaluepositionembeddings kwargs popkeyvaluepositionembeddings none if kwargs raise valueerrorfunexpected arguments kwargs keys if positionembeddings is not none and objectqueries is not none raise valueerror cannot specify both positionembeddings and objectqueries please use just objectqueries if keyvaluepositionembeddings is not none and spatialpositionembeddings is not none raise valueerror cannot specify both keyvaluepositionembeddings and spatialpositionembeddings please use just spatialpositionembeddings if positionembeddings is not none logger warningonce positionembeddings has been deprecated and will be removed in v4 34 please use objectqueries instead objectqueries positionembeddings if keyvaluepositionembeddings is not none logger warningonce keyvaluepositionembeddings has been deprecated and will be removed in v4 34 please use spatialpositionembeddings instead spatialpositionembeddings keyvaluepositionembeddings if keyvaluestates are provided this layer is used as a crossattention layer for the decoder iscrossattention keyvaluestates is not none batchsize targetlen embeddim hiddenstates size add position embeddings to the hidden states before projecting to queries and keys if objectqueries is not none hiddenstatesoriginal hiddenstates hiddenstates self withposembedhiddenstates objectqueries add keyvalue position embeddings to the key value states if spatialpositionembeddings is not none keyvaluestatesoriginal keyvaluestates keyvaluestates self withposembedkeyvaluestates spatialpositionembeddings get query proj querystates self qprojhiddenstates self scaling get key value proj if iscrossattention crossattentions keystates self shapeself kprojkeyvaluestates 1 batchsize valuestates self shapeself vprojkeyvaluestatesoriginal 1 batchsize else selfattention keystates self shapeself kprojhiddenstates 1 batchsize valuestates self shapeself vprojhiddenstatesoriginal 1 batchsize projshape batchsize self numheads 1 self headdim querystates self shapequerystates targetlen batchsize viewprojshape keystates keystates viewprojshape valuestates valuestates viewprojshape sourcelen keystates size1 attnweights torch bmmquerystates keystates transpose1 2 if attnweights size batchsize self numheads targetlen sourcelen raise valueerror fattention weights should be of size batchsize self numheads targetlen sourcelen but is f attnweights size if attentionmask is not none if attentionmask size batchsize 1 targetlen sourcelen raise valueerror fattention mask should be of size batchsize 1 targetlen sourcelen but is f attentionmask size attnweights attnweights viewbatchsize self numheads targetlen sourcelen attentionmask attnweights attnweights viewbatchsize self numheads targetlen sourcelen attnweights nn functional softmaxattnweights dim1 if outputattentions this operation is a bit awkward but it s required to make sure that attnweights keeps its gradient in order to do so attnweights have to reshaped twice and have to be reused in the following attnweightsreshaped attnweights viewbatchsize self numheads targetlen sourcelen attnweights attnweightsreshaped viewbatchsize self numheads targetlen sourcelen else attnweightsreshaped none attnprobs nn functional dropoutattnweights pself dropout trainingself training attnoutput torch bmmattnprobs valuestates if attnoutput size batchsize self numheads targetlen self headdim raise valueerror fattnoutput should be of size batchsize self numheads targetlen self headdim but is f attnoutput size attnoutput attnoutput viewbatchsize self numheads targetlen self headdim attnoutput attnoutput transpose1 2 attnoutput attnoutput reshapebatchsize targetlen embeddim attnoutput self outprojattnoutput return attnoutput attnweightsreshaped copied from transformers models detr modelingdetr detrdecoderlayer class detrdecoderlayernn module def initself config detrconfig super init self embeddim config dmodel self selfattn detrattention embeddimself embeddim numheadsconfig decoderattentionheads dropoutconfig attentiondropout self dropout config dropout self activationfn act2fnconfig activationfunction self activationdropout config activationdropout self selfattnlayernorm nn layernormself embeddim self encoderattn detrattention self embeddim config decoderattentionheads dropoutconfig attentiondropout self encoderattnlayernorm nn layernormself embeddim self fc1 nn linearself embeddim config decoderffndim self fc2 nn linearconfig decoderffndim self embeddim self finallayernorm nn layernormself embeddim def forward self hiddenstates torch tensor attentionmask optionaltorch tensor none objectqueries optionaltorch tensor none querypositionembeddings optionaltorch tensor none encoderhiddenstates optionaltorch tensor none encoderattentionmask optionaltorch tensor none outputattentions optionalbool false kwargs positionembeddings kwargs poppositionembeddings none if kwargs raise valueerrorfunexpected arguments kwargs keys if positionembeddings is not none and objectqueries is not none raise valueerror cannot specify both positionembeddings and objectqueries please use just objectqueries if positionembeddings is not none logger warningonce positionembeddings has been deprecated and will be removed in v4 34 please use objectqueries instead objectqueries positionembeddings residual hiddenstates self attention hiddenstates selfattnweights self selfattn hiddenstateshiddenstates objectqueriesquerypositionembeddings attentionmaskattentionmask outputattentionsoutputattentions hiddenstates nn functional dropouthiddenstates pself dropout trainingself training hiddenstates residual hiddenstates hiddenstates self selfattnlayernormhiddenstates crossattention block crossattnweights none if encoderhiddenstates is not none residual hiddenstates hiddenstates crossattnweights self encoderattn hiddenstateshiddenstates objectqueriesquerypositionembeddings keyvaluestatesencoderhiddenstates attentionmaskencoderattentionmask spatialpositionembeddingsobjectqueries outputattentionsoutputattentions hiddenstates nn functional dropouthiddenstates pself dropout trainingself training hiddenstates residual hiddenstates hiddenstates self encoderattnlayernormhiddenstates fully connected residual hiddenstates hiddenstates self activationfnself fc1hiddenstates hiddenstates nn functional dropouthiddenstates pself activationdropout trainingself training hiddenstates self fc2hiddenstates hiddenstates nn functional dropouthiddenstates pself dropout trainingself training hiddenstates residual hiddenstates hiddenstates self finallayernormhiddenstates outputs hiddenstates if outputattentions outputs selfattnweights crossattnweights return outputs class detrdecodernn module def initself config detrconfig super init self config config self dropout config dropout self layerdrop config decoderlayerdrop self layers nn modulelistdetrdecoderlayerconfig for in rangeconfig decoderlayers in detr the decoder uses layernorm after the last decoder layer output self layernorm nn layernormconfig dmodel self gradientcheckpointing false def forward self inputsembedsnone attentionmasknone encoderhiddenstatesnone encoderattentionmasknone objectqueriesnone querypositionembeddingsnone outputattentionsnone outputhiddenstatesnone returndictnone kwargs r args inputsembeds torch floattensor of shape batchsize sequencelength hiddensize the query embeddings that are passed into the decoder attentionmask torch tensor of shape batchsize sequencelength optional mask to avoid performing attention on certain queries mask values selected in 0 1 1 for queries that are not masked 0 for queries that are masked what are attention masks glossaryattentionmask encoderhiddenstates torch floattensor of shape batchsize encodersequencelength hiddensize optional sequence of hiddenstates at the output of the last layer of the encoder used in the crossattention of the decoder encoderattentionmask torch longtensor of shape batchsize encodersequencelength optional mask to avoid performing crossattention on padding pixelvalues of the encoder mask values selected in 0 1 1 for pixels that are real i e not masked 0 for pixels that are padding i e masked objectqueries torch floattensor of shape batchsize sequencelength hiddensize optional position embeddings that are added to the queries and keys in each crossattention layer querypositionembeddings torch floattensor of shape batchsize numqueries hiddensize optional position embeddings that are added to the queries and keys in each selfattention layer outputattentions bool optional whether or not to return the attentions tensors of all attention layers see attentions under returned tensors for more detail outputhiddenstates bool optional whether or not to return the hidden states of all layers see hiddenstates under returned tensors for more detail returndict bool optional whether or not to return a utils modeloutput instead of a plain tuple expand encoder attention mask bsz seqlen bsz 1 tgtseqlen srcseqlen optional intermediate hidden states decoder layers add layerdrop see https arxiv orgabs1909 11556 for description finally apply layernorm add hidden states from the last decoder layer stack intermediate decoder activations refactored from original implementation this class computes an assignment between the labels and the predictions of the network for efficiency reasons the labels don t include the noobject because of this in general there are more predictions than labels in this case we do a 1to1 matching of the best predictions while the others are unmatched and thus treated as nonobjects creates the matcher params costclass float optional defaults to 1 0 this is the relative weight of the classification error in the matching cost costmask float optional defaults to 1 0 this is the relative weight of the focal loss of the binary mask in the matching cost costdice float optional defaults to 1 0 this is the relative weight of the dice loss of the binary mask in the matching cost performs the matching params masksquerieslogits torch tensor a tensor of dim batchsize numqueries numlabels with the classification logits classquerieslogits torch tensor a tensor of dim batchsize numqueries height width with the predicted masks classlabels torch tensor a tensor of dim numtargetboxes where numtargetboxes is the number of groundtruth objects in the target containing the class labels masklabels torch tensor a tensor of dim numtargetboxes height width containing the target masks returns listtupletensor a list of size batchsize containing tuples of indexi indexj where indexi is the indices of the selected predictions in order indexj is the indices of the corresponding selected labels in order for each batch element it holds lenindexi lenindexj minnumqueries numtargetboxes iterate through batch size downsample the target mask save memory compute the classification cost contrary to the loss we don t use the nll but approximate it in 1 probatarget class the 1 is a constant that doesn t change the matching it can be ommitted flatten spatial dimension q h w q h w same for targetmask c h w c h w compute the focal loss between each mask pairs shape numqueries numlabels compute the dice loss betwen each mask pairs shape numqueries numlabels final cost matrix do the assigmented using the hungarian algorithm in scipy it could be stacked in one tensor copied and adapted from original implementation the maskformer loss the loss is computed very similar to detr the process happens in two steps 1 we compute hungarian assignment between ground truth masks and the outputs of the model 2 we supervise each pair of matched groundtruth prediction supervise class and mask args numlabels int the number of classes matcher maskformerhungarianmatcher a torch module that computes the assigments between the predictions and labels weightdict dictstr float a dictionary of weights to be applied to the different losses eoscoef float weight to apply to the null class get the maximum size in the batch compute finel size get metadata pad the tensors to the size of the biggest one compute the losses related to the labels using cross entropy args classquerieslogits torch tensor a tensor of shape batchsize numqueries numlabels classlabels listtorch tensor list of class labels of shape labels indices tuplenp array the indices computed by the hungarian matcher returns dictstr tensor a dict of torch tensor containing the following key losscrossentropy the loss computed using cross entropy on the predicted and ground truth labels shape batchsize numqueries shape batchsize numqueries targetclasses is a batchsize numlabels numqueries we need to permute predlogits b q c b c q compute the losses related to the masks using focal and dice loss args masksquerieslogits torch tensor a tensor of shape batchsize numqueries height width masklabels torch tensor list of mask labels of shape labels height width indices tuplenp array the indices computed by the hungarian matcher nummasks int the number of masks used for normalization returns dictstr tensor a dict of torch tensor containing two keys lossmask the loss computed using sigmoid focal loss on the predicted and ground truth masks lossdice the loss computed using dice loss on the predicted on the predicted and ground truth masks shape batchsize numqueries height width shape batchsize numqueries height width pad all and stack the targets to the numlabels dimension upsample predictions to the target size we have to add one dim to use interpolate permute predictions following indices permute labels following indices this performs the loss computation args masksquerieslogits torch tensor a tensor of shape batchsize numqueries height width classquerieslogits torch tensor a tensor of shape batchsize numqueries numlabels masklabels torch tensor list of mask labels of shape labels height width classlabels listtorch tensor list of class labels of shape labels auxiliarypredictions dictstr torch tensor optional if useauxiliaryloss was set to true in maskformerconfig then it contains the logits from the inner layers of the detr s decoder returns dictstr tensor a dict of torch tensor containing two keys losscrossentropy the loss computed using cross entropy on the predicted and ground truth labels lossmask the loss computed using sigmoid focal loss on the predicted and ground truth masks lossdice the loss computed using dice loss on the predicted on the predicted and ground truth masks if useauxiliaryloss was set to true in maskformerconfig the dictionary contains addional losses for each auxiliary predictions retrieve the matching between the outputs of the last layer and the labels compute the average number of target masks for normalization purposes get all the losses in case of auxiliary losses we repeat this process with the output of each intermediate layer computes the average number of target masks across the batch for normalization purposes a basic module that executes conv norm in sequence used in maskformer args infeatures int the number of input features channels outfeatures int the number of outputs features channels provide backwards compatibility from when the class inherited from nn sequential in nn sequential subclasses the name given to the layer is its index in the sequence in nn module subclasses they derived from the instance attribute they are assigned to e g self mylayername layer we can t give instance attributes integer names i e self 0 is not permitted and so need to register explicitly a feature pyramid network layer fpn layer it creates a feature map by aggregating features from the previous and backbone layer due to the spatial mismatch the tensor coming from the previous layer is upsampled args infeatures int the number of input features channels lateralfeatures int the number of lateral features channels feature pyramid network given an input tensor and a set of feature map of different featurespatial size it creates a list of feature maps with the same feature size args infeatures int the number of input features channels lateralwidths listint a list with the features channels size of each lateral connection featuresize int optional defaults to 256 the features channels of the resulting feature maps super init self fpn maskformerfpnmodelargs featuresizefeaturesize kwargs self maskprojection nn conv2dfeaturesize maskfeaturesize kernelsize3 padding1 def forward self features listtensor outputhiddenstates bool false returndict bool true maskformerpixeldecoderoutput fpnfeatures self fpnfeatures we use the last feature map lastfeatureprojected self maskprojectionfpnfeatures1 if not returndict return lastfeatureprojected tuplefpnfeatures if outputhiddenstates else lastfeatureprojected return maskformerpixeldecoderoutput lasthiddenstatelastfeatureprojected hiddenstatestuplefpnfeatures if outputhiddenstates else copied and adapted from original implementation also practically equal to detrsinepositionembedding class maskformersinepositionembeddingnn module def init self numposfeats int 64 temperature int 10000 normalize bool false scale optionalfloat none super init if scale is not none and normalize is false raise valueerrornormalize should be true if scale is passed self numposfeats numposfeats self temperature temperature self normalize normalize self scale 2 math pi if scale is none else scale def forwardself x tensor mask optionaltensor none tensor if mask is none mask torch zerosx size0 x size2 x size3 devicex device dtypetorch bool notmask mask tox dtype yembed notmask cumsum1 xembed notmask cumsum2 if self normalize eps 1e6 yembed yembed yembed 1 eps self scale xembed xembed xembed 1 eps self scale dimt torch arangeself numposfeats dtypex dtype devicex device dimt self temperature 2 torch divdimt 2 roundingmodefloor self numposfeats posx xembed none dimt posy yembed none dimt posx torch stackposx 0 2 sin posx 1 2 cos dim4 flatten3 posy torch stackposy 0 2 sin posy 1 2 cos dim4 flatten3 pos torch catposy posx dim3 permute0 3 1 2 return pos class predictionblocknn module def initself indim int outdim int activation nn module none super init self layers nn linearindim outdim activation maintain submodule indexing as if part of a sequential block for i layer in enumerateself layers self addmodulestri layer def forwardself input tensor tensor hiddenstate input for layer in self layers hiddenstate layerhiddenstate return hiddenstate class maskformermlppredictionheadnn module def initself inputdim int hiddendim int outputdim int numlayers int 3 super init indims inputdim hiddendim numlayers 1 outdims hiddendim numlayers 1 outputdim self layers for i indim outdim in enumeratezipindims outdims activation nn relu if i numlayers 1 else nn identity layer predictionblockindim outdim activationactivation self layers appendlayer provide backwards compatibility from when the class inherited from nn sequential in nn sequential subclasses the name given to the layer is its index in the sequence in nn module subclasses they derived from the instance attribute they are assigned to e g self mylayername layer we can t give instance attributes integer names i e self 0 is not permitted and so need to register explicitly self addmodulestri layer def forwardself input tensor tensor hiddenstate input for layer in self layers hiddenstate layerhiddenstate return hiddenstate class maskformerpixellevelmodulenn module def initself config maskformerconfig super init todd add method to load pretrained weights of backbone backboneconfig config backboneconfig if backboneconfig modeltype swin for backwards compatibility backboneconfig maskformerswinconfig fromdictbackboneconfig todict backboneconfig outfeatures stage1 stage2 stage3 stage4 self encoder autobackbone fromconfigbackboneconfig featurechannels self encoder channels self decoder maskformerpixeldecoder infeaturesfeaturechannels1 featuresizeconfig fpnfeaturesize maskfeaturesizeconfig maskfeaturesize lateralwidthsfeaturechannels 1 def forward self pixelvalues tensor outputhiddenstates bool false returndict bool true maskformerpixellevelmoduleoutput features self encoderpixelvalues featuremaps decoderoutput self decoderfeatures outputhiddenstates returndictreturndict if not returndict lasthiddenstate decoderoutput0 outputs features1 lasthiddenstate if outputhiddenstates hiddenstates decoderoutput1 outputs outputs tuplefeatures hiddenstates return outputs return maskformerpixellevelmoduleoutput the last feature is actually the output from the last layer encoderlasthiddenstatefeatures1 decoderlasthiddenstatedecoderoutput lasthiddenstate encoderhiddenstatestuplefeatures if outputhiddenstates else decoderhiddenstatesdecoderoutput hiddenstates if outputhiddenstates else class maskformertransformermodulenn module def initself infeatures int config maskformerconfig super init hiddensize config decoderconfig hiddensize shouldproject infeatures hiddensize self positionembedder maskformersinepositionembeddingnumposfeatshiddensize 2 normalizetrue self queriesembedder nn embeddingconfig decoderconfig numqueries hiddensize self inputprojection nn conv2dinfeatures hiddensize kernelsize1 if shouldproject else none self decoder detrdecoderconfigconfig decoderconfig def forward self imagefeatures tensor outputhiddenstates bool false outputattentions bool false returndict optionalbool none detrdecoderoutput if self inputprojection is not none imagefeatures self inputprojectionimagefeatures objectqueries self positionembedderimagefeatures repeat the queries q c b q c batchsize imagefeatures shape0 queriesembeddings self queriesembedder weight unsqueeze0 repeatbatchsize 1 1 inputsembeds torch zeroslikequeriesembeddings requiresgradtrue batchsize numchannels height width imagefeatures shape rearrange both imagefeatures and objectqueries b c h w b h w c imagefeatures imagefeatures viewbatchsize numchannels height width permute0 2 1 objectqueries objectqueries viewbatchsize numchannels height width permute0 2 1 decoderoutput detrdecoderoutput self decoder inputsembedsinputsembeds attentionmasknone encoderhiddenstatesimagefeatures encoderattentionmasknone objectqueriesobjectqueries querypositionembeddingsqueriesembeddings outputattentionsoutputattentions outputhiddenstatesoutputhiddenstates returndictreturndict return decoderoutput maskformerstartdocstring r this model is a pytorch torch nn modulehttps pytorch orgdocsstablenn htmltorch nn module subclass use it as a regular pytorch module and refer to the pytorch documentation for all matter related to general usage and behavior parameters config maskformerconfig model configuration class with all the parameters of the model initializing with a config file does not load the weights associated with the model only the configuration check out the pretrainedmodel frompretrained method to load the model weights class maskformerpretrainedmodelpretrainedmodel configclass maskformerconfig basemodelprefix model maininputname pixelvalues def initweightsself module nn module xavierstd self config initxavierstd std self config initstd if isinstancemodule maskformertransformermodule if module inputprojection is not none nn init xavieruniformmodule inputprojection weight gainxavierstd nn init constantmodule inputprojection bias 0 fpn elif isinstancemodule maskformerfpnmodel nn init xavieruniformmodule stem getsubmodule0 weight gainxavierstd elif isinstancemodule maskformerfpnlayer nn init xavieruniformmodule proj0 weight gainxavierstd elif isinstancemodule maskformerfpnconvlayer nn init xavieruniformmodule getsubmodule0 weight gainxavierstd the mlp head elif isinstancemodule maskformermlppredictionhead i was not able to find the correct initializer in the original implementation we ll use xavier for submodule in module modules if isinstancesubmodule nn linear nn init xavieruniformsubmodule weight gainxavierstd nn init constantsubmodule bias 0 elif isinstancemodule nn layernorm module bias data zero module weight data fill1 0 copied from detr if isinstancemodule nn linear nn conv2d nn batchnorm2d slightly different from the tf version which uses truncatednormal for initialization cf https github compytorchpytorchpull5617 module weight data normalmean0 0 stdstd if module bias is not none module bias data zero elif isinstancemodule nn embedding module weight data normalmean0 0 stdstd if module paddingidx is not none module weight datamodule paddingidx zero addstartdocstrings the bare maskformer model outputting raw hiddenstates without any specific head on top maskformerstartdocstring class maskformermodelmaskformerpretrainedmodel def initself config maskformerconfig super initconfig self pixellevelmodule maskformerpixellevelmoduleconfig self transformermodule maskformertransformermodule infeaturesself pixellevelmodule encoder channels1 configconfig self postinit addstartdocstringstomodelforwardmaskformerinputsdocstring replacereturndocstringsoutputtypemaskformermodeloutput configclassconfigfordoc def forward self pixelvalues tensor pixelmask optionaltensor none outputhiddenstates optionalbool none outputattentions optionalbool none returndict optionalbool none maskformermodeloutput r returns examples python from transformers import autoimageprocessor maskformermodel from pil import image import requests load maskformer finetuned on ade20k semantic segmentation imageprocessor autoimageprocessor frompretrainedfacebookmaskformerswinbaseade model maskformermodel frompretrainedfacebookmaskformerswinbaseade url http images cocodataset orgval2017000000039769 jpg image image openrequests geturl streamtrue raw inputs imageprocessorimage returntensorspt forward pass outputs modelinputs the decoder of maskformer outputs hidden states of shape batchsize numqueries hiddensize transformerdecoderlasthiddenstate outputs transformerdecoderlasthiddenstate listtransformerdecoderlasthiddenstate shape 1 100 256 if pixelvalues is none raise valueerroryou have to specify pixelvalues outputattentions outputattentions if outputattentions is not none else self config outputattentions outputhiddenstates outputhiddenstates if outputhiddenstates is not none else self config outputhiddenstates returndict returndict if returndict is not none else self config usereturndict batchsize height width pixelvalues shape if pixelmask is none pixelmask torch onesbatchsize height width devicepixelvalues device pixellevelmoduleoutput self pixellevelmodule pixelvalues outputhiddenstates returndictreturndict imagefeatures pixellevelmoduleoutput0 pixelembeddings pixellevelmoduleoutput1 transformermoduleoutput self transformermoduleimagefeatures outputhiddenstates outputattentions queries transformermoduleoutput lasthiddenstate encoderhiddenstates none pixeldecoderhiddenstates none transformerdecoderhiddenstates none hiddenstates none if outputhiddenstates encoderhiddenstates pixellevelmoduleoutput2 pixeldecoderhiddenstates pixellevelmoduleoutput3 transformerdecoderhiddenstates transformermoduleoutput1 hiddenstates encoderhiddenstates pixeldecoderhiddenstates transformerdecoderhiddenstates output maskformermodeloutput encoderlasthiddenstateimagefeatures pixeldecoderlasthiddenstatepixelembeddings transformerdecoderlasthiddenstatequeries encoderhiddenstatesencoderhiddenstates pixeldecoderhiddenstatespixeldecoderhiddenstates transformerdecoderhiddenstatestransformerdecoderhiddenstates hiddenstateshiddenstates attentionstransformermoduleoutput attentions if not returndict output tuplev for v in output values return output class maskformerforinstancesegmentationmaskformerpretrainedmodel def initself config maskformerconfig super initconfig self model maskformermodelconfig hiddensize config decoderconfig hiddensize 1 because we add the null class self classpredictor nn linearhiddensize config numlabels 1 self maskembedder maskformermlppredictionheadhiddensize hiddensize config maskfeaturesize self matcher maskformerhungarianmatcher costclass1 0 costdiceconfig diceweight costmaskconfig maskweight self weightdict dictstr float losscrossentropy config crossentropyweight lossmask config maskweight lossdice config diceweight self criterion maskformerloss config numlabels matcherself matcher weightdictself weightdict eoscoefconfig noobjectweight self postinit def getlossdict self masksquerieslogits tensor classquerieslogits tensor masklabels tensor classlabels tensor auxiliarylogits dictstr tensor dictstr tensor lossdict dictstr tensor self criterion masksquerieslogits classquerieslogits masklabels classlabels auxiliarylogits weight each loss by self weightdictlossname including auxiliary losses for key weight in self weightdict items for losskey loss in lossdict items if key in losskey loss weight return lossdict def getlossself lossdict dictstr tensor tensor return sumlossdict values def getlogitsself outputs maskformermodeloutput tupletensor tensor dictstr tensor pixelembeddings outputs pixeldecoderlasthiddenstate get the auxiliary predictions one for each decoder s layer auxiliarylogits liststr tensor this code is a little bit cumbersome an improvement can be to return a list of predictions if we have auxiliary loss then we are going to return more than one element in the list if self config useauxiliaryloss stackedtransformerdecoderoutputs torch stackoutputs transformerdecoderhiddenstates classes self classpredictorstackedtransformerdecoderoutputs classquerieslogits classes1 get the masks maskembeddings self maskembedderstackedtransformerdecoderoutputs equivalent to einsum lbqc bchw lbqhw but jit friendly numembeddings batchsize numqueries numchannels maskembeddings shape height width pixelembeddings shape binariesmasks torch zeros numembeddings batchsize numqueries height width devicemaskembeddings device for c in rangenumchannels binariesmasks maskembeddings c none none pixelembeddingsnone none c masksquerieslogits binariesmasks1 go til 1 because the last one is always used for auxbinarymasks auxclasses in zipbinariesmasks 1 classes 1 auxiliarylogits append masksquerieslogits auxbinarymasks classquerieslogits auxclasses else transformerdecoderhiddenstates outputs transformerdecoderlasthiddenstate classes self classpredictortransformerdecoderhiddenstates classquerieslogits classes get the masks maskembeddings self maskembeddertransformerdecoderhiddenstates sum up over the channels equivalent to einsum bqc bchw bqhw but jit friendly batchsize numqueries numchannels maskembeddings shape height width pixelembeddings shape masksquerieslogits torch zerosbatchsize numqueries height width devicemaskembeddings device for c in rangenumchannels masksquerieslogits maskembeddings c none none pixelembeddings none c return classquerieslogits masksquerieslogits auxiliarylogits addstartdocstringstomodelforwardmaskformerinputsdocstring replacereturndocstringsoutputtypemaskformerforinstancesegmentationoutput configclassconfigfordoc def forward self pixelvalues tensor masklabels optionallisttensor none classlabels optionallisttensor none pixelmask optionaltensor none outputauxiliarylogits optionalbool none outputhiddenstates optionalbool none outputattentions optionalbool none returndict optionalbool none maskformerforinstancesegmentationoutput r masklabels listtorch tensor optional list of mask labels of shape numlabels height width to be fed to a model classlabels listtorch longtensor optional list of target class labels of shape numlabels height width to be fed to a model they identify the labels of masklabels e g the label of masklabelsij if classlabelsij returns examples semantic segmentation example python from transformers import autoimageprocessor maskformerforinstancesegmentation from pil import image import requests load maskformer finetuned on ade20k semantic segmentation imageprocessor autoimageprocessor frompretrainedfacebookmaskformerswinbaseade model maskformerforinstancesegmentation frompretrainedfacebookmaskformerswinbaseade url https huggingface codatasetshfinternaltestingfixturesade20kresolvemainadeval00000001 jpg image image openrequests geturl streamtrue raw inputs imageprocessorimagesimage returntensorspt outputs modelinputs model predicts classquerieslogits of shape batchsize numqueries and masksquerieslogits of shape batchsize numqueries height width classquerieslogits outputs classquerieslogits masksquerieslogits outputs masksquerieslogits you can pass them to imageprocessor for postprocessing predictedsemanticmap imageprocessor postprocesssemanticsegmentation outputs targetsizesimage size 1 0 we refer to the demo notebooks for visualization see resources section in the maskformer docs listpredictedsemanticmap shape 512 683 panoptic segmentation example python from transformers import autoimageprocessor maskformerforinstancesegmentation from pil import image import requests load maskformer finetuned on coco panoptic segmentation imageprocessor autoimageprocessor frompretrainedfacebookmaskformerswinbasecoco model maskformerforinstancesegmentation frompretrainedfacebookmaskformerswinbasecoco url http images cocodataset orgval2017000000039769 jpg image image openrequests geturl streamtrue raw inputs imageprocessorimagesimage returntensorspt outputs modelinputs model predicts classquerieslogits of shape batchsize numqueries and masksquerieslogits of shape batchsize numqueries height width classquerieslogits outputs classquerieslogits masksquerieslogits outputs masksquerieslogits you can pass them to imageprocessor for postprocessing result imageprocessor postprocesspanopticsegmentationoutputs targetsizesimage size 10 we refer to the demo notebooks for visualization see resources section in the maskformer docs predictedpanopticmap resultsegmentation listpredictedpanopticmap shape 480 640 we need to have rawoutputs optionally be returned as a dict to use torch compile for backwards compatibility we convert to a dataclass for the rest of the model logic coding utf 8 2022 meta platforms inc s 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 pytorch maskformer model see all maskformer models at https huggingface co models filter maskformer copied from transformers models detr modeling_detr detrdecoderoutput base class for outputs of the detr decoder this class adds one attribute to basemodeloutputwithcrossattentions namely an optional stack of intermediate decoder activations i e the output of each decoder layer each of them gone through a layernorm this is useful when training the model with auxiliary decoding losses args last_hidden_state torch floattensor of shape batch_size sequence_length hidden_size sequence of hidden states at the output of the last layer of the model hidden_states tuple torch floattensor optional returned when output_hidden_states true is passed or when config output_hidden_states true tuple of torch floattensor one for the output of the embeddings one for the output of each layer of shape batch_size sequence_length hidden_size hidden states of the model at the output of each layer plus the initial embedding outputs attentions tuple torch floattensor optional returned when output_attentions true is passed or when config output_attentions true tuple of torch floattensor one for each layer of shape batch_size num_heads sequence_length sequence_length attentions weights after the attention softmax used to compute the weighted average in the self attention heads cross_attentions tuple torch floattensor optional returned when output_attentions true and config add_cross_attention true is passed or when config output_attentions true tuple of torch floattensor one for each layer of shape batch_size num_heads sequence_length sequence_length attentions weights of the decoder s cross attention layer after the attention softmax used to compute the weighted average in the cross attention heads intermediate_hidden_states torch floattensor of shape config decoder_layers batch_size num_queries hidden_size optional returned when config auxiliary_loss true intermediate decoder activations i e the output of each decoder layer each of them gone through a layernorm maskformer s pixel level module output it returns both the last and optionally the hidden states from the encoder and decoder by default the encoder is a maskformerswin transformer and the decoder is a feature pyramid network fpn the encoder_last_hidden_state are referred on the paper as images features while decoder_last_hidden_state as pixel embeddings args encoder_last_hidden_state torch floattensor of shape batch_size num_channels height width last hidden states final feature map of the last stage of the encoder encoder_hidden_states tuple torch floattensor optional returned when output_hidden_states true is passed or when config output_hidden_states true tuple of torch floattensor one for the output of the embeddings one for the output of each stage of shape batch_size num_channels height width hidden states also called feature maps of the model at the output of each stage decoder_last_hidden_state torch floattensor of shape batch_size num_channels height width last hidden states final feature map of the last stage of the decoder decoder_hidden_states tuple torch floattensor optional returned when output_hidden_states true is passed or when config output_hidden_states true tuple of torch floattensor one for the output of the embeddings one for the output of each stage of shape batch_size num_channels height width hidden states also called feature maps of the model at the output of each stage maskformer s pixel decoder module output practically a feature pyramid network it returns the last hidden state and optionally the hidden states args last_hidden_state torch floattensor of shape batch_size num_channels height width last hidden states final feature map of the last stage of the model hidden_states tuple torch floattensor optional returned when output_hidden_states true is passed or when config output_hidden_states true tuple of torch floattensor one for the output of the embeddings one for the output of each layer of shape batch_size num_channels height width hidden states of the model at the output of each layer plus the initial embedding outputs attentions tuple torch floattensor optional returned when output_attentions true is passed or when config output_attentions true tuple of torch floattensor one for each layer of shape batch_size num_heads sequence_length sequence_length attentions weights from detr s decoder after the attention softmax used to compute the weighted average in the self attention heads class for outputs of maskformermodel this class returns all the needed hidden states to compute the logits args encoder_last_hidden_state torch floattensor of shape batch_size num_channels height width last hidden states final feature map of the last stage of the encoder model backbone pixel_decoder_last_hidden_state torch floattensor of shape batch_size num_channels height width last hidden states final feature map of the last stage of the pixel decoder model fpn transformer_decoder_last_hidden_state torch floattensor of shape batch_size sequence_length hidden_size last hidden states final feature map of the last stage of the transformer decoder model encoder_hidden_states tuple torch floattensor optional returned when output_hidden_states true is passed or when config output_hidden_states true tuple of torch floattensor one for the output of the embeddings one for the output of each stage of shape batch_size num_channels height width hidden states also called feature maps of the encoder model at the output of each stage pixel_decoder_hidden_states tuple torch floattensor optional returned when output_hidden_states true is passed or when config output_hidden_states true tuple of torch floattensor one for the output of the embeddings one for the output of each stage of shape batch_size num_channels height width hidden states also called feature maps of the pixel decoder model at the output of each stage transformer_decoder_hidden_states tuple torch floattensor optional returned when output_hidden_states true is passed or when config output_hidden_states true tuple of torch floattensor one for the output of the embeddings one for the output of each stage of shape batch_size sequence_length hidden_size hidden states also called feature maps of the transformer decoder at the output of each stage hidden_states tuple torch floattensor optional returned when output_hidden_states true is passed or when config output_hidden_states true tuple of torch floattensor containing encoder_hidden_states pixel_decoder_hidden_states and decoder_hidden_states attentions tuple torch floattensor optional returned when output_attentions true is passed or when config output_attentions true tuple of torch floattensor one for each layer of shape batch_size num_heads sequence_length sequence_length attentions weights from detr s decoder after the attention softmax used to compute the weighted average in the self attention heads class for outputs of maskformerforinstancesegmentation this output can be directly passed to maskformerimageprocessor post_process_semantic_segmentation or or maskformerimageprocessor post_process_instance_segmentation or maskformerimageprocessor post_process_panoptic_segmentation depending on the task please see maskformerimageprocessor for details regarding usage args loss torch tensor optional the computed loss returned when labels are present class_queries_logits torch floattensor a tensor of shape batch_size num_queries num_labels 1 representing the proposed classes for each query note the 1 is needed because we incorporate the null class masks_queries_logits torch floattensor a tensor of shape batch_size num_queries height width representing the proposed masks for each query encoder_last_hidden_state torch floattensor of shape batch_size num_channels height width last hidden states final feature map of the last stage of the encoder model backbone pixel_decoder_last_hidden_state torch floattensor of shape batch_size num_channels height width last hidden states final feature map of the last stage of the pixel decoder model fpn transformer_decoder_last_hidden_state torch floattensor of shape batch_size sequence_length hidden_size last hidden states final feature map of the last stage of the transformer decoder model encoder_hidden_states tuple torch floattensor optional returned when output_hidden_states true is passed or when config output_hidden_states true tuple of torch floattensor one for the output of the embeddings one for the output of each stage of shape batch_size num_channels height width hidden states also called feature maps of the encoder model at the output of each stage pixel_decoder_hidden_states tuple torch floattensor optional returned when output_hidden_states true is passed or when config output_hidden_states true tuple of torch floattensor one for the output of the embeddings one for the output of each stage of shape batch_size num_channels height width hidden states also called feature maps of the pixel decoder model at the output of each stage transformer_decoder_hidden_states tuple torch floattensor optional returned when output_hidden_states true is passed or when config output_hidden_states true tuple of torch floattensor one for the output of the embeddings one for the output of each stage of shape batch_size sequence_length hidden_size hidden states of the transformer decoder at the output of each stage hidden_states tuple torch floattensor optional returned when output_hidden_states true is passed or when config output_hidden_states true tuple of torch floattensor containing encoder_hidden_states pixel_decoder_hidden_states and decoder_hidden_states attentions tuple torch floattensor optional returned when output_attentions true is passed or when config output_attentions true tuple of torch floattensor one for each layer of shape batch_size num_heads sequence_length sequence_length attentions weights from detr s decoder after the attention softmax used to compute the weighted average in the self attention heads an utility function that upsamples pixel_values to match the dimension of like args pixel_values torch tensor the tensor we wish to upsample like torch tensor the tensor we wish to use as size target mode str optional defaults to bilinear the interpolation mode returns torch tensor the upsampled tensor refactored from original implementation compute the dice loss similar to generalized iou for masks as follows mathcal l _ text dice x y 1 frac 2 x cap y x cup y 1 in practice since labels is a binary mask only 0s and 1s dice can be computed as follow mathcal l _ text dice x y 1 frac 2 x y x y 1 args inputs torch tensor a tensor representing a mask labels torch tensor a tensor with the same shape as inputs stores the binary classification labels for each element in inputs 0 for the negative class and 1 for the positive class num_masks int the number of masks present in the current batch used for normalization returns torch tensor the computed loss refactored from original implementation focal loss proposed in focal loss for dense object detection https arxiv org abs 1708 02002 originally used in retinanet the loss is computed as follows mathcal l _ text focal loss 1 p_t gamma log p_t where ce p_t log p_t ce is the standard cross entropy loss please refer to equation 1 2 3 of the paper for a better understanding args inputs torch tensor a float tensor of arbitrary shape labels torch tensor a tensor with the same shape as inputs stores the binary classification labels for each element in inputs 0 for the negative class and 1 for the positive class num_masks int the number of masks present in the current batch used for normalization alpha float optional defaults to 0 25 weighting factor in range 0 1 to balance positive vs negative examples gamma float optional defaults to 2 0 exponent of the modulating factor 1 p_t to balance easy vs hard examples returns torch tensor the computed loss refactored from original implementation a pair wise version of the dice loss see dice_loss for usage args inputs torch tensor a tensor representing a mask labels torch tensor a tensor with the same shape as inputs stores the binary classification labels for each element in inputs 0 for the negative class and 1 for the positive class returns torch tensor the computed loss between each pairs using broadcasting to get a num_queries num_classes matrix refactored from original implementation a pair wise version of the focal loss see sigmoid_focal_loss for usage args inputs torch tensor a tensor representing a mask labels torch tensor a tensor with the same shape as inputs stores the binary classification labels for each element in inputs 0 for the negative class and 1 for the positive class alpha float optional defaults to 0 25 weighting factor in range 0 1 to balance positive vs negative examples gamma float optional defaults to 2 0 exponent of the modulating factor 1 p_t to balance easy vs hard examples returns torch tensor the computed loss between each pairs copied from transformers models detr modeling_detr detrattention multi headed attention from attention is all you need paper here we add position embeddings to the queries and keys as explained in the detr paper input shape batch x time x channel if key_value_states are provided this layer is used as a cross attention layer for the decoder add position embeddings to the hidden states before projecting to queries and keys add key value position embeddings to the key value states get query proj get key value proj cross_attentions self_attention this operation is a bit awkward but it s required to make sure that attn_weights keeps its gradient in order to do so attn_weights have to reshaped twice and have to be reused in the following copied from transformers models detr modeling_detr detrdecoderlayer args hidden_states torch floattensor input to the layer of shape batch seq_len embed_dim attention_mask torch floattensor attention mask of size batch 1 target_len source_len where padding elements are indicated by very large negative values object_queries torch floattensor optional object_queries that are added to the hidden states in the cross attention layer query_position_embeddings torch floattensor optional position embeddings that are added to the queries and keys in the self attention layer encoder_hidden_states torch floattensor cross attention input to the layer of shape batch seq_len embed_dim encoder_attention_mask torch floattensor encoder attention mask of size batch 1 target_len source_len where padding elements are indicated by very large negative values output_attentions bool optional whether or not to return the attentions tensors of all attention layers see attentions under returned tensors for more detail self attention cross attention block fully connected transformer decoder consisting of config decoder_layers layers each layer is a detrdecoderlayer the decoder updates the query embeddings through multiple self attention and cross attention layers some small tweaks for detr object_queries and query_position_embeddings are added to the forward pass if self config auxiliary_loss is set to true also returns a stack of activations from all decoding layers args config detrconfig in detr the decoder uses layernorm after the last decoder layer output args inputs_embeds torch floattensor of shape batch_size sequence_length hidden_size the query embeddings that are passed into the decoder attention_mask torch tensor of shape batch_size sequence_length optional mask to avoid performing attention on certain queries mask values selected in 0 1 1 for queries that are not masked 0 for queries that are masked what are attention masks glossary attention mask encoder_hidden_states torch floattensor of shape batch_size encoder_sequence_length hidden_size optional sequence of hidden states at the output of the last layer of the encoder used in the cross attention of the decoder encoder_attention_mask torch longtensor of shape batch_size encoder_sequence_length optional mask to avoid performing cross attention on padding pixel_values of the encoder mask values selected in 0 1 1 for pixels that are real i e not masked 0 for pixels that are padding i e masked object_queries torch floattensor of shape batch_size sequence_length hidden_size optional position embeddings that are added to the queries and keys in each cross attention layer query_position_embeddings torch floattensor of shape batch_size num_queries hidden_size optional position embeddings that are added to the queries and keys in each self attention layer output_attentions bool optional whether or not to return the attentions tensors of all attention layers see attentions under returned tensors for more detail output_hidden_states bool optional whether or not to return the hidden states of all layers see hidden_states under returned tensors for more detail return_dict bool optional whether or not to return a utils modeloutput instead of a plain tuple expand encoder attention mask bsz seq_len bsz 1 tgt_seq_len src_seq_len optional intermediate hidden states decoder layers add layerdrop see https arxiv org abs 1909 11556 for description finally apply layernorm add hidden states from the last decoder layer stack intermediate decoder activations refactored from original implementation this class computes an assignment between the labels and the predictions of the network for efficiency reasons the labels don t include the no_object because of this in general there are more predictions than labels in this case we do a 1 to 1 matching of the best predictions while the others are un matched and thus treated as non objects creates the matcher params cost_class float optional defaults to 1 0 this is the relative weight of the classification error in the matching cost cost_mask float optional defaults to 1 0 this is the relative weight of the focal loss of the binary mask in the matching cost cost_dice float optional defaults to 1 0 this is the relative weight of the dice loss of the binary mask in the matching cost performs the matching params masks_queries_logits torch tensor a tensor of dim batch_size num_queries num_labels with the classification logits class_queries_logits torch tensor a tensor of dim batch_size num_queries height width with the predicted masks class_labels torch tensor a tensor of dim num_target_boxes where num_target_boxes is the number of ground truth objects in the target containing the class labels mask_labels torch tensor a tensor of dim num_target_boxes height width containing the target masks returns list tuple tensor a list of size batch_size containing tuples of index_i index_j where index_i is the indices of the selected predictions in order index_j is the indices of the corresponding selected labels in order for each batch element it holds len index_i len index_j min num_queries num_target_boxes iterate through batch size downsample the target mask save memory compute the classification cost contrary to the loss we don t use the nll but approximate it in 1 proba target class the 1 is a constant that doesn t change the matching it can be ommitted flatten spatial dimension q h w q h w num_queries height width same for target_mask c h w c h w num_total_labels height width compute the focal loss between each mask pairs shape num_queries num_labels compute the dice loss betwen each mask pairs shape num_queries num_labels final cost matrix do the assigmented using the hungarian algorithm in scipy it could be stacked in one tensor copied and adapted from original implementation the maskformer loss the loss is computed very similar to detr the process happens in two steps 1 we compute hungarian assignment between ground truth masks and the outputs of the model 2 we supervise each pair of matched ground truth prediction supervise class and mask args num_labels int the number of classes matcher maskformerhungarianmatcher a torch module that computes the assigments between the predictions and labels weight_dict dict str float a dictionary of weights to be applied to the different losses eos_coef float weight to apply to the null class get the maximum size in the batch compute finel size get metadata pad the tensors to the size of the biggest one compute the losses related to the labels using cross entropy args class_queries_logits torch tensor a tensor of shape batch_size num_queries num_labels class_labels list torch tensor list of class labels of shape labels indices tuple np array the indices computed by the hungarian matcher returns dict str tensor a dict of torch tensor containing the following key loss_cross_entropy the loss computed using cross entropy on the predicted and ground truth labels shape batch_size num_queries shape batch_size num_queries target_classes is a batch_size num_labels num_queries we need to permute pred_logits b q c b c q compute the losses related to the masks using focal and dice loss args masks_queries_logits torch tensor a tensor of shape batch_size num_queries height width mask_labels torch tensor list of mask labels of shape labels height width indices tuple np array the indices computed by the hungarian matcher num_masks int the number of masks used for normalization returns dict str tensor a dict of torch tensor containing two keys loss_mask the loss computed using sigmoid focal loss on the predicted and ground truth masks loss_dice the loss computed using dice loss on the predicted on the predicted and ground truth masks shape batch_size num_queries height width shape batch_size num_queries height width pad all and stack the targets to the num_labels dimension upsample predictions to the target size we have to add one dim to use interpolate permute predictions following indices permute labels following indices this performs the loss computation args masks_queries_logits torch tensor a tensor of shape batch_size num_queries height width class_queries_logits torch tensor a tensor of shape batch_size num_queries num_labels mask_labels torch tensor list of mask labels of shape labels height width class_labels list torch tensor list of class labels of shape labels auxiliary_predictions dict str torch tensor optional if use_auxiliary_loss was set to true in maskformerconfig then it contains the logits from the inner layers of the detr s decoder returns dict str tensor a dict of torch tensor containing two keys loss_cross_entropy the loss computed using cross entropy on the predicted and ground truth labels loss_mask the loss computed using sigmoid focal loss on the predicted and ground truth masks loss_dice the loss computed using dice loss on the predicted on the predicted and ground truth masks if use_auxiliary_loss was set to true in maskformerconfig the dictionary contains addional losses for each auxiliary predictions retrieve the matching between the outputs of the last layer and the labels compute the average number of target masks for normalization purposes get all the losses in case of auxiliary losses we repeat this process with the output of each intermediate layer computes the average number of target masks across the batch for normalization purposes a basic module that executes conv norm in sequence used in maskformer args in_features int the number of input features channels out_features int the number of outputs features channels provide backwards compatibility from when the class inherited from nn sequential in nn sequential subclasses the name given to the layer is its index in the sequence in nn module subclasses they derived from the instance attribute they are assigned to e g self my_layer_name layer we can t give instance attributes integer names i e self 0 is not permitted and so need to register explicitly a feature pyramid network layer fpn layer it creates a feature map by aggregating features from the previous and backbone layer due to the spatial mismatch the tensor coming from the previous layer is upsampled args in_features int the number of input features channels lateral_features int the number of lateral features channels feature pyramid network given an input tensor and a set of feature map of different feature spatial size it creates a list of feature maps with the same feature size args in_features int the number of input features channels lateral_widths list int a list with the features channels size of each lateral connection feature_size int optional defaults to 256 the features channels of the resulting feature maps pixel decoder module proposed in per pixel classification is not all you need for semantic segmentation https arxiv org abs 2107 06278 it first runs the backbone s features into a feature pyramid network creating a list of feature maps then it projects the last one to the correct mask_size args feature_size int optional defaults to 256 the feature size channel dimension of the fpn feature maps mask_feature_size int optional defaults to 256 the features channels of the target masks size c_ epsilon in the paper we use the last feature map copied and adapted from original implementation also practically equal to detrsinepositionembedding this is a more standard version of the position embedding very similar to the one used by the attention is all you need paper generalized to work on images maintain submodule indexing as if part of a sequential block a classic multi layer perceptron mlp args input_dim int the input dimensions hidden_dim int the hidden dimensions output_dim int the output dimensions num_layers int optional defaults to 3 the number of layers provide backwards compatibility from when the class inherited from nn sequential in nn sequential subclasses the name given to the layer is its index in the sequence in nn module subclasses they derived from the instance attribute they are assigned to e g self my_layer_name layer we can t give instance attributes integer names i e self 0 is not permitted and so need to register explicitly pixel level module proposed in per pixel classification is not all you need for semantic segmentation https arxiv org abs 2107 06278 it runs the input image through a backbone and a pixel decoder generating an image feature map and pixel embeddings args config maskformerconfig the configuration used to instantiate this model todd add method to load pretrained weights of backbone for backwards compatibility the last feature is actually the output from the last layer the maskformer s transformer module repeat the queries q c b q c rearrange both image_features and object_queries b c h w b h w c this model is a pytorch torch nn module https pytorch org docs stable nn html torch nn module sub class use it as a regular pytorch module and refer to the pytorch documentation for all matter related to general usage and behavior parameters config maskformerconfig model configuration class with all the parameters of the model initializing with a config file does not load the weights associated with the model only the configuration check out the pretrainedmodel from_pretrained method to load the model weights args pixel_values torch floattensor of shape batch_size num_channels height width pixel values pixel values can be obtained using autoimageprocessor see maskformerimageprocessor __call__ for details pixel_mask torch longtensor of shape batch_size height width optional mask to avoid performing attention on padding pixel values mask values selected in 0 1 1 for pixels that are real i e not masked 0 for pixels that are padding i e masked what are attention masks glossary attention mask output_hidden_states bool optional whether or not to return the hidden states of all layers see hidden_states under returned tensors for more detail output_attentions bool optional whether or not to return the attentions tensors of detr s decoder attention layers return_dict bool optional whether or not to return a maskformermodeloutput instead of a plain tuple fpn the mlp head i was not able to find the correct initializer in the original implementation we ll use xavier copied from detr slightly different from the tf version which uses truncated_normal for initialization cf https github com pytorch pytorch pull 5617 returns examples python from transformers import autoimageprocessor maskformermodel from pil import image import requests load maskformer fine tuned on ade20k semantic segmentation image_processor autoimageprocessor from_pretrained facebook maskformer swin base ade model maskformermodel from_pretrained facebook maskformer swin base ade url http images cocodataset org val2017 000000039769 jpg image image open requests get url stream true raw inputs image_processor image return_tensors pt forward pass outputs model inputs the decoder of maskformer outputs hidden states of shape batch_size num_queries hidden_size transformer_decoder_last_hidden_state outputs transformer_decoder_last_hidden_state list transformer_decoder_last_hidden_state shape 1 100 256 1 because we add the null class weight each loss by self weight_dict loss_name including auxiliary losses get the auxiliary predictions one for each decoder s layer this code is a little bit cumbersome an improvement can be to return a list of predictions if we have auxiliary loss then we are going to return more than one element in the list get the masks equivalent to einsum lbqc bchw lbqhw but jit friendly go til 1 because the last one is always used get the masks sum up over the channels equivalent to einsum bqc bchw bqhw but jit friendly mask_labels list torch tensor optional list of mask labels of shape num_labels height width to be fed to a model class_labels list torch longtensor optional list of target class labels of shape num_labels height width to be fed to a model they identify the labels of mask_labels e g the label of mask_labels i j if class_labels i j returns examples semantic segmentation example python from transformers import autoimageprocessor maskformerforinstancesegmentation from pil import image import requests load maskformer fine tuned on ade20k semantic segmentation image_processor autoimageprocessor from_pretrained facebook maskformer swin base ade model maskformerforinstancesegmentation from_pretrained facebook maskformer swin base ade url https huggingface co datasets hf internal testing fixtures_ade20k resolve main ade_val_00000001 jpg image image open requests get url stream true raw inputs image_processor images image return_tensors pt outputs model inputs model predicts class_queries_logits of shape batch_size num_queries and masks_queries_logits of shape batch_size num_queries height width class_queries_logits outputs class_queries_logits masks_queries_logits outputs masks_queries_logits you can pass them to image_processor for postprocessing predicted_semantic_map image_processor post_process_semantic_segmentation outputs target_sizes image size 1 0 we refer to the demo notebooks for visualization see resources section in the maskformer docs list predicted_semantic_map shape 512 683 panoptic segmentation example python from transformers import autoimageprocessor maskformerforinstancesegmentation from pil import image import requests load maskformer fine tuned on coco panoptic segmentation image_processor autoimageprocessor from_pretrained facebook maskformer swin base coco model maskformerforinstancesegmentation from_pretrained facebook maskformer swin base coco url http images cocodataset org val2017 000000039769 jpg image image open requests get url stream true raw inputs image_processor images image return_tensors pt outputs model inputs model predicts class_queries_logits of shape batch_size num_queries and masks_queries_logits of shape batch_size num_queries height width class_queries_logits outputs class_queries_logits masks_queries_logits outputs masks_queries_logits you can pass them to image_processor for postprocessing result image_processor post_process_panoptic_segmentation outputs target_sizes image size 1 0 we refer to the demo notebooks for visualization see resources section in the maskformer docs predicted_panoptic_map result segmentation list predicted_panoptic_map shape 480 640 we need to have raw_outputs optionally be returned as a dict to use torch compile for backwards compatibility we convert to a dataclass for the rest of the model logic
import math
from dataclasses import dataclass
from numbers import Number
from typing import Dict, List, Optional, Tuple

import numpy as np
import torch
from torch import Tensor, nn

from ... import AutoBackbone
from ...activations import ACT2FN
from ...modeling_attn_mask_utils import _prepare_4d_attention_mask
from ...modeling_outputs import BaseModelOutputWithCrossAttentions
from ...modeling_utils import PreTrainedModel
from ...utils import (
    ModelOutput,
    add_start_docstrings,
    add_start_docstrings_to_model_forward,
    is_scipy_available,
    logging,
    replace_return_docstrings,
    requires_backends,
)
from ..detr import DetrConfig
from .configuration_maskformer import MaskFormerConfig
from .configuration_maskformer_swin import MaskFormerSwinConfig


if is_scipy_available():
    from scipy.optimize import linear_sum_assignment

logger = logging.get_logger(__name__)


_CONFIG_FOR_DOC = "MaskFormerConfig"
_CHECKPOINT_FOR_DOC = "facebook/maskformer-swin-base-ade"

MASKFORMER_PRETRAINED_MODEL_ARCHIVE_LIST = [
    "facebook/maskformer-swin-base-ade",
    
]


@dataclass

class DetrDecoderOutput(BaseModelOutputWithCrossAttentions):
    

    intermediate_hidden_states: Optional[torch.FloatTensor] = None


@dataclass
class MaskFormerPixelLevelModuleOutput(ModelOutput):
    

    encoder_last_hidden_state: Optional[torch.FloatTensor] = None
    decoder_last_hidden_state: Optional[torch.FloatTensor] = None
    encoder_hidden_states: Optional[Tuple[torch.FloatTensor]] = None
    decoder_hidden_states: Optional[Tuple[torch.FloatTensor]] = None


@dataclass
class MaskFormerPixelDecoderOutput(ModelOutput):
    

    last_hidden_state: torch.FloatTensor = None
    hidden_states: Optional[Tuple[torch.FloatTensor]] = None
    attentions: Optional[Tuple[torch.FloatTensor]] = None


@dataclass
class MaskFormerModelOutput(ModelOutput):
    

    encoder_last_hidden_state: Optional[torch.FloatTensor] = None
    pixel_decoder_last_hidden_state: Optional[torch.FloatTensor] = None
    transformer_decoder_last_hidden_state: Optional[torch.FloatTensor] = None
    encoder_hidden_states: Optional[Tuple[torch.FloatTensor]] = None
    pixel_decoder_hidden_states: Optional[Tuple[torch.FloatTensor]] = None
    transformer_decoder_hidden_states: Optional[Tuple[torch.FloatTensor]] = None
    hidden_states: Optional[Tuple[torch.FloatTensor]] = None
    attentions: Optional[Tuple[torch.FloatTensor]] = None


@dataclass
class MaskFormerForInstanceSegmentationOutput(ModelOutput):
    

    loss: Optional[torch.FloatTensor] = None
    class_queries_logits: torch.FloatTensor = None
    masks_queries_logits: torch.FloatTensor = None
    auxiliary_logits: torch.FloatTensor = None
    encoder_last_hidden_state: Optional[torch.FloatTensor] = None
    pixel_decoder_last_hidden_state: Optional[torch.FloatTensor] = None
    transformer_decoder_last_hidden_state: Optional[torch.FloatTensor] = None
    encoder_hidden_states: Optional[Tuple[torch.FloatTensor]] = None
    pixel_decoder_hidden_states: Optional[Tuple[torch.FloatTensor]] = None
    transformer_decoder_hidden_states: Optional[Tuple[torch.FloatTensor]] = None
    hidden_states: Optional[Tuple[torch.FloatTensor]] = None
    attentions: Optional[Tuple[torch.FloatTensor]] = None


def upsample_like(pixel_values: Tensor, like: Tensor, mode: str = "bilinear") -> Tensor:
    
    _, _, height, width = like.shape
    upsampled = nn.functional.interpolate(pixel_values, size=(height, width), mode=mode, align_corners=False)
    return upsampled



def dice_loss(inputs: Tensor, labels: Tensor, num_masks: int) -> Tensor:
    r
    probs = inputs.sigmoid().flatten(1)
    numerator = 2 * (probs * labels).sum(-1)
    denominator = probs.sum(-1) + labels.sum(-1)
    loss = 1 - (numerator + 1) / (denominator + 1)
    loss = loss.sum() / num_masks
    return loss



def sigmoid_focal_loss(
    inputs: Tensor, labels: Tensor, num_masks: int, alpha: float = 0.25, gamma: float = 2
) -> Tensor:
    r
    criterion = nn.BCEWithLogitsLoss(reduction="none")
    probs = inputs.sigmoid()
    cross_entropy_loss = criterion(inputs, labels)
    p_t = probs * labels + (1 - probs) * (1 - labels)
    loss = cross_entropy_loss * ((1 - p_t) ** gamma)

    if alpha >= 0:
        alpha_t = alpha * labels + (1 - alpha) * (1 - labels)
        loss = alpha_t * loss

    loss = loss.mean(1).sum() / num_masks
    return loss



def pair_wise_dice_loss(inputs: Tensor, labels: Tensor) -> Tensor:
    
    inputs = inputs.sigmoid().flatten(1)
    numerator = 2 * torch.matmul(inputs, labels.T)
    
    denominator = inputs.sum(-1)[:, None] + labels.sum(-1)[None, :]
    loss = 1 - (numerator + 1) / (denominator + 1)
    return loss



def pair_wise_sigmoid_focal_loss(inputs: Tensor, labels: Tensor, alpha: float = 0.25, gamma: float = 2.0) -> Tensor:
    r
    if alpha < 0:
        raise ValueError("alpha must be positive")

    height_and_width = inputs.shape[1]

    criterion = nn.BCEWithLogitsLoss(reduction="none")
    prob = inputs.sigmoid()
    cross_entropy_loss_pos = criterion(inputs, torch.ones_like(inputs))
    focal_pos = ((1 - prob) ** gamma) * cross_entropy_loss_pos
    focal_pos *= alpha

    cross_entropy_loss_neg = criterion(inputs, torch.zeros_like(inputs))

    focal_neg = (prob**gamma) * cross_entropy_loss_neg
    focal_neg *= 1 - alpha

    loss = torch.matmul(focal_pos, labels.T) + torch.matmul(focal_neg, (1 - labels).T)

    return loss / height_and_width



class DetrAttention(nn.Module):
    

    def __init__(
        self,
        embed_dim: int,
        num_heads: int,
        dropout: float = 0.0,
        bias: bool = True,
    ):
        super().__init__()
        self.embed_dim = embed_dim
        self.num_heads = num_heads
        self.dropout = dropout
        self.head_dim = embed_dim // num_heads
        if self.head_dim * num_heads != self.embed_dim:
            raise ValueError(
                f"embed_dim must be divisible by num_heads (got `embed_dim`: {self.embed_dim} and `num_heads`:"
                f" {num_heads})."
            )
        self.scaling = self.head_dim**-0.5

        self.k_proj = nn.Linear(embed_dim, embed_dim, bias=bias)
        self.v_proj = nn.Linear(embed_dim, embed_dim, bias=bias)
        self.q_proj = nn.Linear(embed_dim, embed_dim, bias=bias)
        self.out_proj = nn.Linear(embed_dim, embed_dim, bias=bias)

    def _shape(self, tensor: torch.Tensor, seq_len: int, batch_size: int):
        return tensor.view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2).contiguous()

    def with_pos_embed(self, tensor: torch.Tensor, object_queries: Optional[Tensor], **kwargs):
        position_embeddings = kwargs.pop("position_embeddings", None)

        if kwargs:
            raise ValueError(f"Unexpected arguments {kwargs.keys()}")

        if position_embeddings is not None and object_queries is not None:
            raise ValueError(
                "Cannot specify both position_embeddings and object_queries. Please use just object_queries"
            )

        if position_embeddings is not None:
            logger.warning_once(
                "position_embeddings has been deprecated and will be removed in v4.34. Please use object_queries instead"
            )
            object_queries = position_embeddings

        return tensor if object_queries is None else tensor + object_queries

    def forward(
        self,
        hidden_states: torch.Tensor,
        attention_mask: Optional[torch.Tensor] = None,
        object_queries: Optional[torch.Tensor] = None,
        key_value_states: Optional[torch.Tensor] = None,
        spatial_position_embeddings: Optional[torch.Tensor] = None,
        output_attentions: bool = False,
        **kwargs,
    ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
        

        position_embeddings = kwargs.pop("position_ebmeddings", None)
        key_value_position_embeddings = kwargs.pop("key_value_position_embeddings", None)

        if kwargs:
            raise ValueError(f"Unexpected arguments {kwargs.keys()}")

        if position_embeddings is not None and object_queries is not None:
            raise ValueError(
                "Cannot specify both position_embeddings and object_queries. Please use just object_queries"
            )

        if key_value_position_embeddings is not None and spatial_position_embeddings is not None:
            raise ValueError(
                "Cannot specify both key_value_position_embeddings and spatial_position_embeddings. Please use just spatial_position_embeddings"
            )

        if position_embeddings is not None:
            logger.warning_once(
                "position_embeddings has been deprecated and will be removed in v4.34. Please use object_queries instead"
            )
            object_queries = position_embeddings

        if key_value_position_embeddings is not None:
            logger.warning_once(
                "key_value_position_embeddings has been deprecated and will be removed in v4.34. Please use spatial_position_embeddings instead"
            )
            spatial_position_embeddings = key_value_position_embeddings

        
        
        is_cross_attention = key_value_states is not None
        batch_size, target_len, embed_dim = hidden_states.size()

        
        if object_queries is not None:
            hidden_states_original = hidden_states
            hidden_states = self.with_pos_embed(hidden_states, object_queries)

        
        if spatial_position_embeddings is not None:
            key_value_states_original = key_value_states
            key_value_states = self.with_pos_embed(key_value_states, spatial_position_embeddings)

        
        query_states = self.q_proj(hidden_states) * self.scaling
        
        if is_cross_attention:
            
            key_states = self._shape(self.k_proj(key_value_states), -1, batch_size)
            value_states = self._shape(self.v_proj(key_value_states_original), -1, batch_size)
        else:
            
            key_states = self._shape(self.k_proj(hidden_states), -1, batch_size)
            value_states = self._shape(self.v_proj(hidden_states_original), -1, batch_size)

        proj_shape = (batch_size * self.num_heads, -1, self.head_dim)
        query_states = self._shape(query_states, target_len, batch_size).view(*proj_shape)
        key_states = key_states.view(*proj_shape)
        value_states = value_states.view(*proj_shape)

        source_len = key_states.size(1)

        attn_weights = torch.bmm(query_states, key_states.transpose(1, 2))

        if attn_weights.size() != (batch_size * self.num_heads, target_len, source_len):
            raise ValueError(
                f"Attention weights should be of size {(batch_size * self.num_heads, target_len, source_len)}, but is"
                f" {attn_weights.size()}"
            )

        if attention_mask is not None:
            if attention_mask.size() != (batch_size, 1, target_len, source_len):
                raise ValueError(
                    f"Attention mask should be of size {(batch_size, 1, target_len, source_len)}, but is"
                    f" {attention_mask.size()}"
                )
            attn_weights = attn_weights.view(batch_size, self.num_heads, target_len, source_len) + attention_mask
            attn_weights = attn_weights.view(batch_size * self.num_heads, target_len, source_len)

        attn_weights = nn.functional.softmax(attn_weights, dim=-1)

        if output_attentions:
            
            
            
            
            attn_weights_reshaped = attn_weights.view(batch_size, self.num_heads, target_len, source_len)
            attn_weights = attn_weights_reshaped.view(batch_size * self.num_heads, target_len, source_len)
        else:
            attn_weights_reshaped = None

        attn_probs = nn.functional.dropout(attn_weights, p=self.dropout, training=self.training)

        attn_output = torch.bmm(attn_probs, value_states)

        if attn_output.size() != (batch_size * self.num_heads, target_len, self.head_dim):
            raise ValueError(
                f"`attn_output` should be of size {(batch_size, self.num_heads, target_len, self.head_dim)}, but is"
                f" {attn_output.size()}"
            )

        attn_output = attn_output.view(batch_size, self.num_heads, target_len, self.head_dim)
        attn_output = attn_output.transpose(1, 2)
        attn_output = attn_output.reshape(batch_size, target_len, embed_dim)

        attn_output = self.out_proj(attn_output)

        return attn_output, attn_weights_reshaped



class DetrDecoderLayer(nn.Module):
    def __init__(self, config: DetrConfig):
        super().__init__()
        self.embed_dim = config.d_model

        self.self_attn = DetrAttention(
            embed_dim=self.embed_dim,
            num_heads=config.decoder_attention_heads,
            dropout=config.attention_dropout,
        )
        self.dropout = config.dropout
        self.activation_fn = ACT2FN[config.activation_function]
        self.activation_dropout = config.activation_dropout

        self.self_attn_layer_norm = nn.LayerNorm(self.embed_dim)
        self.encoder_attn = DetrAttention(
            self.embed_dim,
            config.decoder_attention_heads,
            dropout=config.attention_dropout,
        )
        self.encoder_attn_layer_norm = nn.LayerNorm(self.embed_dim)
        self.fc1 = nn.Linear(self.embed_dim, config.decoder_ffn_dim)
        self.fc2 = nn.Linear(config.decoder_ffn_dim, self.embed_dim)
        self.final_layer_norm = nn.LayerNorm(self.embed_dim)

    def forward(
        self,
        hidden_states: torch.Tensor,
        attention_mask: Optional[torch.Tensor] = None,
        object_queries: Optional[torch.Tensor] = None,
        query_position_embeddings: Optional[torch.Tensor] = None,
        encoder_hidden_states: Optional[torch.Tensor] = None,
        encoder_attention_mask: Optional[torch.Tensor] = None,
        output_attentions: Optional[bool] = False,
        **kwargs,
    ):
        
        position_embeddings = kwargs.pop("position_embeddings", None)

        if kwargs:
            raise ValueError(f"Unexpected arguments {kwargs.keys()}")

        if position_embeddings is not None and object_queries is not None:
            raise ValueError(
                "Cannot specify both position_embeddings and object_queries. Please use just object_queries"
            )

        if position_embeddings is not None:
            logger.warning_once(
                "position_embeddings has been deprecated and will be removed in v4.34. Please use object_queries instead"
            )
            object_queries = position_embeddings

        residual = hidden_states

        
        hidden_states, self_attn_weights = self.self_attn(
            hidden_states=hidden_states,
            object_queries=query_position_embeddings,
            attention_mask=attention_mask,
            output_attentions=output_attentions,
        )

        hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training)
        hidden_states = residual + hidden_states
        hidden_states = self.self_attn_layer_norm(hidden_states)

        
        cross_attn_weights = None
        if encoder_hidden_states is not None:
            residual = hidden_states

            hidden_states, cross_attn_weights = self.encoder_attn(
                hidden_states=hidden_states,
                object_queries=query_position_embeddings,
                key_value_states=encoder_hidden_states,
                attention_mask=encoder_attention_mask,
                spatial_position_embeddings=object_queries,
                output_attentions=output_attentions,
            )

            hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training)
            hidden_states = residual + hidden_states
            hidden_states = self.encoder_attn_layer_norm(hidden_states)

        
        residual = hidden_states
        hidden_states = self.activation_fn(self.fc1(hidden_states))
        hidden_states = nn.functional.dropout(hidden_states, p=self.activation_dropout, training=self.training)
        hidden_states = self.fc2(hidden_states)
        hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training)
        hidden_states = residual + hidden_states
        hidden_states = self.final_layer_norm(hidden_states)

        outputs = (hidden_states,)

        if output_attentions:
            outputs += (self_attn_weights, cross_attn_weights)

        return outputs


class DetrDecoder(nn.Module):
    

    def __init__(self, config: DetrConfig):
        super().__init__()
        self.config = config
        self.dropout = config.dropout
        self.layerdrop = config.decoder_layerdrop

        self.layers = nn.ModuleList([DetrDecoderLayer(config) for _ in range(config.decoder_layers)])
        
        self.layernorm = nn.LayerNorm(config.d_model)

        self.gradient_checkpointing = False

    def forward(
        self,
        inputs_embeds=None,
        attention_mask=None,
        encoder_hidden_states=None,
        encoder_attention_mask=None,
        object_queries=None,
        query_position_embeddings=None,
        output_attentions=None,
        output_hidden_states=None,
        return_dict=None,
        **kwargs,
    ):
        r
        position_embeddings = kwargs.pop("position_embeddings", None)
        if kwargs:
            raise ValueError(f"Unexpected arguments {kwargs.keys()}")

        if position_embeddings is not None and object_queries is not None:
            raise ValueError(
                "Cannot specify both position_embeddings and object_queries. Please use just object_queries"
            )

        if position_embeddings is not None:
            logger.warning_once(
                "position_embeddings has been deprecated and will be removed in v4.34. Please use object_queries instead"
            )
            object_queries = position_embeddings

        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
        output_hidden_states = (
            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
        )
        return_dict = return_dict if return_dict is not None else self.config.use_return_dict

        if inputs_embeds is not None:
            hidden_states = inputs_embeds
            input_shape = inputs_embeds.size()[:-1]

        
        if encoder_hidden_states is not None and encoder_attention_mask is not None:
            
            encoder_attention_mask = _prepare_4d_attention_mask(
                encoder_attention_mask, inputs_embeds.dtype, tgt_len=input_shape[-1]
            )

        
        intermediate = () if self.config.auxiliary_loss else None

        
        all_hidden_states = () if output_hidden_states else None
        all_self_attns = () if output_attentions else None
        all_cross_attentions = () if (output_attentions and encoder_hidden_states is not None) else None

        for idx, decoder_layer in enumerate(self.layers):
            
            if output_hidden_states:
                all_hidden_states += (hidden_states,)
            if self.training:
                dropout_probability = torch.rand([])
                if dropout_probability < self.layerdrop:
                    continue

            if self.gradient_checkpointing and self.training:
                layer_outputs = self._gradient_checkpointing_func(
                    decoder_layer.__call__,
                    hidden_states,
                    None,
                    encoder_hidden_states,
                    encoder_attention_mask,
                    None,
                    output_attentions,
                )
            else:
                layer_outputs = decoder_layer(
                    hidden_states,
                    attention_mask=None,
                    object_queries=object_queries,
                    query_position_embeddings=query_position_embeddings,
                    encoder_hidden_states=encoder_hidden_states,
                    encoder_attention_mask=encoder_attention_mask,
                    output_attentions=output_attentions,
                )

            hidden_states = layer_outputs[0]

            if self.config.auxiliary_loss:
                hidden_states = self.layernorm(hidden_states)
                intermediate += (hidden_states,)

            if output_attentions:
                all_self_attns += (layer_outputs[1],)

                if encoder_hidden_states is not None:
                    all_cross_attentions += (layer_outputs[2],)

        
        hidden_states = self.layernorm(hidden_states)

        
        if output_hidden_states:
            all_hidden_states += (hidden_states,)

        
        if self.config.auxiliary_loss:
            intermediate = torch.stack(intermediate)

        if not return_dict:
            return tuple(
                v
                for v in [hidden_states, all_hidden_states, all_self_attns, all_cross_attentions, intermediate]
                if v is not None
            )
        return DetrDecoderOutput(
            last_hidden_state=hidden_states,
            hidden_states=all_hidden_states,
            attentions=all_self_attns,
            cross_attentions=all_cross_attentions,
            intermediate_hidden_states=intermediate,
        )



class MaskFormerHungarianMatcher(nn.Module):
    

    def __init__(self, cost_class: float = 1.0, cost_mask: float = 1.0, cost_dice: float = 1.0):
        
        super().__init__()
        if cost_class == 0 and cost_mask == 0 and cost_dice == 0:
            raise ValueError("All costs cant be 0")
        self.cost_class = cost_class
        self.cost_mask = cost_mask
        self.cost_dice = cost_dice

    @torch.no_grad()
    def forward(self, masks_queries_logits, class_queries_logits, mask_labels, class_labels) -> List[Tuple[Tensor]]:
        
        indices: List[Tuple[np.array]] = []

        preds_masks = masks_queries_logits
        preds_probs = class_queries_logits
        
        for pred_probs, pred_mask, target_mask, labels in zip(preds_probs, preds_masks, mask_labels, class_labels):
            
            target_mask = nn.functional.interpolate(target_mask[:, None], size=pred_mask.shape[-2:], mode="nearest")
            pred_probs = pred_probs.softmax(-1)
            
            
            
            cost_class = -pred_probs[:, labels]
            
            pred_mask_flat = pred_mask.flatten(1)  
            
            target_mask_flat = target_mask[:, 0].flatten(1)  
            
            cost_mask = pair_wise_sigmoid_focal_loss(pred_mask_flat, target_mask_flat)
            
            cost_dice = pair_wise_dice_loss(pred_mask_flat, target_mask_flat)
            
            cost_matrix = self.cost_mask * cost_mask + self.cost_class * cost_class + self.cost_dice * cost_dice
            
            assigned_indices: Tuple[np.array] = linear_sum_assignment(cost_matrix.cpu())
            indices.append(assigned_indices)

        
        matched_indices = [
            (torch.as_tensor(i, dtype=torch.int64), torch.as_tensor(j, dtype=torch.int64)) for i, j in indices
        ]
        return matched_indices

    def __repr__(self):
        head = "Matcher " + self.__class__.__name__
        body = [
            f"cost_class: {self.cost_class}",
            f"cost_mask: {self.cost_mask}",
            f"cost_dice: {self.cost_dice}",
        ]
        _repr_indent = 4
        lines = [head] + [" " * _repr_indent + line for line in body]
        return "\n".join(lines)



class MaskFormerLoss(nn.Module):
    def __init__(
        self,
        num_labels: int,
        matcher: MaskFormerHungarianMatcher,
        weight_dict: Dict[str, float],
        eos_coef: float,
    ):
        

        super().__init__()
        requires_backends(self, ["scipy"])
        self.num_labels = num_labels
        self.matcher = matcher
        self.weight_dict = weight_dict
        self.eos_coef = eos_coef
        empty_weight = torch.ones(self.num_labels + 1)
        empty_weight[-1] = self.eos_coef
        self.register_buffer("empty_weight", empty_weight)

    def _max_by_axis(self, the_list: List[List[int]]) -> List[int]:
        maxes = the_list[0]
        for sublist in the_list[1:]:
            for index, item in enumerate(sublist):
                maxes[index] = max(maxes[index], item)
        return maxes

    def _pad_images_to_max_in_batch(self, tensors: List[Tensor]) -> Tuple[Tensor, Tensor]:
        
        max_size = self._max_by_axis([list(tensor.shape) for tensor in tensors])
        batch_size = len(tensors)
        
        batch_shape = [batch_size] + max_size
        b, _, h, w = batch_shape
        
        dtype = tensors[0].dtype
        device = tensors[0].device
        padded_tensors = torch.zeros(batch_shape, dtype=dtype, device=device)
        padding_masks = torch.ones((b, h, w), dtype=torch.bool, device=device)
        
        for tensor, padded_tensor, padding_mask in zip(tensors, padded_tensors, padding_masks):
            padded_tensor[: tensor.shape[0], : tensor.shape[1], : tensor.shape[2]].copy_(tensor)
            padding_mask[: tensor.shape[1], : tensor.shape[2]] = False

        return padded_tensors, padding_masks

    def loss_labels(
        self, class_queries_logits: Tensor, class_labels: List[Tensor], indices: Tuple[np.array]
    ) -> Dict[str, Tensor]:
        

        pred_logits = class_queries_logits
        batch_size, num_queries, _ = pred_logits.shape
        criterion = nn.CrossEntropyLoss(weight=self.empty_weight)
        idx = self._get_predictions_permutation_indices(indices)
        
        target_classes_o = torch.cat([target[j] for target, (_, j) in zip(class_labels, indices)])
        
        target_classes = torch.full(
            (batch_size, num_queries), fill_value=self.num_labels, dtype=torch.int64, device=pred_logits.device
        )
        target_classes[idx] = target_classes_o
        
        pred_logits_transposed = pred_logits.transpose(1, 2)
        loss_ce = criterion(pred_logits_transposed, target_classes)
        losses = {"loss_cross_entropy": loss_ce}
        return losses

    def loss_masks(
        self, masks_queries_logits: Tensor, mask_labels: List[Tensor], indices: Tuple[np.array], num_masks: int
    ) -> Dict[str, Tensor]:
        
        src_idx = self._get_predictions_permutation_indices(indices)
        tgt_idx = self._get_targets_permutation_indices(indices)
        
        pred_masks = masks_queries_logits[src_idx]
        
        
        target_masks, _ = self._pad_images_to_max_in_batch(mask_labels)
        target_masks = target_masks[tgt_idx]
        
        pred_masks = nn.functional.interpolate(
            pred_masks[:, None], size=target_masks.shape[-2:], mode="bilinear", align_corners=False
        )
        pred_masks = pred_masks[:, 0].flatten(1)

        target_masks = target_masks.flatten(1)
        losses = {
            "loss_mask": sigmoid_focal_loss(pred_masks, target_masks, num_masks),
            "loss_dice": dice_loss(pred_masks, target_masks, num_masks),
        }
        return losses

    def _get_predictions_permutation_indices(self, indices):
        
        batch_indices = torch.cat([torch.full_like(src, i) for i, (src, _) in enumerate(indices)])
        predictions_indices = torch.cat([src for (src, _) in indices])
        return batch_indices, predictions_indices

    def _get_targets_permutation_indices(self, indices):
        
        batch_indices = torch.cat([torch.full_like(tgt, i) for i, (_, tgt) in enumerate(indices)])
        target_indices = torch.cat([tgt for (_, tgt) in indices])
        return batch_indices, target_indices

    def forward(
        self,
        masks_queries_logits: Tensor,
        class_queries_logits: Tensor,
        mask_labels: List[Tensor],
        class_labels: List[Tensor],
        auxiliary_predictions: Optional[Dict[str, Tensor]] = None,
    ) -> Dict[str, Tensor]:
        

        
        indices = self.matcher(masks_queries_logits, class_queries_logits, mask_labels, class_labels)
        
        num_masks: Number = self.get_num_masks(class_labels, device=class_labels[0].device)
        
        losses: Dict[str, Tensor] = {
            **self.loss_masks(masks_queries_logits, mask_labels, indices, num_masks),
            **self.loss_labels(class_queries_logits, class_labels, indices),
        }
        
        if auxiliary_predictions is not None:
            for idx, aux_outputs in enumerate(auxiliary_predictions):
                masks_queries_logits = aux_outputs["masks_queries_logits"]
                class_queries_logits = aux_outputs["class_queries_logits"]
                loss_dict = self.forward(masks_queries_logits, class_queries_logits, mask_labels, class_labels)
                loss_dict = {f"{key}_{idx}": value for key, value in loss_dict.items()}
                losses.update(loss_dict)

        return losses

    def get_num_masks(self, class_labels: torch.Tensor, device: torch.device) -> torch.Tensor:
        
        num_masks = sum([len(classes) for classes in class_labels])
        num_masks_pt = torch.as_tensor([num_masks], dtype=torch.float, device=device)
        return num_masks_pt


class MaskFormerFPNConvLayer(nn.Module):
    def __init__(self, in_features: int, out_features: int, kernel_size: int = 3, padding: int = 1):
        
        super().__init__()
        self.layers = [
            nn.Conv2d(in_features, out_features, kernel_size=kernel_size, padding=padding, bias=False),
            nn.GroupNorm(32, out_features),
            nn.ReLU(inplace=True),
        ]
        for i, layer in enumerate(self.layers):
            
            
            
            
            
            
            self.add_module(str(i), layer)

    def forward(self, input: Tensor) -> Tensor:
        hidden_state = input
        for layer in self.layers:
            hidden_state = layer(hidden_state)
        return hidden_state


class MaskFormerFPNLayer(nn.Module):
    def __init__(self, in_features: int, lateral_features: int):
        
        super().__init__()
        self.proj = nn.Sequential(
            nn.Conv2d(lateral_features, in_features, kernel_size=1, padding=0, bias=False),
            nn.GroupNorm(32, in_features),
        )

        self.block = MaskFormerFPNConvLayer(in_features, in_features)

    def forward(self, down: Tensor, left: Tensor) -> Tensor:
        left = self.proj(left)
        down = nn.functional.interpolate(down, size=left.shape[-2:], mode="nearest")
        down += left
        down = self.block(down)
        return down


class MaskFormerFPNModel(nn.Module):
    def __init__(self, in_features: int, lateral_widths: List[int], feature_size: int = 256):
        
        super().__init__()
        self.stem = MaskFormerFPNConvLayer(in_features, feature_size)
        self.layers = nn.Sequential(
            *[MaskFormerFPNLayer(feature_size, lateral_width) for lateral_width in lateral_widths[::-1]]
        )

    def forward(self, features: List[Tensor]) -> List[Tensor]:
        fpn_features = []
        last_feature = features[-1]
        other_features = features[:-1]
        output = self.stem(last_feature)
        for layer, left in zip(self.layers, other_features[::-1]):
            output = layer(output, left)
            fpn_features.append(output)
        return fpn_features


class MaskFormerPixelDecoder(nn.Module):
    def __init__(self, *args, feature_size: int = 256, mask_feature_size: int = 256, **kwargs):
        r
        super().__init__()

        self.fpn = MaskFormerFPNModel(*args, feature_size=feature_size, **kwargs)
        self.mask_projection = nn.Conv2d(feature_size, mask_feature_size, kernel_size=3, padding=1)

    def forward(
        self, features: List[Tensor], output_hidden_states: bool = False, return_dict: bool = True
    ) -> MaskFormerPixelDecoderOutput:
        fpn_features = self.fpn(features)
        
        last_feature_projected = self.mask_projection(fpn_features[-1])

        if not return_dict:
            return (last_feature_projected, tuple(fpn_features)) if output_hidden_states else (last_feature_projected,)

        return MaskFormerPixelDecoderOutput(
            last_hidden_state=last_feature_projected, hidden_states=tuple(fpn_features) if output_hidden_states else ()
        )



class MaskFormerSinePositionEmbedding(nn.Module):
    

    def __init__(
        self, num_pos_feats: int = 64, temperature: int = 10000, normalize: bool = False, scale: Optional[float] = None
    ):
        super().__init__()
        if scale is not None and normalize is False:
            raise ValueError("normalize should be True if scale is passed")
        self.num_pos_feats = num_pos_feats
        self.temperature = temperature
        self.normalize = normalize
        self.scale = 2 * math.pi if scale is None else scale

    def forward(self, x: Tensor, mask: Optional[Tensor] = None) -> Tensor:
        if mask is None:
            mask = torch.zeros((x.size(0), x.size(2), x.size(3)), device=x.device, dtype=torch.bool)
        not_mask = (~mask).to(x.dtype)
        y_embed = not_mask.cumsum(1)
        x_embed = not_mask.cumsum(2)
        if self.normalize:
            eps = 1e-6
            y_embed = y_embed / (y_embed[:, -1:, :] + eps) * self.scale
            x_embed = x_embed / (x_embed[:, :, -1:] + eps) * self.scale

        dim_t = torch.arange(self.num_pos_feats, dtype=x.dtype, device=x.device)
        dim_t = self.temperature ** (2 * torch.div(dim_t, 2, rounding_mode="floor") / self.num_pos_feats)

        pos_x = x_embed[:, :, :, None] / dim_t
        pos_y = y_embed[:, :, :, None] / dim_t
        pos_x = torch.stack((pos_x[:, :, :, 0::2].sin(), pos_x[:, :, :, 1::2].cos()), dim=4).flatten(3)
        pos_y = torch.stack((pos_y[:, :, :, 0::2].sin(), pos_y[:, :, :, 1::2].cos()), dim=4).flatten(3)
        pos = torch.cat((pos_y, pos_x), dim=3).permute(0, 3, 1, 2)
        return pos


class PredictionBlock(nn.Module):
    def __init__(self, in_dim: int, out_dim: int, activation: nn.Module) -> None:
        super().__init__()
        self.layers = [nn.Linear(in_dim, out_dim), activation]
        
        for i, layer in enumerate(self.layers):
            self.add_module(str(i), layer)

    def forward(self, input: Tensor) -> Tensor:
        hidden_state = input
        for layer in self.layers:
            hidden_state = layer(hidden_state)
        return hidden_state


class MaskformerMLPPredictionHead(nn.Module):
    def __init__(self, input_dim: int, hidden_dim: int, output_dim: int, num_layers: int = 3):
        
        super().__init__()
        in_dims = [input_dim] + [hidden_dim] * (num_layers - 1)
        out_dims = [hidden_dim] * (num_layers - 1) + [output_dim]

        self.layers = []
        for i, (in_dim, out_dim) in enumerate(zip(in_dims, out_dims)):
            activation = nn.ReLU() if i < num_layers - 1 else nn.Identity()
            layer = PredictionBlock(in_dim, out_dim, activation=activation)
            self.layers.append(layer)
            
            
            
            
            
            
            self.add_module(str(i), layer)

    def forward(self, input: Tensor) -> Tensor:
        hidden_state = input
        for layer in self.layers:
            hidden_state = layer(hidden_state)
        return hidden_state


class MaskFormerPixelLevelModule(nn.Module):
    def __init__(self, config: MaskFormerConfig):
        
        super().__init__()

        
        backbone_config = config.backbone_config
        if backbone_config.model_type == "swin":
            
            backbone_config = MaskFormerSwinConfig.from_dict(backbone_config.to_dict())
            backbone_config.out_features = ["stage1", "stage2", "stage3", "stage4"]
        self.encoder = AutoBackbone.from_config(backbone_config)

        feature_channels = self.encoder.channels
        self.decoder = MaskFormerPixelDecoder(
            in_features=feature_channels[-1],
            feature_size=config.fpn_feature_size,
            mask_feature_size=config.mask_feature_size,
            lateral_widths=feature_channels[:-1],
        )

    def forward(
        self, pixel_values: Tensor, output_hidden_states: bool = False, return_dict: bool = True
    ) -> MaskFormerPixelLevelModuleOutput:
        features = self.encoder(pixel_values).feature_maps
        decoder_output = self.decoder(features, output_hidden_states, return_dict=return_dict)

        if not return_dict:
            last_hidden_state = decoder_output[0]
            outputs = (features[-1], last_hidden_state)
            if output_hidden_states:
                hidden_states = decoder_output[1]
                outputs = outputs + (tuple(features),) + (hidden_states,)
            return outputs

        return MaskFormerPixelLevelModuleOutput(
            
            encoder_last_hidden_state=features[-1],
            decoder_last_hidden_state=decoder_output.last_hidden_state,
            encoder_hidden_states=tuple(features) if output_hidden_states else (),
            decoder_hidden_states=decoder_output.hidden_states if output_hidden_states else (),
        )


class MaskFormerTransformerModule(nn.Module):
    

    def __init__(self, in_features: int, config: MaskFormerConfig):
        super().__init__()
        hidden_size = config.decoder_config.hidden_size
        should_project = in_features != hidden_size
        self.position_embedder = MaskFormerSinePositionEmbedding(num_pos_feats=hidden_size // 2, normalize=True)
        self.queries_embedder = nn.Embedding(config.decoder_config.num_queries, hidden_size)
        self.input_projection = nn.Conv2d(in_features, hidden_size, kernel_size=1) if should_project else None
        self.decoder = DetrDecoder(config=config.decoder_config)

    def forward(
        self,
        image_features: Tensor,
        output_hidden_states: bool = False,
        output_attentions: bool = False,
        return_dict: Optional[bool] = None,
    ) -> DetrDecoderOutput:
        if self.input_projection is not None:
            image_features = self.input_projection(image_features)
        object_queries = self.position_embedder(image_features)
        
        batch_size = image_features.shape[0]
        queries_embeddings = self.queries_embedder.weight.unsqueeze(0).repeat(batch_size, 1, 1)
        inputs_embeds = torch.zeros_like(queries_embeddings, requires_grad=True)

        batch_size, num_channels, height, width = image_features.shape
        
        image_features = image_features.view(batch_size, num_channels, height * width).permute(0, 2, 1)
        object_queries = object_queries.view(batch_size, num_channels, height * width).permute(0, 2, 1)

        decoder_output: DetrDecoderOutput = self.decoder(
            inputs_embeds=inputs_embeds,
            attention_mask=None,
            encoder_hidden_states=image_features,
            encoder_attention_mask=None,
            object_queries=object_queries,
            query_position_embeddings=queries_embeddings,
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
            return_dict=return_dict,
        )
        return decoder_output


MASKFORMER_START_DOCSTRING = r

MASKFORMER_INPUTS_DOCSTRING = r


class MaskFormerPreTrainedModel(PreTrainedModel):
    config_class = MaskFormerConfig
    base_model_prefix = "model"
    main_input_name = "pixel_values"

    def _init_weights(self, module: nn.Module):
        xavier_std = self.config.init_xavier_std
        std = self.config.init_std
        if isinstance(module, MaskFormerTransformerModule):
            if module.input_projection is not None:
                nn.init.xavier_uniform_(module.input_projection.weight, gain=xavier_std)
                nn.init.constant_(module.input_projection.bias, 0)
        
        elif isinstance(module, MaskFormerFPNModel):
            nn.init.xavier_uniform_(module.stem.get_submodule("0").weight, gain=xavier_std)

        elif isinstance(module, MaskFormerFPNLayer):
            nn.init.xavier_uniform_(module.proj[0].weight, gain=xavier_std)

        elif isinstance(module, MaskFormerFPNConvLayer):
            nn.init.xavier_uniform_(module.get_submodule("0").weight, gain=xavier_std)
        
        elif isinstance(module, MaskformerMLPPredictionHead):
            
            
            for submodule in module.modules():
                if isinstance(submodule, nn.Linear):
                    nn.init.xavier_uniform_(submodule.weight, gain=xavier_std)
                    nn.init.constant_(submodule.bias, 0)
        elif isinstance(module, nn.LayerNorm):
            module.bias.data.zero_()
            module.weight.data.fill_(1.0)
        
        if isinstance(module, (nn.Linear, nn.Conv2d, nn.BatchNorm2d)):
            
            
            module.weight.data.normal_(mean=0.0, std=std)
            if module.bias is not None:
                module.bias.data.zero_()
        elif isinstance(module, nn.Embedding):
            module.weight.data.normal_(mean=0.0, std=std)
            if module.padding_idx is not None:
                module.weight.data[module.padding_idx].zero_()


@add_start_docstrings(
    "The bare MaskFormer Model outputting raw hidden-states without any specific head on top.",
    MASKFORMER_START_DOCSTRING,
)
class MaskFormerModel(MaskFormerPreTrainedModel):
    def __init__(self, config: MaskFormerConfig):
        super().__init__(config)
        self.pixel_level_module = MaskFormerPixelLevelModule(config)
        self.transformer_module = MaskFormerTransformerModule(
            in_features=self.pixel_level_module.encoder.channels[-1], config=config
        )

        self.post_init()

    @add_start_docstrings_to_model_forward(MASKFORMER_INPUTS_DOCSTRING)
    @replace_return_docstrings(output_type=MaskFormerModelOutput, config_class=_CONFIG_FOR_DOC)
    def forward(
        self,
        pixel_values: Tensor,
        pixel_mask: Optional[Tensor] = None,
        output_hidden_states: Optional[bool] = None,
        output_attentions: Optional[bool] = None,
        return_dict: Optional[bool] = None,
    ) -> MaskFormerModelOutput:
        r

        if pixel_values is None:
            raise ValueError("You have to specify pixel_values")

        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
        output_hidden_states = (
            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
        )
        return_dict = return_dict if return_dict is not None else self.config.use_return_dict

        batch_size, _, height, width = pixel_values.shape

        if pixel_mask is None:
            pixel_mask = torch.ones((batch_size, height, width), device=pixel_values.device)

        pixel_level_module_output = self.pixel_level_module(
            pixel_values, output_hidden_states, return_dict=return_dict
        )
        image_features = pixel_level_module_output[0]
        pixel_embeddings = pixel_level_module_output[1]

        transformer_module_output = self.transformer_module(image_features, output_hidden_states, output_attentions)
        queries = transformer_module_output.last_hidden_state

        encoder_hidden_states = None
        pixel_decoder_hidden_states = None
        transformer_decoder_hidden_states = None
        hidden_states = None

        if output_hidden_states:
            encoder_hidden_states = pixel_level_module_output[2]
            pixel_decoder_hidden_states = pixel_level_module_output[3]
            transformer_decoder_hidden_states = transformer_module_output[1]
            hidden_states = encoder_hidden_states + pixel_decoder_hidden_states + transformer_decoder_hidden_states

        output = MaskFormerModelOutput(
            encoder_last_hidden_state=image_features,
            pixel_decoder_last_hidden_state=pixel_embeddings,
            transformer_decoder_last_hidden_state=queries,
            encoder_hidden_states=encoder_hidden_states,
            pixel_decoder_hidden_states=pixel_decoder_hidden_states,
            transformer_decoder_hidden_states=transformer_decoder_hidden_states,
            hidden_states=hidden_states,
            attentions=transformer_module_output.attentions,
        )

        if not return_dict:
            output = tuple(v for v in output.values())

        return output


class MaskFormerForInstanceSegmentation(MaskFormerPreTrainedModel):
    def __init__(self, config: MaskFormerConfig):
        super().__init__(config)
        self.model = MaskFormerModel(config)
        hidden_size = config.decoder_config.hidden_size
        
        self.class_predictor = nn.Linear(hidden_size, config.num_labels + 1)
        self.mask_embedder = MaskformerMLPPredictionHead(hidden_size, hidden_size, config.mask_feature_size)

        self.matcher = MaskFormerHungarianMatcher(
            cost_class=1.0, cost_dice=config.dice_weight, cost_mask=config.mask_weight
        )

        self.weight_dict: Dict[str, float] = {
            "loss_cross_entropy": config.cross_entropy_weight,
            "loss_mask": config.mask_weight,
            "loss_dice": config.dice_weight,
        }

        self.criterion = MaskFormerLoss(
            config.num_labels,
            matcher=self.matcher,
            weight_dict=self.weight_dict,
            eos_coef=config.no_object_weight,
        )

        self.post_init()

    def get_loss_dict(
        self,
        masks_queries_logits: Tensor,
        class_queries_logits: Tensor,
        mask_labels: Tensor,
        class_labels: Tensor,
        auxiliary_logits: Dict[str, Tensor],
    ) -> Dict[str, Tensor]:
        loss_dict: Dict[str, Tensor] = self.criterion(
            masks_queries_logits, class_queries_logits, mask_labels, class_labels, auxiliary_logits
        )
        
        for key, weight in self.weight_dict.items():
            for loss_key, loss in loss_dict.items():
                if key in loss_key:
                    loss *= weight

        return loss_dict

    def get_loss(self, loss_dict: Dict[str, Tensor]) -> Tensor:
        return sum(loss_dict.values())

    def get_logits(self, outputs: MaskFormerModelOutput) -> Tuple[Tensor, Tensor, Dict[str, Tensor]]:
        pixel_embeddings = outputs.pixel_decoder_last_hidden_state
        
        auxiliary_logits: List[str, Tensor] = []
        
        if self.config.use_auxiliary_loss:
            stacked_transformer_decoder_outputs = torch.stack(outputs.transformer_decoder_hidden_states)
            classes = self.class_predictor(stacked_transformer_decoder_outputs)
            class_queries_logits = classes[-1]
            
            mask_embeddings = self.mask_embedder(stacked_transformer_decoder_outputs)

            
            num_embeddings, batch_size, num_queries, num_channels = mask_embeddings.shape
            _, _, height, width = pixel_embeddings.shape
            binaries_masks = torch.zeros(
                (num_embeddings, batch_size, num_queries, height, width), device=mask_embeddings.device
            )
            for c in range(num_channels):
                binaries_masks += mask_embeddings[..., c][..., None, None] * pixel_embeddings[None, :, None, c]

            masks_queries_logits = binaries_masks[-1]
            
            for aux_binary_masks, aux_classes in zip(binaries_masks[:-1], classes[:-1]):
                auxiliary_logits.append(
                    {"masks_queries_logits": aux_binary_masks, "class_queries_logits": aux_classes}
                )

        else:
            transformer_decoder_hidden_states = outputs.transformer_decoder_last_hidden_state
            classes = self.class_predictor(transformer_decoder_hidden_states)
            class_queries_logits = classes
            
            mask_embeddings = self.mask_embedder(transformer_decoder_hidden_states)
            

            
            batch_size, num_queries, num_channels = mask_embeddings.shape
            _, _, height, width = pixel_embeddings.shape
            masks_queries_logits = torch.zeros((batch_size, num_queries, height, width), device=mask_embeddings.device)
            for c in range(num_channels):
                masks_queries_logits += mask_embeddings[..., c][..., None, None] * pixel_embeddings[:, None, c]

        return class_queries_logits, masks_queries_logits, auxiliary_logits

    @add_start_docstrings_to_model_forward(MASKFORMER_INPUTS_DOCSTRING)
    @replace_return_docstrings(output_type=MaskFormerForInstanceSegmentationOutput, config_class=_CONFIG_FOR_DOC)
    def forward(
        self,
        pixel_values: Tensor,
        mask_labels: Optional[List[Tensor]] = None,
        class_labels: Optional[List[Tensor]] = None,
        pixel_mask: Optional[Tensor] = None,
        output_auxiliary_logits: Optional[bool] = None,
        output_hidden_states: Optional[bool] = None,
        output_attentions: Optional[bool] = None,
        return_dict: Optional[bool] = None,
    ) -> MaskFormerForInstanceSegmentationOutput:
        r

        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
        output_hidden_states = (
            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
        )
        return_dict = return_dict if return_dict is not None else self.config.use_return_dict

        raw_outputs = self.model(
            pixel_values,
            pixel_mask,
            output_hidden_states=output_hidden_states or self.config.use_auxiliary_loss,
            return_dict=return_dict,
            output_attentions=output_attentions,
        )
        
        
        outputs = MaskFormerModelOutput(
            encoder_last_hidden_state=raw_outputs[0],
            pixel_decoder_last_hidden_state=raw_outputs[1],
            transformer_decoder_last_hidden_state=raw_outputs[2],
            encoder_hidden_states=raw_outputs[3] if output_hidden_states else None,
            pixel_decoder_hidden_states=raw_outputs[4] if output_hidden_states else None,
            transformer_decoder_hidden_states=raw_outputs[5] if output_hidden_states else None,
            hidden_states=raw_outputs[6] if output_hidden_states else None,
            attentions=raw_outputs[-1] if output_attentions else None,
        )

        loss, loss_dict, auxiliary_logits = None, None, None

        class_queries_logits, masks_queries_logits, auxiliary_logits = self.get_logits(outputs)

        if mask_labels is not None and class_labels is not None:
            loss_dict: Dict[str, Tensor] = self.get_loss_dict(
                masks_queries_logits, class_queries_logits, mask_labels, class_labels, auxiliary_logits
            )
            loss = self.get_loss(loss_dict)

        output_auxiliary_logits = (
            self.config.output_auxiliary_logits if output_auxiliary_logits is None else output_auxiliary_logits
        )
        if not output_auxiliary_logits:
            auxiliary_logits = None

        if not return_dict:
            output = tuple(
                v
                for v in (loss, class_queries_logits, masks_queries_logits, auxiliary_logits, *outputs.values())
                if v is not None
            )
            return output

        return MaskFormerForInstanceSegmentationOutput(
            loss=loss,
            **outputs,
            class_queries_logits=class_queries_logits,
            masks_queries_logits=masks_queries_logits,
            auxiliary_logits=auxiliary_logits,
        )